API Reference: DynSpec Engine¶

class FormulaEvaluator(Interpreter):
    """
    An interpreter that evaluates mathematical formulas as defined by the grammar.

    This class can evaluate:
    - Basic arithmetic operations (add, sub, mul, div, pow, neg)
    - Numbers and symbols (constants, values, variables)
    - Function calls
    - Assignments and equations
    """

    def __init__(
        self,
        context: Dict[str, Any] = {},
        function_table: Dict[str, Callable] = function_table_0,
        unknown_as_nan=True,
        raise_on_nan=False,
        steady_state: bool = False,
    ):
        """
        Initialize the evaluator.

        Args:
            symbol_table: Dictionary mapping symbol names to their values
            function_table: Dictionary mapping function names to callable functions
            steady_state: If True, evaluates variables at their steady state (only the name of the symbol is taken into account)
            raise_on_nan: If True, raise a DefinitionError as soon as any node evaluates to NaN,
                pinpointing the offending subtree instead of letting NaN propagate silently.
        """
        super().__init__()

        self.function_table: Dict[str, Callable[..., Any]] = dict(function_table or {})
        self.unknown_as_nan = unknown_as_nan
        self.raise_on_nan = raise_on_nan
        self.steady_state = steady_state

        self.constants = context.get("constants", {})
        self.processes = context.get("processes", {})
        self.values = context.get("values", {})
        self.variables = context.get("variables", {})
        self.steady_states = context.get("steady_states", {})
        self.metadata = context.get("metadata", {})

        self.time = None  # None or integer
        self.errors: List[Any] = []

        # Add default mathematical functions
        from .autodiff import MATH_FUNCTIONS

        self.function_table.update(MATH_FUNCTIONS)
        self.function_table.update({"N": _normal_distribution})

    def _undefined(self, message: str, tree):
        """Report an undefined value: raise if `unknown_as_nan` is False, else return NaN."""
        if not self.unknown_as_nan:
            raise DefinitionError(message, tree=tree)
        err = DefinitionError(message, tree=tree)
        self.errors.append(err)
        import warnings
        from dyno.errors import UndefinedSymbolWarning

        if getattr(tree, "data", None) == "constant":
            try:
                loc = f" at ({tree.meta.line}, {tree.meta.column})"
            except Exception:
                loc = ""
            warnings.warn(
                f"{message}{loc}. Defaulting to NaN.",
                UndefinedSymbolWarning,
                stacklevel=3,
            )
        return math.nan

    def visit(self, tree):
        """Visit a node, optionally raising as soon as its result is NaN.

        Centralizing the check here (rather than in every node evaluator) means it
        applies uniformly and reports the innermost subtree where the NaN first appears.
        """
        result = super().visit(tree)
        if self.raise_on_nan and isinstance(result, float) and math.isnan(result):
            raise DefinitionError(
                f"NaN encountered while evaluating `{tree.data}`", tree=tree
            )
        return result

    # Arithmetic operations
    def add(self, tree):
        """Handle addition: a + b"""
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        return left + right

    def sub(self, tree):
        """Handle subtraction: a - b"""
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        return left - right

    def mul(self, tree):
        """Handle multiplication: a * b"""
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        return left * right

    def div(self, tree):
        """Handle division: a / b"""
        left = self.visit(tree.children[0])
        right = self.visit(tree.children[1])
        # TODO: maybe add a safe division?
        # if right == 0:
        #     raise ZeroDivisionError("Division by zero")
        return left / right

    def pow(self, tree):
        """Handle exponentiation: a ^ b or a ** b"""
        base = self.visit(tree.children[0])
        exponent = self.visit(tree.children[1])
        return base**exponent

    def neg(self, tree):
        """Handle negation: -a"""
        value = self.visit(tree.children[0])
        return -value

    # Numbers and literals
    def number(self, tree):
        """Handle numeric literals"""
        value = tree.children[0].value
        # Try to parse as int first, then float
        try:
            return int(value)
        except ValueError:
            return float(value)

    # Symbols
    def constant(self, tree):
        """Handle constants (symbols without time indexing)"""
        name = str(tree.children[0].children[0])
        if name in self.constants:
            return self.constants[name]
        else:
            return self._undefined(f"Undefined value: {name}", tree)

    def value(self, tree):
        """Handle values with specific time: name[time]"""
        name = str(tree.children[0].children[0])
        time = int(tree.children[1].children[0])

        # Create a key for the symbol table
        if self.steady_state:
            if name not in self.steady_states:
                return self._undefined(
                    f"Undefined steady state for value {name}[~]", tree
                )
            return self.steady_states[name]
        else:
            if name not in self.values:
                return self._undefined(f"Undefined value {name}[~]", tree)
            else:
                vvs = self.values[name]
                if time not in vvs:
                    return self._undefined(f"Undefined value {name}[{time}]", tree)
                else:
                    return vvs[time]

    def variable(self, tree):
        """Handle variables with time indexing: name[t+shift]"""
        name = str(tree.children[0].children[0])
        index = str(tree.children[1].children[0])  # Usually 't'
        shift = int(tree.children[2].children[0])

        # TODO deal with index ~ #### this should disappear
        if name not in self.variables:
            self.variables[name] = {}

        if self.time is not None:
            time = self.time + shift
            vvs = self.values.get(name, {})
            if time not in vvs:
                return self._undefined(f"Undefined value {name}[{time}]", tree)
            return vvs[time]
        elif self.steady_state or (index == "~"):
            # from rich import print
            if name not in self.steady_states:
                return self._undefined(
                    f"Undefined steady state for variable {name}[~]", tree
                )
            return self.steady_states[name]
        else:
            if shift not in self.variables[name]:
                return self._undefined(f"Undefined variable {name}[{shift}]", tree)
            return self.variables[name][shift]

    # Function calls
    def call(self, tree):
        """Handle function calls: func_name(arg)"""
        func_name = str(tree.children[0].children[0])

        if func_name in self.function_table:
            args = [self.visit(c) for c in tree.children[1:]]
            return self.function_table[func_name](*args)
        elif func_name == "steady_state":
            arg = tree.children[1]
            if getattr(arg, "data", None) == "variable":
                name = str(arg.children[0].children[0])
                if name not in self.steady_states:
                    return self._undefined(
                        f"Undefined steady state for variable {name}[~]", tree
                    )
                return self.steady_states[name]
            else:
                prev_ss = self.steady_state
                try:
                    self.steady_state = True
                    return self.visit(arg)
                finally:
                    self.steady_state = prev_ss
        else:
            raise ValueError(f"Undefined function: {func_name}")

    def bare_formula(self, tree):
        """Handle a standalone formula equation (implicitly `formula = 0`)"""
        return self.visit(tree.children[0])

