ENH: Air brakes controller functions now support 8-parameter signature (#854)

* solved environment parameter global usage, now can be used locally

* f string placeholder changes in test file

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

* make lint

* solved environment parameter global usage, now can be used locally

* f string placeholder changes in test file

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

* make lint

* fix tests

---------

Co-authored-by: Nitheesh Krishna <nitheeshhs@gmail.com>
Co-authored-by: Gui-FernandesBR <guilherme_fernandes@usp.br>
Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com>
Co-authored-by: Mateus Stano Junqueira <69485049+MateusStano@users.noreply.github.com>

Author: 4 people

CHANGELOG.md

@@ -31,6 +31,8 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 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)
 - TST: Add acceptance tests for 3DOF flight simulation based on Bella Lui rocket [#914] (https://github.com/RocketPy-Team/RocketPy/pull/914_
 - ENH: Add background map auto download functionality to Monte Carlo plots [#896](https://github.com/RocketPy-Team/RocketPy/pull/896)
 - MNT: net thrust addition to 3 dof in flight class [#907] (https://github.com/RocketPy-Team/RocketPy/pull/907)

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.

rocketpy/control/controller.py

@@ -60,7 +60,11 @@ def __init__(
             7. `sensors` (list): A list of sensors that are attached to the
                 rocket. The most recent measurements of the sensors are provided
                 with the ``sensor.measurement`` attribute. The sensors are
-                listed in the same order as they are added to the rocket
+                listed in the same order as they are added to the rocket.
+            8. `environment` (Environment): The environment object containing
+                atmospheric conditions, wind data, gravity, and other
+                environmental parameters. This allows the controller to access
+                environmental data locally without relying on global variables.
 
             This function will be called during the simulation at the specified
             sampling rate. The function should evaluate and change the interactive
@@ -102,7 +106,7 @@ def __init__(
     def __init_controller_function(self, controller_function):
         """Checks number of arguments of the controller function and initializes
         it with the correct number of arguments. This is a workaround to allow
-        the controller function to receive sensors without breaking changes"""
+        the controller function to receive sensors and environment without breaking changes"""
         sig = signature(controller_function)
         if len(sig.parameters) == 6:
             # pylint: disable=unused-argument
@@ -114,6 +118,7 @@ def new_controller_function(
                 observed_variables,
                 interactive_objects,
                 sensors,
+                environment,
             ):
                 return controller_function(
                     time,
@@ -125,18 +130,43 @@ def new_controller_function(
                 )
 
         elif len(sig.parameters) == 7:
+            # pylint: disable=unused-argument
+            def new_controller_function(
+                time,
+                sampling_rate,
+                state_vector,
+                state_history,
+                observed_variables,
+                interactive_objects,
+                sensors,
+                environment,
+            ):
+                return controller_function(
+                    time,
+                    sampling_rate,
+                    state_vector,
+                    state_history,
+                    observed_variables,
+                    interactive_objects,
+                    sensors,
+                )
+
+        elif len(sig.parameters) == 8:
             new_controller_function = controller_function
         else:
             raise ValueError(
-                "The controller function must have 6 or 7 arguments. "
+                "The controller function must have 6, 7, or 8 arguments. "
                 "The arguments must be in the following order: "
                 "(time, sampling_rate, state_vector, state_history, "
-                "observed_variables, interactive_objects, sensors)."
-                "Sensors argument is optional."
+                "observed_variables, interactive_objects, sensors, environment). "
+                "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
 
-    def __call__(self, time, state_vector, state_history, sensors):
+    def __call__(self, time, state_vector, state_history, sensors, environment):
         """Call the controller function. This is used by the simulation class.
 
         Parameters
@@ -157,6 +187,9 @@ def __call__(self, time, state_vector, state_history, sensors):
             measurements of the sensors are provided with the
             ``sensor.measurement`` attribute. The sensors are listed in the same
             order as they are added to the rocket.
+        environment : Environment
+            The environment object containing atmospheric conditions, wind data,
+            gravity, and other environmental parameters.
 
         Returns
         -------
@@ -170,6 +203,7 @@ def __call__(self, time, state_vector, state_history, sensors):
             self.observed_variables,
             self.interactive_objects,
             sensors,
+            environment,
         )
         if observed_variables is not None:
             self.observed_variables.append(observed_variables)

rocketpy/rocket/rocket.py

@@ -1696,9 +1696,13 @@ def add_air_brakes(
                listed in the same order as they are provided in the
                `interactive_objects` argument.
             7. `sensors` (list): A list of sensors that are attached to the
-               rocket. The most recent measurements of the sensors are provided
-               with the ``sensor.measurement`` attribute. The sensors are
-               listed in the same order as they are added to the rocket.
+                rocket. The most recent measurements of the sensors are provided
+                with the ``sensor.measurement`` attribute. The sensors are
+                listed in the same order as they are added to the rocket.
+            8. `environment` (Environment): The environment object containing
+                atmospheric conditions, wind data, gravity, and other
+                environmental parameters. This allows the controller to access
+                environmental data locally without relying on global variables.
 
             This function will be called during the simulation at the specified
             sampling rate. The function should evaluate and change the observed

rocketpy/simulation/flight.py

@@ -716,6 +716,7 @@ def __simulate(self, verbose):
                         self.y_sol,
                         self.solution,
                         self.sensors,
+                        self.env,
                     )
 
                 for parachute in node.parachutes:
@@ -874,6 +875,7 @@ def __process_sensors_and_controllers_at_current_node(self, node, phase):
                 self.y_sol,
                 self.solution,
                 self.sensors,
+                self.env,
             )
 
     def __measure_sensors(self, component_sensors, u_dot, t=None, y_sol=None):
@@ -1295,6 +1297,7 @@ def __process_overshootable_nodes(self, phase, phase_index, node_index):
                     overshootable_node.y_sol,
                     self.solution,
                     self.sensors,
+                    self.env,
                 )
 
             # Process sensors at overshootable node

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",

tests/fixtures/controller/__init__.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

tests/fixtures/controller/controller_fixtures.py

@@ -84,47 +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 lambda_quad_func():
     """Create a lambda function based on a string.