|
| 1 | +"""Characterization tests for ``Flight.step_simulation()``. |
| 2 | +
|
| 3 | +The fork's stepped-simulation API (``run_simulation=False`` plus repeated |
| 4 | +``step_simulation()`` calls) must reproduce a one-shot ``simulate()`` so that |
| 5 | +callers driving the flight one node at a time -- e.g. the Balloon Popping |
| 6 | +Challenge environment, which steps it every timestep -- obtain the same |
| 7 | +trajectory. A parachute-free rocket (``flight_calisto``) is used because |
| 8 | +parachute triggers are not yet migrated into the stepping path. |
| 9 | +
|
| 10 | +Scope: ``TestStepSimulation`` guards that *uncontrolled* stepping matches |
| 11 | +``simulate()``. ``TestControlledStepSimulation`` guards the controlled use case |
| 12 | +the Balloon Popping Challenge actually relies on: injecting an actuator command |
| 13 | +between ``step_simulation()`` calls must change the trajectory. |
| 14 | +""" |
| 15 | + |
| 16 | +import copy |
| 17 | + |
| 18 | +import numpy as np |
| 19 | + |
| 20 | +from rocketpy import Flight |
| 21 | + |
| 22 | + |
| 23 | +def _stepped_twin(reference_flight): |
| 24 | + """A non-simulated ``Flight`` twin of ``reference_flight``, to step by hand. |
| 25 | +
|
| 26 | + Launch parameters are read back from the reference so the twin cannot drift |
| 27 | + from it. |
| 28 | + """ |
| 29 | + return Flight( |
| 30 | + environment=reference_flight.env, |
| 31 | + rocket=reference_flight.rocket, |
| 32 | + rail_length=reference_flight.rail_length, |
| 33 | + inclination=reference_flight.inclination, |
| 34 | + heading=reference_flight.heading, |
| 35 | + terminate_on_apogee=reference_flight.terminate_on_apogee, |
| 36 | + run_simulation=False, |
| 37 | + ) |
| 38 | + |
| 39 | + |
| 40 | +def _run_stepped(flight, max_steps=100000): |
| 41 | + """Drive ``step_simulation`` to completion. |
| 42 | +
|
| 43 | + Returns the number of calls made and the set of phase indices visited. |
| 44 | + """ |
| 45 | + steps = 0 |
| 46 | + phases_seen = {flight._step_state["phase_index"]} |
| 47 | + while not flight._step_state["finished"]: |
| 48 | + flight.step_simulation() |
| 49 | + phases_seen.add(flight._step_state["phase_index"]) |
| 50 | + steps += 1 |
| 51 | + assert steps < max_steps, "stepped simulation did not terminate" |
| 52 | + return steps, phases_seen |
| 53 | + |
| 54 | + |
| 55 | +class TestStepSimulation: |
| 56 | + """Stepping must match a one-shot ``simulate()`` and finalise correctly.""" |
| 57 | + |
| 58 | + def test_initial_state_is_unfinished_at_first_phase(self, flight_calisto): |
| 59 | + stepped = _stepped_twin(flight_calisto) |
| 60 | + assert stepped._step_state["finished"] is False |
| 61 | + assert stepped._step_state["phase_index"] == 0 |
| 62 | + assert stepped._step_state["node_index"] == 0 |
| 63 | + |
| 64 | + def test_stepping_visits_multiple_phases_then_finishes(self, flight_calisto): |
| 65 | + stepped = _stepped_twin(flight_calisto) |
| 66 | + _, phases_seen = _run_stepped(stepped) |
| 67 | + assert stepped._step_state["finished"] is True |
| 68 | + assert len(phases_seen) > 1 # at least a rail phase and a flight phase |
| 69 | + |
| 70 | + def test_stepped_trajectory_matches_simulate(self, flight_calisto): |
| 71 | + stepped = _stepped_twin(flight_calisto) |
| 72 | + _run_stepped(stepped) |
| 73 | + # Stepping replays the same solver nodes as simulate(). A tight tolerance |
| 74 | + # (rather than exact equality) keeps the guard robust to last-bit noise in |
| 75 | + # the LSODA Fortran solver across platforms, while still catching any real |
| 76 | + # trajectory divergence. |
| 77 | + np.testing.assert_allclose(stepped.t, flight_calisto.t, rtol=1e-8, atol=1e-10) |
| 78 | + np.testing.assert_allclose( |
| 79 | + np.array(stepped.solution), |
| 80 | + np.array(flight_calisto.solution), |
| 81 | + rtol=1e-8, |
| 82 | + atol=1e-10, |
| 83 | + ) |
| 84 | + |
| 85 | + def test_post_process_artifacts_exist_after_stepping(self, flight_calisto): |
| 86 | + stepped = _stepped_twin(flight_calisto) |
| 87 | + _run_stepped(stepped) |
| 88 | + # post_process_simulation() and initialize_prints_plots() fire on finish. |
| 89 | + assert stepped.t_final == stepped.t |
| 90 | + assert stepped.prints is not None |
| 91 | + assert stepped.plots is not None |
| 92 | + |
| 93 | + def test_stepping_after_finished_is_a_noop(self, flight_calisto): |
| 94 | + stepped = _stepped_twin(flight_calisto) |
| 95 | + _run_stepped(stepped) |
| 96 | + t_final, y_final = stepped.t, np.array(stepped.y_sol) |
| 97 | + stepped.step_simulation() # already finished -> must return immediately |
| 98 | + assert stepped.t == t_final |
| 99 | + np.testing.assert_array_equal(stepped.y_sol, y_final) |
| 100 | + |
| 101 | + |
| 102 | +def _no_op_roll_logger( |
| 103 | + time, |
| 104 | + sampling_rate, |
| 105 | + state, |
| 106 | + state_history, |
| 107 | + observed_variables, |
| 108 | + roll_control, |
| 109 | + sensors, |
| 110 | + environment, |
| 111 | +): # pylint: disable=unused-argument |
| 112 | + """Controller that only *logs* the roll torque, never sets it. |
| 113 | +
|
| 114 | + This mirrors the Balloon Popping Challenge wiring, where the controller is a |
| 115 | + passive logger and the command is injected externally between steps. It must |
| 116 | + not override the externally set ``roll_torque``. |
| 117 | + """ |
| 118 | + return (time, roll_control.roll_torque) |
| 119 | + |
| 120 | + |
| 121 | +def _controlled_twin(calisto, max_roll_torque=100.0): |
| 122 | + """A deep copy of ``calisto`` fitted with a no-op-logger roll actuator. |
| 123 | +
|
| 124 | + A copy is used so the original fixture rocket stays actuator-free and can |
| 125 | + serve as the uncontrolled baseline within the same test. Roll control is |
| 126 | + chosen because its torque is added straight onto the body-axis moment |
| 127 | + (``M3 += roll_control.roll_torque``), giving a thrust-independent effect that |
| 128 | + shows up cleanly in the roll rate ``w3``. ``max_roll_torque`` is left well |
| 129 | + above the commanded value so the command passes through unclamped, and no |
| 130 | + rate limit / time constant is set so the injected value is applied verbatim. |
| 131 | + """ |
| 132 | + twin = copy.deepcopy(calisto) |
| 133 | + twin.add_roll_control( |
| 134 | + controller_function=_no_op_roll_logger, |
| 135 | + sampling_rate=10.0, |
| 136 | + max_roll_torque=max_roll_torque, |
| 137 | + ) |
| 138 | + return twin |
| 139 | + |
| 140 | + |
| 141 | +def _step_with_roll(env, rocket, command, max_steps=100000): |
| 142 | + """Drive a stepped flight, injecting a roll command before every step. |
| 143 | +
|
| 144 | + Mirrors the BPC ``step()`` loop, which sets |
| 145 | + ``rocket.roll_control.roll_torque`` ahead of each ``step_simulation()``. |
| 146 | + ``command`` is either a constant torque (N.m) or a callable ``command(t)`` |
| 147 | + returning the torque to inject at the current flight time. |
| 148 | +
|
| 149 | + ``time_overshoot=False`` is essential and is exactly what BPC uses: it turns |
| 150 | + the controller's sampling nodes into real solver time nodes, so each |
| 151 | + ``step_simulation()`` advances one small node and the command is injected |
| 152 | + every timestep. With the default ``time_overshoot=True`` each call would |
| 153 | + instead integrate a whole flight phase, injecting the command only a handful |
| 154 | + of times -- not the per-timestep control loop under test. The on-rail phase |
| 155 | + is roll-constrained, so commanding through it is harmless; the body only |
| 156 | + starts spinning up once the rocket leaves the rail. |
| 157 | +
|
| 158 | + Returns the finished flight and the number of injected steps. |
| 159 | + """ |
| 160 | + command_fn = command if callable(command) else (lambda t: command) |
| 161 | + flight = Flight( |
| 162 | + environment=env, |
| 163 | + rocket=rocket, |
| 164 | + rail_length=5.2, |
| 165 | + inclination=85, |
| 166 | + heading=0, |
| 167 | + terminate_on_apogee=True, |
| 168 | + run_simulation=False, |
| 169 | + time_overshoot=False, |
| 170 | + ) |
| 171 | + steps = 0 |
| 172 | + while not flight._step_state["finished"]: |
| 173 | + rocket.roll_control.roll_torque = command_fn(flight.t) |
| 174 | + flight.step_simulation() |
| 175 | + steps += 1 |
| 176 | + assert steps < max_steps, "stepped simulation did not terminate" |
| 177 | + return flight, steps |
| 178 | + |
| 179 | + |
| 180 | +class TestControlledStepSimulation: |
| 181 | + """Injecting an actuator command between steps must move the trajectory. |
| 182 | +
|
| 183 | + This is the fork's actual use case -- the one the Balloon Popping Challenge |
| 184 | + relies on: a passive logger controller plus a roll command set externally |
| 185 | + between ``step_simulation()`` calls. |
| 186 | + """ |
| 187 | + |
| 188 | + # Roll torque commanded between steps, in N.m. Well inside the actuator |
| 189 | + # range so it is applied verbatim; large relative to the tiny roll inertia |
| 190 | + # so the angular response is unmistakable. |
| 191 | + COMMAND = 5.0 |
| 192 | + # Flight time (s) at which the reversal test flips the command sign. Comfor- |
| 193 | + # tably between rail departure and the (~25 s) apogee of this rocket, and |
| 194 | + # early enough that the post-reversal torque decisively reverses the spin. |
| 195 | + REVERSAL_TIME = 10.0 |
| 196 | + |
| 197 | + def test_injected_roll_command_changes_angular_state( |
| 198 | + self, calisto, example_plain_env |
| 199 | + ): |
| 200 | + rocket = _controlled_twin(calisto) |
| 201 | + commanded, commanded_steps = _step_with_roll( |
| 202 | + example_plain_env, rocket, self.COMMAND |
| 203 | + ) |
| 204 | + neutral, _ = _step_with_roll(example_plain_env, rocket, 0.0) |
| 205 | + |
| 206 | + # The command really was injected per timestep, not once per phase: the |
| 207 | + # fine-grained loop runs for many steps (hundreds for this rocket). |
| 208 | + assert commanded_steps > 50 |
| 209 | + |
| 210 | + # Sample the roll rate off the rail, up to just before apogee, on a grid |
| 211 | + # common to both flights (their nodes need not coincide). |
| 212 | + grid = np.linspace(2.0, min(commanded.t, neutral.t) - 0.5, 12) |
| 213 | + commanded_w3 = np.array([commanded.w3(t) for t in grid]) |
| 214 | + neutral_w3 = np.array([neutral.w3(t) for t in grid]) |
| 215 | + |
| 216 | + # Neutral command -> the actuator contributes no moment -> no roll. |
| 217 | + assert np.max(np.abs(neutral_w3)) < 1e-6 |
| 218 | + |
| 219 | + # Commanded -> the body spins up, the roll rate growing while the torque |
| 220 | + # is sustained. Orders of magnitude above any solver noise. |
| 221 | + assert np.all(np.diff(commanded_w3) > 0) |
| 222 | + assert np.abs(commanded_w3[-1]) > 100.0 |
| 223 | + |
| 224 | + # The two angular trajectories diverge: the injected command had a real |
| 225 | + # dynamical effect. |
| 226 | + divergence = np.max(np.abs(commanded_w3 - neutral_w3)) |
| 227 | + assert divergence > 100.0 |
| 228 | + |
| 229 | + def test_mid_flight_command_reversal_is_tracked(self, calisto, example_plain_env): |
| 230 | + # Spin one way, then -- partway up -- command the opposite torque. This |
| 231 | + # can only register if the command is re-read on every step: a command |
| 232 | + # applied once (or latched) could never make the roll rate turn around. |
| 233 | + rocket = _controlled_twin(calisto) |
| 234 | + |
| 235 | + def command(t): |
| 236 | + return self.COMMAND if t < self.REVERSAL_TIME else -self.COMMAND |
| 237 | + |
| 238 | + reversed_flight, _ = _step_with_roll(example_plain_env, rocket, command) |
| 239 | + |
| 240 | + grid = np.linspace(2.0, reversed_flight.t - 0.5, 20) |
| 241 | + w3 = np.array([reversed_flight.w3(t) for t in grid]) |
| 242 | + |
| 243 | + # Roll rate climbs, peaks in the interior (the sign flip), then comes |
| 244 | + # back down through zero -- the body is now spinning the other way. |
| 245 | + assert w3.max() > 100.0 |
| 246 | + assert 0 < int(np.argmax(w3)) < len(w3) - 1 |
| 247 | + assert w3[-1] < 0 |
| 248 | + |
| 249 | + def test_zero_command_matches_uncontrolled(self, calisto, example_plain_env): |
| 250 | + # Build the controlled twin before flying the baseline so both share an |
| 251 | + # identical, simulation-untouched starting rocket. |
| 252 | + rocket = _controlled_twin(calisto) |
| 253 | + |
| 254 | + # Uncontrolled baseline: the untouched fixture rocket, run one-shot. |
| 255 | + baseline = Flight( |
| 256 | + environment=example_plain_env, |
| 257 | + rocket=calisto, |
| 258 | + rail_length=5.2, |
| 259 | + inclination=85, |
| 260 | + heading=0, |
| 261 | + terminate_on_apogee=True, |
| 262 | + ) |
| 263 | + |
| 264 | + # Same rocket plus a no-op actuator, stepped with a zero command. |
| 265 | + neutral, _ = _step_with_roll(example_plain_env, rocket, 0.0) |
| 266 | + |
| 267 | + grid = np.linspace(0.5, min(baseline.t, neutral.t) - 0.2, 60) |
| 268 | + |
| 269 | + # The actuator acts only on the roll axis, so commanding zero must leave |
| 270 | + # the angular state bit-for-bit identical to the uncontrolled flight: |
| 271 | + # the divergence in the tests above is the *command*, never the mere |
| 272 | + # presence of the actuator. |
| 273 | + for state in ("w1", "w2", "w3"): |
| 274 | + baseline_state = getattr(baseline, state) |
| 275 | + neutral_state = getattr(neutral, state) |
| 276 | + np.testing.assert_allclose( |
| 277 | + [neutral_state(t) for t in grid], |
| 278 | + [baseline_state(t) for t in grid], |
| 279 | + rtol=0, |
| 280 | + atol=1e-9, |
| 281 | + ) |
| 282 | + |
| 283 | + # The translational trajectory also tracks the uncontrolled run. It is |
| 284 | + # not bit-identical: ``time_overshoot=False`` forces an LSODA restart at |
| 285 | + # every controller node, and those restarts accumulate sub-metre path |
| 286 | + # noise over the multi-kilometre flight (well within the integrator's own |
| 287 | + # tolerance). This is pure numerical path noise from the inert actuator, |
| 288 | + # not a dynamical effect -- the roll axis it acts on is left exactly |
| 289 | + # unchanged (asserted above) -- so it cannot be confused with a command. |
| 290 | + for state in ("x", "y", "z", "vx", "vy", "vz"): |
| 291 | + baseline_state = getattr(baseline, state) |
| 292 | + neutral_state = getattr(neutral, state) |
| 293 | + np.testing.assert_allclose( |
| 294 | + [neutral_state(t) for t in grid], |
| 295 | + [baseline_state(t) for t in grid], |
| 296 | + rtol=2e-3, |
| 297 | + atol=1e-1, |
| 298 | + ) |
0 commit comments