Merge pull request #861 from Ipuch/master

catching when reading
pyomeca · Mar 15, 2024 · a377e94 · a377e94
2 parents 95fd84e + 6a6f2b0
commit a377e94
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/bioptim/dynamics/integrator.py b/bioptim/dynamics/integrator.py
@@ -1,5 +1,5 @@
-from casadi import Function, vertcat, horzcat, collocation_points, tangent, rootfinder, DM, MX, SX, linspace
 import numpy as np
+from casadi import Function, vertcat, horzcat, collocation_points, tangent, rootfinder, DM, MX, SX, linspace
 
 from ..misc.enums import ControlType, DefectType
 from ..models.protocols.biomodel import BioModel
@@ -170,7 +170,7 @@ def get_u(self, u: np.ndarray, t: float) -> np.ndarray:
         The control at a given time
         """
 
-        if self.control_type == ControlType.CONSTANT or self.control_type == ControlType.CONSTANT_WITH_LAST_NODE:
+        if self.control_type in (ControlType.CONSTANT, ControlType.CONSTANT_WITH_LAST_NODE):
             return u
         elif self.control_type == ControlType.LINEAR_CONTINUOUS:
             dt_norm = (t - self.t_span_sym[0]) / self.t_span_sym[1]
@@ -595,7 +595,7 @@ def get_u(self, u: np.ndarray, t: float | MX | SX) -> np.ndarray:
         The control at a given time
         """
 
-        if self.control_type == ControlType.CONSTANT or self.control_type == ControlType.CONSTANT_WITH_LAST_NODE:
+        if self.control_type in (ControlType.CONSTANT, ControlType.CONSTANT_WITH_LAST_NODE):
             return super(COLLOCATION, self).get_u(u, t)
         else:
             raise NotImplementedError(f"{self.control_type} ControlType not implemented yet with COLLOCATION")