From 6a6f2b09457b062accc264792c679525e04b3b62 Mon Sep 17 00:00:00 2001 From: Ipuch Date: Fri, 15 Mar 2024 14:59:08 -0400 Subject: [PATCH] catching when reading --- bioptim/dynamics/integrator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/bioptim/dynamics/integrator.py b/bioptim/dynamics/integrator.py index f98d617b9..fcf8a7922 100644 --- 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")