[RTR/RTM] Algebraic states are states

bioptim/dynamics/configure_new_variable.py

@@ -1,5 +1,5 @@
-from casadi import MX, SX, vertcat
 import numpy as np
+from casadi import MX, SX, vertcat
 
 from .fatigue.fatigue_dynamics import FatigueList, MultiFatigueInterface
 from ..gui.plot import CustomPlot
@@ -291,7 +291,8 @@ def _declare_cx_and_plot(self):
             for node_index in range(
                 self.nlp.n_controls_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
             ):
-                cx_scaled = self.define_cx_scaled(n_col=3, node_index=node_index)
+                n_cx = 3
+                cx_scaled = self.define_cx_scaled(n_col=n_cx, node_index=node_index)
                 cx = self.define_cx_unscaled(cx_scaled, self.nlp.u_scaling[self.name].scaling)
                 self.nlp.controls.append(
                     self.name,
@@ -338,10 +339,9 @@ def _declare_cx_and_plot(self):
             for node_index in range(
                 self.nlp.n_states_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
             ):
-                n_cx = 2
+                n_cx = self.nlp.ode_solver.n_required_cx + 2
                 cx_scaled = self.define_cx_scaled(n_col=n_cx, node_index=node_index)
                 cx = self.define_cx_unscaled(cx_scaled, self.nlp.a_scaling[self.name].scaling)
-
                 self.nlp.algebraic_states.append(
                     self.name,
                     cx,
@@ -350,18 +350,16 @@ def _declare_cx_and_plot(self):
                     node_index,
                 )
                 if not self.skip_plot:
-                    all_variables_in_one_subplot = True if self.name in ["m", "cov", "k"] else False
                     self.nlp.plot[f"{self.name}_algebraic"] = CustomPlot(
                         lambda t0, phases_dt, node_idx, x, u, p, a, d: (
                             a[self.nlp.algebraic_states.key_index(self.name), :]
                             if a.any()
                             else np.ndarray((cx[0].shape[0], 1)) * np.nan
                         ),
-                        plot_type=PlotType.STEP,
+                        plot_type=PlotType.INTEGRATED,
                         axes_idx=self.axes_idx,
                         legend=self.legend,
                         combine_to=self.combine_name,
-                        all_variables_in_one_subplot=all_variables_in_one_subplot,
                     )
 
 

bioptim/dynamics/configure_problem.py

@@ -401,12 +401,7 @@
             ocp, nlp, n_noised_controls=n_noised_tau, n_references=nlp.model.n_references
         )
         ConfigureProblem.configure_stochastic_ref(ocp, nlp, n_references=nlp.model.n_references)
-        n_collocation_points = 1
-        if isinstance(problem_type, SocpType.COLLOCATION):
-            n_collocation_points += problem_type.polynomial_degree
-        ConfigureProblem.configure_stochastic_m(
-            ocp, nlp, n_noised_states=n_noised_states, n_collocation_points=n_collocation_points
-        )
+        ConfigureProblem.configure_stochastic_m(ocp, nlp, n_noised_states=n_noised_states)
 
         if isinstance(problem_type, SocpType.TRAPEZOIDAL_EXPLICIT):
             if initial_matrix is None:
@@ -479,12 +474,7 @@
             ocp, nlp, n_noised_controls=n_noised_tau, n_references=nlp.model.n_references
         )
         ConfigureProblem.configure_stochastic_ref(ocp, nlp, n_references=nlp.model.n_references)
-        n_collocation_points = 1
-        if isinstance(problem_type, SocpType.COLLOCATION):
-            n_collocation_points += problem_type.polynomial_degree
-        ConfigureProblem.configure_stochastic_m(
-            ocp, nlp, n_noised_states=n_noised_states, n_collocation_points=n_collocation_points
-        )
+        ConfigureProblem.configure_stochastic_m(ocp, nlp, n_noised_states=n_noised_states)
 
         if isinstance(problem_type, SocpType.TRAPEZOIDAL_EXPLICIT):
             if initial_matrix is None:
@@ -1515,9 +1505,9 @@
             ocp,
             nlp,
             as_states=False,
-            as_controls=False,
+            as_controls=True,
             as_states_dot=False,
-            as_algebraic_states=True,
+            as_algebraic_states=False,
         )
 
     @staticmethod
@@ -1548,9 +1538,9 @@
             ocp,
             nlp,
             as_states=False,
-            as_controls=False,
+            as_controls=True,
             as_states_dot=False,
-            as_algebraic_states=True,
+            as_algebraic_states=False,
             skip_plot=True,
         )
 
@@ -1580,9 +1570,9 @@
             ocp,
             nlp,
             as_states=False,
-            as_controls=False,
+            as_controls=True,
             as_states_dot=False,
-            as_algebraic_states=True,
+            as_algebraic_states=False,
             skip_plot=True,
         )
 
