diff --git a/agentlib_mpc/models/casadi_ml_model.py b/agentlib_mpc/models/casadi_ml_model.py index 81ef136..fbcd613 100644 --- a/agentlib_mpc/models/casadi_ml_model.py +++ b/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()]) diff --git a/agentlib_mpc/models/casadi_model.py b/agentlib_mpc/models/casadi_model.py index 6ac7048..f425986 100644 --- a/agentlib_mpc/models/casadi_model.py +++ b/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) diff --git a/agentlib_mpc/optimization_backends/casadi_/basic.py b/agentlib_mpc/optimization_backends/casadi_/basic.py index 8710395..0f45252 100644 --- a/agentlib_mpc/optimization_backends/casadi_/basic.py +++ b/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( diff --git a/agentlib_mpc/optimization_backends/casadi_/casadi_admm_ml.py b/agentlib_mpc/optimization_backends/casadi_/casadi_admm_ml.py index a99ed22..1f7122a 100644 --- a/agentlib_mpc/optimization_backends/casadi_/casadi_admm_ml.py +++ b/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( diff --git a/agentlib_mpc/optimization_backends/casadi_/casadi_ml.py b/agentlib_mpc/optimization_backends/casadi_/casadi_ml.py index 6921501..4069bcf 100644 --- a/agentlib_mpc/optimization_backends/casadi_/casadi_ml.py +++ b/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: diff --git a/agentlib_mpc/optimization_backends/casadi_/mhe.py b/agentlib_mpc/optimization_backends/casadi_/mhe.py index ef88e9d..323e3a6 100644 --- a/agentlib_mpc/optimization_backends/casadi_/mhe.py +++ b/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(