allow algebraics again

agentlib_mpc/models/casadi_ml_model.py

@@ -282,7 +282,7 @@ def _fixed_during_integration(
             bb_results: The results of the evaluation of the blackbox functions
         """
         all_inputs = self._all_inputs()
-        exclude = [v.name for v in self.differentials + self.outputs]
+        exclude = [v.name for v in self.differentials + self.outputs + self.algebraics]
         # take the mean of start/finish values of variables that have already been
         # integrated by a discrete blackbox function
         if bb_results:
@@ -311,12 +311,16 @@ def _make_integrator(self, ignore_algebraics: bool = False) -> ca.Function:
         par = ca.vertcat(*integration_params.values())
 
         # if we have no differentials and no algebraics, this function should do nothing
-        if (not self.differentials) and (ignore_algebraics or not self.outputs):
+        if (not self.differentials) and (
+            ignore_algebraics or not self.outputs + self.algebraics
+        ):
             return ca.Function("empty", [[], par], [[], []], ["x0", "p"], ["xf", "zf"])
 
         x = ca.vertcat(*[sta.sym for sta in self.differentials])
         # if we have a pure ode, we can use an ode solver which is more efficient
-        if self.differentials and (ignore_algebraics or not self.outputs):
+        if self.differentials and (
+            ignore_algebraics or not self.outputs + self.algebraics
+        ):
             ode = {
                 "x": x,
                 "p": par,
@@ -329,8 +333,8 @@ def _make_integrator(self, ignore_algebraics: bool = False) -> ca.Function:
             "x": x,
             "p": par,
             "ode": self.system,
-            "z": ca.vertcat(*[var.sym for var in self.outputs]),
-            "alg": ca.vertcat(*self.output_equations),
+            "z": ca.vertcat(*[var.sym for var in self.outputs + self.algebraics]),
+            "alg": ca.vertcat(*self.algebraic_equations),
         }
         # if there are no differential values, we create a dummy to make integrator
         # callable
@@ -507,7 +511,7 @@ def _make_unified_predict_function(
         # with keywords names
         differentials_dict = {var.name: var.sym for var in self.differentials}
         if not ignore_algebraics:
-            alg_dict = {var.name: var.sym for var in self.outputs}
+            alg_dict = {var.name: var.sym for var in self.outputs + self.algebraics}
         else:
             alg_dict = {}
         stacked_alg = ca.vertcat(*[mx for mx in alg_dict.values()])

agentlib_mpc/models/casadi_model.py

@@ -166,6 +166,7 @@ class CasadiState(CasadiVariable):
     Class that stores various attributes of CasADi differential variables.
     """
 
+    _alg: Optional[CasadiTypes] = attrs.field(default=None, alias="_alg")
     _ode: Optional[CasadiTypes] = attrs.field(default=None, alias="_ode")
 
     def __attrs_post_init__(self):
@@ -175,22 +176,12 @@ def __attrs_post_init__(self):
 
     @property
     def alg(self) -> CasadiTypes:
-        raise AttributeError(
-            "Casadi States should not have .alg assignments. If you wish to provide "
-            "algebraic relationships to states, add them in the constraints."
-        )
-        return -1
+        return self._alg
 
     @alg.setter
     def alg(self, equation: Union[CasadiTypes, CasadiVariable]):
-        raise AttributeError(
-            "Casadi States should not have .alg assignments. Consider the following: \n"
-            " 1. If you need equality constraints in your MPC, please add them in the "
-            "constraints. \n"
-            " 2. If you use this to bundle an expression, consider using a regular "
-            "Python variable. \n"
-            " 3. Implicit algebraic equations are currently not supported."
-        )
+        self._alg = get_symbolic(equation)
+        self.exclusive_ode_or_alg()
 
     @property
     def ode(self) -> CasadiTypes:
@@ -199,6 +190,7 @@ def ode(self) -> CasadiTypes:
     @ode.setter
     def ode(self, equation: Union[CasadiTypes, CasadiVariable]):
         self._ode = get_symbolic(equation)
+        self.exclusive_ode_or_alg()
 
     def json(self, indent: int = 2, **kwargs):
         data = self.dict(**kwargs)
@@ -208,6 +200,18 @@ def json(self, indent: int = 2, **kwargs):
         data.pop("_alg")
         json.dumps(data, indent=indent)
 
+    def exclusive_ode_or_alg(self):
+        """Ensures a CasadiState cannot have both a differential and an
+        algebraic equation."""
+        if (self.ode is not None) and (self.alg is not None):
+            raise ValueError(
+                f"CasadiStates can be exclusively differential or algebraic "
+                f"variables. However, CasadiState '{self.name}' was "
+                f"assigned both a differential equation and an algebraic "
+                f"equation."
+            )
+        return self
+
 
 @attrs.define(slots=True, weakref_slot=False, kw_only=True)
 class CasadiInput(CasadiVariable):
@@ -322,6 +326,13 @@ def _assert_outputs_are_defined(self):
                     f"sure you specify '{out.name}.alg' in 'setup_system()'."
                 )
 
+    def _setup_algebraic_equations(self) -> list[ca.MX]:
+        """Collects the algebraic equations defined in setup_system and also adds them
+        to constraints."""
+        equality_constraints = [(0, alg, 0) for alg in self.algebraic_equations]
+        self.constraints = self.constraints + equality_constraints
+        return self.constraints + equality_constraints
+
     def do_step(self, *, t_start, t_sample=None):
         if t_sample is None:
             t_sample = self.dt
@@ -337,8 +348,8 @@ def do_step(self, *, t_start, t_sample=None):
             self.set_differential_values(np.array(result["xf"]).flatten())
         else:
             result = self.integrator(p=pars)