class DNumber:

    def __init__(self, value, derivatives=None):
        self.value = value
        self.derivatives = derivatives if derivatives is not None else {}

    def __add__(self, other):
        if isinstance(other, DNumber):
            new_value = self.value + other.value
            new_derivatives = self.derivatives.copy()
            for var, deriv in other.derivatives.items():
                if var in new_derivatives:
                    new_derivatives[var] += deriv
                else:
                    new_derivatives[var] = deriv
            return DNumber(new_value, new_derivatives)
        else:
            return DNumber(self.value + other, self.derivatives.copy())

    def __radd__(self, other):
        return self.__add__(other)

    def __sub__(self, other):
        if isinstance(other, DNumber):
            new_value = self.value - other.value
            new_derivatives = self.derivatives.copy()
            for var, deriv in other.derivatives.items():
                if var in new_derivatives:
                    new_derivatives[var] -= deriv
                else:
                    new_derivatives[var] = -deriv
            return DNumber(new_value, new_derivatives)
        else:
            return DNumber(self.value - other, self.derivatives.copy())

    def __rsub__(self, other):
        if isinstance(other, DNumber):
            return other.__sub__(self)
        else:
            new_derivatives = {var: -deriv for var, deriv in self.derivatives.items()}
            return DNumber(other - self.value, new_derivatives)

    def __mul__(self, other):
        if isinstance(other, DNumber):
            new_value = self.value * other.value
            new_derivatives = {}
            for var in set(self.derivatives.keys()).union(other.derivatives.keys()):
                deriv1 = self.derivatives.get(var, 0)
                deriv2 = other.derivatives.get(var, 0)
                new_derivatives[var] = deriv1 * other.value + deriv2 * self.value
            return DNumber(new_value, new_derivatives)
        else:
            new_derivatives = {
                var: deriv * other for var, deriv in self.derivatives.items()
            }
            return DNumber(self.value * other, new_derivatives)

    def __rmul__(self, other):
        return self.__mul__(other)

    def __truediv__(self, other):
        if isinstance(other, DNumber):
            new_value = self.value / other.value
            new_derivatives = {}
            for var in set(self.derivatives.keys()).union(other.derivatives.keys()):
                deriv1 = self.derivatives.get(var, 0)
                deriv2 = other.derivatives.get(var, 0)
                new_derivatives[var] = (deriv1 * other.value - deriv2 * self.value) / (
                    other.value**2
                )
            return DNumber(new_value, new_derivatives)
        else:
            new_derivatives = {
                var: deriv / other for var, deriv in self.derivatives.items()
            }
            return DNumber(self.value / other, new_derivatives)

    def __rtruediv__(self, other):

        if isinstance(other, DNumber):
            return other.__truediv__(self)
        else:
            new_derivatives = {}
            for var, deriv in self.derivatives.items():
                new_derivatives[var] = -deriv * other / (self.value**2)
            return DNumber(other / self.value, new_derivatives)

    def __pow__(self, power):

        if isinstance(power, DNumber):
            new_value = self.value**power.value
            new_derivatives = {}
            m = _get_math_module(self.value)
            for var in set(self.derivatives.keys()).union(power.derivatives.keys()):
                deriv1 = self.derivatives.get(var, 0)
                deriv2 = power.derivatives.get(var, 0)
                new_derivatives[var] = new_value * (
                    deriv1 * power.value / self.value + deriv2 * m.log(self.value)
                )
            return DNumber(new_value, new_derivatives)
        else:
            new_value = self.value**power
            new_derivatives = {
                var: deriv * power * (self.value ** (power - 1))
                for var, deriv in self.derivatives.items()
            }
            return DNumber(new_value, new_derivatives)

    def __rpow__(self, base):
        if isinstance(base, DNumber):
            return base.__pow__(self)
        else:
            new_value = base**self.value
            m = _get_math_module(self.value)
            new_derivatives = {
                var: deriv * new_value * m.log(base)
                for var, deriv in self.derivatives.items()
            }
            return DNumber(new_value, new_derivatives)

    def __neg__(self):
        new_derivatives = {var: -deriv for var, deriv in self.derivatives.items()}
        return DNumber(-self.value, new_derivatives)

    def __repr__(self):
        return f"DNumber(value={self.value}, derivatives={self.derivatives})"

