diff --git a/bioptim/dynamics/configure_new_variable.py b/bioptim/dynamics/configure_new_variable.py index 5531598ae..8aeaaa495 100644 --- a/bioptim/dynamics/configure_new_variable.py +++ b/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, diff --git a/bioptim/dynamics/ode_solver_base.py b/bioptim/dynamics/ode_solver_base.py index e095ca218..6606cb706 100644 --- a/bioptim/dynamics/ode_solver_base.py +++ b/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: """ diff --git a/bioptim/dynamics/rk_base.py b/bioptim/dynamics/rk_base.py index 88cb6ebe3..9ecc38193 100644 --- a/bioptim/dynamics/rk_base.py +++ b/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 diff --git a/tests/shard2/test_global_getting_started.py b/tests/shard2/test_global_getting_started.py index 073325241..b08209353 100644 --- a/tests/shard2/test_global_getting_started.py +++ b/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])