On hold — discussion with @EveCharbie required — Reduce unused RK state symbols

bioptim/dynamics/configure_new_variable.py

@@ -166,6 +166,15 @@ def define_cx_unscaled(self, _cx_scaled: CXList, scaling: NpArray) -> CXList:
             _cx[j] = _cx_scaled[j] * scaling
         return _cx
 
+    def define_cx_scaled_for_state(self, n_cx: Int, node_index: Int) -> CXList:
+        """Define state-like CX without a middle symbol when the integrator cannot consume it."""
+        if (
+            self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE
+            and self.nlp.dynamics_type.ode_solver.uses_only_cx_start
+        ):
+            n_cx -= 1
+        return self.define_cx_scaled(n_col=n_cx, node_index=node_index)
+
     def _declare_auto_variable_mapping(self) -> None:
         """Declare the mapping of the new variable if not already declared"""
         if self.name not in self.nlp.variable_mappings:
@@ -228,7 +237,7 @@ def _declare_cx_and_plot(self) -> None:
                 self.nlp.n_states_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
             ):
                 n_cx = self.nlp.dynamics_type.ode_solver.n_required_cx + 2
-                cx_scaled = self.define_cx_scaled(n_col=n_cx, node_index=node_index)
+                cx_scaled = self.define_cx_scaled_for_state(n_cx=n_cx, node_index=node_index)
                 cx = self.define_cx_unscaled(cx_scaled, self.nlp.x_scaling[self.name].scaling)
                 self.nlp.states.append(
                     self.name,
@@ -260,7 +269,7 @@ def _declare_cx_and_plot(self) -> None:
                 self.nlp.n_states_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
             ):
                 n_cx = self.nlp.dynamics_type.ode_solver.n_required_cx + 2
-                cx_scaled = self.define_cx_scaled(n_col=n_cx, node_index=node_index)
+                cx_scaled = self.define_cx_scaled_for_state(n_cx=n_cx, node_index=node_index)
                 cx = self.define_cx_unscaled(cx_scaled, np.ones_like(self.nlp.x_scaling[self.name].scaling))
                 self.nlp.states_dot.append(
                     self.name,
@@ -308,7 +317,7 @@ def _declare_cx_and_plot(self) -> None:
                 self.nlp.n_states_nodes if self.nlp.phase_dynamics == PhaseDynamics.ONE_PER_NODE else 1
             ):
                 n_cx = self.nlp.dynamics_type.ode_solver.n_required_cx + 2
-                cx_scaled = self.define_cx_scaled(n_col=n_cx, node_index=node_index)
+                cx_scaled = self.define_cx_scaled_for_state(n_cx=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,

bioptim/dynamics/ode_solver_base.py

@@ -77,6 +77,11 @@ def n_required_cx(self) -> Int:
         """
         raise RuntimeError("This method should be implemented in the child class")
 
+    @property
+    def uses_only_cx_start(self) -> Bool:
+        """Whether the integrator only consumes the first state and algebraic-state CX."""
+        return False
+
     @property
     def defects_type(self) -> DefectType:
         """

bioptim/dynamics/rk_base.py

@@ -37,6 +37,10 @@ def is_direct_shooting(self) -> Bool:
     def n_required_cx(self) -> Int:
         return 1
 
+    @property
+    def uses_only_cx_start(self) -> Bool:
+        return True
+
     @property
     def defects_type(self) -> DefectType:
         return DefectType.NOT_APPLICABLE

tests/shard2/test_global_getting_started.py

@@ -36,6 +36,30 @@
 test_memory = {}
 
 
+def test_rk_one_per_node_does_not_declare_unused_state_cx():
+    from bioptim.examples.getting_started import basic_ocp as ocp_module
+
+    ocp = ocp_module.prepare_ocp(
+        biorbd_model_path=TestUtils.bioptim_folder() + "/examples/models/pendulum.bioMod",
+        final_time=1,
+        n_shooting=3,
+        ode_solver=OdeSolver.RK4(),
+        phase_dynamics=PhaseDynamics.ONE_PER_NODE,
+        expand_dynamics=False,
+    )
+    nlp = ocp.nlp[0]
+
+    for node_index in range(nlp.n_states_nodes):
+        nlp.states.node_index = node_index
+        nlp.states_dot.node_index = node_index
+        assert len(nlp.states["q"].original_cx) == 2
+        assert len(nlp.states_dot["q"].original_cx) == 2
+
+    # Controls are deliberately outside this optimization and retain their existing CX layout.
+    nlp.controls.node_index = 0
+    assert len(nlp.controls["tau"].original_cx) == 3
+
+
 @pytest.mark.parametrize("phase_dynamics", [PhaseDynamics.SHARED_DURING_THE_PHASE, PhaseDynamics.ONE_PER_NODE])
 @pytest.mark.parametrize("n_threads", [1, 2])
 @pytest.mark.parametrize("use_sx", [False, True])