diff --git a/bioptim/dynamics/configure_variables.py b/bioptim/dynamics/configure_variables.py index 97beb3983..508653184 100644 --- a/bioptim/dynamics/configure_variables.py +++ b/bioptim/dynamics/configure_variables.py @@ -1106,7 +1106,7 @@ def configure_rigid_contact_function(ocp, nlp, **extra_params) -> None: else: contact_names_in_phase = [name for name in nlp.model.rigid_contact_names] axes_idx = BiMapping( - to_first=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], + to_first=list(range(len(contact_names_in_phase))), to_second=[i for i, c in enumerate(all_contact_names) if c in contact_names_in_phase], ) diff --git a/tests/shard5/test_configure_problem.py b/tests/shard5/test_configure_problem.py index 46210fe71..c2aafd167 100644 --- a/tests/shard5/test_configure_problem.py +++ b/tests/shard5/test_configure_problem.py @@ -1,3 +1,5 @@ +from types import SimpleNamespace + import numpy.testing as npt import pytest from casadi import MX, SX @@ -180,6 +182,41 @@ def test_configure_soft_contacts(cx): npt.assert_equal(nlp.states_dot.keys(), ["soft_contact_forces"]) +def test_rigid_contact_plot_mapping_when_first_contact_is_removed(): + class Model: + def __init__(self, rigid_contact_names): + self.rigid_contact_names = rigid_contact_names + + def get_rigid_contact_forces(self, *args, **kwargs): + return MX.zeros(len(self.rigid_contact_names), 1) + + def make_nlp(contact_names, phase_idx): + empty_variable = SimpleNamespace(scaled=SimpleNamespace(cx=MX.zeros(0, 1))) + return SimpleNamespace( + model=Model(contact_names), + time_cx=MX.sym(f"time_{phase_idx}", 1, 1), + dt=MX.sym(f"dt_{phase_idx}", 1, 1), + states=empty_variable, + controls=empty_variable, + parameters=empty_variable, + algebraic_states=empty_variable, + numerical_timeseries=SimpleNamespace(cx=MX.zeros(0, 1)), + plot_mapping={}, + plot={}, + phase_idx=phase_idx, + ) + + phase_0 = make_nlp(["heel", "toe"], 0) + phase_1 = make_nlp(["toe"], 1) + ocp = SimpleNamespace(nlp=[phase_0, phase_1]) + + ConfigureVariables.configure_rigid_contact_function(ocp, phase_1) + + mapping = phase_1.plot["rigid_contact_forces"].phase_mappings + npt.assert_equal(mapping.to_first.map_idx, [0]) + npt.assert_equal(mapping.to_second.map_idx, [1]) + + @pytest.mark.parametrize("cx", [MX, SX]) def test_configure_muscles(cx):