From b628c19d69b8451c2567be295ca53f0adbf221bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tiberiu=20Sab=C4=83u?= Date: Tue, 16 Dec 2025 12:50:59 +0100 Subject: [PATCH 1/5] ENH: replace ifelif else chains with match statement --- pyproject.toml | 2 +- rocketpy/environment/environment.py | 142 ++++----- rocketpy/mathutils/function.py | 356 ++++++++++++---------- rocketpy/motors/motor.py | 19 +- rocketpy/plots/rocket_plots.py | 58 ++-- rocketpy/rocket/aero_surface/nose_cone.py | 141 ++++----- rocketpy/rocket/rocket.py | 19 +- rocketpy/units.py | 11 +- 8 files changed, 394 insertions(+), 354 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 306299b2b..1442ed015 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -78,7 +78,7 @@ exclude_also = [ [tool.ruff] -target-version = "py39" +target-version = "py310" line-length = 88 indent-width = 4 diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index eb2eacd5a..1dd984ff9 100644 --- a/rocketpy/environment/environment.py +++ b/rocketpy/environment/environment.py @@ -1277,31 +1277,32 @@ def set_atmospheric_model( # pylint: disable=too-many-statements self.atmospheric_model_type = type type = type.lower() - # Handle each case # TODO: use match case when python 3.9 is no longer supported - if type == "standard_atmosphere": - self.process_standard_atmosphere() - elif type == "wyoming_sounding": - self.process_wyoming_sounding(file) - elif type == "custom_atmosphere": - self.process_custom_atmosphere(pressure, temperature, wind_u, wind_v) - elif type == "windy": - self.process_windy_atmosphere(file) - elif type in ["forecast", "reanalysis", "ensemble"]: - dictionary = self.__validate_dictionary(file, dictionary) - try: - fetch_function = self.__atm_type_file_to_function_map[type][file] - except KeyError: - fetch_function = None - - # Fetches the dataset using OpenDAP protocol or uses the file path - dataset = fetch_function() if fetch_function is not None else file - - if type in ["forecast", "reanalysis"]: - self.process_forecast_reanalysis(dataset, dictionary) - else: - self.process_ensemble(dataset, dictionary) - else: # pragma: no cover - raise ValueError(f"Unknown model type '{type}'.") + # Handle each case using match statement + match type: + case "standard_atmosphere": + self.process_standard_atmosphere() + case "wyoming_sounding": + self.process_wyoming_sounding(file) + case "custom_atmosphere": + self.process_custom_atmosphere(pressure, temperature, wind_u, wind_v) + case "windy": + self.process_windy_atmosphere(file) + case "forecast" | "reanalysis" | "ensemble": + dictionary = self.__validate_dictionary(file, dictionary) + try: + fetch_function = self.__atm_type_file_to_function_map[type][file] + except KeyError: + fetch_function = None + + # Fetches the dataset using OpenDAP protocol or uses the file path + dataset = fetch_function() if fetch_function is not None else file + + if type in ["forecast", "reanalysis"]: + self.process_forecast_reanalysis(dataset, dictionary) + else: + self.process_ensemble(dataset, dictionary) + case _: # pragma: no cover + raise ValueError(f"Unknown model type '{type}'.") if type not in ["ensemble"]: # Ensemble already computed these values @@ -2686,51 +2687,52 @@ def from_dict(cls, data): # pylint: disable=too-many-statements ) atmospheric_model = data["atmospheric_model_type"] - if atmospheric_model == "standard_atmosphere": - env.set_atmospheric_model("standard_atmosphere") - elif atmospheric_model == "custom_atmosphere": - env.set_atmospheric_model( - type="custom_atmosphere", - pressure=data["pressure"], - temperature=data["temperature"], - wind_u=data["wind_velocity_x"], - wind_v=data["wind_velocity_y"], - ) - else: - env.__set_pressure_function(data["pressure"]) - env.__set_temperature_function(data["temperature"]) - env.__set_wind_velocity_x_function(data["wind_velocity_x"]) - env.__set_wind_velocity_y_function(data["wind_velocity_y"]) - env.__set_wind_heading_function(data["wind_heading"]) - env.__set_wind_direction_function(data["wind_direction"]) - env.__set_wind_speed_function(data["wind_speed"]) - env.elevation = data["elevation"] - env.max_expected_height = data["max_expected_height"] - - if atmospheric_model in ("windy", "forecast", "reanalysis", "ensemble"): - env.atmospheric_model_init_date = data["atmospheric_model_init_date"] - env.atmospheric_model_end_date = data["atmospheric_model_end_date"] - env.atmospheric_model_interval = data["atmospheric_model_interval"] - env.atmospheric_model_init_lat = data["atmospheric_model_init_lat"] - env.atmospheric_model_end_lat = data["atmospheric_model_end_lat"] - env.atmospheric_model_init_lon = data["atmospheric_model_init_lon"] - env.atmospheric_model_end_lon = data["atmospheric_model_end_lon"] - - if atmospheric_model == "ensemble": - env.level_ensemble = data["level_ensemble"] - env.height_ensemble = data["height_ensemble"] - env.temperature_ensemble = data["temperature_ensemble"] - env.wind_u_ensemble = data["wind_u_ensemble"] - env.wind_v_ensemble = data["wind_v_ensemble"] - env.wind_heading_ensemble = data["wind_heading_ensemble"] - env.wind_direction_ensemble = data["wind_direction_ensemble"] - env.wind_speed_ensemble = data["wind_speed_ensemble"] - env.num_ensemble_members = data["num_ensemble_members"] - - env.__reset_barometric_height_function() - env.calculate_density_profile() - env.calculate_speed_of_sound_profile() - env.calculate_dynamic_viscosity() + match atmospheric_model: + case "standard_atmosphere": + env.set_atmospheric_model("standard_atmosphere") + case "custom_atmosphere": + env.set_atmospheric_model( + type="custom_atmosphere", + pressure=data["pressure"], + temperature=data["temperature"], + wind_u=data["wind_velocity_x"], + wind_v=data["wind_velocity_y"], + ) + case _: + env.__set_pressure_function(data["pressure"]) + env.__set_temperature_function(data["temperature"]) + env.__set_wind_velocity_x_function(data["wind_velocity_x"]) + env.__set_wind_velocity_y_function(data["wind_velocity_y"]) + env.__set_wind_heading_function(data["wind_heading"]) + env.__set_wind_direction_function(data["wind_direction"]) + env.__set_wind_speed_function(data["wind_speed"]) + env.elevation = data["elevation"] + env.max_expected_height = data["max_expected_height"] + + if atmospheric_model in ("windy", "forecast", "reanalysis", "ensemble"): + env.atmospheric_model_init_date = data["atmospheric_model_init_date"] + env.atmospheric_model_end_date = data["atmospheric_model_end_date"] + env.atmospheric_model_interval = data["atmospheric_model_interval"] + env.atmospheric_model_init_lat = data["atmospheric_model_init_lat"] + env.atmospheric_model_end_lat = data["atmospheric_model_end_lat"] + env.atmospheric_model_init_lon = data["atmospheric_model_init_lon"] + env.atmospheric_model_end_lon = data["atmospheric_model_end_lon"] + + if atmospheric_model == "ensemble": + env.level_ensemble = data["level_ensemble"] + env.height_ensemble = data["height_ensemble"] + env.temperature_ensemble = data["temperature_ensemble"] + env.wind_u_ensemble = data["wind_u_ensemble"] + env.wind_v_ensemble = data["wind_v_ensemble"] + env.wind_heading_ensemble = data["wind_heading_ensemble"] + env.wind_direction_ensemble = data["wind_direction_ensemble"] + env.wind_speed_ensemble = data["wind_speed_ensemble"] + env.num_ensemble_members = data["num_ensemble_members"] + + env.__reset_barometric_height_function() + env.calculate_density_profile() + env.calculate_speed_of_sound_profile() + env.calculate_dynamic_viscosity() return env diff --git a/rocketpy/mathutils/function.py b/rocketpy/mathutils/function.py index 9fe343687..ca5ebe54d 100644 --- a/rocketpy/mathutils/function.py +++ b/rocketpy/mathutils/function.py @@ -327,17 +327,18 @@ def __update_interpolation_coefficients(self, method): """Update interpolation coefficients for the given method.""" # Spline, akima and polynomial need data processing # Shepard, and linear do not - if method == "polynomial": - self.__interpolate_polynomial__() - self._coeffs = self.__polynomial_coefficients__ - elif method == "akima": - self.__interpolate_akima__() - self._coeffs = self.__akima_coefficients__ - elif method == "spline" or method is None: - self.__interpolate_spline__() - self._coeffs = self.__spline_coefficients__ - else: - self._coeffs = [] + match method: + case "polynomial": + self.__interpolate_polynomial__() + self._coeffs = self.__polynomial_coefficients__ + case "akima": + self.__interpolate_akima__() + self._coeffs = self.__akima_coefficients__ + case "spline" | None: + self.__interpolate_spline__() + self._coeffs = self.__spline_coefficients__ + case _: + self._coeffs = [] def set_extrapolation(self, method="constant"): """Set extrapolation behavior of data set. @@ -370,107 +371,12 @@ def __set_interpolation_func(self): # pylint: disable=too-many-statements interpolation method has its own function`. The function is stored in the attribute _interpolation_func.""" interpolation = INTERPOLATION_TYPES[self.__interpolation__] - if interpolation == 0: # linear - if self.__dom_dim__ == 1: - - def linear_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - x_interval = bisect_left(x_data, x) - x_left = x_data[x_interval - 1] - y_left = y_data[x_interval - 1] - dx = float(x_data[x_interval] - x_left) - dy = float(y_data[x_interval] - y_left) - return (x - x_left) * (dy / dx) + y_left - - else: - interpolator = LinearNDInterpolator(self._domain, self._image) - - def linear_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - return interpolator(x) - - self._interpolation_func = linear_interpolation - - elif interpolation == 1: # polynomial - - def polynomial_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - return np.sum(coeffs * x ** np.arange(len(coeffs))) - - self._interpolation_func = polynomial_interpolation - - elif interpolation == 2: # akima - - def akima_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - x_interval = bisect_left(x_data, x) - x_interval = x_interval if x_interval != 0 else 1 - a = coeffs[4 * x_interval - 4 : 4 * x_interval] - return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] - - self._interpolation_func = akima_interpolation - - elif interpolation == 3: # spline - - def spline_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - x_interval = bisect_left(x_data, x) - x_interval = max(x_interval, 1) - a = coeffs[:, x_interval - 1] - x = x - x_data[x_interval - 1] - return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] - - self._interpolation_func = spline_interpolation - - elif interpolation == 4: # shepard - # pylint: disable=unused-argument - def shepard_interpolation(x, x_min, x_max, x_data, y_data, _): - arg_qty, arg_dim = x.shape - result = np.empty(arg_qty) - x = x.reshape((arg_qty, 1, arg_dim)) - sub_matrix = x_data - x - distances_squared = np.sum(sub_matrix**2, axis=2) - - # Remove zero distances from further calculations - zero_distances = np.where(distances_squared == 0) - valid_indexes = np.ones(arg_qty, dtype=bool) - valid_indexes[zero_distances[0]] = False - - weights = distances_squared[valid_indexes] ** (-1.5) - numerator_sum = np.sum(y_data * weights, axis=1) - denominator_sum = np.sum(weights, axis=1) - result[valid_indexes] = numerator_sum / denominator_sum - result[~valid_indexes] = y_data[zero_distances[1]] - - return result - - self._interpolation_func = shepard_interpolation - - elif interpolation == 5: # RBF - interpolator = RBFInterpolator(self._domain, self._image, neighbors=100) - - def rbf_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - return interpolator(x) - - self._interpolation_func = rbf_interpolation - - else: - raise ValueError(f"Interpolation {interpolation} method not recognized.") - - def __set_extrapolation_func(self): # pylint: disable=too-many-statements - """Defines extrapolation function used by the Function. Each - extrapolation method has its own function. The function is stored in - the attribute _extrapolation_func.""" - interpolation = INTERPOLATION_TYPES[self.__interpolation__] - extrapolation = EXTRAPOLATION_TYPES[self.__extrapolation__] - - if extrapolation == 0: # zero - - def zero_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - return 0 - - self._extrapolation_func = zero_extrapolation - elif extrapolation == 1: # natural - if interpolation == 0: # linear + match interpolation: + case 0: # linear if self.__dom_dim__ == 1: - def natural_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - x_interval = 1 if x < x_min else -1 + def linear_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + x_interval = bisect_left(x_data, x) x_left = x_data[x_interval - 1] y_left = y_data[x_interval - 1] dx = float(x_data[x_interval] - x_left) @@ -478,38 +384,44 @@ def natural_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: d return (x - x_left) * (dy / dx) + y_left else: - interpolator = RBFInterpolator( - self._domain, self._image, neighbors=100 - ) + interpolator = LinearNDInterpolator(self._domain, self._image) - def natural_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + def linear_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument return interpolator(x) - elif interpolation == 1: # polynomial + self._interpolation_func = linear_interpolation - def natural_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + case 1: # polynomial + + def polynomial_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument return np.sum(coeffs * x ** np.arange(len(coeffs))) - elif interpolation == 2: # akima + self._interpolation_func = polynomial_interpolation + + case 2: # akima - def natural_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - a = coeffs[:4] if x < x_min else coeffs[-4:] + def akima_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + x_interval = bisect_left(x_data, x) + x_interval = x_interval if x_interval != 0 else 1 + a = coeffs[4 * x_interval - 4 : 4 * x_interval] return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] - elif interpolation == 3: # spline + self._interpolation_func = akima_interpolation + + case 3: # spline - def natural_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - if x < x_min: - a = coeffs[:, 0] - x = x - x_data[0] - else: - a = coeffs[:, -1] - x = x - x_data[-2] + def spline_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + x_interval = bisect_left(x_data, x) + x_interval = max(x_interval, 1) + a = coeffs[:, x_interval - 1] + x = x - x_data[x_interval - 1] return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] - elif interpolation == 4: # shepard + self._interpolation_func = spline_interpolation + + case 4: # shepard # pylint: disable=unused-argument - def natural_extrapolation(x, x_min, x_max, x_data, y_data, _): + def shepard_interpolation(x, x_min, x_max, x_data, y_data, _): arg_qty, arg_dim = x.shape result = np.empty(arg_qty) x = x.reshape((arg_qty, 1, arg_dim)) @@ -529,34 +441,144 @@ def natural_extrapolation(x, x_min, x_max, x_data, y_data, _): return result - elif interpolation == 5: # RBF + self._interpolation_func = shepard_interpolation + + case 5: # RBF interpolator = RBFInterpolator(self._domain, self._image, neighbors=100) - def natural_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + def rbf_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument return interpolator(x) - else: + self._interpolation_func = rbf_interpolation + + case _: raise ValueError( - f"Natural extrapolation not defined for {interpolation}." + f"Interpolation {interpolation} method not recognized." ) - self._extrapolation_func = natural_extrapolation - elif extrapolation == 2: # constant - if self.__dom_dim__ == 1: + def __set_extrapolation_func(self): # pylint: disable=too-many-statements + """Defines extrapolation function used by the Function. Each + extrapolation method has its own function. The function is stored in + the attribute _extrapolation_func.""" + interpolation = INTERPOLATION_TYPES[self.__interpolation__] + extrapolation = EXTRAPOLATION_TYPES[self.__extrapolation__] - def constant_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument - return y_data[0] if x < x_min else y_data[-1] + match extrapolation: + case 0: # zero + + def zero_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + return 0 + + self._extrapolation_func = zero_extrapolation + case 1: # natural + match interpolation: + case 0: # linear + if self.__dom_dim__ == 1: + + def natural_extrapolation( + x, x_min, x_max, x_data, y_data, coeffs + ): # pylint: disable=unused-argument + x_interval = 1 if x < x_min else -1 + x_left = x_data[x_interval - 1] + y_left = y_data[x_interval - 1] + dx = float(x_data[x_interval] - x_left) + dy = float(y_data[x_interval] - y_left) + return (x - x_left) * (dy / dx) + y_left + + else: + interpolator = RBFInterpolator( + self._domain, self._image, neighbors=100 + ) + + def natural_extrapolation( + x, x_min, x_max, x_data, y_data, coeffs + ): # pylint: disable=unused-argument + return interpolator(x) + + case 1: # polynomial + + def natural_extrapolation( + x, x_min, x_max, x_data, y_data, coeffs + ): # pylint: disable=unused-argument + return np.sum(coeffs * x ** np.arange(len(coeffs))) + + case 2: # akima + + def natural_extrapolation( + x, x_min, x_max, x_data, y_data, coeffs + ): # pylint: disable=unused-argument + a = coeffs[:4] if x < x_min else coeffs[-4:] + return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] + + case 3: # spline + + def natural_extrapolation( + x, x_min, x_max, x_data, y_data, coeffs + ): # pylint: disable=unused-argument + if x < x_min: + a = coeffs[:, 0] + x = x - x_data[0] + else: + a = coeffs[:, -1] + x = x - x_data[-2] + return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] + + case 4: # shepard + # pylint: disable=unused-argument + def natural_extrapolation(x, x_min, x_max, x_data, y_data, _): + arg_qty, arg_dim = x.shape + result = np.empty(arg_qty) + x = x.reshape((arg_qty, 1, arg_dim)) + sub_matrix = x_data - x + distances_squared = np.sum(sub_matrix**2, axis=2) + + # Remove zero distances from further calculations + zero_distances = np.where(distances_squared == 0) + valid_indexes = np.ones(arg_qty, dtype=bool) + valid_indexes[zero_distances[0]] = False + + weights = distances_squared[valid_indexes] ** (-1.5) + numerator_sum = np.sum(y_data * weights, axis=1) + denominator_sum = np.sum(weights, axis=1) + result[valid_indexes] = numerator_sum / denominator_sum + result[~valid_indexes] = y_data[zero_distances[1]] + + return result + + case 5: # RBF + interpolator = RBFInterpolator( + self._domain, self._image, neighbors=100 + ) + + def natural_extrapolation( + x, x_min, x_max, x_data, y_data, coeffs + ): # pylint: disable=unused-argument + return interpolator(x) + + case _: + raise ValueError( + f"Natural extrapolation not defined for {interpolation}." + ) + + self._extrapolation_func = natural_extrapolation + case 2: # constant + if self.__dom_dim__ == 1: - else: - extrapolator = NearestNDInterpolator(self._domain, self._image) + def constant_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + return y_data[0] if x < x_min else y_data[-1] - def constant_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): - # pylint: disable=unused-argument - return extrapolator(x) + else: + extrapolator = NearestNDInterpolator(self._domain, self._image) - self._extrapolation_func = constant_extrapolation - else: - raise ValueError(f"Extrapolation {extrapolation} method not recognized.") + def constant_extrapolation(x, x_min, x_max, x_data, y_data, coeffs): + # pylint: disable=unused-argument + return extrapolator(x) + + self._extrapolation_func = constant_extrapolation + case _: + raise ValueError( + f"Extrapolation {extrapolation} method not recognized." + ) def set_get_value_opt(self): """Defines a method that evaluates interpolations. @@ -2038,17 +2060,18 @@ def plot_2d( # pylint: disable=too-many-statements vmax=z_max, ) figure.colorbar(surf) - elif disp_type == "wireframe": - axes.plot_wireframe(mesh_x, mesh_y, z, rstride=1, cstride=1) - elif disp_type == "contour": - figure.clf() - contour_set = plt.contour(mesh_x, mesh_y, z) - plt.clabel(contour_set, inline=1, fontsize=10) - elif disp_type == "contourf": - figure.clf() - contour_set = plt.contour(mesh_x, mesh_y, z) - plt.contourf(mesh_x, mesh_y, z) - plt.clabel(contour_set, inline=1, fontsize=10) + match disp_type: + case "wireframe": + axes.plot_wireframe(mesh_x, mesh_y, z, rstride=1, cstride=1) + case "contour": + figure.clf() + contour_set = plt.contour(mesh_x, mesh_y, z) + plt.clabel(contour_set, inline=1, fontsize=10) + case "contourf": + figure.clf() + contour_set = plt.contour(mesh_x, mesh_y, z) + plt.contourf(mesh_x, mesh_y, z) + plt.clabel(contour_set, inline=1, fontsize=10) plt.title(self.title) axes.set_xlabel(self.__inputs__[0].title()) axes.set_ylabel(self.__inputs__[1].title()) @@ -3235,14 +3258,17 @@ def differentiate(self, x, dx=1e-6, order=1): ans : float Evaluated derivative. """ - if order == 1: - return (self.get_value_opt(x + dx) - self.get_value_opt(x - dx)) / (2 * dx) - elif order == 2: - return ( - self.get_value_opt(x + dx) - - 2 * self.get_value_opt(x) - + self.get_value_opt(x - dx) - ) / dx**2 + match order: + case 1: + return (self.get_value_opt(x + dx) - self.get_value_opt(x - dx)) / ( + 2 * dx + ) + case 2: + return ( + self.get_value_opt(x + dx) + - 2 * self.get_value_opt(x) + + self.get_value_opt(x - dx) + ) / dx**2 def differentiate_complex_step(self, x, dx=1e-200, order=1): """Differentiate a Function object at a given point using the complex diff --git a/rocketpy/motors/motor.py b/rocketpy/motors/motor.py index 7930ed52b..796acee49 100644 --- a/rocketpy/motors/motor.py +++ b/rocketpy/motors/motor.py @@ -259,15 +259,16 @@ class Function. Thrust units are Newtons. """ # Define coordinate system orientation self.coordinate_system_orientation = coordinate_system_orientation - if coordinate_system_orientation == "nozzle_to_combustion_chamber": - self._csys = 1 - elif coordinate_system_orientation == "combustion_chamber_to_nozzle": - self._csys = -1 - else: # pragma: no cover - raise ValueError( - "Invalid coordinate system orientation. Options are " - "'nozzle_to_combustion_chamber' and 'combustion_chamber_to_nozzle'." - ) + match coordinate_system_orientation: + case "nozzle_to_combustion_chamber": + self._csys = 1 + case "combustion_chamber_to_nozzle": + self._csys = -1 + case _: # pragma: no cover + raise ValueError( + "Invalid coordinate system orientation. Options are " + "'nozzle_to_combustion_chamber' and 'combustion_chamber_to_nozzle'." + ) # Motor parameters self.interpolate = interpolation_method diff --git a/rocketpy/plots/rocket_plots.py b/rocketpy/plots/rocket_plots.py index 8eaaded16..0f087a5d6 100644 --- a/rocketpy/plots/rocket_plots.py +++ b/rocketpy/plots/rocket_plots.py @@ -348,18 +348,19 @@ def _draw_generic_surface( ): """Draws the generic surface and saves the position of the points of interest for the tubes.""" - if plane == "xz": - # z position of the sensor is the x position in the plot - x_pos = position[2] - # x position of the surface is the y position in the plot - y_pos = position[0] - elif plane == "yz": - # z position of the surface is the x position in the plot - x_pos = position[2] - # y position of the surface is the y position in the plot - y_pos = position[1] - else: # pragma: no cover - raise ValueError("Plane must be 'xz' or 'yz'.") + match plane: + case "xz": + # z position of the sensor is the x position in the plot + x_pos = position[2] + # x position of the surface is the y position in the plot + y_pos = position[0] + case "yz": + # z position of the surface is the x position in the plot + x_pos = position[2] + # y position of the surface is the y position in the plot + y_pos = position[1] + case _: # pragma: no cover + raise ValueError("Plane must be 'xz' or 'yz'.") ax.scatter( x_pos, @@ -592,22 +593,23 @@ def _draw_sensors(self, ax, sensors, plane): for i, sensor_pos in enumerate(sensors): sensor = sensor_pos[0] pos = sensor_pos[1] - if plane == "xz": - # z position of the sensor is the x position in the plot - x_pos = pos[2] - normal_x = sensor.normal_vector.z - # x position of the sensor is the y position in the plot - y_pos = pos[0] - normal_y = sensor.normal_vector.x - elif plane == "yz": - # z position of the sensor is the x position in the plot - x_pos = pos[2] - normal_x = sensor.normal_vector.z - # y position of the sensor is the y position in the plot - y_pos = pos[1] - normal_y = sensor.normal_vector.y - else: # pragma: no cover - raise ValueError("Plane must be 'xz' or 'yz'.") + match plane: + case "xz": + # z position of the sensor is the x position in the plot + x_pos = pos[2] + normal_x = sensor.normal_vector.z + # x position of the sensor is the y position in the plot + y_pos = pos[0] + normal_y = sensor.normal_vector.x + case "yz": + # z position of the sensor is the x position in the plot + x_pos = pos[2] + normal_x = sensor.normal_vector.z + # y position of the sensor is the y position in the plot + y_pos = pos[1] + normal_y = sensor.normal_vector.y + case _: # pragma: no cover + raise ValueError("Plane must be 'xz' or 'yz'.") # line length is 2/5 of the rocket radius line_length = self.rocket.radius / 2.5 diff --git a/rocketpy/rocket/aero_surface/nose_cone.py b/rocketpy/rocket/aero_surface/nose_cone.py index 4e6ea8e95..ddf157586 100644 --- a/rocketpy/rocket/aero_surface/nose_cone.py +++ b/rocketpy/rocket/aero_surface/nose_cone.py @@ -226,79 +226,84 @@ def kind(self, value): # pylint: disable=too-many-statements self._kind = value value = (value.replace(" ", "")).lower() - if value == "conical": - self.k = 2 / 3 - self.y_nosecone = Function(lambda x: x * self.base_radius / self.length) - - elif value == "lvhaack": - self.k = 0.563 - - def theta(x): - return np.arccos(1 - 2 * max(min(x / self.length, 1), 0)) - - self.y_nosecone = Function( - lambda x: self.base_radius - * (theta(x) - np.sin(2 * theta(x)) / 2 + (np.sin(theta(x)) ** 3) / 3) - ** (0.5) - / (np.pi**0.5) - ) + match value: + case "conical": + self.k = 2 / 3 + self.y_nosecone = Function(lambda x: x * self.base_radius / self.length) + + case "lvhaack": + self.k = 0.563 + + def theta(x): + return np.arccos(1 - 2 * max(min(x / self.length, 1), 0)) + + self.y_nosecone = Function( + lambda x: self.base_radius + * ( + theta(x) + - np.sin(2 * theta(x)) / 2 + + (np.sin(theta(x)) ** 3) / 3 + ) + ** (0.5) + / (np.pi**0.5) + ) - elif value in ["tangent", "tangentogive", "ogive"]: - rho = (self.base_radius**2 + self.length**2) / (2 * self.base_radius) - volume = np.pi * ( - self.length * rho**2 - - (self.length**3) / 3 - - (rho - self.base_radius) * rho**2 * np.arcsin(self.length / rho) - ) - area = np.pi * self.base_radius**2 - self.k = 1 - volume / (area * self.length) - self.y_nosecone = Function( - lambda x: np.sqrt(rho**2 - (min(x - self.length, 0)) ** 2) - + (self.base_radius - rho) - ) + case "tangent" | "tangentogive" | "ogive": + rho = (self.base_radius**2 + self.length**2) / (2 * self.base_radius) + volume = np.pi * ( + self.length * rho**2 + - (self.length**3) / 3 + - (rho - self.base_radius) * rho**2 * np.arcsin(self.length / rho) + ) + area = np.pi * self.base_radius**2 + self.k = 1 - volume / (area * self.length) + self.y_nosecone = Function( + lambda x: np.sqrt(rho**2 - (min(x - self.length, 0)) ** 2) + + (self.base_radius - rho) + ) - elif value == "elliptical": - self.k = 1 / 3 - self.y_nosecone = Function( - lambda x: self.base_radius - * np.sqrt(1 - ((x - self.length) / self.length) ** 2) - ) + case "elliptical": + self.k = 1 / 3 + self.y_nosecone = Function( + lambda x: self.base_radius + * np.sqrt(1 - ((x - self.length) / self.length) ** 2) + ) - elif value == "vonkarman": - self.k = 0.5 + case "vonkarman": + self.k = 0.5 - def theta(x): - return np.arccos(1 - 2 * max(min(x / self.length, 1), 0)) + def theta(x): + return np.arccos(1 - 2 * max(min(x / self.length, 1), 0)) - self.y_nosecone = Function( - lambda x: self.base_radius - * (theta(x) - np.sin(2 * theta(x)) / 2) ** (0.5) - / (np.pi**0.5) - ) - elif value == "parabolic": - self.k = 0.5 - self.y_nosecone = Function( - lambda x: self.base_radius - * ((2 * x / self.length - (x / self.length) ** 2) / (2 - 1)) - ) - elif value == "powerseries": - self.k = (2 * self.power) / ((2 * self.power) + 1) - self.y_nosecone = Function( - lambda x: self.base_radius * np.power(x / self.length, self.power) - ) - else: # pragma: no cover - raise ValueError( - f"Nose Cone kind '{self.kind}' not found, " - + "please use one of the following Nose Cone kinds:" - + '\n\t"conical"' - + '\n\t"ogive"' - + '\n\t"lvhaack"' - + '\n\t"tangent"' - + '\n\t"vonkarman"' - + '\n\t"elliptical"' - + '\n\t"powerseries"' - + '\n\t"parabolic"\n' - ) + self.y_nosecone = Function( + lambda x: self.base_radius + * (theta(x) - np.sin(2 * theta(x)) / 2) ** (0.5) + / (np.pi**0.5) + ) + case "parabolic": + self.k = 0.5 + self.y_nosecone = Function( + lambda x: self.base_radius + * ((2 * x / self.length - (x / self.length) ** 2) / (2 - 1)) + ) + case "powerseries": + self.k = (2 * self.power) / ((2 * self.power) + 1) + self.y_nosecone = Function( + lambda x: self.base_radius * np.power(x / self.length, self.power) + ) + case _: # pragma: no cover + raise ValueError( + f"Nose Cone kind '{self.kind}' not found, " + + "please use one of the following Nose Cone kinds:" + + '\n\t"conical"' + + '\n\t"ogive"' + + '\n\t"lvhaack"' + + '\n\t"tangent"' + + '\n\t"vonkarman"' + + '\n\t"elliptical"' + + '\n\t"powerseries"' + + '\n\t"parabolic"\n' + ) self.evaluate_center_of_pressure() self.evaluate_geometrical_parameters() diff --git a/rocketpy/rocket/rocket.py b/rocketpy/rocket/rocket.py index 1112a98f3..95b51d9f2 100644 --- a/rocketpy/rocket/rocket.py +++ b/rocketpy/rocket/rocket.py @@ -279,15 +279,16 @@ def __init__( # pylint: disable=too-many-statements """ # Define coordinate system orientation self.coordinate_system_orientation = coordinate_system_orientation - if coordinate_system_orientation == "tail_to_nose": - self._csys = 1 - elif coordinate_system_orientation == "nose_to_tail": - self._csys = -1 - else: # pragma: no cover - raise TypeError( - "Invalid coordinate system orientation. Please choose between " - + '"tail_to_nose" and "nose_to_tail".' - ) + match coordinate_system_orientation: + case "tail_to_nose": + self._csys = 1 + case "nose_to_tail": + self._csys = -1 + case _: # pragma: no cover + raise TypeError( + "Invalid coordinate system orientation. Please choose between " + + '"tail_to_nose" and "nose_to_tail".' + ) # Define rocket inertia attributes in SI units self.mass = mass diff --git a/rocketpy/units.py b/rocketpy/units.py index 4d70806b9..b8e4c199c 100644 --- a/rocketpy/units.py +++ b/rocketpy/units.py @@ -77,10 +77,13 @@ def convert_units_Functions(variable, from_unit, to_unit, axis=1): else: variable_source[:, axis] *= conversion_factor(from_unit, to_unit) # Rename axis labels - if axis == 0: - variable.__inputs__[0] = variable.__inputs__[0].replace(from_unit, to_unit) - elif axis == 1: - variable.__outputs__[0] = variable.__outputs__[0].replace(from_unit, to_unit) + match axis: + case 0: + variable.__inputs__[0] = variable.__inputs__[0].replace(from_unit, to_unit) + case 1: + variable.__outputs__[0] = variable.__outputs__[0].replace( + from_unit, to_unit + ) # Create new Function instance with converted data return Function( source=variable_source, From e98f9697b59c54c6147ba4b5109a9c318bb51c38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tiberiu=20Sab=C4=83u?= <96194994+tibisabau@users.noreply.github.com> Date: Tue, 16 Dec 2025 15:27:50 +0100 Subject: [PATCH 2/5] Update rocketpy/environment/environment.py Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com> --- rocketpy/environment/environment.py | 1 - 1 file changed, 1 deletion(-) diff --git a/rocketpy/environment/environment.py b/rocketpy/environment/environment.py index dbda5acff..6b6d399b4 100644 --- a/rocketpy/environment/environment.py +++ b/rocketpy/environment/environment.py @@ -1277,7 +1277,6 @@ def set_atmospheric_model( # pylint: disable=too-many-statements self.atmospheric_model_type = type type = type.lower() - # Handle each case using match statement match type: case "standard_atmosphere": self.process_standard_atmosphere() From cdf12afb4c1ade3c4a945d474873388aee54f4bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tiberiu=20Sab=C4=83u?= Date: Tue, 16 Dec 2025 15:49:42 +0100 Subject: [PATCH 3/5] fix CI tests and add changes to CHANGELOG.md --- CHANGELOG.md | 1 + rocketpy/mathutils/function.py | 47 +++++++++++++++++++++++++++++----- 2 files changed, 42 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 70051de22..56738d804 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -49,6 +49,7 @@ Attention: The newest changes should be on top --> ### Changed +- ENH: replace if elif else chains with match statement - ENH: Refactor Flight class to improve time node handling and sensor/controllers [#843](https://github.com/RocketPy-Team/RocketPy/pull/843) ### Fixed diff --git a/rocketpy/mathutils/function.py b/rocketpy/mathutils/function.py index 72446a695..5529b25ea 100644 --- a/rocketpy/mathutils/function.py +++ b/rocketpy/mathutils/function.py @@ -465,6 +465,41 @@ def rbf_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disab self._interpolation_func = rbf_interpolation + case 6: # regular_grid (RegularGridInterpolator) + # For grid interpolation, the actual interpolator is stored separately + # This function is a placeholder that should not be called directly + # since __get_value_opt_grid is used instead + if hasattr(self, "_grid_interpolator"): + + def grid_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disable=unused-argument + return self._grid_interpolator(x) + + self._interpolation_func = grid_interpolation + else: + # Fallback to shepard if grid interpolator not available + warnings.warn( + "Grid interpolator not found, falling back to shepard interpolation" + ) + + def shepard_fallback(x, x_min, x_max, x_data, y_data, _): + # pylint: disable=unused-argument + arg_qty, arg_dim = x.shape + result = np.empty(arg_qty) + x = x.reshape((arg_qty, 1, arg_dim)) + sub_matrix = x_data - x + distances_squared = np.sum(sub_matrix**2, axis=2) + zero_distances = np.where(distances_squared == 0) + valid_indexes = np.ones(arg_qty, dtype=bool) + valid_indexes[zero_distances[0]] = False + weights = distances_squared[valid_indexes] ** (-1.5) + numerator_sum = np.sum(y_data * weights, axis=1) + denominator_sum = np.sum(weights, axis=1) + result[valid_indexes] = numerator_sum / denominator_sum + result[~valid_indexes] = y_data[zero_distances[1]] + return result + + self._interpolation_func = shepard_fallback + case _: raise ValueError( f"Interpolation {interpolation} method not recognized." @@ -511,14 +546,14 @@ def natural_extrapolation( case 1: # polynomial - def natural_extrapolation( + def natural_extrapolation( # pylint: disable=function-redefined x, x_min, x_max, x_data, y_data, coeffs ): # pylint: disable=unused-argument return np.sum(coeffs * x ** np.arange(len(coeffs))) case 2: # akima - def natural_extrapolation( + def natural_extrapolation( # pylint: disable=function-redefined x, x_min, x_max, x_data, y_data, coeffs ): # pylint: disable=unused-argument a = coeffs[:4] if x < x_min else coeffs[-4:] @@ -526,19 +561,19 @@ def natural_extrapolation( case 3: # spline - def natural_extrapolation( + def natural_extrapolation( # pylint: disable=function-redefined x, x_min, x_max, x_data, y_data, coeffs ): # pylint: disable=unused-argument if x < x_min: a = coeffs[:, 0] - x = x - x_data[0] + x = x_data[0] else: a = coeffs[:, -1] x = x - x_data[-2] return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] case 4: # shepard - # pylint: disable=unused-argument + # pylint: disable=unused-argument,function-redefined def natural_extrapolation(x, x_min, x_max, x_data, y_data, _): arg_qty, arg_dim = x.shape result = np.empty(arg_qty) @@ -564,7 +599,7 @@ def natural_extrapolation(x, x_min, x_max, x_data, y_data, _): self._domain, self._image, neighbors=100 ) - def natural_extrapolation( + def natural_extrapolation( # pylint: disable=function-redefined x, x_min, x_max, x_data, y_data, coeffs ): # pylint: disable=unused-argument return interpolator(x) From 44d6e8bc2357b41ed8e84dbb404680dc4e54ec3d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tiberiu=20Sab=C4=83u?= Date: Tue, 16 Dec 2025 17:29:49 +0100 Subject: [PATCH 4/5] update CHANGELOG.md --- CHANGELOG.md | 2 +- rocketpy/mathutils/function.py | 22 +++++++++++----------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 56738d804..04b2f4d59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -49,7 +49,7 @@ Attention: The newest changes should be on top --> ### Changed -- ENH: replace if elif else chains with match statement +- ENH: replace if elif else chains with match statement [#921](https://github.com/RocketPy-Team/RocketPy/pull/921/changes) - ENH: Refactor Flight class to improve time node handling and sensor/controllers [#843](https://github.com/RocketPy-Team/RocketPy/pull/843) ### Fixed diff --git a/rocketpy/mathutils/function.py b/rocketpy/mathutils/function.py index 5529b25ea..3a46a3116 100644 --- a/rocketpy/mathutils/function.py +++ b/rocketpy/mathutils/function.py @@ -566,11 +566,11 @@ def natural_extrapolation( # pylint: disable=function-redefined ): # pylint: disable=unused-argument if x < x_min: a = coeffs[:, 0] - x = x_data[0] + x_offset = x - x_data[0] else: a = coeffs[:, -1] - x = x - x_data[-2] - return a[3] * x**3 + a[2] * x**2 + a[1] * x + a[0] + x_offset = x - x_data[-2] + return a[3] * x_offset**3 + a[2] * x_offset**2 + a[1] * x_offset + a[0] case 4: # shepard # pylint: disable=unused-argument,function-redefined @@ -3254,20 +3254,20 @@ def integral(self, a, b, numerical=False): # pylint: disable=too-many-statement ans += y_data[0] * (min(x_data[0], b) - a) elif self.__extrapolation__ == "natural": c = coeffs[:, 0] - sub_b = a - x_data[0] - sub_a = min(b, x_data[0]) - x_data[0] + sub_a = a - x_data[0] + sub_b = min(b, x_data[0]) - x_data[0] ans += ( - (c[3] * sub_a**4) / 4 - + (c[2] * sub_a**3 / 3) - + (c[1] * sub_a**2 / 2) - + c[0] * sub_a - ) - ans -= ( (c[3] * sub_b**4) / 4 + (c[2] * sub_b**3 / 3) + (c[1] * sub_b**2 / 2) + c[0] * sub_b ) + ans -= ( + (c[3] * sub_a**4) / 4 + + (c[2] * sub_a**3 / 3) + + (c[1] * sub_a**2 / 2) + + c[0] * sub_a + ) else: # self.__extrapolation__ = 'zero' pass From c92dacc19ea6467c0c6a65699d522682dabd5570 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tiberiu=20Sab=C4=83u?= Date: Tue, 16 Dec 2025 19:20:58 +0100 Subject: [PATCH 5/5] fix linters --- rocketpy/mathutils/function.py | 7 ++++++- rocketpy/rocket/aero_surface/nose_cone.py | 12 ++++++------ 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/rocketpy/mathutils/function.py b/rocketpy/mathutils/function.py index 3a46a3116..fbece42dd 100644 --- a/rocketpy/mathutils/function.py +++ b/rocketpy/mathutils/function.py @@ -570,7 +570,12 @@ def natural_extrapolation( # pylint: disable=function-redefined else: a = coeffs[:, -1] x_offset = x - x_data[-2] - return a[3] * x_offset**3 + a[2] * x_offset**2 + a[1] * x_offset + a[0] + return ( + a[3] * x_offset**3 + + a[2] * x_offset**2 + + a[1] * x_offset + + a[0] + ) case 4: # shepard # pylint: disable=unused-argument,function-redefined diff --git a/rocketpy/rocket/aero_surface/nose_cone.py b/rocketpy/rocket/aero_surface/nose_cone.py index ddf157586..6ebfdd2ab 100644 --- a/rocketpy/rocket/aero_surface/nose_cone.py +++ b/rocketpy/rocket/aero_surface/nose_cone.py @@ -234,15 +234,15 @@ def kind(self, value): # pylint: disable=too-many-statements case "lvhaack": self.k = 0.563 - def theta(x): + def theta_lvhaack(x): return np.arccos(1 - 2 * max(min(x / self.length, 1), 0)) self.y_nosecone = Function( lambda x: self.base_radius * ( - theta(x) - - np.sin(2 * theta(x)) / 2 - + (np.sin(theta(x)) ** 3) / 3 + theta_lvhaack(x) + - np.sin(2 * theta_lvhaack(x)) / 2 + + (np.sin(theta_lvhaack(x)) ** 3) / 3 ) ** (0.5) / (np.pi**0.5) @@ -272,12 +272,12 @@ def theta(x): case "vonkarman": self.k = 0.5 - def theta(x): + def theta_vonkarman(x): return np.arccos(1 - 2 * max(min(x / self.length, 1), 0)) self.y_nosecone = Function( lambda x: self.base_radius - * (theta(x) - np.sin(2 * theta(x)) / 2) ** (0.5) + * (theta_vonkarman(x) - np.sin(2 * theta_vonkarman(x)) / 2) ** (0.5) / (np.pi**0.5) ) case "parabolic":