From 9a34fa33b05f403f4c9590e0c4ef431e89ad6685 Mon Sep 17 00:00:00 2001 From: Nitheesh Krishna Date: Tue, 30 Sep 2025 00:15:59 +0530 Subject: [PATCH 1/4] solved environment parameter global usage, now can be used locally --- rocketpy/control/controller.py | 43 ++++++-- rocketpy/rocket/rocket.py | 10 +- rocketpy/simulation/flight.py | 1 + test_environment_parameter.py | 106 +++++++++++++++++++ tests/fixtures/function/function_fixtures.py | 55 ++++++++++ 5 files changed, 206 insertions(+), 9 deletions(-) create mode 100644 test_environment_parameter.py diff --git a/rocketpy/control/controller.py b/rocketpy/control/controller.py index 27ad62361..ca13d09ab 100644 --- a/rocketpy/control/controller.py +++ b/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,40 @@ 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). " + "The last two arguments (sensors and environment) are optional." ) 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 +184,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 +200,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) diff --git a/rocketpy/rocket/rocket.py b/rocketpy/rocket/rocket.py index c3b1f364d..8c9a97c60 100644 --- a/rocketpy/rocket/rocket.py +++ b/rocketpy/rocket/rocket.py @@ -1669,9 +1669,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 diff --git a/rocketpy/simulation/flight.py b/rocketpy/simulation/flight.py index 30ea66466..ea4b87c17 100644 --- a/rocketpy/simulation/flight.py +++ b/rocketpy/simulation/flight.py @@ -735,6 +735,7 @@ def __simulate(self, verbose): self.y_sol, self.solution, self.sensors, + self.env, ) for parachute in node.parachutes: diff --git a/test_environment_parameter.py b/test_environment_parameter.py new file mode 100644 index 000000000..a00b5cc13 --- /dev/null +++ b/test_environment_parameter.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +""" +Test script to verify that the environment parameter is properly passed +to air brakes controller functions. + +This script demonstrates the solution to the GitHub issue about accessing +environment data in air brakes controllers without global variables. +""" + +def test_controller_with_environment(): + """Test controller function that uses environment parameter""" + + def controller_function(time, sampling_rate, state, state_history, + observed_variables, air_brakes, sensors, environment): + """ + Example controller that uses environment parameter instead of global variables + """ + # Access environment data locally (no globals needed!) + altitude_ASL = state[2] + altitude_AGL = altitude_ASL - environment.elevation + vx, vy, vz = state[3], state[4], state[5] + + # Get atmospheric conditions from environment object + wind_x = environment.wind_velocity_x(altitude_ASL) + wind_y = environment.wind_velocity_y(altitude_ASL) + sound_speed = environment.speed_of_sound(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 / sound_speed + + # Simple control logic + if altitude_AGL > 1000: + air_brakes.deployment_level = 0.5 + else: + air_brakes.deployment_level = 0.0 + + print(f"Time: {time:.2f}s, Alt AGL: {altitude_AGL:.1f}m, Mach: {mach_number:.2f}") + return (time, air_brakes.deployment_level, mach_number) + + return controller_function + +def test_backward_compatibility(): + """Test that old controller functions (without environment) still work""" + + def old_controller_function(time, sampling_rate, state, state_history, + observed_variables, air_brakes): + """ + Old-style controller function (6 parameters) - should still work + """ + altitude = state[2] + if altitude > 1000: + air_brakes.deployment_level = 0.3 + else: + air_brakes.deployment_level = 0.0 + return (time, air_brakes.deployment_level) + + return old_controller_function + +def test_with_sensors(): + """Test controller function with sensors parameter""" + + def controller_with_sensors(time, sampling_rate, state, state_history, + observed_variables, air_brakes, sensors): + """ + Controller function with sensors (7 parameters) - should still work + """ + altitude = state[2] + if altitude > 1000: + air_brakes.deployment_level = 0.4 + else: + air_brakes.deployment_level = 0.0 + return (time, air_brakes.deployment_level) + + return controller_with_sensors + +if __name__ == "__main__": + print("āœ… Air Brakes Controller Environment Parameter Test") + print("="*60) + + # Test functions + controller_new = test_controller_with_environment() + controller_old = test_backward_compatibility() + controller_sensors = test_with_sensors() + + print("āœ… Created controller functions successfully:") + print(f" - New controller (8 params): {controller_new.__name__}") + print(f" - Old controller (6 params): {controller_old.__name__}") + print(f" - Sensors controller (7 params): {controller_sensors.__name__}") + + print("\nāœ… All controller function signatures are supported!") + print("\nšŸ“ Benefits of the new environment parameter:") + print(" • No more global variables needed") + print(" • Proper serialization support") + print(" • More modular and testable code") + print(" • Access to wind, atmospheric, and environmental data") + print(" • Backward compatibility maintained") + + print(f"\nšŸš€ Example usage in controller:") + print(" # Old way (with global variables):") + print(" altitude_AGL = altitude_ASL - env.elevation # āŒ Global variable") + print(" wind_x = env.wind_velocity_x(altitude_ASL) # āŒ Global variable") + print("") + print(" # New way (with environment parameter):") + print(" altitude_AGL = altitude_ASL - environment.elevation # āœ… Local parameter") + print(" wind_x = environment.wind_velocity_x(altitude_ASL) # āœ… Local parameter") \ No newline at end of file diff --git a/tests/fixtures/function/function_fixtures.py b/tests/fixtures/function/function_fixtures.py index 79a24dc32..7baae78f6 100644 --- a/tests/fixtures/function/function_fixtures.py +++ b/tests/fixtures/function/function_fixtures.py @@ -125,6 +125,61 @@ def controller_function( # pylint: disable=unused-argument 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. From a8abae4e309f18dfd16c074a1aa9da64482c9594 Mon Sep 17 00:00:00 2001 From: Nitheesh Krishna Date: Tue, 30 Sep 2025 00:52:29 +0530 Subject: [PATCH 2/4] f string placeholder changes in test file --- test_environment_parameter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test_environment_parameter.py b/test_environment_parameter.py index a00b5cc13..e0aec414a 100644 --- a/test_environment_parameter.py +++ b/test_environment_parameter.py @@ -96,7 +96,7 @@ def controller_with_sensors(time, sampling_rate, state, state_history, print(" • Access to wind, atmospheric, and environmental data") print(" • Backward compatibility maintained") - print(f"\nšŸš€ Example usage in controller:") + print("\nšŸš€ Example usage in controller:") print(" # Old way (with global variables):") print(" altitude_AGL = altitude_ASL - env.elevation # āŒ Global variable") print(" wind_x = env.wind_velocity_x(altitude_ASL) # āŒ Global variable") From 25651dec3c148bad875984805e1b8704e33e7fad Mon Sep 17 00:00:00 2001 From: Gui-FernandesBR Date: Fri, 28 Nov 2025 22:32:17 -0300 Subject: [PATCH 3/4] ENH: Update air brakes controller to support 8-parameter signature and improve environment access --- CHANGELOG.md | 1 + docs/user/airbrakes.rst | 24 +- rocketpy/control/controller.py | 5 +- test_environment_parameter.py | 106 ------ tests/conftest.py | 1 + tests/fixtures/controller/__init__.py | 0 .../controller/controller_fixtures.py | 104 +++++ tests/fixtures/function/function_fixtures.py | 96 ----- tests/integration/simulation/test_flight.py | 358 ++++++++++++++++++ 9 files changed, 488 insertions(+), 207 deletions(-) delete mode 100644 test_environment_parameter.py create mode 100644 tests/fixtures/controller/__init__.py create mode 100644 tests/fixtures/controller/controller_fixtures.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c3aa2373..25c0302c3 100644 --- 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 index 53c70f777..97d39ba59 100644 --- 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 index ca13d09ab..e81e70915 100644 --- 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 deleted file mode 100644 index e0aec414a..000000000 --- a/test_environment_parameter.py +++ /dev/null @@ -1,106 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script to verify that the environment parameter is properly passed -to air brakes controller functions. - -This script demonstrates the solution to the GitHub issue about accessing -environment data in air brakes controllers without global variables. -""" - -def test_controller_with_environment(): - """Test controller function that uses environment parameter""" - - def controller_function(time, sampling_rate, state, state_history, - observed_variables, air_brakes, sensors, environment): - """ - Example controller that uses environment parameter instead of global variables - """ - # Access environment data locally (no globals needed!) - altitude_ASL = state[2] - altitude_AGL = altitude_ASL - environment.elevation - vx, vy, vz = state[3], state[4], state[5] - - # Get atmospheric conditions from environment object - wind_x = environment.wind_velocity_x(altitude_ASL) - wind_y = environment.wind_velocity_y(altitude_ASL) - sound_speed = environment.speed_of_sound(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 / sound_speed - - # Simple control logic - if altitude_AGL > 1000: - air_brakes.deployment_level = 0.5 - else: - air_brakes.deployment_level = 0.0 - - print(f"Time: {time:.2f}s, Alt AGL: {altitude_AGL:.1f}m, Mach: {mach_number:.2f}") - return (time, air_brakes.deployment_level, mach_number) - - return controller_function - -def test_backward_compatibility(): - """Test that old controller functions (without environment) still work""" - - def old_controller_function(time, sampling_rate, state, state_history, - observed_variables, air_brakes): - """ - Old-style controller function (6 parameters) - should still work - """ - altitude = state[2] - if altitude > 1000: - air_brakes.deployment_level = 0.3 - else: - air_brakes.deployment_level = 0.0 - return (time, air_brakes.deployment_level) - - return old_controller_function - -def test_with_sensors(): - """Test controller function with sensors parameter""" - - def controller_with_sensors(time, sampling_rate, state, state_history, - observed_variables, air_brakes, sensors): - """ - Controller function with sensors (7 parameters) - should still work - """ - altitude = state[2] - if altitude > 1000: - air_brakes.deployment_level = 0.4 - else: - air_brakes.deployment_level = 0.0 - return (time, air_brakes.deployment_level) - - return controller_with_sensors - -if __name__ == "__main__": - print("āœ… Air Brakes Controller Environment Parameter Test") - print("="*60) - - # Test functions - controller_new = test_controller_with_environment() - controller_old = test_backward_compatibility() - controller_sensors = test_with_sensors() - - print("āœ… Created controller functions successfully:") - print(f" - New controller (8 params): {controller_new.__name__}") - print(f" - Old controller (6 params): {controller_old.__name__}") - print(f" - Sensors controller (7 params): {controller_sensors.__name__}") - - print("\nāœ… All controller function signatures are supported!") - print("\nšŸ“ Benefits of the new environment parameter:") - print(" • No more global variables needed") - print(" • Proper serialization support") - print(" • More modular and testable code") - print(" • Access to wind, atmospheric, and environmental data") - print(" • Backward compatibility maintained") - - print("\nšŸš€ Example usage in controller:") - print(" # Old way (with global variables):") - print(" altitude_AGL = altitude_ASL - env.elevation # āŒ Global variable") - print(" wind_x = env.wind_velocity_x(altitude_ASL) # āŒ Global variable") - print("") - print(" # New way (with environment parameter):") - print(" altitude_AGL = altitude_ASL - environment.elevation # āœ… Local parameter") - print(" wind_x = environment.wind_velocity_x(altitude_ASL) # āœ… Local parameter") \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index 12d07c334..4ce1b45a6 100644 --- 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 new file mode 100644 index 000000000..e69de29bb diff --git a/tests/fixtures/controller/controller_fixtures.py b/tests/fixtures/controller/controller_fixtures.py new file mode 100644 index 000000000..515addd38 --- /dev/null +++ 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 index 7baae78f6..70fa9e9b9 100644 --- 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 index f40eb6b27..3a31a8aee 100644 --- a/tests/integration/simulation/test_flight.py +++ b/tests/integration/simulation/test_flight.py @@ -438,3 +438,361 @@ def test_rocket_csys_equivalence( flight_calisto_robust.initial_solution, flight_calisto_nose_to_tail_robust.initial_solution, ) + + +def test_air_brakes_with_environment_parameter( + calisto_robust, controller_function_with_environment, example_plain_env +): + """Test that air brakes controller can access environment parameter during flight. + + This test verifies that: + - The 8-parameter controller signature works correctly + - Environment data is accessible within the controller + - The flight simulation completes successfully + - Controller observed variables are properly stored + + This addresses issue #853 where environment had to be accessed via global variables. + + Parameters + ---------- + calisto_robust : rocketpy.Rocket + Calisto rocket without air brakes + controller_function_with_environment : function + Controller function using the new 8-parameter signature + example_plain_env : rocketpy.Environment + Environment object for the simulation + """ + # Add air brakes with 8-parameter controller + calisto_robust.parachutes = [] # Remove parachutes for cleaner test + calisto_robust.add_air_brakes( + drag_coefficient_curve="data/rockets/calisto/air_brakes_cd.csv", + controller_function=controller_function_with_environment, + sampling_rate=10, + clamp=True, + ) + + # Run flight simulation + flight = Flight( + rocket=calisto_robust, + environment=example_plain_env, + rail_length=5.2, + inclination=85, + heading=0, + terminate_on_apogee=True, + ) + + # Verify flight completed successfully + assert flight.t_final > 0 + assert flight.apogee > 0 + + # Verify controller was called and observed variables were stored + # Controller is attached to the rocket, not the air brakes object + controllers = [c for c in calisto_robust._controllers if "AirBrakes" in c.name] + assert len(controllers) > 0 + controller = controllers[0] + assert len(controller.observed_variables) > 0 + + # Verify observed variables contain expected data (time, deployment_level, mach_number) + for observed in controller.observed_variables: + if observed is not None: + assert len(observed) == 3 + time, deployment_level, mach_number = observed + assert time >= 0 + assert 0 <= deployment_level <= 1 # Should be clamped + assert mach_number >= 0 + + +def test_air_brakes_serialization_with_environment( + calisto_robust, controller_function_with_environment, example_plain_env +): + """Test that rockets with air brakes using environment parameter can be serialized. + + This test specifically addresses issue #853 - serialization of rockets with + air brakes that use controllers should work without relying on global variables. + + Parameters + ---------- + calisto_robust : rocketpy.Rocket + Calisto rocket without air brakes + controller_function_with_environment : function + Controller function using the new 8-parameter signature + example_plain_env : rocketpy.Environment + Environment object for the simulation + """ + # Add air brakes with 8-parameter controller + calisto_robust.parachutes = [] + calisto_robust.add_air_brakes( + drag_coefficient_curve="data/rockets/calisto/air_brakes_cd.csv", + controller_function=controller_function_with_environment, + sampling_rate=10, + clamp=True, + ) + + # Serialize the rocket + rocket_dict = calisto_robust.to_dict() + + # Verify serialization succeeded and contains air brakes data + assert "air_brakes" in rocket_dict + assert len(rocket_dict["air_brakes"]) > 0 + + # Run a flight with the original rocket + flight_original = Flight( + rocket=calisto_robust, + environment=example_plain_env, + rail_length=5.2, + inclination=85, + heading=0, + terminate_on_apogee=True, + ) + + # Verify flight completed + assert flight_original.t_final > 0 + assert flight_original.apogee > 0 + + +def test_backward_compatibility_6_parameter_controller( + calisto_robust, controller_function, example_plain_env +): + """Test that old 6-parameter controllers still work (backward compatibility). + + Parameters + ---------- + calisto_robust : rocketpy.Rocket + Calisto rocket without air brakes + controller_function : function + Controller function using the old 6-parameter signature + example_plain_env : rocketpy.Environment + Environment object for the simulation + """ + # Add air brakes with old-style 6-parameter controller + calisto_robust.parachutes = [] + calisto_robust.add_air_brakes( + drag_coefficient_curve="data/rockets/calisto/air_brakes_cd.csv", + controller_function=controller_function, + sampling_rate=10, + clamp=True, + ) + + # Run flight simulation + flight = Flight( + rocket=calisto_robust, + environment=example_plain_env, + rail_length=5.2, + inclination=85, + heading=0, + terminate_on_apogee=True, + ) + + # Verify flight completed successfully + assert flight.t_final > 0 + assert flight.apogee > 0 + + # Verify controller exists + controllers = [c for c in calisto_robust._controllers if "AirBrakes" in c.name] + assert len(controllers) > 0 + + +def test_7_parameter_controller_with_sensors(calisto_robust, example_plain_env): + """Test that 7-parameter controllers (with sensors, no environment) work correctly. + + Parameters + ---------- + calisto_robust : rocketpy.Rocket + Calisto rocket without air brakes + example_plain_env : rocketpy.Environment + Environment object for the simulation + """ + + # Define a 7-parameter controller + def controller_7_params( + time, + sampling_rate, + state, + state_history, + observed_variables, + air_brakes, + sensors, + ): + """Controller with 7 parameters (includes sensors, but not environment).""" + altitude = state[2] + vz = state[5] + + if time < 3.9: + return None + + if altitude < 1500: + air_brakes.deployment_level = 0 + else: + # Simple proportional control + air_brakes.deployment_level = min(0.5, max(0, vz / 100)) + + return (time, air_brakes.deployment_level) + + # Add air brakes with 7-parameter controller + calisto_robust.parachutes = [] + calisto_robust.add_air_brakes( + drag_coefficient_curve="data/rockets/calisto/air_brakes_cd.csv", + controller_function=controller_7_params, + sampling_rate=10, + clamp=True, + ) + + # Run flight simulation + flight = Flight( + rocket=calisto_robust, + environment=example_plain_env, + rail_length=5.2, + inclination=85, + heading=0, + terminate_on_apogee=True, + ) + + # Verify flight completed successfully + assert flight.t_final > 0 + assert flight.apogee > 0 + + +def test_invalid_controller_parameter_count(calisto_robust): + """Test that controllers with invalid parameter counts raise ValueError. + + Parameters + ---------- + calisto_robust : rocketpy.Rocket + Calisto rocket without air brakes + """ + + # Define controller with wrong number of parameters (5) + def invalid_controller_5_params( + time, sampling_rate, state, state_history, observed_variables + ): + """Invalid controller with only 5 parameters.""" + return None + + # Define controller with wrong number of parameters (9) + def invalid_controller_9_params( + time, + sampling_rate, + state, + state_history, + observed_variables, + air_brakes, + sensors, + environment, + extra_param, + ): + """Invalid controller with 9 parameters.""" + return None + + calisto_robust.parachutes = [] + + # Test that 5-parameter controller raises ValueError + with pytest.raises(ValueError, match="must have 6, 7, or 8 arguments"): + calisto_robust.add_air_brakes( + drag_coefficient_curve="data/rockets/calisto/air_brakes_cd.csv", + controller_function=invalid_controller_5_params, + sampling_rate=10, + clamp=True, + ) + + # Test that 9-parameter controller raises ValueError + with pytest.raises(ValueError, match="must have 6, 7, or 8 arguments"): + calisto_robust.add_air_brakes( + drag_coefficient_curve="data/rockets/calisto/air_brakes_cd.csv", + controller_function=invalid_controller_9_params, + sampling_rate=10, + clamp=True, + ) + + +def test_environment_methods_accessible_in_controller( + calisto_robust, example_plain_env +): + """Test that all environment methods are accessible within the controller. + + This test verifies that the environment object passed to the controller + provides access to all necessary atmospheric and environmental data. + + Parameters + ---------- + calisto_robust : rocketpy.Rocket + Calisto rocket without air brakes + example_plain_env : rocketpy.Environment + Environment object for the simulation + """ + # Track which environment methods were successfully called + methods_called = { + "elevation": False, + "wind_velocity_x": False, + "wind_velocity_y": False, + "speed_of_sound": False, + "pressure": False, + "temperature": False, + } + + def controller_test_environment_access( + time, + sampling_rate, + state, + state_history, + observed_variables, + air_brakes, + sensors, + environment, + ): + """Controller that tests access to various environment methods.""" + altitude_ASL = state[2] + + if time < 3.9: + return None + + # Test accessing various environment methods + try: + _ = environment.elevation + methods_called["elevation"] = True + + _ = environment.wind_velocity_x(altitude_ASL) + methods_called["wind_velocity_x"] = True + + _ = environment.wind_velocity_y(altitude_ASL) + methods_called["wind_velocity_y"] = True + + _ = environment.speed_of_sound(altitude_ASL) + methods_called["speed_of_sound"] = True + + _ = environment.pressure(altitude_ASL) + methods_called["pressure"] = True + + _ = environment.temperature(altitude_ASL) + methods_called["temperature"] = True + + air_brakes.deployment_level = 0.3 + except AttributeError as e: + # If any method is not accessible, the test should fail + raise AssertionError(f"Environment method not accessible: {e}") + + return (time, air_brakes.deployment_level) + + # Add air brakes with environment-testing controller + calisto_robust.parachutes = [] + calisto_robust.add_air_brakes( + drag_coefficient_curve="data/rockets/calisto/air_brakes_cd.csv", + controller_function=controller_test_environment_access, + sampling_rate=10, + clamp=True, + ) + + # Run flight simulation + flight = Flight( + rocket=calisto_robust, + environment=example_plain_env, + rail_length=5.2, + inclination=85, + heading=0, + terminate_on_apogee=True, + ) + + # Verify flight completed + assert flight.t_final > 0 + + # Verify all environment methods were successfully called + assert all(methods_called.values()), f"Not all methods called: {methods_called}" From fc10614b311326874c158cab583328877700f033 Mon Sep 17 00:00:00 2001 From: Gui-FernandesBR Date: Fri, 28 Nov 2025 22:36:58 -0300 Subject: [PATCH 4/4] make lint --- .../controller/controller_fixtures.py | 14 ++++++------ tests/integration/simulation/test_flight.py | 22 +++++++++---------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/tests/fixtures/controller/controller_fixtures.py b/tests/fixtures/controller/controller_fixtures.py index 515addd38..d00899d29 100644 --- a/tests/fixtures/controller/controller_fixtures.py +++ b/tests/fixtures/controller/controller_fixtures.py @@ -65,24 +65,24 @@ def controller_function( # pylint: disable=unused-argument 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_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) + 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) + mach_number = free_stream_speed / environment.speed_of_sound(altitude_asl) if time < 3.9: return None - if altitude_AGL < 1500: + if altitude_agl < 1500: air_brakes.deployment_level = 0 else: previous_vz = state_history[-1][5] if state_history else vz diff --git a/tests/integration/simulation/test_flight.py b/tests/integration/simulation/test_flight.py index 3a31a8aee..61a28ede2 100644 --- a/tests/integration/simulation/test_flight.py +++ b/tests/integration/simulation/test_flight.py @@ -604,7 +604,7 @@ def test_7_parameter_controller_with_sensors(calisto_robust, example_plain_env): """ # Define a 7-parameter controller - def controller_7_params( + def controller_7_params( # pylint: disable=unused-argument time, sampling_rate, state, @@ -662,14 +662,14 @@ def test_invalid_controller_parameter_count(calisto_robust): """ # Define controller with wrong number of parameters (5) - def invalid_controller_5_params( + def invalid_controller_5_params( # pylint: disable=unused-argument time, sampling_rate, state, state_history, observed_variables ): """Invalid controller with only 5 parameters.""" return None # Define controller with wrong number of parameters (9) - def invalid_controller_9_params( + def invalid_controller_9_params( # pylint: disable=unused-argument time, sampling_rate, state, @@ -729,7 +729,7 @@ def test_environment_methods_accessible_in_controller( "temperature": False, } - def controller_test_environment_access( + def controller_test_environment_access( # pylint: disable=unused-argument time, sampling_rate, state, @@ -740,7 +740,7 @@ def controller_test_environment_access( environment, ): """Controller that tests access to various environment methods.""" - altitude_ASL = state[2] + altitude_asl = state[2] if time < 3.9: return None @@ -750,25 +750,25 @@ def controller_test_environment_access( _ = environment.elevation methods_called["elevation"] = True - _ = environment.wind_velocity_x(altitude_ASL) + _ = environment.wind_velocity_x(altitude_asl) methods_called["wind_velocity_x"] = True - _ = environment.wind_velocity_y(altitude_ASL) + _ = environment.wind_velocity_y(altitude_asl) methods_called["wind_velocity_y"] = True - _ = environment.speed_of_sound(altitude_ASL) + _ = environment.speed_of_sound(altitude_asl) methods_called["speed_of_sound"] = True - _ = environment.pressure(altitude_ASL) + _ = environment.pressure(altitude_asl) methods_called["pressure"] = True - _ = environment.temperature(altitude_ASL) + _ = environment.temperature(altitude_asl) methods_called["temperature"] = True air_brakes.deployment_level = 0.3 except AttributeError as e: # If any method is not accessible, the test should fail - raise AssertionError(f"Environment method not accessible: {e}") + raise AssertionError(f"Environment method not accessible: {e}") from e return (time, air_brakes.deployment_level)