diff --git a/rocketpy/mathutils/function.py b/rocketpy/mathutils/function.py index 9fe343687..a1fbc1595 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, + "linear_grid": 6, } EXTRAPOLATION_TYPES = {"zero": 0, "natural": 1, "constant": 2} @@ -449,6 +451,41 @@ def rbf_interpolation(x, x_min, x_max, x_data, y_data, coeffs): # pylint: disab self._interpolation_func = rbf_interpolation + elif interpolation == 6: # linear_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 +672,66 @@ def __get_value_opt_nd(self, *args): return result + def __get_value_opt_grid(self, *args): + """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 +3988,11 @@ 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", "linear_grid"]: warnings.warn( ( "Interpolation method set to 'shepard'. The methods " - "'linear', 'shepard' and 'rbf' are supported for " + "'linear', 'shepard', 'rbf' and 'linear_grid' are supported for " "multiple dimensions." ), ) @@ -3950,6 +4047,181 @@ 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="linear_grid", + extrapolation="constant", + **kwargs, + ): + """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 'linear_grid'. + Currently only 'linear_grid' is supported for grid data. + extrapolation : str, optional + Extrapolation behavior. Default is 'constant', which clamps to edge values. + 'constant': Use nearest edge value for out-of-bounds points. + 'zero': Return zero for out-of-bounds points. + **kwargs : dict, optional + Additional arguments passed to the Function constructor. + + Returns + ------- + Function + A Function object using RegularGridInterpolator for evaluation. + + 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) + + 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. + """ + # 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 + 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" + ) + + # 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) + + # 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 + # This flattens the grid for any code expecting these attributes + 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, etc.) + func.source = np.column_stack([domain_points, func._image]) + + # 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] + func.y_array = func._image[: len(axes[0])] # Placeholder + func.y_initial, func.y_final = func._image[0], func._image[-1] + 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() + 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 1112a98f3..f5b3f13c1 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 a38be7d93..47ffa9af4 100644 --- a/rocketpy/simulation/flight.py +++ b/rocketpy/simulation/flight.py @@ -1348,6 +1348,76 @@ 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 + """ + # Check if drag function is multi-dimensional by examining its inputs + if hasattr(drag_function, "__inputs__") and len(drag_function.__inputs__) > 1: + # 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 + reynolds = rho * freestream_speed * characteristic_length / mu + + # 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) + else: + # 1D drag function - only mach number + return drag_function.get_value_opt(mach) + 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. @@ -1384,7 +1454,36 @@ 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) + + vx_b = a11 * vx + a21 * vy + a31 * vz + vy_b = a12 * vx + a22 * vy + a32 * vz + vz_b = a13 * vx + a23 * vy + a33 * vz + + # 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) @@ -1552,12 +1651,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: @@ -1579,10 +1704,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]) @@ -1831,6 +1952,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( @@ -1838,10 +1965,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: @@ -1859,8 +1996,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/integration/test_multidim_drag.py b/tests/integration/test_multidim_drag.py new file mode 100644 index 000000000..642a2a263 --- /dev/null +++ b/tests/integration/test_multidim_drag.py @@ -0,0 +1,126 @@ +"""Integration tests for multi-dimensional drag coefficient support.""" + +import numpy as np + +from rocketpy import Environment, Flight, Function, Rocket, SolidMotor + + +def test_flight_with_1d_drag(example_plain_env, calisto): + """Test that flights with 1D drag curves still work (backward compatibility).""" + + flight = Flight( + rocket=calisto, + environment=example_plain_env, + rail_length=5.2, + inclination=85, + heading=0, + ) + + # 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(): + """Test that a simple 3D drag function works.""" + # Create environment + env = Environment(gravity=9.81) + env.set_atmospheric_model(type="standard_atmosphere") + + # Create motor with simple constant thrust + motor = SolidMotor( + thrust_source=lambda t: 2000 if t < 3 else 0, + burn_time=3.0, + nozzle_radius=0.033, + dry_mass=1.815, + dry_inertia=(0.125, 0.125, 0.002), + center_of_dry_mass_position=0.317, + grains_center_of_mass_position=0.397, + grain_number=5, + grain_separation=0.005, + grain_density=1815, + grain_outer_radius=0.033, + grain_initial_inner_radius=0.015, + grain_initial_height=0.120, + nozzle_position=0, + throat_radius=0.011, + ) + + # 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.""" + # 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, Re, 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 + # 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 diff --git a/tests/unit/mathutils/test_function_grid.py b/tests/unit/mathutils/test_function_grid.py new file mode 100644 index 000000000..994689544 --- /dev/null +++ b/tests/unit/mathutils/test_function_grid.py @@ -0,0 +1,139 @@ +"""Unit tests for Function.from_grid() method and grid interpolation.""" + +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