ENH: Update air brakes controller to support 8-parameter signature and improve environment access

Gui-FernandesBR · Gui-FernandesBR · commit 25651dec3c14 · 2025-11-28T22:32:17.000-03:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -32,6 +32,7 @@ Attention: The newest changes should be on top -->
 
 ### Added
 
+- ENH: Air brakes controller functions now support 8-parameter signature [#854](https://github.com/RocketPy-Team/RocketPy/pull/854)
 - ENH: Add save functionality to `_MonteCarloPlots.all` method [#848](https://github.com/RocketPy-Team/RocketPy/pull/848)
 - ENH: Add persistent caching for ThrustCurve API [#881](https://github.com/RocketPy-Team/RocketPy/pull/881)
 - ENH: Compatibility with MERRA-2 atmosphere reanalysis files [#825](https://github.com/RocketPy-Team/RocketPy/pull/825)
diff --git a/docs/user/airbrakes.rst b/docs/user/airbrakes.rst
@@ -167,21 +167,21 @@ Lets define the controller function:
 .. jupyter-execute::
 
     def controller_function(
-        time, sampling_rate, state, state_history, observed_variables, air_brakes
+        time, sampling_rate, state, state_history, observed_variables, air_brakes, sensors, environment
     ):
         # state = [x, y, z, vx, vy, vz, e0, e1, e2, e3, wx, wy, wz]
         altitude_ASL = state[2]
-        altitude_AGL = altitude_ASL - env.elevation
+        altitude_AGL = altitude_ASL - environment.elevation
         vx, vy, vz = state[3], state[4], state[5]
 
         # Get winds in x and y directions
-        wind_x, wind_y = env.wind_velocity_x(altitude_ASL), env.wind_velocity_y(altitude_ASL)
+        wind_x, wind_y = environment.wind_velocity_x(altitude_ASL), environment.wind_velocity_y(altitude_ASL)
 
         # Calculate Mach number
         free_stream_speed = (
             (wind_x - vx) ** 2 + (wind_y - vy) ** 2 + (vz) ** 2
         ) ** 0.5
-        mach_number = free_stream_speed / env.speed_of_sound(altitude_ASL)
+        mach_number = free_stream_speed / environment.speed_of_sound(altitude_ASL)
 
         # Get previous state from state_history
         previous_state = state_history[-1]
@@ -224,6 +224,22 @@ Lets define the controller function:
 
 .. note::
 
+    - The ``controller_function`` accepts 6, 7, or 8 parameters for backward
+      compatibility:
+
+      * **6 parameters** (original): ``time``, ``sampling_rate``, ``state``,
+        ``state_history``, ``observed_variables``, ``air_brakes``
+      * **7 parameters** (with sensors): adds ``sensors`` as the 7th parameter
+      * **8 parameters** (with environment): adds ``sensors`` and ``environment``
+        as the 7th and 8th parameters
+
+    - The **environment parameter** provides access to atmospheric conditions
+      (wind, temperature, pressure, elevation) without relying on global variables.
+      This enables proper serialization of rockets with air brakes and improves
+      code modularity. Available methods include ``environment.elevation``,
+      ``environment.wind_velocity_x(altitude)``, ``environment.wind_velocity_y(altitude)``,
+      ``environment.speed_of_sound(altitude)``, and others.
+
     - The code inside the ``controller_function`` can be as complex as needed.
       Anything can be implemented inside the function, including filters,
       apogee prediction, and any controller logic.
diff --git a/rocketpy/control/controller.py b/rocketpy/control/controller.py
@@ -159,7 +159,10 @@ def new_controller_function(
                 "The arguments must be in the following order: "
                 "(time, sampling_rate, state_vector, state_history, "
                 "observed_variables, interactive_objects, sensors, environment). "
-                "The last two arguments (sensors and environment) are optional."
+                "Supported signatures: "
+                "6 parameters (no sensors, no environment), "
+                "7 parameters (with sensors, no environment), or "
+                "8 parameters (with sensors and environment)."
             )
         return new_controller_function
 
diff --git a/test_environment_parameter.py b/test_environment_parameter.py
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -11,6 +11,7 @@
     "tests.fixtures.environment.environment_fixtures",
     "tests.fixtures.flight.flight_fixtures",
     "tests.fixtures.function.function_fixtures",
+    "tests.fixtures.controller.controller_fixtures",
     "tests.fixtures.motor.liquid_fixtures",
     "tests.fixtures.motor.hybrid_fixtures",
     "tests.fixtures.motor.solid_motor_fixtures",
diff --git a/tests/fixtures/controller/__init__.py b/tests/fixtures/controller/__init__.py
diff --git a/tests/fixtures/controller/controller_fixtures.py b/tests/fixtures/controller/controller_fixtures.py
@@ -0,0 +1,104 @@
+import pytest
+
+
+## Controller
+@pytest.fixture
+def controller_function():
+    """Create a controller function that updates the air brakes deployment level
+    based on the altitude and vertical velocity of the rocket. This is the same
+    controller function that is used in the air brakes example in the
+    documentation.
+
+    Returns
+    -------
+    function
+        A controller function
+    """
+
+    def controller_function(  # pylint: disable=unused-argument
+        time, sampling_rate, state, state_history, observed_variables, air_brakes
+    ):
+        z = state[2]
+        vz = state[5]
+        previous_vz = state_history[-1][5]
+        if time < 3.9:
+            return None
+        if z < 1500:
+            air_brakes.deployment_level = 0
+        else:
+            new_deployment_level = (
+                air_brakes.deployment_level + 0.1 * vz + 0.01 * previous_vz**2
+            )
+            if new_deployment_level > air_brakes.deployment_level + 0.2 / sampling_rate:
+                new_deployment_level = air_brakes.deployment_level + 0.2 / sampling_rate
+            elif (
+                new_deployment_level < air_brakes.deployment_level - 0.2 / sampling_rate
+            ):
+                new_deployment_level = air_brakes.deployment_level - 0.2 / sampling_rate
+            else:
+                new_deployment_level = air_brakes.deployment_level
+            air_brakes.deployment_level = new_deployment_level
+
+    return controller_function
+
+
+@pytest.fixture
+def controller_function_with_environment():
+    """Create a controller function that uses the environment parameter to access
+    atmospheric conditions without relying on global variables. This demonstrates
+    the new environment parameter feature for air brakes controllers.
+
+    Returns
+    -------
+    function
+        A controller function that uses environment parameter
+    """
+
+    def controller_function(  # pylint: disable=unused-argument
+        time,
+        sampling_rate,
+        state,
+        state_history,
+        observed_variables,
+        air_brakes,
+        sensors,
+        environment,
+    ):
+        # state = [x, y, z, vx, vy, vz, e0, e1, e2, e3, wx, wy, wz]
+        altitude_ASL = state[2]  # altitude above sea level
+        altitude_AGL = (
+            altitude_ASL - environment.elevation
+        )  # altitude above ground level
+        vx, vy, vz = state[3], state[4], state[5]
+
+        # Use environment parameter instead of global variable
+        wind_x = environment.wind_velocity_x(altitude_ASL)
+        wind_y = environment.wind_velocity_y(altitude_ASL)
+
+        # Calculate Mach number using environment data
+        free_stream_speed = ((wind_x - vx) ** 2 + (wind_y - vy) ** 2 + (vz) ** 2) ** 0.5
+        mach_number = free_stream_speed / environment.speed_of_sound(altitude_ASL)
+
+        if time < 3.9:
+            return None
+
+        if altitude_AGL < 1500:
+            air_brakes.deployment_level = 0
+        else:
+            previous_vz = state_history[-1][5] if state_history else vz
+            new_deployment_level = (
+                air_brakes.deployment_level + 0.1 * vz + 0.01 * previous_vz**2
+            )
+            # Rate limiting
+            max_change = 0.2 / sampling_rate
+            if new_deployment_level > air_brakes.deployment_level + max_change:
+                new_deployment_level = air_brakes.deployment_level + max_change
+            elif new_deployment_level < air_brakes.deployment_level - max_change:
+                new_deployment_level = air_brakes.deployment_level - max_change
+
+            air_brakes.deployment_level = new_deployment_level
+
+        # Return observed variables including Mach number
+        return (time, air_brakes.deployment_level, mach_number)
+
+    return controller_function
diff --git a/tests/fixtures/function/function_fixtures.py b/tests/fixtures/function/function_fixtures.py
@@ -84,102 +84,6 @@ def func_2d_from_csv():
     return func
 
 
-## Controller
-@pytest.fixture
-def controller_function():
-    """Create a controller function that updates the air brakes deployment level
-    based on the altitude and vertical velocity of the rocket. This is the same
-    controller function that is used in the air brakes example in the
-    documentation.
-
-    Returns
-    -------
-    function
-        A controller function
-    """
-
-    def controller_function(  # pylint: disable=unused-argument
-        time, sampling_rate, state, state_history, observed_variables, air_brakes
-    ):
-        z = state[2]
-        vz = state[5]
-        previous_vz = state_history[-1][5]
-        if time < 3.9:
-            return None
-        if z < 1500:
-            air_brakes.deployment_level = 0
-        else:
-            new_deployment_level = (
-                air_brakes.deployment_level + 0.1 * vz + 0.01 * previous_vz**2
-            )
-            if new_deployment_level > air_brakes.deployment_level + 0.2 / sampling_rate:
-                new_deployment_level = air_brakes.deployment_level + 0.2 / sampling_rate
-            elif (
-                new_deployment_level < air_brakes.deployment_level - 0.2 / sampling_rate
-            ):
-                new_deployment_level = air_brakes.deployment_level - 0.2 / sampling_rate
-            else:
-                new_deployment_level = air_brakes.deployment_level
-            air_brakes.deployment_level = new_deployment_level
-
-    return controller_function
-
-
-@pytest.fixture
-def controller_function_with_environment():
-    """Create a controller function that uses the environment parameter to access
-    atmospheric conditions without relying on global variables. This demonstrates
-    the new environment parameter feature for air brakes controllers.
-
-    Returns
-    -------
-    function
-        A controller function that uses environment parameter
-    """
-
-    def controller_function(  # pylint: disable=unused-argument
-        time, sampling_rate, state, state_history, observed_variables, air_brakes, sensors, environment
-    ):
-        # state = [x, y, z, vx, vy, vz, e0, e1, e2, e3, wx, wy, wz]
-        altitude_ASL = state[2]  # altitude above sea level
-        altitude_AGL = altitude_ASL - environment.elevation  # altitude above ground level
-        vx, vy, vz = state[3], state[4], state[5]
-        
-        # Use environment parameter instead of global variable
-        wind_x = environment.wind_velocity_x(altitude_ASL)
-        wind_y = environment.wind_velocity_y(altitude_ASL)
-        
-        # Calculate Mach number using environment data
-        free_stream_speed = (
-            (wind_x - vx) ** 2 + (wind_y - vy) ** 2 + (vz) ** 2
-        ) ** 0.5
-        mach_number = free_stream_speed / environment.speed_of_sound(altitude_ASL)
-        
-        if time < 3.9:
-            return None
-            
-        if altitude_AGL < 1500:
-            air_brakes.deployment_level = 0
-        else:
-            previous_vz = state_history[-1][5] if state_history else vz
-            new_deployment_level = (
-                air_brakes.deployment_level + 0.1 * vz + 0.01 * previous_vz**2
-            )
-            # Rate limiting
-            max_change = 0.2 / sampling_rate
-            if new_deployment_level > air_brakes.deployment_level + max_change:
-                new_deployment_level = air_brakes.deployment_level + max_change
-            elif new_deployment_level < air_brakes.deployment_level - max_change:
-                new_deployment_level = air_brakes.deployment_level - max_change
-            
-            air_brakes.deployment_level = new_deployment_level
-            
-        # Return observed variables including Mach number
-        return (time, air_brakes.deployment_level, mach_number)
-
-    return controller_function
-
-
 @pytest.fixture
 def lambda_quad_func():
     """Create a lambda function based on a string.
diff --git a/tests/integration/simulation/test_flight.py b/tests/integration/simulation/test_flight.py