@@ -1638,9 +1628,9 @@
             ocp,
             nlp,
             as_states=False,
-            as_controls=False,
+            as_controls=True,
             as_states_dot=False,
-            as_algebraic_states=True,
+            as_algebraic_states=False,
         )
 
     @staticmethod
@@ -1669,9 +1659,9 @@
             ocp,
             nlp,
             as_states=False,
-            as_controls=False,
+            as_controls=True,
             as_states_dot=False,
-            as_algebraic_states=True,
+            as_algebraic_states=False,
         )
 
     @staticmethod
@@ -1697,13 +1687,13 @@
             ocp,
             nlp,
             as_states=False,
-            as_controls=False,
+            as_controls=True,
             as_states_dot=False,
-            as_algebraic_states=True,
+            as_algebraic_states=False,
         )
 
     @staticmethod
-    def configure_stochastic_m(ocp, nlp, n_noised_states: int, n_collocation_points: int = 1):
+    def configure_stochastic_m(ocp, nlp, n_noised_states: int):
         """
         Configure the helper matrix M (from Gillis 2013 : https://doi.org/10.1109/CDC.2013.6761121).
 
@@ -1719,11 +1709,11 @@
 
         name_m = []
         for name_1 in [f"X_{i}" for i in range(n_noised_states)]:
-            for name_2 in [f"X_{i}" for i in range(n_noised_states * n_collocation_points)]:
+            for name_2 in [f"X_{i}" for i in range(n_noised_states)]:
                 name_m += [name_1 + "_&_" + name_2]
         nlp.variable_mappings[name] = BiMapping(
-            list(range(n_noised_states * n_noised_states * n_collocation_points)),
-            list(range(n_noised_states * n_noised_states * n_collocation_points)),
+            list(range(n_noised_states * n_noised_states)),
+            list(range(n_noised_states * n_noised_states)),
         )
         ConfigureProblem.configure_new_variable(
             name,

bioptim/dynamics/integrator.py

@@ -87,7 +87,7 @@ def __init__(self, ode: dict, ode_opt: dict):
                 self._x_sym_modified,
                 self.u_sym,
                 self.param_sym,
-                self.a_sym,
+                self._a_sym_modified,
                 self.numerical_timeseries_sym,
             ],
             self.dxdt(
@@ -130,6 +130,10 @@ def _time_xall_from_dt_func(self) -> Function:
     def _x_sym_modified(self):
         return self.x_sym
 
+    @property
+    def _a_sym_modified(self):
+        return self.a_sym
+
     @property
     def _input_names(self):
         return ["t_span", "x0", "u", "p", "a", "d"]
@@ -585,6 +589,10 @@ def _initialize(self, ode: dict, ode_opt: dict):
     def _x_sym_modified(self):
         return horzcat(*self.x_sym) if self.duplicate_starting_point else horzcat(*self.x_sym[1:])
 
+    @property
+    def _a_sym_modified(self):
+        return horzcat(*self.a_sym) if self.duplicate_starting_point else horzcat(*self.a_sym[1:])
+
     @property
     def _output_names(self):
         return ["xf", "xall", "defects"]
@@ -664,7 +672,7 @@ def dxdt(
                     states[j + 1],
                     self.get_u(controls, self._integration_time[j]),
                     params,
-                    algebraic_states,
+                    algebraic_states[j + 1],
                     numerical_timeseries,
                 )[:, self.ode_idx]
                 defects.append(xp_j - f_j * self.h)
@@ -676,7 +684,7 @@ def dxdt(
                         states[j + 1],
                         self.get_u(controls, self._integration_time[j]),
                         params,
-                        algebraic_states,
+                        algebraic_states[j + 1],
                         numerical_timeseries,
                         xp_j / self.h,
                     )
@@ -741,6 +749,9 @@ def dxdt(
 
         # Root-finding function, implicitly defines x_collocation_points as a function of x0 and p
         collocation_states = vertcat(*states[1:]) if self.duplicate_starting_point else vertcat(*states[2:])
+        algebraic_states = (
+            vertcat(*algebraic_states[1:]) if self.duplicate_starting_point else vertcat(*algebraic_states[2:])
+        )
         vfcn = Function(
             "vfcn",
             [collocation_states, self.t_span_sym, states[0], controls, params, algebraic_states, numerical_timeseries],

bioptim/dynamics/ode_solvers.py

@@ -167,7 +167,11 @@ def p_ode(self, nlp):
             return nlp.controls.scaled.cx_start
 
         def a_ode(self, nlp):
-            return nlp.algebraic_states.scaled.cx_start
+            out = [nlp.algebraic_states.scaled.cx_start]
+            if not self.duplicate_starting_point:
+                out += [nlp.algebraic_states.scaled.cx_start]
+            out += nlp.algebraic_states.scaled.cx_intermediates_list
+            return out
 
         def d_ode(self, nlp):
             return nlp.numerical_timeseries.cx_start