diff --git a/.pylintrc b/.pylintrc index e417e0b11..991d9445f 100644 --- a/.pylintrc +++ b/.pylintrc @@ -229,6 +229,7 @@ good-names=FlightPhases, center_of_mass_without_motor_to_CDM, motor_center_of_dry_mass_to_CDM, generic_motor_cesaroni_M1520, + Re, # Reynolds number # Good variable names regexes, separated by a comma. If names match any regex, # they will always be accepted diff --git a/CHANGELOG.md b/CHANGELOG.md index cd40ae8d0..73596780f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -32,6 +32,8 @@ Attention: The newest changes should be on top --> ### Added +- ENH: Add multi-dimensional drag coefficient support (Cd as function of M, Re, α) [#875](https://github.com/RocketPy-Team/RocketPy/pull/875) +- ENH: Add save functionality to `_MonteCarloPlots.all` method [#848](https://github.com/RocketPy-Team/RocketPy/pull/848) - ENH: add animations for motor propellant mass and tank fluid volumes [#894](https://github.com/RocketPy-Team/RocketPy/pull/894) - ENH: Rail button bending moments calculation in Flight class [#893](https://github.com/RocketPy-Team/RocketPy/pull/893) - ENH: Implement Bootstrapping for Confidence Interval Estimation [#891](https://github.com/RocketPy-Team/RocketPy/pull/897) diff --git a/rocketpy/mathutils/function.py b/rocketpy/mathutils/function.py index 9fe343687..3ebac65d6 100644 --- a/rocketpy/mathutils/function.py +++ b/rocketpy/mathutils/function.py @@ -22,6 +22,7 @@ LinearNDInterpolator, NearestNDInterpolator, RBFInterpolator, + RegularGridInterpolator, ) from rocketpy.plots.plot_helpers import show_or_save_plot @@ -43,6 +44,7 @@ "spline": 3, "shepard": 4, "rbf": 5, + "regular_grid": 6, } EXTRAPOLATION_TYPES = {"zero": 0, "natural": 1, "constant": 2} @@ -365,6 +367,18 @@ def set_extrapolation(self, method="constant"): self.__set_extrapolation_func() return self + @property + def is_multidimensional(self): + """Return True when the Function has domain dimension greater than 1. + + This abstracts checks for multi-dimensionality so callers don't need + to inspect internal attributes like ``__inputs__`` or ``__dom_dim__``. + """ + try: + return int(self.__dom_dim__) > 1 + except (AttributeError, TypeError): + return False + def __set_interpolation_func(self): # pylint: disable=too-many-statements """Defines interpolation function used by the Function. Each interpolation method has its own function`. @@ -449,6 +463,41 @@ def rbf_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disab self._interpolation_func = rbf_interpolation + elif interpolation == 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 + else: raise ValueError(f"Interpolation {interpolation} method not recognized.") @@ -635,6 +684,66 @@ def __get_value_opt_nd(self, *args): return result + def __get_value_opt_grid(self, *args): # pylint: disable=unused-private-member + """Evaluate the Function using RegularGridInterpolator for structured grids. + + This method is dynamically assigned in from_grid() class method. + + Parameters + ---------- + args : tuple + Values where the Function is to be evaluated. Must match the number + of dimensions of the grid. + + Returns + ------- + result : scalar or ndarray + Value of the Function at the specified points. + """ + # Check if we have the grid interpolator + if not hasattr(self, "_grid_interpolator"): + raise RuntimeError( + "Grid interpolator not initialized. Use from_grid() to create " + "a Function with grid interpolation." + ) + + # Convert args to appropriate format for RegularGridInterpolator + # RegularGridInterpolator expects points as (N, ndim) array + if len(args) != self.__dom_dim__: + raise ValueError( + f"Expected {self.__dom_dim__} arguments but got {len(args)}" + ) + + # Handle single point evaluation + point = np.array(args).reshape(1, -1) + + # Handle extrapolation based on the extrapolation setting + if self.__extrapolation__ == "constant": + # Clamp point to grid boundaries for constant extrapolation + for i, axis in enumerate(self._grid_axes): + point[0, i] = np.clip(point[0, i], axis[0], axis[-1]) + result = self._grid_interpolator(point) + elif self.__extrapolation__ == "zero": + # Check if point is outside bounds + outside_bounds = False + for i, axis in enumerate(self._grid_axes): + if point[0, i] < axis[0] or point[0, i] > axis[-1]: + outside_bounds = True + break + if outside_bounds: + result = np.array([0.0]) + else: + result = self._grid_interpolator(point) + else: + # Natural or other extrapolation - use interpolator directly + result = self._grid_interpolator(point) + + # Return scalar for single evaluation + if result.size == 1: + return float(result[0]) + + return result + def __determine_1d_domain_bounds(self, lower, upper): """Determine domain bounds for 1-D function discretization. @@ -3891,11 +4000,16 @@ def __validate_interpolation(self, interpolation): elif self.__dom_dim__ > 1: if interpolation is None: interpolation = "shepard" - if interpolation.lower() not in ["shepard", "linear", "rbf"]: + if interpolation.lower() not in [ + "shepard", + "linear", + "rbf", + "regular_grid", + ]: warnings.warn( ( "Interpolation method set to 'shepard'. The methods " - "'linear', 'shepard' and 'rbf' are supported for " + "'linear', 'shepard', 'rbf' and 'regular_grid' are supported for " "multiple dimensions." ), ) @@ -3950,6 +4064,250 @@ def to_dict(self, **kwargs): # pylint: disable=unused-argument "extrapolation": self.__extrapolation__, } + @classmethod + def from_grid( + cls, + grid_data, + axes, + inputs=None, + outputs=None, + interpolation="regular_grid", + extrapolation="constant", + flatten_for_compatibility=True, + **kwargs, + ): # pylint: disable=too-many-statements #TODO: Refactor this method into smaller methods + """Creates a Function from N-dimensional grid data. + + This method is designed for structured grid data, such as CFD simulation + results where values are computed on a regular grid. It uses + scipy.interpolate.RegularGridInterpolator for efficient interpolation. + + Parameters + ---------- + grid_data : ndarray + N-dimensional array containing the function values on the grid. + For example, for a 3D function Cd(M, Re, α), this would be a 3D array + where grid_data[i, j, k] = Cd(M[i], Re[j], α[k]). + axes : list of ndarray + List of 1D arrays defining the grid points along each axis. + Each array should be sorted in ascending order. + For example: [M_axis, Re_axis, alpha_axis]. + inputs : list of str, optional + Names of the input variables. If None, generic names will be used. + For example: ['Mach', 'Reynolds', 'Alpha']. + outputs : str, optional + Name of the output variable. For example: 'Cd'. + interpolation : str, optional + Interpolation method. Default is 'regular_grid'. + Currently only 'regular_grid' is supported for grid data. + extrapolation : str, optional + Extrapolation behavior. Default is ``'constant'`` which clamps to + edge values. Supported options are:: + + 'constant' + Use nearest edge value for out-of-bounds points (clamp). + 'zero' + Return zero for out-of-bounds points. + 'natural' + Use the interpolator's natural behavior: when the + underlying ``RegularGridInterpolator`` is created with + ``fill_value=None`` and ``method='linear'``, this results + in linear extrapolation based on the edge gradients. + + If an unsupported extrapolation value is supplied a ``ValueError`` + is raised. + flatten_for_compatibility : bool, optional + If True (default), creates flattened ``_domain``, ``_image``, and + ``source`` arrays for backward compatibility with existing Function + methods and serialization. For large N-dimensional grids (e.g., + 100x100x100 points), this requires O(n^d) additional memory where n + is the typical axis length and d is the number of dimensions. + Set to False to skip this flattening and reduce memory usage if + compatibility with legacy code paths is not required. + **kwargs : dict, optional + Additional arguments passed to the Function constructor. + + Returns + ------- + Function + A Function object using RegularGridInterpolator for evaluation. + + Notes + ----- + - Grid data must be on a regular (structured) grid. + - For unstructured data, use the regular Function constructor with + scattered points. + - Extrapolation with 'constant' mode uses the nearest edge values, + which is appropriate for aerodynamic coefficients where extrapolation + beyond the data range should be avoided. + + Examples + -------- + >>> import numpy as np + >>> # Create 3D drag coefficient data + >>> mach = np.array([0.0, 0.5, 1.0, 1.5, 2.0]) + >>> reynolds = np.array([1e5, 5e5, 1e6]) + >>> alpha = np.array([0.0, 2.0, 4.0, 6.0]) + >>> # Create a simple drag coefficient function + >>> M, Re, A = np.meshgrid(mach, reynolds, alpha, indexing='ij') + >>> cd_data = 0.3 + 0.1 * M + 1e-7 * Re + 0.01 * A + >>> # Create Function object + >>> cd_func = Function.from_grid( + ... cd_data, + ... [mach, reynolds, alpha], + ... inputs=['Mach', 'Reynolds', 'Alpha'], + ... outputs='Cd' + ... ) + >>> # Evaluate at a point + >>> cd_func(1.2, 3e5, 3.0) + 0.48000000000000004 + + """ + # Validate inputs + if not isinstance(grid_data, np.ndarray): + grid_data = np.array(grid_data) + + if not isinstance(axes, (list, tuple)): + raise ValueError("axes must be a list or tuple of 1D arrays") + + # Ensure all axes are numpy arrays + axes = [ + np.array(axis) if not isinstance(axis, np.ndarray) else axis + for axis in axes + ] + + # Check dimensions match + if len(axes) != grid_data.ndim: + raise ValueError( + f"Number of axes ({len(axes)}) must match grid_data dimensions " + f"({grid_data.ndim})" + ) + + # Check each axis matches corresponding grid dimension and is sorted + for i, axis in enumerate(axes): + if len(axis) != grid_data.shape[i]: + raise ValueError( + f"Axis {i} has {len(axis)} points but grid dimension {i} " + f"has {grid_data.shape[i]} points" + ) + # Check if axis is sorted in ascending order + if not np.all(np.diff(axis) > 0): + warnings.warn( + f"Axis {i} is not strictly sorted in ascending order. " + "RegularGridInterpolator requires sorted axes. " + "This may cause unexpected interpolation results.", + UserWarning, + ) + + # Set default inputs if not provided + if inputs is None: + inputs = [f"x{i}" for i in range(len(axes))] + elif len(inputs) != len(axes): + raise ValueError( + f"Number of inputs ({len(inputs)}) must match number of axes ({len(axes)})" + ) + + # Create a new Function instance + func = cls.__new__(cls) + + # Validate extrapolation option for grid-based interpolation + allowed_extrap = ("constant", "zero", "natural") + if extrapolation not in allowed_extrap: + raise ValueError( + "Unsupported extrapolation for grid interpolation. " + f"Supported values: {allowed_extrap}" + ) + + # Store grid-specific data first + func._grid_axes = axes + func._grid_data = grid_data + + # Create RegularGridInterpolator + # We handle extrapolation manually in __get_value_opt_grid, + # so we set bounds_error=False and let it extrapolate linearly + # (which we'll override when needed) + func._grid_interpolator = RegularGridInterpolator( + axes, + grid_data, + method="linear", + bounds_error=False, + fill_value=None, # Linear extrapolation (will be overridden by manual handling) + ) + + # Create placeholder domain and image for compatibility. + # For large grids this requires O(n^d) memory; set flatten_for_compatibility=False + # to skip this if legacy code compatibility is not required. + if flatten_for_compatibility: + mesh = np.meshgrid(*axes, indexing="ij") + domain_points = np.column_stack([m.ravel() for m in mesh]) + func._domain = domain_points + func._image = grid_data.ravel() + # Set source as flattened data array (for compatibility with serialization) + func.source = np.column_stack([domain_points, func._image]) + else: + # Minimal placeholders - grid interpolator is the primary data source + func._domain = None + func._image = None + func.source = None + + # Initialize basic attributes + func.__inputs__ = inputs + func.__outputs__ = outputs if outputs is not None else "f" + func.__interpolation__ = interpolation + func.__extrapolation__ = extrapolation + func.title = kwargs.get("title", None) + func.__img_dim__ = 1 + func.__cropped_domain__ = (None, None) + func._source_type = SourceType.ARRAY + func.__dom_dim__ = len(axes) + + # Set basic array attributes for compatibility + func.x_array = axes[0] + func.x_initial, func.x_final = axes[0][0], axes[0][-1] + if flatten_for_compatibility: + # For grid-based (N-D) functions, a 1-D `y_array` is not a meaningful + # representation of the function values. Some legacy code paths and + # serialization expect a `y_array` attribute to exist, so provide the + # full flattened image for compatibility rather than a truncated slice. + # Callers should avoid relying on `y_array` for multidimensional + # Functions; use the interpolator / `get_value_opt` instead. + func.y_array = func._image + # Use the global min/max of the flattened image as a sensible + # `y_initial`/`y_final` for compatibility with code that inspects + # scalar bounds. These describe the image range, not an ordering + # along any particular axis. + func.y_initial, func.y_final = ( + float(func._image.min()), + float(func._image.max()), + ) + else: + # Minimal placeholders when flattening is disabled + func.y_array = None + func.y_initial, func.y_final = ( + float(grid_data.min()), + float(grid_data.max()), + ) + if len(axes) > 2: + func.z_array = axes[2] + func.z_initial, func.z_final = axes[2][0], axes[2][-1] + + # Set get_value_opt to use grid interpolation + func.get_value_opt = func.__get_value_opt_grid + + # Set interpolation and extrapolation functions + func.__set_interpolation_func() + # Only set extrapolation function if we have flattened data, otherwise + # extrapolation is handled by __get_value_opt_grid directly + if flatten_for_compatibility: + func.__set_extrapolation_func() + + # Set inputs and outputs properly + func.set_inputs(inputs) + func.set_outputs(outputs) + func.set_title(func.title) + + return func + @classmethod def from_dict(cls, func_dict): """Creates a Function instance from a dictionary. diff --git a/rocketpy/rocket/rocket.py b/rocketpy/rocket/rocket.py index c3b1f364d..9ac1bb86e 100644 --- a/rocketpy/rocket/rocket.py +++ b/rocketpy/rocket/rocket.py @@ -341,20 +341,28 @@ def __init__( # pylint: disable=too-many-statements ) # Define aerodynamic drag coefficients - self.power_off_drag = Function( - power_off_drag, - "Mach Number", - "Drag Coefficient with Power Off", - "linear", - "constant", - ) - self.power_on_drag = Function( - power_on_drag, - "Mach Number", - "Drag Coefficient with Power On", - "linear", - "constant", - ) + # If already a Function, use it directly (preserves multi-dimensional drag) + if isinstance(power_off_drag, Function): + self.power_off_drag = power_off_drag + else: + self.power_off_drag = Function( + power_off_drag, + "Mach Number", + "Drag Coefficient with Power Off", + "linear", + "constant", + ) + + if isinstance(power_on_drag, Function): + self.power_on_drag = power_on_drag + else: + self.power_on_drag = Function( + power_on_drag, + "Mach Number", + "Drag Coefficient with Power On", + "linear", + "constant", + ) # Create a, possibly, temporary empty motor # self.motors = Components() # currently unused, only 1 motor is supported diff --git a/rocketpy/simulation/flight.py b/rocketpy/simulation/flight.py index 05d746b20..90a2c8fd8 100644 --- a/rocketpy/simulation/flight.py +++ b/rocketpy/simulation/flight.py @@ -1391,6 +1391,90 @@ def lateral_surface_wind(self): return -wind_u * np.cos(heading_rad) + wind_v * np.sin(heading_rad) + def __get_drag_coefficient(self, drag_function, mach, z, freestream_velocity_body): + """Calculate drag coefficient, handling both 1D and multi-dimensional functions. + + Parameters + ---------- + drag_function : Function + The drag coefficient function (power_on_drag or power_off_drag) + mach : float + Mach number + z : float + Altitude in meters + freestream_velocity_body : Vector or array-like + Freestream velocity in body frame [stream_vx_b, stream_vy_b, stream_vz_b] + + Returns + ------- + float + Drag coefficient value + """ + # Early return for 1D drag functions (only mach number) + if not isinstance(drag_function, Function) or not getattr( + drag_function, "is_multidimensional", False + ): + return drag_function.get_value_opt(mach) + + # Multi-dimensional drag function - calculate additional parameters + + # Calculate Reynolds number: Re = rho * V * L / mu + # where L is characteristic length (rocket diameter) + rho = self.env.density.get_value_opt(z) + mu = self.env.dynamic_viscosity.get_value_opt(z) + freestream_speed = np.linalg.norm(freestream_velocity_body) + characteristic_length = 2 * self.rocket.radius # Diameter + # Defensive: avoid division by zero or non-finite viscosity values. + # Use a small epsilon fallback if `mu` is zero, negative, NaN or infinite. + try: + mu_val = float(mu) + except (TypeError, ValueError, OverflowError): + # Only catch errors related to invalid numeric conversion. + # Avoid catching broad Exception to satisfy linters and + # allow other unexpected errors to surface. + mu_val = 0.0 + if not np.isfinite(mu_val) or mu_val <= 0.0: + mu_safe = 1e-10 + else: + mu_safe = mu_val + + reynolds = rho * freestream_speed * characteristic_length / mu_safe + + # Calculate angle of attack + # Angle between freestream velocity and rocket axis (z-axis in body frame) + # The z component of freestream velocity in body frame + if hasattr(freestream_velocity_body, "z"): + stream_vz_b = -freestream_velocity_body.z + else: + stream_vz_b = -freestream_velocity_body[2] + + # Normalize and calculate angle + if freestream_speed > 1e-6: + cos_alpha = stream_vz_b / freestream_speed + # Clamp to [-1, 1] to avoid numerical issues + cos_alpha = np.clip(cos_alpha, -1.0, 1.0) + alpha_rad = np.arccos(cos_alpha) + alpha_deg = np.rad2deg(alpha_rad) + else: + alpha_deg = 0.0 + + # Determine which parameters to pass based on input names + input_names = [name.lower() for name in drag_function.__inputs__] + args = [] + + for name in input_names: + if "mach" in name or name == "m": + args.append(mach) + elif "reynolds" in name or name == "re": + args.append(reynolds) + elif "alpha" in name or name == "a" or "attack" in name: + args.append(alpha_deg) + else: + # Unknown parameter, default to mach + args.append(mach) + + return drag_function.get_value_opt(*args) + def udot_rail1(self, t, u, post_processing=False): """Calculates derivative of u state vector with respect to time when rocket is flying in 1 DOF motion in the rail. @@ -1427,7 +1511,32 @@ def udot_rail1(self, t, u, post_processing=False): + (vz) ** 2 ) ** 0.5 free_stream_mach = free_stream_speed / self.env.speed_of_sound.get_value_opt(z) - drag_coeff = self.rocket.power_on_drag.get_value_opt(free_stream_mach) + + # For rail motion, rocket is constrained - velocity mostly along z-axis in body frame + # Calculate velocity in body frame (simplified for rail) + a11 = 1 - 2 * (e2**2 + e3**2) + a12 = 2 * (e1 * e2 - e0 * e3) + a13 = 2 * (e1 * e3 + e0 * e2) + a21 = 2 * (e1 * e2 + e0 * e3) + a22 = 1 - 2 * (e1**2 + e3**2) + a23 = 2 * (e2 * e3 - e0 * e1) + a31 = 2 * (e1 * e3 - e0 * e2) + a32 = 2 * (e2 * e3 + e0 * e1) + a33 = 1 - 2 * (e1**2 + e2**2) + + # Freestream velocity in body frame + wind_vx = self.env.wind_velocity_x.get_value_opt(z) + wind_vy = self.env.wind_velocity_y.get_value_opt(z) + stream_vx_b = a11 * (wind_vx - vx) + a21 * (wind_vy - vy) + a31 * (-vz) + stream_vy_b = a12 * (wind_vx - vx) + a22 * (wind_vy - vy) + a32 * (-vz) + stream_vz_b = a13 * (wind_vx - vx) + a23 * (wind_vy - vy) + a33 * (-vz) + + drag_coeff = self.__get_drag_coefficient( + self.rocket.power_on_drag, + free_stream_mach, + z, + [stream_vx_b, stream_vy_b, stream_vz_b], + ) # Calculate Forces pressure = self.env.pressure.get_value_opt(z) @@ -1597,12 +1706,38 @@ def u_dot(self, t, u, post_processing=False): # pylint: disable=too-many-locals ) ** 0.5 free_stream_mach = free_stream_speed / speed_of_sound + # Get rocket velocity in body frame (needed for drag calculation) + vx_b = a11 * vx + a21 * vy + a31 * vz + vy_b = a12 * vx + a22 * vy + a32 * vz + vz_b = a13 * vx + a23 * vy + a33 * vz + + # Calculate freestream velocity in body frame + stream_vx_b = ( + a11 * (wind_velocity_x - vx) + a21 * (wind_velocity_y - vy) + a31 * (-vz) + ) + stream_vy_b = ( + a12 * (wind_velocity_x - vx) + a22 * (wind_velocity_y - vy) + a32 * (-vz) + ) + stream_vz_b = ( + a13 * (wind_velocity_x - vx) + a23 * (wind_velocity_y - vy) + a33 * (-vz) + ) + # Determine aerodynamics forces # Determine Drag Force if t < self.rocket.motor.burn_out_time: - drag_coeff = self.rocket.power_on_drag.get_value_opt(free_stream_mach) + drag_coeff = self.__get_drag_coefficient( + self.rocket.power_on_drag, + free_stream_mach, + z, + [stream_vx_b, stream_vy_b, stream_vz_b], + ) else: - drag_coeff = self.rocket.power_off_drag.get_value_opt(free_stream_mach) + drag_coeff = self.__get_drag_coefficient( + self.rocket.power_off_drag, + free_stream_mach, + z, + [stream_vx_b, stream_vy_b, stream_vz_b], + ) rho = self.env.density.get_value_opt(z) R3 = -0.5 * rho * (free_stream_speed**2) * self.rocket.area * drag_coeff for air_brakes in self.rocket.air_brakes: @@ -1624,10 +1759,6 @@ def u_dot(self, t, u, post_processing=False): # pylint: disable=too-many-locals # Off center moment M1 += self.rocket.cp_eccentricity_y * R3 M2 -= self.rocket.cp_eccentricity_x * R3 - # Get rocket velocity in body frame - vx_b = a11 * vx + a21 * vy + a31 * vz - vy_b = a12 * vx + a22 * vy + a32 * vz - vz_b = a13 * vx + a23 * vy + a33 * vz # Calculate lift and moment for each component of the rocket velocity_in_body_frame = Vector([vx_b, vy_b, vz_b]) w = Vector([omega1, omega2, omega3]) @@ -1992,6 +2123,12 @@ def u_dot_generalized(self, t, u, post_processing=False): # pylint: disable=too speed_of_sound = self.env.speed_of_sound.get_value_opt(z) free_stream_mach = free_stream_speed / speed_of_sound + # Get rocket velocity in body frame (needed for drag calculation) + velocity_in_body_frame = Kt @ v + # Calculate freestream velocity in body frame + freestream_velocity = wind_velocity - v + freestream_velocity_body = Kt @ freestream_velocity + if self.rocket.motor.burn_start_time < t < self.rocket.motor.burn_out_time: pressure = self.env.pressure.get_value_opt(z) net_thrust = max( @@ -1999,10 +2136,20 @@ def u_dot_generalized(self, t, u, post_processing=False): # pylint: disable=too + self.rocket.motor.pressure_thrust(pressure), 0, ) - drag_coeff = self.rocket.power_on_drag.get_value_opt(free_stream_mach) + drag_coeff = self.__get_drag_coefficient( + self.rocket.power_on_drag, + free_stream_mach, + z, + freestream_velocity_body, + ) else: net_thrust = 0 - drag_coeff = self.rocket.power_off_drag.get_value_opt(free_stream_mach) + drag_coeff = self.__get_drag_coefficient( + self.rocket.power_off_drag, + free_stream_mach, + z, + freestream_velocity_body, + ) R3 += -0.5 * rho * (free_stream_speed**2) * self.rocket.area * drag_coeff for air_brakes in self.rocket.air_brakes: if air_brakes.deployment_level > 0: @@ -2020,8 +2167,6 @@ def u_dot_generalized(self, t, u, post_processing=False): # pylint: disable=too R3 = air_brakes_force # Substitutes rocket drag coefficient else: R3 += air_brakes_force - # Get rocket velocity in body frame - velocity_in_body_frame = Kt @ v # Calculate lift and moment for each component of the rocket for aero_surface, _ in self.rocket.aerodynamic_surfaces: # Component cp relative to CDM in body frame diff --git a/tests/fixtures/flight/flight_fixtures.py b/tests/fixtures/flight/flight_fixtures.py index da47e6cdb..e0e960317 100644 --- a/tests/fixtures/flight/flight_fixtures.py +++ b/tests/fixtures/flight/flight_fixtures.py @@ -1,6 +1,7 @@ +import numpy as np import pytest -from rocketpy import Flight +from rocketpy import Flight, Function, Rocket @pytest.fixture @@ -273,3 +274,90 @@ def flight_calisto_with_sensors(calisto_with_sensors, example_plain_env): time_overshoot=False, terminate_on_apogee=True, ) + + +@pytest.fixture +def flight_alpha(example_plain_env, cesaroni_m1670): + """Fixture that returns a Flight using an alpha-dependent 3D Cd function.""" + # Create grid data + mach = np.array([0.0, 0.5, 1.0, 1.5]) + reynolds = np.array([1e5, 1e6]) + alpha = np.array([0.0, 5.0, 10.0, 15.0]) + M, _, A = np.meshgrid(mach, reynolds, alpha, indexing="ij") + cd_data = 0.3 + 0.05 * M + 0.03 * A + cd_data = np.clip(cd_data, 0.2, 2.0) + + drag_func = Function.from_grid( + cd_data, + [mach, reynolds, alpha], + inputs=["Mach", "Reynolds", "Alpha"], + outputs="Cd", + ) + + env = example_plain_env + env.set_atmospheric_model(type="standard_atmosphere") + + # Build rocket and flight + rocket = Rocket( + radius=0.0635, + mass=16.24, + inertia=(6.321, 6.321, 0.034), + power_off_drag=drag_func, + power_on_drag=drag_func, + center_of_mass_without_motor=0, + coordinate_system_orientation="tail_to_nose", + ) + rocket.set_rail_buttons(0.2, -0.5, 30) + rocket.add_motor(cesaroni_m1670, position=-1.255) + + return Flight( + rocket=rocket, + environment=env, + rail_length=5.2, + inclination=85, + heading=0, + ) + + +@pytest.fixture +def flight_flat(example_plain_env, cesaroni_m1670): + """Fixture that returns a Flight using an alpha-averaged (flat) Cd function.""" + # Create grid data + mach = np.array([0.0, 0.5, 1.0, 1.5]) + reynolds = np.array([1e5, 1e6]) + alpha = np.array([0.0, 5.0, 10.0, 15.0]) + M, _, A = np.meshgrid(mach, reynolds, alpha, indexing="ij") + cd_data = 0.3 + 0.05 * M + 0.03 * A + cd_data = np.clip(cd_data, 0.2, 2.0) + + cd_flat = cd_data.mean(axis=2) + drag_flat = Function.from_grid( + cd_flat, + [mach, reynolds], + inputs=["Mach", "Reynolds"], + outputs="Cd", + ) + + env = example_plain_env + env.set_atmospheric_model(type="standard_atmosphere") + + # Build rocket and flight + rocket = Rocket( + radius=0.0635, + mass=16.24, + inertia=(6.321, 6.321, 0.034), + power_off_drag=drag_flat, + power_on_drag=drag_flat, + center_of_mass_without_motor=0, + coordinate_system_orientation="tail_to_nose", + ) + rocket.set_rail_buttons(0.2, -0.5, 30) + rocket.add_motor(cesaroni_m1670, position=-1.255) + + return Flight( + rocket=rocket, + environment=env, + rail_length=5.2, + inclination=85, + heading=0, + ) diff --git a/tests/integration/test_multidim_drag.py b/tests/integration/test_multidim_drag.py new file mode 100644 index 000000000..a6280e53b --- /dev/null +++ b/tests/integration/test_multidim_drag.py @@ -0,0 +1,157 @@ +"""Integration tests for multi-dimensional drag coefficient support.""" + +import numpy as np + +from rocketpy import Flight, Function, Rocket + + +def test_flight_with_1d_drag(flight_calisto): + """Test that flights with 1D drag curves still work (backward compatibility).""" + + # `flight_calisto` is a fixture that already runs the simulation + flight = flight_calisto + + # Check that flight completed successfully + assert flight.t_final > 0 + assert flight.apogee > 0 + assert flight.apogee_time > 0 + + +def test_flight_with_3d_drag_basic(example_plain_env, cesaroni_m1670): + """Test that a simple 3D drag function works.""" + # Use fixtures for environment and motor + env = example_plain_env + env.set_atmospheric_model(type="standard_atmosphere") + motor = cesaroni_m1670 + + # Create 3D drag + mach = np.array([0.0, 0.5, 1.0, 1.5, 2.0]) + reynolds = np.array([1e5, 5e5, 1e6]) + alpha = np.array([0.0, 2.0, 4.0, 6.0]) + + M, Re, A = np.meshgrid(mach, reynolds, alpha, indexing="ij") + cd_data = 0.3 + 0.1 * M - 1e-7 * Re + 0.01 * A + cd_data = np.clip(cd_data, 0.2, 1.0) + + power_off_drag = Function.from_grid( + cd_data, + [mach, reynolds, alpha], + inputs=["Mach", "Reynolds", "Alpha"], + outputs="Cd", + ) + power_on_drag = Function.from_grid( + cd_data * 1.1, + [mach, reynolds, alpha], + inputs=["Mach", "Reynolds", "Alpha"], + outputs="Cd", + ) + + # Create rocket + rocket = Rocket( + radius=0.0635, + mass=16.24, + inertia=(6.321, 6.321, 0.034), + power_off_drag=power_off_drag, + power_on_drag=power_on_drag, + center_of_mass_without_motor=0, + coordinate_system_orientation="tail_to_nose", + ) + rocket.set_rail_buttons(0.2, -0.5, 30) + rocket.add_motor(motor, position=-1.255) + + # Run flight + flight = Flight( + rocket=rocket, + environment=env, + rail_length=5.2, + inclination=85, + heading=0, + ) + + # Check results - should launch and have non-zero apogee + assert flight.apogee > 100, f"Apogee too low: {flight.apogee}m" + assert flight.apogee < 5000, f"Apogee too high: {flight.apogee}m" + assert hasattr(flight, "angle_of_attack") + + +def test_3d_drag_with_varying_alpha(): + """Test that 3D drag responds to angle of attack changes. + + This test only verifies the Function mapping from alpha -> Cd. The + integration-level comparison is placed in a separate test to keep each + test function small and easier to lint/maintain. + """ + # Create drag function with strong alpha dependency + mach = np.array([0.0, 0.5, 1.0, 1.5]) + reynolds = np.array([1e5, 1e6]) + alpha = np.array([0.0, 5.0, 10.0, 15.0]) + + M, _, A = np.meshgrid(mach, reynolds, alpha, indexing="ij") + # Strong alpha dependency: Cd increases significantly with alpha + cd_data = 0.3 + 0.05 * M + 0.03 * A + cd_data = np.clip(cd_data, 0.2, 2.0) + + drag_func = Function.from_grid( + cd_data, + [mach, reynolds, alpha], + inputs=["Mach", "Reynolds", "Alpha"], + outputs="Cd", + ) + + # Test at different angles of attack (direct function call) + # At zero alpha, Cd should be lower + cd_0 = drag_func(0.8, 5e5, 0.0) + cd_10 = drag_func(0.8, 5e5, 10.0) + + # Cd should increase with alpha + assert cd_10 > cd_0 + assert cd_10 - cd_0 > 0.2 # Should show significant difference + + +def test_flight_apogee_diff(flight_alpha, flight_flat): + """Run paired flights (fixtures) and assert their apogees differ.""" + + # Flights should both launch + assert flight_alpha.apogee > 100 + assert flight_flat.apogee > 100 + + # Apogees should differ + assert flight_alpha.apogee != flight_flat.apogee + + +def test_flight_cd_sample_consistency(flight_alpha, flight_flat): + """Sample Cd during a flight and ensure Cd difference matches apogee ordering. + + Uses the `flight_alpha` and `flight_flat` fixtures which provide paired + flights constructed with alpha-dependent and alpha-averaged Cd functions. + """ + + # Sample a mid-ascent time and compare Cd evaluations + speeds = flight_alpha.free_stream_speed[:, 1] + idx_candidates = np.where(speeds > 5)[0] + assert idx_candidates.size > 0 + idx = idx_candidates[len(idx_candidates) // 2] + t_sample = flight_alpha.time[idx] + + mach_sample = flight_alpha.mach_number.get_value_opt(t_sample) + v_sample = flight_alpha.free_stream_speed.get_value_opt(t_sample) + reynolds_sample = ( + flight_alpha.density.get_value_opt(t_sample) + * v_sample + * (2 * flight_alpha.rocket.radius) + / flight_alpha.dynamic_viscosity.get_value_opt(t_sample) + ) + alpha_sample = flight_alpha.angle_of_attack.get_value_opt(t_sample) + + cd_alpha_sample = flight_alpha.rocket.power_on_drag.get_value_opt( + mach_sample, reynolds_sample, alpha_sample + ) + cd_flat_sample = flight_flat.rocket.power_on_drag.get_value_opt( + mach_sample, reynolds_sample + ) + + assert cd_alpha_sample != cd_flat_sample + if cd_alpha_sample > cd_flat_sample: + assert flight_alpha.apogee < flight_flat.apogee + else: + assert flight_alpha.apogee > flight_flat.apogee diff --git a/tests/unit/mathutils/test_function_from_grid.py b/tests/unit/mathutils/test_function_from_grid.py new file mode 100644 index 000000000..cc0e214c5 --- /dev/null +++ b/tests/unit/mathutils/test_function_from_grid.py @@ -0,0 +1,40 @@ +import numpy as np +import pytest + +from rocketpy.mathutils.function import Function + + +def test_from_grid_unsupported_extrapolation_raises(): + """from_grid should reject unsupported extrapolation names with ValueError.""" + mach = np.array([0.0, 1.0]) + reynolds = np.array([1e5, 2e5]) + grid = np.zeros((mach.size, reynolds.size)) + + with pytest.raises(ValueError): + Function.from_grid(grid, [mach, reynolds], extrapolation="unsupported_mode") + + +def test_from_grid_is_multidimensional_property(): + """Ensure `is_multidimensional` is True for ND grid Functions and False for 1D.""" + mach = np.array([0.0, 0.5, 1.0]) + reynolds = np.array([1e5, 2e5, 3e5]) + alpha = np.array([0.0, 2.0]) + + M, R, A = np.meshgrid(mach, reynolds, alpha, indexing="ij") + cd_data = 0.1 + 0.2 * M + 1e-7 * R + 0.01 * A + + func_nd = Function.from_grid( + cd_data, + [mach, reynolds, alpha], + inputs=["Mach", "Reynolds", "Alpha"], + outputs="Cd", + ) + + assert hasattr(func_nd, "is_multidimensional") + assert func_nd.is_multidimensional is True + + # 1D Function constructed from a two-column array should not be multidimensional + src = np.column_stack((mach, 0.5 + 0.1 * mach)) + func_1d = Function(src, inputs=["Mach"], outputs="Cd") + assert hasattr(func_1d, "is_multidimensional") + assert func_1d.is_multidimensional is False diff --git a/tests/unit/mathutils/test_function_grid.py b/tests/unit/mathutils/test_function_grid.py new file mode 100644 index 000000000..3752d4de0 --- /dev/null +++ b/tests/unit/mathutils/test_function_grid.py @@ -0,0 +1,352 @@ +"""Unit tests for Function.from_grid() method and grid interpolation.""" + +import warnings + +import numpy as np +import pytest + +from rocketpy import Function + + +def test_from_grid_1d(): + """Test from_grid with 1D data (edge case).""" + x = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) + y_data = np.array([0.0, 1.0, 4.0, 9.0, 16.0]) # y = x^2 + + func = Function.from_grid(y_data, [x], inputs=["x"], outputs="y") + + # Test interpolation + assert abs(func(1.5) - 2.25) < 0.5 # Should be close to 1.5^2 + + +def test_from_grid_2d(): + """Test from_grid with 2D data.""" + x = np.array([0.0, 1.0, 2.0]) + y = np.array([0.0, 1.0, 2.0]) + + # Create grid: f(x, y) = x + 2*y + X, Y = np.meshgrid(x, y, indexing="ij") + z_data = X + 2 * Y + + func = Function.from_grid(z_data, [x, y], inputs=["x", "y"], outputs="z") + + # Test exact points + assert func(0.0, 0.0) == 0.0 + assert func(1.0, 1.0) == 3.0 + assert func(2.0, 2.0) == 6.0 + + # Test interpolation + result = func(1.0, 0.5) + expected = 1.0 + 2 * 0.5 # = 2.0 + assert abs(result - expected) < 0.01 + + +def test_from_grid_3d_drag_coefficient(): + """Test from_grid with 3D drag coefficient data (Mach, Reynolds, Alpha).""" + # Create sample aerodynamic data + mach = np.array([0.0, 0.5, 1.0, 1.5, 2.0]) + reynolds = np.array([1e5, 5e5, 1e6]) + alpha = np.array([0.0, 2.0, 4.0, 6.0]) + + # Create a simple drag coefficient model + # Cd increases with Mach and alpha, slight dependency on Reynolds + M, Re, A = np.meshgrid(mach, reynolds, alpha, indexing="ij") + cd_data = 0.3 + 0.1 * M - 1e-7 * Re + 0.01 * A + + cd_func = Function.from_grid( + cd_data, + [mach, reynolds, alpha], + inputs=["Mach", "Reynolds", "Alpha"], + outputs="Cd", + ) + + # Test at grid points + assert abs(cd_func(0.0, 1e5, 0.0) - 0.29) < 0.01 # 0.3 - 1e-7*1e5 + assert abs(cd_func(1.0, 5e5, 0.0) - 0.35) < 0.01 # 0.3 + 0.1*1 - 1e-7*5e5 + + # Test interpolation between points + result = cd_func(0.5, 3e5, 1.0) + # Expected roughly: 0.3 + 0.1*0.5 - 1e-7*3e5 + 0.01*1.0 = 0.32 + assert 0.31 < result < 0.34 + + +def test_from_grid_extrapolation_constant(): + """Test that constant extrapolation clamps to edge values.""" + x = np.array([0.0, 1.0, 2.0]) + y = np.array([0.0, 1.0, 4.0]) # y = x^2 + + func = Function.from_grid( + y, [x], inputs=["x"], outputs="y", extrapolation="constant" + ) + + # Test below lower bound - should return value at x=0 + assert func(-1.0) == 0.0 + + # Test above upper bound - should return value at x=2 + assert func(3.0) == 4.0 + + +def test_from_grid_validation_errors(): + """Test that from_grid raises appropriate errors for invalid inputs.""" + x = np.array([0.0, 1.0, 2.0]) + y = np.array([0.0, 1.0, 2.0]) + + # Mismatched dimensions + X, Y = np.meshgrid(x, y, indexing="ij") + z_data = X + Y + + # Wrong number of axes + with pytest.raises(ValueError, match="Number of axes"): + Function.from_grid(z_data, [x], inputs=["x"], outputs="z") + + # Wrong axis length + with pytest.raises(ValueError, match="Axis 1 has"): + Function.from_grid( + z_data, [x, np.array([0.0, 1.0])], inputs=["x", "y"], outputs="z" + ) + + # Wrong number of inputs + with pytest.raises(ValueError, match="Number of inputs"): + Function.from_grid(z_data, [x, y], inputs=["x"], outputs="z") + + +def test_from_grid_default_inputs(): + """Test that from_grid uses default input names when not provided.""" + x = np.array([0.0, 1.0, 2.0]) + y = np.array([0.0, 1.0, 2.0]) + + X, Y = np.meshgrid(x, y, indexing="ij") + z_data = X + Y + + func = Function.from_grid(z_data, [x, y]) + + # Should use default names + assert "x0" in func.__inputs__ + assert "x1" in func.__inputs__ + + +def test_from_grid_backward_compatibility(): + """Test that regular Function creation still works after adding from_grid.""" + # Test 1D function from list + func1 = Function([[0, 0], [1, 1], [2, 4], [3, 9]]) + assert func1(1.5) > 0 # Should interpolate + + # Test 2D function from array + data = np.array([[0, 0, 0], [1, 0, 1], [0, 1, 2], [1, 1, 3]]) + func2 = Function(data) + assert func2(0.5, 0.5) > 0 # Should interpolate + + # Test callable function + func3 = Function(lambda x: x**2) + assert func3(2) == 4 + + +def test_regular_grid_without_grid_interpolator_warns(): + """Test that setting `regular_grid` without a grid interpolator warns. + + This test constructs a Function from scattered points (no structured + grid). If `regular_grid` interpolation is later selected without a + grid interpolator being configured, the implementation currently + falls back to shepard interpolation and should emit a warning. The + test ensures a warning is raised in this scenario. + """ + # Create a 2D function with scattered points (not structured grid) + source = [(0, 0, 0), (1, 0, 1), (0, 1, 2), (1, 1, 3)] + func = Function( + source=source, inputs=["x", "y"], outputs="z", interpolation="shepard" + ) + + # Now manually change interpolation to regular_grid without setting up the grid + # This simulates the fallback scenario + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + func.set_interpolation("regular_grid") + + # Check that a warning was issued + assert len(w) == 1 + assert "falling back to shepard interpolation" in str(w[0].message) + + +def test_shepard_fallback_2d_interpolation(): + """Test that shepard_fallback produces correct interpolation for 2D data. + + This test verifies the fallback interpolation works correctly when + regular_grid is set without a grid interpolator. + """ + # Create a 2D function: z = x + y + source = [ + (0, 0, 0), # f(0, 0) = 0 + (1, 0, 1), # f(1, 0) = 1 + (0, 1, 1), # f(0, 1) = 1 + (1, 1, 2), # f(1, 1) = 2 + ] + + # First, create with shepard to get baseline results + func_shepard = Function( + source=source, inputs=["x", "y"], outputs="z", interpolation="shepard" + ) + + # Create another function and trigger the fallback + func_fallback = Function( + source=source, inputs=["x", "y"], outputs="z", interpolation="shepard" + ) + + # Trigger fallback + with warnings.catch_warnings(): + warnings.simplefilter("ignore") # Suppress warnings for this test + func_fallback.set_interpolation("regular_grid") + + # Test that both produce the same results at exact points + assert func_fallback(0, 0) == func_shepard(0, 0) + assert func_fallback(1, 1) == func_shepard(1, 1) + + # Test interpolation at an intermediate point + result_fallback = func_fallback(0.5, 0.5) + result_shepard = func_shepard(0.5, 0.5) + assert np.isclose(result_fallback, result_shepard, atol=1e-6) + + +def test_shepard_fallback_3d_interpolation(): + """Test that shepard_fallback produces correct interpolation for 3D data. + + This test verifies the fallback interpolation works correctly for + 3-dimensional input data. + """ + # Create a 3D function: w = x + y + z + source = [ + (0, 0, 0, 0), # f(0, 0, 0) = 0 + (1, 0, 0, 1), # f(1, 0, 0) = 1 + (0, 1, 0, 1), # f(0, 1, 0) = 1 + (0, 0, 1, 1), # f(0, 0, 1) = 1 + (1, 1, 1, 3), # f(1, 1, 1) = 3 + ] + + # Create with shepard to get baseline results + func_shepard = Function( + source=source, + inputs=["x", "y", "z"], + outputs="w", + interpolation="shepard", + ) + + # Create another function and trigger the fallback + func_fallback = Function( + source=source, + inputs=["x", "y", "z"], + outputs="w", + interpolation="shepard", + ) + + # Trigger fallback + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + func_fallback.set_interpolation("regular_grid") + + # Test that both produce the same results at exact points + assert func_fallback(0, 0, 0) == func_shepard(0, 0, 0) + assert func_fallback(1, 1, 1) == func_shepard(1, 1, 1) + + # Test interpolation at an intermediate point + result_fallback = func_fallback(0.5, 0.5, 0.5) + result_shepard = func_shepard(0.5, 0.5, 0.5) + assert np.isclose(result_fallback, result_shepard, atol=1e-6) + + +def test_shepard_fallback_at_exact_data_points(): + """Test that shepard_fallback returns exact values at data points. + + When querying at exact data points, the fallback should return the + exact value stored at that point. + """ + # Create a 2D function + source = [ + (0, 0, 10), + (1, 0, 20), + (0, 1, 30), + (1, 1, 40), + ] + + func = Function( + source=source, inputs=["x", "y"], outputs="z", interpolation="shepard" + ) + + # Trigger fallback + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + func.set_interpolation("regular_grid") + + # Test exact data points - should return exact values + assert func(0, 0) == 10 + assert func(1, 0) == 20 + assert func(0, 1) == 30 + assert func(1, 1) == 40 + + +def test_from_grid_unsorted_axis_warns(): + """Test that from_grid warns when axes are not sorted in ascending order.""" + y_data = np.array([0.0, 1.0, 4.0]) + + # Test with unsorted axis (descending order) + unsorted_axis = np.array([2.0, 1.0, 0.0]) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + Function.from_grid(y_data, [unsorted_axis], inputs=["x"], outputs="y") + + # Check that a warning was issued + assert len(w) == 1 + assert "not strictly sorted in ascending order" in str(w[0].message) + + +def test_from_grid_repeated_values_warns(): + """Test that from_grid warns when axes have repeated values. + + Note: RegularGridInterpolator requires strictly ascending or descending + axes. Repeated values will cause scipy to raise a ValueError after our + warning is issued. + """ + y_data = np.array([0.0, 1.0, 4.0]) + + # Test with repeated values (not strictly ascending) + repeated_axis = np.array([0.0, 1.0, 1.0]) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + # Scipy will raise ValueError after our warning, so we expect both + try: + Function.from_grid(y_data, [repeated_axis], inputs=["x"], outputs="y") + except ValueError as e: + # scipy raises this error for non-strictly-sorted axes + assert "strictly ascending" in str(e).lower() or "dimension 0" in str(e) + + # Check that a warning was issued before the error + assert len(w) == 1 + assert "not strictly sorted in ascending order" in str(w[0].message) + + +def test_from_grid_flatten_for_compatibility_false(): + """Test that flatten_for_compatibility=False skips flattening.""" + x = np.array([0.0, 1.0, 2.0]) + y = np.array([0.0, 1.0]) + + X, Y = np.meshgrid(x, y, indexing="ij") + z_data = X + Y + + func = Function.from_grid( + z_data, + [x, y], + inputs=["x", "y"], + outputs="z", + flatten_for_compatibility=False, + ) + + # Check that flattened attributes are None + assert func._domain is None + assert func._image is None + assert func.source is None + assert func.y_array is None + + # But the function should still work correctly + assert func(0.0, 0.0) == 0.0 + assert func(1.0, 1.0) == 2.0 + assert func(2.0, 1.0) == 3.0