diff --git a/CHANGELOG.md b/CHANGELOG.md index 75e20416e..373cdce77 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: 3-dof lateral motion improvement [#883](https://github.com/RocketPy-Team/RocketPy/pull/883) - ENH: Add multi-dimensional drag coefficient support (Cd as function of M, Re, α) [#875](https://github.com/RocketPy-Team/RocketPy/pull/875) - ENH: Add save functionality to `_MonteCarloPlots.all` method [#848](https://github.com/RocketPy-Team/RocketPy/pull/848) 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 27ad62361..e81e70915 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,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) diff --git a/rocketpy/rocket/rocket.py b/rocketpy/rocket/rocket.py index 9ac1bb86e..fdea1ac47 100644 --- a/rocketpy/rocket/rocket.py +++ b/rocketpy/rocket/rocket.py @@ -1677,9 +1677,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 84ab880dc..fa68862f9 100644 --- a/rocketpy/simulation/flight.py +++ b/rocketpy/simulation/flight.py @@ -757,6 +757,7 @@ def __simulate(self, verbose): self.y_sol, self.solution, self.sensors, + self.env, ) for parachute in node.parachutes: diff --git a/tests/conftest.py b/tests/conftest.py index 456de43ca..5a8ca93fa 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..d00899d29 --- /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 79a24dc32..70fa9e9b9 100644 --- a/tests/fixtures/function/function_fixtures.py +++ b/tests/fixtures/function/function_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. diff --git a/tests/integration/simulation/test_flight.py b/tests/integration/simulation/test_flight.py index f40eb6b27..61a28ede2 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( # pylint: disable=unused-argument + 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( # 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( # pylint: disable=unused-argument + 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( # pylint: disable=unused-argument + 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}") from 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}"