Add initial equation support.

Signed-off-by: James Goppert <james.goppert@gmail.com>

Author: jgoppert

cyecca/dsl/ROADMAP.md

@@ -137,7 +137,28 @@
     submodel,
     var,
 )
-from cyecca.dsl.operators import abs, acos, asin, atan, atan2, cos, exp, log, sin, sqrt, tan
+from cyecca.dsl.operators import (
+    abs,
+    acos,
+    asin,
+    atan,
+    atan2,
+    ceil,
+    cos,
+    cosh,
+    exp,
+    floor,
+    log,
+    log10,
+    max,
+    min,
+    sign,
+    sin,
+    sinh,
+    sqrt,
+    tan,
+    tanh,
+)
 from cyecca.dsl.simulation import SimulationResult, Simulator
 from cyecca.dsl.types import DType, Indices, Shape, Var, VarKind
 
@@ -193,5 +214,14 @@
     "sqrt",
     "exp",
     "log",
+    "log10",
     "abs",
+    "sign",
+    "floor",
+    "ceil",
+    "sinh",
+    "cosh",
+    "tanh",
+    "min",
+    "max",
 ]

cyecca/dsl/__init__.py

@@ -197,9 +197,36 @@ def expr_to_casadi(expr: Expr) -> ca.SX:
             elif expr.kind == ExprKind.LOG:
                 return ca.log(expr_to_casadi(expr.children[0]))
 
+            elif expr.kind == ExprKind.LOG10:
+                return ca.log10(expr_to_casadi(expr.children[0]))
+
             elif expr.kind == ExprKind.ABS:
                 return ca.fabs(expr_to_casadi(expr.children[0]))
 
+            elif expr.kind == ExprKind.SIGN:
+                return ca.sign(expr_to_casadi(expr.children[0]))
+
+            elif expr.kind == ExprKind.FLOOR:
+                return ca.floor(expr_to_casadi(expr.children[0]))
+
+            elif expr.kind == ExprKind.CEIL:
+                return ca.ceil(expr_to_casadi(expr.children[0]))
+
+            elif expr.kind == ExprKind.SINH:
+                return ca.sinh(expr_to_casadi(expr.children[0]))
+
+            elif expr.kind == ExprKind.COSH:
+                return ca.cosh(expr_to_casadi(expr.children[0]))
+
+            elif expr.kind == ExprKind.TANH:
+                return ca.tanh(expr_to_casadi(expr.children[0]))
+
+            elif expr.kind == ExprKind.MIN:
+                return ca.fmin(expr_to_casadi(expr.children[0]), expr_to_casadi(expr.children[1]))
+
+            elif expr.kind == ExprKind.MAX:
+                return ca.fmax(expr_to_casadi(expr.children[0]), expr_to_casadi(expr.children[1]))
+
             # Relational operators
             elif expr.kind == ExprKind.LT:
                 return expr_to_casadi(expr.children[0]) < expr_to_casadi(expr.children[1])
@@ -451,9 +478,36 @@ def expr_to_casadi_mx(expr: Expr) -> ca.MX:
             elif expr.kind == ExprKind.LOG:
                 return ca.log(expr_to_casadi_mx(expr.children[0]))
 
+            elif expr.kind == ExprKind.LOG10:
+                return ca.log10(expr_to_casadi_mx(expr.children[0]))
+
             elif expr.kind == ExprKind.ABS:
                 return ca.fabs(expr_to_casadi_mx(expr.children[0]))
 
+            elif expr.kind == ExprKind.SIGN:
+                return ca.sign(expr_to_casadi_mx(expr.children[0]))
+
+            elif expr.kind == ExprKind.FLOOR:
+                return ca.floor(expr_to_casadi_mx(expr.children[0]))
+
+            elif expr.kind == ExprKind.CEIL:
+                return ca.ceil(expr_to_casadi_mx(expr.children[0]))
+
+            elif expr.kind == ExprKind.SINH:
+                return ca.sinh(expr_to_casadi_mx(expr.children[0]))
+
+            elif expr.kind == ExprKind.COSH:
+                return ca.cosh(expr_to_casadi_mx(expr.children[0]))
+
+            elif expr.kind == ExprKind.TANH:
+                return ca.tanh(expr_to_casadi_mx(expr.children[0]))
+
+            elif expr.kind == ExprKind.MIN:
+                return ca.fmin(expr_to_casadi_mx(expr.children[0]), expr_to_casadi_mx(expr.children[1]))
+
+            elif expr.kind == ExprKind.MAX:
+                return ca.fmax(expr_to_casadi_mx(expr.children[0]), expr_to_casadi_mx(expr.children[1]))
+
             # Relational operators
             elif expr.kind == ExprKind.LT:
                 return expr_to_casadi_mx(expr.children[0]) < expr_to_casadi_mx(expr.children[1])

cyecca/dsl/backends/casadi.py

@@ -136,7 +136,16 @@ class ExprKind(Enum):
     SQRT = auto()
     EXP = auto()
     LOG = auto()
+    LOG10 = auto()  # Base-10 logarithm
     ABS = auto()
+    SIGN = auto()  # Sign function (-1, 0, or 1)
+    FLOOR = auto()  # Floor function
+    CEIL = auto()  # Ceiling function
+    SINH = auto()  # Hyperbolic sine
+    COSH = auto()  # Hyperbolic cosine
+    TANH = auto()  # Hyperbolic tangent
+    MIN = auto()  # Minimum of two values
+    MAX = auto()  # Maximum of two values
 
 
 @dataclass(frozen=True)
@@ -197,11 +206,22 @@ def __repr__(self) -> str:
             ExprKind.SQRT,
             ExprKind.EXP,
             ExprKind.LOG,
+            ExprKind.LOG10,
             ExprKind.ABS,
