ENH: continue a Monte Carlo study from the root its rows were drawn with

CHANGELOG.md

@@ -32,6 +32,8 @@ Attention: The newest changes should be on top -->
 
 ### Added
 
+- ENH: Continue a Monte Carlo study from the root its rows were drawn with [#1187](https://github.com/RocketPy-Team/RocketPy/pull/1187) [#1075](https://github.com/RocketPy-Team/RocketPy/issues/1075)
+- ENH: Reproducible Monte Carlo through per-simulation-index seeding [#1054](https://github.com/RocketPy-Team/RocketPy/pull/1054) [#1053](https://github.com/RocketPy-Team/RocketPy/issues/1053)
 - ENH: Support fixed-time parachute deployment triggers [#1133](https://github.com/RocketPy-Team/RocketPy/pull/1133) [#437](https://github.com/RocketPy-Team/RocketPy/issues/437)
 - DOC: Add SIL parachute ejection integration example [#1131](https://github.com/RocketPy-Team/RocketPy/pull/1131) [#524](https://github.com/RocketPy-Team/RocketPy/issues/524)
 - ENH: List NOAA atmosphere datasets and fetch latest [#1136](https://github.com/RocketPy-Team/RocketPy/pull/1136) [#660](https://github.com/RocketPy-Team/RocketPy/issues/660)

docs/user/stochastic.rst

@@ -289,3 +289,12 @@ better reflecting the inherent uncertainties in rocketry.
 .. note::
     See the ``MonteCarlo`` class documentation for more information on how to run \
     Monte Carlo simulations with stochastic objects.
+
+.. note::
+    A whole run is fixed by ``MonteCarlo.simulate(random_seed=...)`` rather than
+    by seeding these models yourself. Each simulation takes its seed from its
+    own index, so simulation 7 draws the same inputs whether the run was serial
+    or split over any number of workers. Every input row records the root it
+    came from, so appending carries that study on whether or not the seed is
+    given again, and a different one is refused rather than mixed in. Without
+    a seed a run draws fresh entropy and reproduces nothing.

rocketpy/simulation/monte_carlo.py

@@ -18,6 +18,7 @@
 import os
 import traceback
 import warnings
+from copy import deepcopy
 from numbers import Real
 from pathlib import Path
 from time import time
@@ -31,6 +32,7 @@
 from rocketpy.prints.monte_carlo_prints import _MonteCarloPrints
 from rocketpy.simulation.flight import Flight
 from rocketpy.tools import (
+    _seed_sequence_to_int,
     generate_monte_carlo_ellipses,
     generate_monte_carlo_ellipses_coordinates,
     import_optional_dependency,
@@ -43,6 +45,125 @@
 # this is the only format it can both resume from and overwrite safely.
 _SIMULATION_LOG_SUFFIX = ".txt"
 
+# Which root drew a row. An append reads it to continue the same stream.
+_SIMULATION_ROOT_KEY = "run_root"
+
+# Told apart from a row that carries ``None``, which no run writes.
+_NOTHING_READ_YET = object()
+
+
+def _root_seed_sequence(random_seed):
+    """The immutable root a run derives every simulation's seed from.
+
+    A ``SeedSequence`` is rebuilt from its full state rather than used as
+    given, since ``spawn`` advances a counter the caller still holds. A
+    ``Generator`` is refused rather than read, because using a consume-on-use
+    object as an immutable seed cannot mean what it says.
+    """
+    if isinstance(random_seed, np.random.SeedSequence):
+        return np.random.SeedSequence(**random_seed.state)
+    if isinstance(random_seed, (np.random.Generator, np.random.BitGenerator)):
+        raise TypeError(
+            f"random_seed must be an int, a sequence of non-negative integers, "
+            f"or a numpy.random.SeedSequence, not a "
+            f"{type(random_seed).__name__}. Pass the seed the generator was "
+            f"built from."
+        )
+    return np.random.SeedSequence(random_seed)
+
+
+def _jsonable_entropy(entropy):
+    """``SeedSequence`` entropy as something ``json`` will take.
+
+    It may be an int, a sequence or an ndarray; only the first survives.
+    """
+    if entropy is None or isinstance(entropy, (int, np.integer)):
+        return None if entropy is None else int(entropy)
+    return [int(part) for part in np.asarray(entropy).ravel()]
+
+
+def _root_written_into_a_row(root_state):
+    """The run's root as one JSON value, carried by every input row.
+
+    In the rows because a file beside a log cannot be shown to belong to it.
+    """
+    entropy, spawn_key, pool_size, base = root_state
+    return {
+        "entropy": _jsonable_entropy(entropy),
+        "spawn_key": [int(key) for key in spawn_key],
+        "pool_size": int(pool_size),
+        "n_children_spawned": int(base),
+    }
+
+
+def _root_a_log_was_written_with(path):
+    """The root every row of a log agrees on, or ``None`` if it holds none.
+
+    ``None`` means the log holds no rows, and nothing else does. Rows that
+    carry no root are refused instead: they cannot be shown to be one study,
+    and reading them as an empty log would start a second one in the file.
+    A log whose rows disagree is refused for the same reason.
+    """
+    first = _NOTHING_READ_YET
+    with open(path, "r", encoding="utf-8") as recorded:
+        for line in recorded:
+            if not line.strip():
+                continue
+            try:
+                row = json.loads(line)
+            except ValueError as error:
+                raise ValueError(
+                    f"cannot continue {path}: a row cannot be read, so what "
+                    f"produced it cannot be established."
+                ) from error
+            if _SIMULATION_ROOT_KEY not in row:
+                raise ValueError(
+                    f"cannot continue {path}: a row does not say which root "
+                    f"drew it, which is how a study written before this "
+                    f"release looks. Start a new one rather than continuing "
+                    f"one whose rows cannot be checked."
+                )
+            root = row[_SIMULATION_ROOT_KEY]
+            if first is _NOTHING_READ_YET:
+                first = root
+            elif root != first:
+                raise ValueError(
+                    f"cannot continue {path}: its rows were not all drawn "
+                    f"from one root, so it holds more than one study."
+                )
+    return None if first is _NOTHING_READ_YET else first
+
+
+def _root_state_of(root):
+    """A root as the four picklable values a worker can rebuild it from.
+
+    Sent to each worker instead of the object, and instead of the list of
+    children, so a run of a million simulations costs four values. The entropy
+    is copied because a sequence one is kept by reference all the way from the
+    caller, who could otherwise still move every child by editing their list.
+    """
+    return (
+        deepcopy(root.entropy),
+        tuple(root.spawn_key),
+        root.pool_size,
+        root.n_children_spawned,
+    )
+
+
+def _seed_of_simulation(root_state, sim_idx):
+    """The seed for one simulation index, without spawning the ones before it.
+
+    ``spawn`` derives child ``i`` by appending ``n_children_spawned + i`` to
+    the parent spawn key, so rebuilding that one child directly reproduces it
+    and any index can be reached from the four values above alone.
+    """
+    entropy, spawn_key, pool_size, base = root_state
+    return np.random.SeedSequence(
+        entropy=entropy,
+        spawn_key=(*spawn_key, base + sim_idx),
+        pool_size=pool_size,
+    )
+
 
 def _refuse_logs_this_run_cannot_write(
     input_file, output_file, error_file, export_config=None
@@ -265,6 +386,8 @@ def simulate(
         append=False,
         parallel=False,
         n_workers=None,
+        *,
+        random_seed=None,
         **kwargs,
     ):
         """
@@ -284,6 +407,19 @@ def simulate(
             number of workers will be equal to the number of CPUs available.
             A minimum of 2 workers is required for parallel mode.
             Default is None.
+        random_seed : int, sequence of int or numpy.random.SeedSequence, optional
+            Fixes what every simulation draws. Simulation ``i`` takes the same
+            inputs whichever way the run was split up, so serial and parallel
+            results agree and the number of workers does not reach the
+            sampling. Keyword-only. Default is None, which draws fresh entropy
+            and reproduces nothing.
+
+            Every input row carries the root it was drawn from, so an append
+            carries on from the study already in the file whether or not the
+            seed is given again. A different one is refused, not mixed in.
+
+            A ``Generator`` or ``BitGenerator`` is refused rather than read.
+            Pass the seed it was built from.
         kwargs : dict
             Custom arguments for simulation export of the ``inputs`` file. Options
             are:
@@ -317,12 +453,21 @@ def simulate(
         self._export_config = kwargs
         self.number_of_simulations = number_of_simulations
         self._initial_sim_idx = self.num_of_loaded_sims if append else 0
-
+        # Validated here, before __setup_files truncates anything, so an
+        # unusable seed cannot cost a previous run its results. Kept as four
+        # picklable values rather than as the object, since a worker rebuilds
+        # any index from them.
+        self.__root_state = _root_state_of(_root_seed_sequence(random_seed))
         # Before anything is opened: __setup_files truncates for append=False.
         _refuse_logs_this_run_cannot_write(
             self.input_file, self.output_file, self.error_file, kwargs
         )
 
+        # After that one, which says plainly that a .csv cannot be a working
+        # log. Reaching this first would report it as a row that cannot be read.
+        if append:
+            self.__continue_the_root_the_rows_carry(random_seed)
+
         print("Starting Monte Carlo analysis")
 
         self.__setup_files(append)
@@ -413,14 +558,19 @@ def __run_in_serial(self):
             n_simulations=self.number_of_simulations,
             start_time=time(),
         )
+        sim_idx = sim_monitor.count
         try:
             while sim_monitor.keep_simulating():
-                sim_monitor.increment()
+                # Counted from zero, as the parallel path already does. The two
+                # named the same simulation differently: three of them wrote
+                # 1, 2, 3 here and 0, 1, 2 there.
+                sim_idx = sim_monitor.increment() - 1
                 inputs_json, outputs_json = "", ""
 
+                self.__seed_this_simulation(sim_idx)
                 flight = self.__run_single_simulation()
-                inputs_json = self.__evaluate_flight_inputs(sim_monitor.count)
-                outputs_json = self.__evaluate_flight_outputs(flight, sim_monitor.count)
+                inputs_json = self.__evaluate_flight_inputs(sim_idx)
+                outputs_json = self.__evaluate_flight_outputs(flight, sim_idx)
 
                 self._append_simulation_record(inputs_json, outputs_json)
 
@@ -434,7 +584,7 @@ def __run_in_serial(self):
                 f.write(inputs_json)
 
         except Exception as error:
-            print(f"Error on iteration {sim_monitor.count}: {error}")
+            print(f"Error on iteration {sim_idx}: {error}")
             with open(self._error_file, "a", encoding="utf-8") as f:
                 f.write(inputs_json)
             raise error
@@ -469,13 +619,14 @@ def __run_in_parallel(self, n_workers=None):
             )
 
             processes = []
-            seeds = np.random.SeedSequence().spawn(n_workers)
 
-            for seed in seeds:
+            # No seed per worker any more: every simulation takes its own from
+            # its index, so the workers are interchangeable and how many there
+            # are does not reach the sampling.
+            for _ in range(n_workers):
                 sim_producer = multiprocess.Process(
                     target=self.__sim_producer,
                     args=(
-                        seed,
                         sim_monitor,
                         mutex,
                         simulation_error_event,
@@ -517,13 +668,11 @@ def __validate_number_of_workers(self, n_workers):
             raise ValueError("Number of workers must be at least 2 for parallel mode.")
         return n_workers
 
-    def __sim_producer(self, seed, sim_monitor, mutex, error_event):  # pylint: disable=too-many-statements
+    def __sim_producer(self, sim_monitor, mutex, error_event):
         """Simulation producer to be used in parallel by multiprocessing.
 
         Parameters
         ----------
-        seed : int
-            The seed to set the random number generator.
         sim_monitor : _SimMonitor
             The simulation monitor object to keep track of the simulations.
         mutex : multiprocess.Lock
@@ -532,15 +681,11 @@ def __sim_producer(self, seed, sim_monitor, mutex, error_event):  # pylint: disa
             Event signaling an error occurred during the simulation.
         """
         try:
-            # Ensure Processes generate different random numbers
-            self.environment._set_stochastic(seed)
-            self.rocket._set_stochastic(seed)
-            self.flight._set_stochastic(seed)
-
             while sim_monitor.keep_simulating():
                 sim_idx = sim_monitor.increment() - 1
                 inputs_json, outputs_json = "", ""
 
+                self.__seed_this_simulation(sim_idx)
                 flight = self.__run_single_simulation()
                 inputs_json = self.__evaluate_flight_inputs(sim_idx)
                 outputs_json = self.__evaluate_flight_outputs(flight, sim_idx)
@@ -580,6 +725,47 @@ def __sim_producer(self, seed, sim_monitor, mutex, error_event):  # pylint: disa
             error_event.set()
             mutex.release()
 
+    def __continue_the_root_the_rows_carry(self, random_seed):
+        """Take the root from the rows being appended to, or refuse to.
+
+        Without this an append draws a second root into one file and nothing
+        afterwards can tell which simulation came from which. Reading it back
+        also means a fresh object can continue a study, which is the ordinary
+        way of resuming one.
+        """
+        recorded = _root_a_log_was_written_with(self.input_file)
+        if recorded is None:
+            return
+        if random_seed is None:
+            self.__root_state = (
+                recorded["entropy"],
+                tuple(recorded["spawn_key"]),
+                recorded["pool_size"],
+                recorded["n_children_spawned"],
+            )
+            return
+        if _root_written_into_a_row(self.__root_state) != recorded:
+            raise ValueError(
+                f"cannot append to {self.input_file}: its rows were drawn from "
+                f"a different root than random_seed gives. Continuing would put "
+                f"two studies in one file. Pass the seed the run started with, "
+                f"or leave random_seed out to carry on from the rows."
+            )
+
+    def __seed_this_simulation(self, sim_idx):
+        """Reseed the three models from this index's own child of the root.
+
+        Per index rather than per worker, which is what makes a simulation's
+        inputs the same however the run was split up. The child is split three
+        ways so the environment, rocket and flight draw independently instead
+        of sharing one stream.
+        """
+        child = _seed_of_simulation(self.__root_state, sim_idx)
+        environment, rocket, flight = child.spawn(3)
+        self.environment._set_stochastic(_seed_sequence_to_int(environment))
+        self.rocket._set_stochastic(_seed_sequence_to_int(rocket))
+        self.flight._set_stochastic(_seed_sequence_to_int(flight))
+
     def __run_single_simulation(self):
         """Runs a single simulation and returns the inputs and outputs.
 
@@ -800,6 +986,7 @@ def __evaluate_flight_inputs(self, sim_idx):
             for item in d.items()
         )
         inputs_dict["index"] = sim_idx
+        inputs_dict[_SIMULATION_ROOT_KEY] = _root_written_into_a_row(self.__root_state)
         return (
             json.dumps(inputs_dict, cls=RocketPyEncoder, **self._export_config) + "\n"
         )

rocketpy/stochastic/stochastic_model.py

@@ -8,7 +8,7 @@
 from rocketpy.mathutils.function import Function
 from rocketpy.stochastic.custom_sampler import CustomSampler
 
-from ..tools import get_distribution
+from ..tools import _seed_sequence_to_int, get_distribution
 
 
 def _names_as_spawn_key(input_names):
@@ -41,6 +41,18 @@ def _format_number(value):
     return f"array of shape {np.shape(value)}"
 
 
+def _seed_as_entropy(seed):
+    """A seed as something ``SeedSequence`` will take as entropy.
+
+    A parallel run is handed a ``SeedSequence``, which it will not take. Any
+    other seed goes through untouched, so the stream an int reaches stays where
+    it was.
+    """
+    if not isinstance(seed, np.random.SeedSequence):
+        return seed
+    return _seed_sequence_to_int(seed)
+
+
 def _sampler_seed(seed, input_names):
     """Derive a seed for one sampler, or for one group that shares a generator.
 
@@ -54,10 +66,10 @@ def _sampler_seed(seed, input_names):
     # Sorted here rather than trusting the caller, so a future call site cannot
     # give one group two different seeds by listing its members another way.
     root = np.random.SeedSequence(
-        entropy=seed, spawn_key=_names_as_spawn_key(tuple(sorted(input_names)))
+        entropy=_seed_as_entropy(seed),
+        spawn_key=_names_as_spawn_key(tuple(sorted(input_names))),
     )
-    words = root.generate_state(4, dtype=np.uint32)
-    return sum(int(word) << (32 * position) for position, word in enumerate(words))
+    return _seed_sequence_to_int(root)
 
 
 # TODO: Stop using assert in production code. Use exceptions instead.