ENH: add force convention to LinearGenericSurface

Author: MateusStano

rocketpy/rocket/aero_surface/generic_surface.py

@@ -242,15 +242,13 @@ def __init__(
         self.force_convention = self._resolve_force_convention(
             coefficients, force_convention
         )
-        # The wind->body conversion only applies to surfaces whose coefficients
-        # are the full body-frame forces (cN/cY/cA). The linear model uses
-        # coefficient derivatives (cN_alpha, ...) whose frame is fixed by name.
+        # Wind-frame force input (cL/cQ/cD) is converted once to the canonical
+        # body-frame coefficients before validation. Each surface supplies the
+        # conversion appropriate to its coefficients: the generic surface rotates
+        # the full force coefficients, while the linear model recombines the
+        # coefficient derivatives (see LinearGenericSurface._wind_input_to_body).
         # A non-dict input falls through to _check_coefficients, which rejects it.
-        if (
-            self.force_convention == "wind"
-            and "cN" in default_coefficients
-            and isinstance(coefficients, dict)
-        ):
+        if self.force_convention == "wind" and isinstance(coefficients, dict):
             coefficients = self._wind_input_to_body(coefficients)
         self._check_coefficients(coefficients, default_coefficients)
         coefficients = self._complete_coefficients(coefficients, default_coefficients)
@@ -505,23 +503,35 @@ def _coefficient_option(option, coeff_name):
     _WIND_FORCE_NAMES = ("cL", "cQ", "cD")
     _BODY_FORCE_NAMES = ("cN", "cY", "cA")
 
+    def _force_frames_present(self, coefficients):
+        """Report which force frames the input coefficient names belong to, as
+        ``(has_wind, has_body)``.
+
+        A generic surface matches the plain force names (``cL``/``cQ``/``cD`` for
+        wind, ``cN``/``cY``/``cA`` for body). The linear model overrides this to
+        match those same names as derivative prefixes (``cL_alpha`` ...).
+        """
+        keys = set(coefficients)
+        has_wind = bool(keys & set(self._WIND_FORCE_NAMES))
+        has_body = bool(keys & set(self._BODY_FORCE_NAMES))
+        return has_wind, has_body
+
     def _resolve_force_convention(self, coefficients, force_convention):
         """Decide whether the input force coefficients are given in the wind
         frame (``cL``/``cQ``/``cD``) or the body frame (``cN``/``cY``/``cA``).
 
         When ``force_convention`` is ``None`` the frame is inferred from the
-        coefficient names; mixing the two frames is rejected.
+        coefficient names; mixing the two frames is rejected. With no force
+        coefficients to infer from, the canonical body frame is assumed.
         """
-        keys = set(coefficients)
-        has_wind = bool(keys & set(self._WIND_FORCE_NAMES))
-        has_body = bool(keys & set(self._BODY_FORCE_NAMES))
+        has_wind, has_body = self._force_frames_present(coefficients)
         if force_convention is None:
             if has_wind and has_body:
                 raise ValueError(
                     "Mixed wind (cL/cQ/cD) and body (cN/cY/cA) force "
                     "coefficients; pass force_convention='wind' or 'body'."
                 )
-            return "body" if has_body else "wind"
+            return "wind" if has_wind else "body"
         if force_convention not in ("wind", "body"):
             raise ValueError(
                 f"force_convention must be 'wind' or 'body', got {force_convention!r}."

rocketpy/rocket/aero_surface/linear_generic_surface.py

@@ -1,6 +1,9 @@
+import inspect
+
 from rocketpy.mathutils import Function
 from rocketpy.plots.aero_surface_plots import _LinearGenericSurfacePlots
 from rocketpy.prints.aero_surface_prints import _LinearGenericSurfacePrints
+from rocketpy.rocket.aero_surface.aero_coefficient import AeroCoefficient
 from rocketpy.rocket.aero_surface.generic_surface import GenericSurface
 
 
@@ -22,6 +25,7 @@ def __init__(
         name="Generic Linear Surface",
         interpolation=None,
         extrapolation=None,
+        force_convention=None,
     ):
         """Create a generic linear aerodynamic surface, defined by its
         aerodynamic coefficients derivatives. This surface is used to model any
@@ -35,6 +39,13 @@ def __init__(
         contain at least one of the following: "alpha", "beta", "mach",
         "reynolds", "pitch_rate", "yaw_rate" and "roll_rate".
 
+        By default the force-coefficient derivatives are the body-frame ones
+        (``cN_*`` normal, ``cY_*`` side, ``cA_*`` axial; see
+        ``force_convention``). You may instead give the wind-frame derivatives
+        ``cL_*`` (lift), ``cQ_*`` (side) and ``cD_*`` (drag) -- for example
+        ``cL_alpha`` in place of ``cN_alpha``; they are converted once to the
+        body-frame set at construction.
+
         See Also
         --------
         :ref:`genericsurfaces`.
@@ -57,7 +68,10 @@ def __init__(
             yaw moment ``cn`` or roll moment ``cl``; the variable is ``0`` (the
             value at zero angle of attack, zero sideslip and zero rates),
             ``alpha``, ``beta``, ``p`` (roll rate), ``q`` (pitch rate) or ``r``
-            (yaw rate). The full list is:\n
+            (yaw rate). With ``force_convention="wind"`` the force derivatives are
+            named after the wind-frame coefficients instead (lift ``cL``, side
+            ``cQ``, drag ``cD`` -- e.g. ``cL_alpha``, ``cD_0``, ``cQ_beta``); the
+            moment names are unchanged. The full (body-frame) list is:\n
             cN_0: callable, str, optional
                 Coefficient of normal force at zero angle of attack. Default is 0.\n
             cN_alpha: callable, str, optional
@@ -193,6 +207,21 @@ def __init__(
             uses ``"constant"`` for tables built here and keeps whatever a
             pre-built ``Function`` already carries. Only affects tabulated
             sources (constants and callables are evaluated directly).
+        force_convention : str, optional
+            The frame your force-coefficient derivatives are given in. ``"body"``
+            for the body-frame derivatives ``cN_*`` (normal), ``cY_*`` (side) and
+            ``cA_*`` (axial); ``"wind"`` for the aerodynamic-frame derivatives
+            ``cL_*`` (lift), ``cQ_*`` (side) and ``cD_*`` (drag). The moment
+            derivatives (``cm_*``, ``cn_*``, ``cl_*``) are the same in both.
+            ``None`` (the default) infers the frame from the coefficient names you
+            pass and assumes body when none are given. A wind-frame input is
+            converted once to the body-frame derivatives the surface stores, by
+            linearizing the angle-of-attack/sideslip rotation about zero: the
+            straight renames ``cN_0 = cL_0``, ``cN_beta = cL_beta``, the rate
+            derivatives, and the cross terms ``cN_alpha = cL_alpha + cD_0``,
+            ``cY_beta = cQ_beta - cD_0``, ``cA_alpha = cD_alpha - cL_0`` and
+            ``cA_beta = cD_beta + cQ_0``. At zero angle this reduces to
+            ``cN = cL``, ``cY = cQ``, ``cA = cD``.
         """
 
         super().__init__(
@@ -203,6 +232,7 @@ def __init__(
             name=name,
             extrapolation=extrapolation,
             interpolation=interpolation,
+            force_convention=force_convention,
         )
 
         self.compute_all_coefficients()
@@ -271,6 +301,126 @@ def _get_default_coefficients(self):
         }
         return default_coefficients
 
+    # Body force-coefficient prefix -> wind force-coefficient prefix, used to
+    # name the accepted wind-frame inputs. The per-plane suffixes (_0, _alpha,
+    # _beta, _p, _q, _r) and the moment coefficients (cm/cn/cl) are frame-shared.
+    _BODY_TO_WIND_PREFIX = {"cN": "cL", "cY": "cQ", "cA": "cD"}
+
+    def _force_frames_present(self, coefficients):
+        """Detect the force frame from the derivative-name prefixes: a wind key
+        looks like ``cL_alpha``/``cD_0``/``cQ_beta`` and a body key like
+        ``cN_alpha``/``cA_0``/``cY_beta``. The moment derivatives (``cm_*``,
+        ``cn_*``, ``cl_*``) are frame-shared and ignored here.
+        """
+        prefixes = {key.split("_", 1)[0] for key in coefficients}
+        has_wind = bool(prefixes & set(self._WIND_FORCE_NAMES))
+        has_body = bool(prefixes & set(self._BODY_FORCE_NAMES))
+        return has_wind, has_body
+
+    def _wind_default_coefficient_names(self):
+        """The valid wind-frame input names: the body defaults with the force
+        prefixes swapped to wind (``cN_* -> cL_*``, ``cY_* -> cQ_*``,
+        ``cA_* -> cD_*``); the moment names are unchanged.
+        """
+        names = set()
+        for key in self._get_default_coefficients():
+            prefix, sep, suffix = key.partition("_")
+            wind_prefix = self._BODY_TO_WIND_PREFIX.get(prefix, prefix)
+            names.add(f"{wind_prefix}{sep}{suffix}")
+        return names
+
+    def _wind_input_to_body(self, coefficients):
+        """Convert wind-frame coefficient derivatives (``cL_*``/``cD_*``/``cQ_*``)
+        into the canonical body-frame derivatives (``cN_*``/``cY_*``/``cA_*``).
+
+        The full body-frame force coefficients are the wind ones rotated by the
+        angle of attack and sideslip; linearizing that rotation about
+        ``alpha = beta = 0`` gives, to first order, a coefficient-derivative map
+        with four cross-frame terms::
+
+            cN_alpha = cL_alpha + cD_0      cA_alpha = cD_alpha - cL_0
+            cY_beta  = cQ_beta  - cD_0      cA_beta  = cD_beta  + cQ_0
+
+        Every other derivative is a straight rename (``cN_0 = cL_0``,
+        ``cN_beta = cL_beta``, the rate derivatives ``cN_p = cL_p`` ..., and the
+        wind side/axial analogues). At zero angle this reduces to ``cN = cL``,
+        ``cY = cQ``, ``cA = cD``, matching the generic surface. The moment
+        derivatives (``cm_*``/``cn_*``/``cl_*``) are frame-shared and pass
+        through unchanged.
+        """
+        invalid = set(coefficients) - self._wind_default_coefficient_names()
+        if invalid:
+            raise ValueError(
+                f"Invalid coefficient name(s) used in key(s): {', '.join(invalid)}. "
+                "Check the documentation for valid names."
+            )
+
+        def wind(name):
+            return coefficients.get(name, 0)
+
+        body = {
+            "cN_0": wind("cL_0"),
+            "cN_alpha": self._combine(wind("cL_alpha"), wind("cD_0"), 1.0, "cN_alpha"),
+            "cN_beta": wind("cL_beta"),
+            "cN_p": wind("cL_p"),
+            "cN_q": wind("cL_q"),
+            "cN_r": wind("cL_r"),
+            "cY_0": wind("cQ_0"),
+            "cY_alpha": wind("cQ_alpha"),
+            "cY_beta": self._combine(wind("cQ_beta"), wind("cD_0"), -1.0, "cY_beta"),
+            "cY_p": wind("cQ_p"),
+            "cY_q": wind("cQ_q"),
+            "cY_r": wind("cQ_r"),
+            "cA_0": wind("cD_0"),
+            "cA_alpha": self._combine(wind("cD_alpha"), wind("cL_0"), -1.0, "cA_alpha"),
+            "cA_beta": self._combine(wind("cD_beta"), wind("cQ_0"), 1.0, "cA_beta"),
+            "cA_p": wind("cD_p"),
+            "cA_q": wind("cD_q"),
+            "cA_r": wind("cD_r"),
+        }
+        # Moment derivatives are the same in both frames; pass them through.
+        for name, value in coefficients.items():
+            if name.split("_", 1)[0] not in self._WIND_FORCE_NAMES:
+                body[name] = value
+        return body
+
+    def _as_coefficient(self, source, name):
+        """Wrap a raw coefficient input as an :class:`AeroCoefficient` over this
+        surface's variables (used when recombining wind-frame derivatives).
+        """
+        return AeroCoefficient(
+            source,
+            unsteady_aero=self._unsteady_aero,
+            control_variables=self.control_variables,
+            name=name,
+        )
+
+    def _combine(self, first, second, sign, name):
+        """Return a coefficient equal to ``first + sign * second``.
+
+        When one term is identically zero the other is returned directly (as a
+        renamed coefficient), so a derivative that is really just a rename keeps
+        its original, low-dimensional form. Otherwise the two are summed by a
+        small wrapper evaluated over the full variable tuple.
+        """
+        coeff_first = self._as_coefficient(first, name)
+        coeff_second = self._as_coefficient(second, name)
+        if coeff_second.is_zero:
+            return coeff_first
+        if coeff_first.is_zero:
+            return coeff_second if sign > 0 else coeff_second * -1.0
+        first_opt = coeff_first.get_value_opt
+        second_opt = coeff_second.get_value_opt
+
+        def combined(*args):
+            return first_opt(*args) + sign * second_opt(*args)
+
+        combined.__signature__ = inspect.Signature(
+            inspect.Parameter(var, inspect.Parameter.POSITIONAL_OR_KEYWORD)
+            for var in self.independent_vars
+        )
+        return self._as_coefficient(combined, name)
+
     _COEFFICIENT_INPUTS = [
         "alpha",
         "beta",

tests/unit/rocket/aero_surface/test_linear_generic_surfaces.py

@@ -1,11 +1,14 @@
 import pytest
 
-from rocketpy import Function, LinearGenericSurface
+from rocketpy import Function, GenericSurface, LinearGenericSurface
 from rocketpy.mathutils import Vector
 
 REFERENCE_AREA = 1
 REFERENCE_LENGTH = 1
 
+# (alpha, beta, mach, reynolds, pitch_rate, yaw_rate, roll_rate)
+_ARGS = (0.0, 0.0, 0.5, 1e6, 0.0, 0.0, 0.0)
+
 
 @pytest.mark.parametrize(
     "coefficients",
@@ -121,3 +124,108 @@ def test_roll_damping_uses_reduced_rate():
     # Old (raw-rate) formulation -- identical value:
     old_damping_scaling = 0.5 * rho * speed * ref_area * ref_length**2 / 2
     assert roll_moment == pytest.approx(old_damping_scaling * cl_p * raw_roll)
+
+
+def test_force_convention_inference():
+    """The input frame is inferred from the derivative names when
+    ``force_convention`` is not given, defaulting to body when there are no
+    force derivatives to infer from."""
+    assert LinearGenericSurface(1, 1, {"cN_alpha": 2.0}).force_convention == "body"
+    assert LinearGenericSurface(1, 1, {"cL_alpha": 2.0}).force_convention == "wind"
+    # Moment-only / empty input carries no force frame -> body (canonical).
+    assert LinearGenericSurface(1, 1, {"cm_alpha": 1.0}).force_convention == "body"
+    assert LinearGenericSurface(1, 1, {}).force_convention == "body"
+
+
+def test_mixed_frame_input_raises():
+    """Supplying both wind (cL_*) and body (cN_*) force derivatives without
+    declaring the frame is rejected."""
+    with pytest.raises(ValueError, match="[Mm]ixed"):
+        LinearGenericSurface(1, 1, {"cL_alpha": 1.0, "cN_0": 1.0})
+
+
+def test_invalid_wind_coefficient_name_raises():
+    """A wind-frame input with an unknown derivative name is rejected."""
+    with pytest.raises(ValueError, match="Invalid coefficient name"):
+        LinearGenericSurface(1, 1, {"cL_gamma": 1.0}, force_convention="wind")
+
+
+def test_wind_derivatives_convert_to_body_first_order():
+    """Wind-frame derivatives are converted to the body-frame set by linearizing
+    the wind/body rotation about zero: the four cross terms plus straight
+    renames. A nonzero base drag ``cD_0`` couples into the normal- and
+    axial-force slopes."""
+    surface = LinearGenericSurface(
+        reference_area=0.01,
+        reference_length=0.1,
+        coefficients={
+            "cL_0": 0.1,
+            "cL_alpha": 5.0,
+            "cD_0": 0.3,
+            "cD_alpha": 0.2,
+            "cQ_beta": -4.0,
+            "cQ_0": 0.05,
+            "cm_alpha": -2.0,  # frame-shared moment, must pass through unchanged
+        },
+        force_convention="wind",
+    )
+    assert surface.force_convention == "wind"
+    assert surface.cN_0.get_value_opt(*_ARGS) == pytest.approx(0.1)  # cL_0
+    assert surface.cN_alpha.get_value_opt(*_ARGS) == pytest.approx(5.3)  # cL_alpha+cD_0
+    assert surface.cY_beta.get_value_opt(*_ARGS) == pytest.approx(-4.3)  # cQ_beta-cD_0
+    assert surface.cA_0.get_value_opt(*_ARGS) == pytest.approx(0.3)  # cD_0
+    assert surface.cA_alpha.get_value_opt(*_ARGS) == pytest.approx(0.1)  # cD_alpha-cL_0
+    assert surface.cA_beta.get_value_opt(*_ARGS) == pytest.approx(0.05)  # cD_beta+cQ_0
+    assert surface.cm_alpha.get_value_opt(*_ARGS) == pytest.approx(-2.0)  # passthrough
+
+
+def test_wind_input_matches_body_input():
+    """A wind-frame surface equals the body-frame surface built from the
+    hand-converted derivatives, at an arbitrary angle."""
+    wind = LinearGenericSurface(
+        1,
+        1,
+        coefficients={"cL_0": 0.1, "cL_alpha": 5.0, "cD_0": 0.3, "cQ_beta": -4.0},
+        force_convention="wind",
+    )
+    body = LinearGenericSurface(
+        1,
+        1,
+        coefficients={
+            "cN_0": 0.1,
+            "cN_alpha": 5.3,
+            "cA_0": 0.3,
+            "cA_alpha": -0.1,  # cD_alpha - cL_0
+            "cY_beta": -4.3,
+        },
+    )
+    args = (0.02, 0.01, 0.5, 1e6, 0.0, 0.0, 0.0)
+    for coeff in ("cN", "cY", "cA"):
+        assert getattr(wind, coeff).get_value_opt(*args) == pytest.approx(
+            getattr(body, coeff).get_value_opt(*args)
+        )
+
+
+def test_wind_linear_matches_generic_surface_to_first_order():
+    """The body-frame forces of a wind-input linear surface agree with the
+    exact rotation used by GenericSurface to first order in the flow angles."""
+    coeffs = {"cL_0": 0.1, "cL_alpha": 5.0, "cD_0": 0.3, "cQ_beta": -4.0}
+    linear = LinearGenericSurface(0.01, 0.1, coeffs, force_convention="wind")
+    generic = GenericSurface(
+        0.01,
+        0.1,
+        coefficients={
+            "cL": lambda a, b, m, re, p, q, r: 0.1 + 5.0 * a,
+            "cD": lambda a, b, m, re, p, q, r: 0.3,
+            "cQ": lambda a, b, m, re, p, q, r: -4.0 * b,
+        },
+        force_convention="wind",
+    )
+    eps = 1e-3
+    for alpha, beta in [(eps, 0.0), (0.0, eps), (eps, eps)]:
+        args = (alpha, beta, 0.5, 1e6, 0.0, 0.0, 0.0)
+        for coeff in ("cN", "cY", "cA"):
+            # Difference is second order in the angle (~1e-6 at eps=1e-3).
+            assert getattr(linear, coeff).get_value_opt(*args) == pytest.approx(
+                getattr(generic, coeff).get_value_opt(*args), abs=1e-5
+            )