+            ExprKind.SIGN,
+            ExprKind.FLOOR,
+            ExprKind.CEIL,
+            ExprKind.SINH,
+            ExprKind.COSH,
+            ExprKind.TANH,
         ):
             return f"{self.kind.name.lower()}({self.children[0]})"
         elif self.kind == ExprKind.ATAN2:
             return f"atan2({self.children[0]}, {self.children[1]})"
+        elif self.kind == ExprKind.MIN:
+            return f"min({self.children[0]}, {self.children[1]})"
+        elif self.kind == ExprKind.MAX:
+            return f"max({self.children[0]}, {self.children[1]})"
         elif self.kind == ExprKind.PRE:
             return f"pre({self.name})"
         elif self.kind == ExprKind.EDGE:
@@ -1683,6 +1703,46 @@ def algorithm(m) -> Generator[Assignment, None, None]:
         return
         yield  # Make this a generator
 
+    def initial_equations(self) -> Generator[Equation, None, None]:
+        """
+        Override this method to define initial equations.
+
+        Initial equations (Modelica: `initial equation` section) are used to
+        specify initial conditions for simulation. They are solved once at
+        t=0 to determine initial values of states and algebraic variables.
+
+        This provides more flexibility than just using `start` values:
+        - Can specify relationships between initial values
+        - Can use equations rather than just fixed values
+        - Can leave some variables to be computed from others
+
+        Example
+        -------
+        >>> @model
+        ... class Pendulum:
+        ...     theta = var()
+        ...     omega = var()
+        ...
+        ...     def equations(m):
+        ...         yield der(m.theta) == m.omega
+        ...         yield der(m.omega) == -9.81 * sin(m.theta)
+        ...
+        ...     def initial_equations(m):
+        ...         yield m.theta == 0.5  # Initial angle
+        ...         yield m.omega == 0.0  # Start at rest
+
+        Notes
+        -----
+        Modelica Spec: Section 8.6 - Initialization, Initial Equation, and Initial Algorithm
+
+        In Modelica, initial equations form a system that is solved to find
+        consistent initial values. Variables with `fixed=True` use their `start`
+        values as fixed constraints. Other variables are determined by the
+        initial equation system.
+        """
+        return
+        yield  # Make this a generator
+
     def flatten(self, expand_arrays: bool = True) -> "FlatModel":
         """
         Flatten the model into a backend-agnostic representation.
@@ -1763,6 +1823,34 @@ def flatten(self, expand_arrays: bool = True) -> "FlatModel":
             else:
                 raise TypeError(f"Expected Assignment in algorithm(), got {type(assign)}")
 
+        # Collect initial equations
+        raw_initial_equations = self.initial_equations()
+        initial_equations_list: List[Equation] = []
+
+        for eq in raw_initial_equations:
+            if isinstance(eq, Equation):
+                initial_equations_list.append(eq)
+            elif isinstance(eq, ArrayEquation):
+                if expand_arrays:
+                    initial_equations_list.extend(eq.expand())
+                # For non-expanded case, we could add array initial equations
+                # but for now just expand them
+                else:
+                    initial_equations_list.extend(eq.expand())
+            else:
+                raise TypeError(f"Expected Equation in initial_equations(), got {type(eq)}")
+
+        # Collect initial equations from submodels
+        for sub_name, sub_instance in self._submodels.items():
+            for eq in sub_instance.initial_equations():
+                if isinstance(eq, Equation):
+                    prefixed_eq = eq._prefix_names(sub_name)
+                    initial_equations_list.append(prefixed_eq)
+                elif isinstance(eq, ArrayEquation):
+                    for scalar_eq in eq.expand():
+                        prefixed_eq = scalar_eq._prefix_names(sub_name)
+                        initial_equations_list.append(prefixed_eq)
+
         # Find all derivatives (der(x)) used in equations to identify states
         derivatives_used: set[str] = set()
         for eq in equations:
@@ -1991,6 +2079,7 @@ def flatten(self, expand_arrays: bool = True) -> "FlatModel":
             input_defaults=input_defaults,
             discrete_defaults=discrete_defaults,
             param_defaults=param_defaults,
+            initial_equations=initial_equations_list,
             algorithm_assignments=algorithm_assignments,
             algorithm_locals=algorithm_locals,
             expand_arrays=expand_arrays,
@@ -2046,6 +2135,7 @@ def model(cls: Type[Any]) -> Type[Any]:
     # non-conformant. Output variables are just vars with output=True flag.
     original_equations = getattr(cls, "equations", None)
     original_algorithm = getattr(cls, "algorithm", None)
+    original_initial_equations = getattr(cls, "initial_equations", None)
 
     # Deprecation check: warn if user defines output_equations (non-Modelica)
     if hasattr(cls, "output_equations"):
@@ -2077,6 +2167,10 @@ def algorithm(self) -> Generator[Assignment, None, None]:
             if original_algorithm is not None:
                 yield from original_algorithm(self)
 
+        def initial_equations(self) -> Generator[Equation, None, None]:
+            if original_initial_equations is not None:
+                yield from original_initial_equations(self)
+
     # Copy the metadata reference
     ModelClass._dsl_metadata = metadata
 
@@ -2396,6 +2490,10 @@ class FlatModel:
     discrete_defaults: Dict[str, Any]
     param_defaults: Dict[str, Any]
 
+    # Initial equations (Modelica: initial equation section)
+    # These are solved once at t=0 to determine initial values
+    initial_equations: List[Equation] = field(default_factory=list)
+
     # Array derivative equations (when expand_arrays=False)
     # For CasADi MX backend: keeps array structure for efficient matrix operations
     # Key is base variable name (e.g., 'pos'), value is {'shape': (3,), 'rhs': SymbolicVar}