@dataclass
class Recipe:
    """A recipe defining the structure of a model type.

    Attributes:
        name: Recipe name (e.g. ``"dtcc"``).
        variable_groups: Ordered list of variable group names.
        equation_groups: Specifications for each equation group.
    """

    name: str
    variable_groups: list[str] = field(default_factory=list)
    equation_groups: list[EquationGroupSpec] = field(default_factory=list)

    def get_group_spec(self, name: str) -> EquationGroupSpec | None:
        """Return the equation group spec with the given name, or ``None``."""
        for spec in self.equation_groups:
            if spec.name == name:
                return spec
        return None

@dataclass
class EquationGroupSpec:
    """Specification for one equation group within a recipe.

    Attributes:
        name: Equation group name (e.g. ``"transition"``, ``"arbitrage"``).
        allowed: Which ``(variable_group, shift)`` pairs are permitted.
        target: If set, the variable group that this block *defines*
            (e.g. ``"states"`` for transitions).
        recursive: If ``True``, the equations must form a DAG when
            interpreted as assignments to target variables.
        n_equations: If set, the variable group whose size must equal
            the number of equations (e.g. ``"states"``).
        optional: If ``True``, the equation group may be absent from the model.
    """

    name: str
    allowed: list[VariableSpec] = field(default_factory=list)
    target: str | None = None
    recursive: bool = True
    n_equations: str | None = None
    optional: bool = False

@dataclass(frozen=True)
class VariableSpec:
    """A specification for a variable occurrence: which group and at which time shift."""

    group: str
    shift: int

    def __repr__(self) -> str:
        sign = "+" if self.shift > 0 else ""
        return (
            f"{self.group}[t{sign}{self.shift}]" if self.shift else f"{self.group}[t]"
        )

@dataclass
class ConformityResult:
    """Result of checking an equation group against a recipe spec."""

    group_name: str
    ok: bool
    violations: list[Violation] = field(default_factory=list)
    dag_order: list[str] | None = None

    def __str__(self) -> str:
        if self.ok:
            msg = f"✓ {self.group_name}: OK"
            if self.dag_order is not None:
                msg += f" (order: {' → '.join(self.dag_order)})"
            return msg
        lines = [f"✗ {self.group_name}: {len(self.violations)} violation(s)"]
        for v in self.violations:
            lines.append(f"  - {v}")
        return "\n".join(lines)

    def __bool__(self) -> bool:
        return self.ok

@dataclass
class Violation:
    """A single conformity violation."""

    equation_index: int
    variable_name: str
    shift: int
    message: str

    def __str__(self) -> str:
        sign = "+" if self.shift > 0 else ""
        shift_str = f"[t{sign}{self.shift}]" if self.shift else "[t]"
        return f"Equation {self.equation_index}: {self.variable_name}{shift_str} — {self.message}"