From d6ef06c91902170787a18127fafcdc85efd5e141 Mon Sep 17 00:00:00 2001
From: Virgile Andreani <armavica@ulminfo.fr>
Date: Thu, 9 Nov 2023 20:49:38 -0500
Subject: [PATCH] Make it possible to specify a seed from Python

---
 CHANGELOG.md             | 4 ++++
 python/rebop/__init__.py | 5 ++++-
 src/lib.rs               | 9 +++++++--
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a0acb04..8ffe864 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+* It is now possible to specify a random seed from Python to run the simulator.
+
 ### Changed
 
 * Upgraded PyO3 to v0.21.2.
diff --git a/python/rebop/__init__.py b/python/rebop/__init__.py
index 4b4e068..6d38a3b 100644
--- a/python/rebop/__init__.py
+++ b/python/rebop/__init__.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import xarray as xr
 
 from .rebop import Gillespie, __version__
@@ -12,13 +14,14 @@ def run_xarray(
     init: dict[str, int],
     tmax: float,
     nb_steps: int,
+    seed: int | None = None,
 ) -> xr.Dataset:
     """Run the system until `tmax` with `nb_steps` steps.
 
     The initial configuration is specified in the dictionary `init`.
     Returns an xarray Dataset.
     """
-    times, result = og_run(self, init, tmax, nb_steps)
+    times, result = og_run(self, init, tmax, nb_steps, seed)
     ds = xr.Dataset(
         data_vars={
             name: xr.DataArray(values, dims="time", coords={"time": times})
diff --git a/src/lib.rs b/src/lib.rs
index ea8c42d..d8c2cbd 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -299,12 +299,13 @@ impl Gillespie {
     /// The initial configuration is specified in the dictionary `init`.
     /// Returns `times, vars` where `times` is an array of `nb_steps + 1` uniformly spaced time
     /// points between `0` and `tmax`, and `vars` is a dictionary of species name to array of
-    /// values at the given time points.
+    /// values at the given time points.  One can specify a random `seed` for reproducibility.
     fn run(
         &self,
         init: HashMap<String, usize>,
         tmax: f64,
         nb_steps: usize,
+        seed: Option<u64>,
     ) -> PyResult<(Vec<f64>, HashMap<String, Vec<isize>>)> {
         let mut x0 = vec![0; self.species.len()];
         for (name, &value) in &init {
@@ -312,7 +313,11 @@ impl Gillespie {
                 x0[id] = value as isize;
             }
         }
-        let mut g = gillespie::Gillespie::new(x0);
+        let mut g = match seed {
+            Some(seed) => gillespie::Gillespie::new_with_seed(x0, seed),
+            None => gillespie::Gillespie::new(x0),
+        };
+
         for (rate, reactants, products) in self.reactions.iter() {
             let mut vreactants = vec![0; self.species.len()];
             for reactant in reactants {