diff --git a/rocketpy/stochastic/stochastic_model.py b/rocketpy/stochastic/stochastic_model.py index 333a6d891..d42fb76c5 100644 --- a/rocketpy/stochastic/stochastic_model.py +++ b/rocketpy/stochastic/stochastic_model.py @@ -124,6 +124,22 @@ def __init__(self, obj, seed=None, **kwargs): self.__stochastic_dict = kwargs self._set_stochastic(seed) + def _declare_stochastic_input(self, input_name, input_value): + """Declare an input that an ``add_*`` method installs after ``__init__``. + + ``dict_generator`` walks the inputs a model declared rather than every + attribute on it (#1109), and that list is built in ``__init__``. Anything + added afterwards is set on the instance and never drawn from unless it + says so here. + + The value is the argument as given, not the validated form, because + ``_set_stochastic`` validates it again on every reseed and binds the + distribution to the generator that is live then. + """ + if input_value is None: + return + self.__stochastic_dict[input_name] = input_value + def _set_stochastic(self, seed=None): """Set the stochastic attributes from the input dictionary. This method is useful to reset or reseed the attributes of the instance. diff --git a/rocketpy/stochastic/stochastic_rocket.py b/rocketpy/stochastic/stochastic_rocket.py index 895e9a2a4..515439f14 100644 --- a/rocketpy/stochastic/stochastic_rocket.py +++ b/rocketpy/stochastic/stochastic_rocket.py @@ -457,7 +457,9 @@ def add_cp_eccentricity(self, x=None, y=None): Object of the StochasticRocket class. """ self.cp_eccentricity_x = self._validate_eccentricity("cp_eccentricity_x", x) + self._declare_stochastic_input("cp_eccentricity_x", x) self.cp_eccentricity_y = self._validate_eccentricity("cp_eccentricity_y", y) + self._declare_stochastic_input("cp_eccentricity_y", y) return self def add_thrust_eccentricity(self, x=None, y=None): @@ -485,9 +487,11 @@ def add_thrust_eccentricity(self, x=None, y=None): self.thrust_eccentricity_x = self._validate_eccentricity( "thrust_eccentricity_x", x ) + self._declare_stochastic_input("thrust_eccentricity_x", x) self.thrust_eccentricity_y = self._validate_eccentricity( "thrust_eccentricity_y", y ) + self._declare_stochastic_input("thrust_eccentricity_y", y) return self def _validate_eccentricity(self, eccentricity, position): diff --git a/tests/unit/stochastic/test_stochastic_rocket.py b/tests/unit/stochastic/test_stochastic_rocket.py index d15eb0bb6..dcf94df36 100644 --- a/tests/unit/stochastic/test_stochastic_rocket.py +++ b/tests/unit/stochastic/test_stochastic_rocket.py @@ -1,3 +1,5 @@ +from numbers import Real + import numpy as np import pytest @@ -191,3 +193,46 @@ def test_add_free_form_fins_wraps_a_deterministic_fin_set(calisto_robust): added = stochastic.aerodynamic_surfaces.get_tuple_by_type(StochasticFreeFormFins) assert len(added) == 1 assert added[0].component.obj is fins + + +@pytest.mark.parametrize( + "add_them, names", + [ + ( + "add_cp_eccentricity", + ("cp_eccentricity_x", "cp_eccentricity_y"), + ), + ( + "add_thrust_eccentricity", + ("thrust_eccentricity_x", "thrust_eccentricity_y"), + ), + ], +) +def test_an_eccentricity_added_after_init_is_still_drawn(calisto, add_them, names): + """``dict_generator`` walks the declared inputs, and these arrive later. + + The list is built in ``__init__``, so a distribution installed by an + ``add_*`` method afterwards was set on the instance and never drawn from: + every simulation used the same value, with nothing to say so. + """ + stochastic = StochasticRocket(rocket=calisto, radius=0.0127 / 2) + getattr(stochastic, add_them)(x=(0.0, 0.001), y=(0.0, 0.001)) + stochastic._set_stochastic(42) + + generated = next(stochastic.dict_generator()) + + assert set(names) <= set(generated), f"{add_them} was set but never sampled" + assert all(isinstance(generated[name], Real) for name in names) + + +def test_two_seeds_move_an_eccentricity_that_was_added_late(calisto): + """Being present is not enough; it has to follow the seed.""" + stochastic = StochasticRocket(rocket=calisto, radius=0.0127 / 2) + stochastic.add_cp_eccentricity(x=(0.0, 0.01), y=(0.0, 0.01)) + + def drawn(seed): + stochastic._set_stochastic(seed) + return next(stochastic.dict_generator())["cp_eccentricity_x"] + + assert drawn(7) == drawn(7) + assert drawn(7) != drawn(8)