From 83aa20e2e3012948b291a1c7dfe63a8839f0f280 Mon Sep 17 00:00:00 2001 From: Pedro Bressan Date: Sat, 7 Dec 2024 12:11:51 +0100 Subject: [PATCH] MNT: Move ode solver validation to separate method. --- rocketpy/simulation/flight.py | 35 +++++++++++++++++++++++------------ 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/rocketpy/simulation/flight.py b/rocketpy/simulation/flight.py index e274b7f63..e4be64f5a 100644 --- a/rocketpy/simulation/flight.py +++ b/rocketpy/simulation/flight.py @@ -1213,18 +1213,7 @@ def __init_solver_monitors(self): self.t = self.solution[-1][0] self.y_sol = self.solution[-1][1:] - if isinstance(self.ode_solver, OdeSolver): - self._solver = self.ode_solver - else: - try: - self._solver = ODE_SOLVER_MAP[self.ode_solver] - except KeyError as e: - raise ValueError( - f"Invalid ``ode_solver`` input: {self.ode_solver}. " - f"Available options are: {', '.join(ODE_SOLVER_MAP.keys())}" - ) from e - - self.__is_lsoda = hasattr(self._solver, "_lsoda_solver") + self.__set_ode_solver(self.ode_solver) def __init_equations_of_motion(self): """Initialize equations of motion.""" @@ -1263,6 +1252,28 @@ def __cache_sensor_data(self): sensor_data[sensor] = sensor.measured_data[:] self.sensor_data = sensor_data + def __set_ode_solver(self, solver): + """Sets the ODE solver to be used in the simulation. + + Parameters + ---------- + solver : str, ``scipy.integrate.OdeSolver`` + Integration method to use to solve the equations of motion ODE, + or a custom ``scipy.integrate.OdeSolver``. + """ + if isinstance(solver, OdeSolver): + self._solver = solver + else: + try: + self._solver = ODE_SOLVER_MAP[solver] + except KeyError as e: + raise ValueError( + f"Invalid ``ode_solver`` input: {solver}. " + f"Available options are: {', '.join(ODE_SOLVER_MAP.keys())}" + ) from e + + self.__is_lsoda = hasattr(self._solver, "_lsoda_solver") + @cached_property def effective_1rl(self): """Original rail length minus the distance measured from nozzle exit