From 6209411dfbebb39736e2afc6281190a208337890 Mon Sep 17 00:00:00 2001 From: Ipuch Date: Fri, 11 Oct 2024 12:18:51 -0400 Subject: [PATCH] refactor : very small but remove if, in forward dynamics to handle and switch between cases. --- bioptim/dynamics/dynamics_functions.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/bioptim/dynamics/dynamics_functions.py b/bioptim/dynamics/dynamics_functions.py index 6b17fadad..7f522c267 100644 --- a/bioptim/dynamics/dynamics_functions.py +++ b/bioptim/dynamics/dynamics_functions.py @@ -1129,19 +1129,17 @@ def forward_dynamics( else: qdot_var_mapping = BiMapping([i for i in range(qdot.shape[0])], [i for i in range(qdot.shape[0])]).to_first - if external_forces is None: - if with_contact: - qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau) - else: - qddot = nlp.model.forward_dynamics(q, qdot, tau) + dynamics_functions = { + (True, True): nlp.model.constrained_forward_dynamics, + (True, False): nlp.model.forward_dynamics, + (False, True): lambda q, qdot, tau: nlp.model.constrained_forward_dynamics(q, qdot, tau, external_forces), + (False, False): lambda q, qdot, tau: nlp.model.forward_dynamics(q, qdot, tau, external_forces), + } - return qdot_var_mapping.map(qddot) - else: - if with_contact: - qddot = nlp.model.constrained_forward_dynamics(q, qdot, tau, external_forces) - else: - qddot = nlp.model.forward_dynamics(q, qdot, tau, external_forces) - return qdot_var_mapping.map(qddot) + dynamics_function = dynamics_functions[(external_forces is None, with_contact)] + + qddot = dynamics_function(q, qdot, tau) + return qdot_var_mapping.map(qddot) @staticmethod def inverse_dynamics(