-        if self.outputs:
-            self.set_output_values(np.array(result["zf"]).flatten())
+        if self.algebraics + self.outputs:
+            self.set_algebraic_values(np.array(result["zf"]).flatten())
 
     def _make_integrator(self) -> ca.Function:
         """Creates the integrator to be used in do_step(). The integrator takes the
@@ -349,8 +360,8 @@ def _make_integrator(self) -> ca.Function:
             *[inp.sym for inp in chain.from_iterable([self.inputs, self.parameters])]
         )
         x = ca.vertcat(*[sta.sym for sta in self.differentials])
-        z = ca.vertcat(*[var.sym for var in self.outputs])
-        algebraic_equations = ca.vertcat(*self.output_equations)
+        z = ca.vertcat(*[var.sym for var in self.outputs + self.algebraics])
+        algebraic_equations = ca.vertcat(*self.algebraic_equations)
 
         if not algebraic_equations.shape[0] and self.differentials:
             # case of pure ode
@@ -398,7 +409,7 @@ def get_constraints(self) -> List[ModelConstraint]:
             for lb, func, ub in self.constraints
         ]
         equality_constraints = [
-            ModelConstraint(0, alg, 0) for alg in self.output_equations
+            ModelConstraint(0, alg, 0) for alg in self.algebraic_equations
         ]
         return base_constraints + equality_constraints
 
@@ -423,10 +434,15 @@ def parameters(self) -> list[CasadiParameter]:
         return list(self._parameters.values())
 
     @property
-    def output_equations(self) -> List[CasadiTypes]:
+    def algebraic_equations(self) -> List[CasadiTypes]:
         """List of algebraic equations RHS in the form
         0 = z - g(x, z, p, ... )"""
-        return [alg_var - alg_var.alg for alg_var in self.outputs]
+        return [alg_var - alg_var.alg for alg_var in self.algebraics + self.outputs]
+
+    @property
+    def algebraics(self) -> List[CasadiState]:
+        """List of all CasadiStates with an associated algebraic equation."""
+        return [var for var in self.states if var.alg is not None]
 
     @property
     def differentials(self) -> List[CasadiState]:
@@ -438,7 +454,7 @@ def auxiliaries(self) -> List[CasadiState]:
         """List of all CasadiStates without an associated equation. Common
         uses for this are slack variables that appear in cost functions and
         constraints of optimization models."""
-        return [var for var in self.states if var.ode is None]
+        return [var for var in self.states if (var.alg is None and var.ode is None)]
 
     @abc.abstractmethod
     def setup_system(self):
@@ -460,11 +476,18 @@ def set_differential_values(self, values: Union[List, np.ndarray]):
         for state, value in zip(self.differentials, values):
             self._states[state.name].value = value
 
-    def set_output_values(self, values: Union[List, np.ndarray]):
-        """Sets the values for all outputs. Provided values list MUST match the order
-        in which outputs are saved, there is no check."""
-        for var, value in zip(self.outputs, values):
-            self._outputs[var.name].value = value
+    def set_algebraic_values(self, values: Union[List, np.ndarray]):
+        """
+        Sets the values for all algebraic states and outputs. Provided
+        values list MUST match the correct order.
+        Correct order is:
+        All variables in self.algebraics, then all variables in self.outputs
+        """
+        for var, value in zip(self.algebraics + self.outputs, values):
+            try:
+                self._outputs[var.name].value = value
+            except KeyError:
+                self._states[var.name].value = value
 
     def get(self, name: str) -> CasadiVariable:
         return super().get(name)

agentlib_mpc/optimization_backends/casadi_/basic.py

@@ -58,7 +58,7 @@ def initialize(self, model: CasadiModel, var_ref: VariableReference):
         )
         self.algebraics = OptimizationVariable.declare(
             denotation="z",
-            variables=model.auxiliaries,
+            variables=model.algebraics + model.auxiliaries,
             ref_list=[],
         )
         self.outputs = OptimizationVariable.declare(

agentlib_mpc/optimization_backends/casadi_/casadi_admm_ml.py

@@ -59,7 +59,7 @@ def initialize(
         )
         self.algebraics = OptimizationVariable.declare(
             denotation="z",
-            variables=model.auxiliaries,
+            variables=model.auxiliaries + model.algebraics,
             ref_list=[],
         )
         self.outputs = OptimizationVariable.declare(

agentlib_mpc/optimization_backends/casadi_/casadi_ml.py

@@ -51,7 +51,7 @@ def initialize(self, model: CasadiMLModel, var_ref: FullVariableReference):
         )
         self.algebraics = OptimizationVariable.declare(
             denotation="z",
-            variables=model.auxiliaries,
+            variables=model.auxiliaries + model.algebraics,
             ref_list=[],
         )
         self.outputs = OptimizationVariable.declare(
@@ -147,11 +147,6 @@ def _discretize(self, sys: CasadiMLSystem):
                 sys.states, lb=x_past, ub=x_past, guess=x_past
             )
             mx_dict[time][sys.initial_state.name] = x_past
-            y_past = self.add_opt_par(sys.initial_output)
-            mx_dict[time][sys.outputs.name] = self.add_opt_var(
-                sys.outputs, lb=y_past, ub=y_past, guess=y_past
-            )
-            mx_dict[time][sys.initial_output.name] = y_past
 
         # add past inputs
         for time in pre_grid_inputs:

agentlib_mpc/optimization_backends/casadi_/mhe.py

@@ -70,7 +70,7 @@ def initialize(self, model: CasadiModel, var_ref: MHEVariableReference):
         )
         self.algebraics = OptimizationVariable.declare(
             denotation="algebraics",
-            variables=model.auxiliaries,
+            variables=model.auxiliaries + model.algebraics,
             ref_list=[],
         )
         self.outputs = OptimizationVariable.declare(