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

Co-authored-by: aZira371 <99824864+aZira371@users.noreply.github.com>

Author: Copilot

rocketpy/mathutils/function.py

@@ -4073,6 +4073,7 @@ def from_grid(
         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.
@@ -4115,6 +4116,14 @@ def from_grid(
 
             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.
 
@@ -4174,13 +4183,21 @@ def from_grid(
                 f"({grid_data.ndim})"
             )
 
-        # Check each axis matches corresponding grid dimension
+        # 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:
@@ -4217,15 +4234,21 @@ def from_grid(
             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])
+        # 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
@@ -4241,21 +4264,29 @@ def from_grid(
         # Set basic array attributes for compatibility
         func.x_array = axes[0]
         func.x_initial, func.x_final = axes[0][0], axes[0][-1]
-        # 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()),
-        )
+        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]
@@ -4265,7 +4296,10 @@ def from_grid(
 
         # Set interpolation and extrapolation functions
         func.__set_interpolation_func()
-        func.__set_extrapolation_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)

rocketpy/simulation/flight.py

@@ -1398,72 +1398,70 @@ def __get_drag_coefficient(self, drag_function, mach, z, freestream_velocity_bod
         float
             Drag coefficient value
         """
-        # Check if drag function is multi-dimensional using Function API
-        if isinstance(drag_function, Function) and getattr(
+        # Early return for 1D drag functions (only mach number)
+        if not isinstance(drag_function, Function) or not getattr(
             drag_function, "is_multidimensional", False
         ):
-            # 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
+            return drag_function.get_value_opt(mach)
 
-            reynolds = rho * freestream_speed * characteristic_length / mu_safe
+        # Multi-dimensional drag function - calculate additional parameters
 
-            # 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)
+        # 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
 
-            return drag_function.get_value_opt(*args)
+        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:
-            # 1D drag function - only mach number
-            return drag_function.get_value_opt(mach)
+            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

tests/unit/mathutils/test_function_grid.py

@@ -280,3 +280,73 @@ def test_shepard_fallback_at_exact_data_points():
     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