ENH: Add multi-dimensional drag coefficient support (Cd as function of M, Re, α) (#875)

* Add grid interpolation support to Function class with from_grid() method

Co-authored-by: Gui-FernandesBR <[email protected]>

* Add multi-dimensional drag coefficient support to Flight class and integration tests

Co-authored-by: Gui-FernandesBR <[email protected]>

* Run ruff format on modified files

Co-authored-by: Gui-FernandesBR <[email protected]>

* MNt: refactoring get_drag_coefficient in flight.py

- MNT: velocity_body was not being used in get_drag_coefficient, removed it as an input

* MNT: refactoring in flight.py and lint corrections to function.py and test_multidim_drag.py

-MNT:  removed unused velocity in body frame parameter requirement from all instances of get_drag_coefficient in flight.py

- MNT: corrected docstring for get_value_opt_grid in function.py

- MNT: shifted import of classes before the definition of functions in test_multidim_drag.py

* MNT: refactoring flight.py to remove unused parameters

* MNT: correction of docstring function.py

- MNT: rearranged the docstring of from_grid in function.py to match the expected output of doctest

* MNT: make format and lint corrections to function.py

- MNT: reran make format and lint on function.py to correct after previous changes to from_grid

* MNT: pylint adjustments for new methods in function.py

- MNT: disables pylint unused private member for get_value_opt_grid as it is called upon dynamically by from_grid

- MNT: disabled pylint too many statement for from_grid for now and added a to-do to refactor it into smaller methods/helper functions

- MNT: updated .pylintrc to record Re as good name

* MNt: make format after previous change to function.py

* MNT: removed Re where unused in test_multidim_drag.py

- MNT: Re variable was unused in test_3d_drag_with_varying_alpha thus replaced it

* TST: Add tests for shepard_fallback in test_function_grid.py (#879)

* Add tests for shepard_fallback in test_function_grid.py

Co-authored-by: aZira371 <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: aZira371 <[email protected]>

* TST: test_multidim_drag.py

- TST: Added integration-level check to verify the flight simulation actually uses alpha when evaluating multi dim drag coeff.
- TST: utilized pytest fixtures where possible.

* MNT: addition of is_multidimensional to function.py

- MNT: Function.is_multidimensional property in function.py. It returns True when the function's internal domain dimension is greater than 1 and safely returns False on errors.

- MNT: Replaced the ad-hoc hasattr/len check in flight.py with a clearer check:
if isinstance and drag_function.is_multidimensional

* MNT: Added validation in from_grid in function.py to raise a ValueError when an unsupported extrapolation value is provided.

* ENH: Added alpha-sensitive flight fixtures to flight_fixtures.py

ENH: Used  shared flight fixtures and simplify multi-drag integration tests in test_multidim_drag.py

ENH: Exposes multidimensionality and validated grid extrapolation in function.py

* MNT: renamed linear_grid to regular_grid for easy to understand nomenclature

- MNT: added a fallback to reynolds calculation in flight.py to avoid mu is zero case

- MNT: updated test_shepard_fallback_warning name and docstring to match implementation in test_function_grid.py

* MNT:  replaced the broad except Exception: with except (TypeError, ValueError, OverflowError): in get_drag_coefficient of flight.py

* TST: added from_grid unit tests to cover constructor-level validation and the explicit API.

* MNT: format and lint update to test_function_from_grid.py

* DOC: changelog.md update for multidim drag

* Address review feedback: add unsorted axis warning, flatten_for_compatibility parameter, and early return guard clause

Co-authored-by: aZira371 <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: Gui-FernandesBR <[email protected]>
Co-authored-by: Ishan <[email protected]>
Co-authored-by: Ishan <[email protected]>

Author: 3 people

.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

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)

rocketpy/mathutils/function.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

rocketpy/rocket/rocket.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,17 +2123,33 @@ 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(
                 self.rocket.motor.thrust.get_value_opt(t)
                 + 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

rocketpy/simulation/flight.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,
+    )