feat: fix and expose sparsity

In `pub fn add_reaction` (gillespie.rs:190), the rate was added sparsely but the jump was not.

I added a `sparse: bool` to the Gillespie struct, and add every reaction (rate and jump) as sparse (or not) according to its value.

In the Python code, I exposed it as a keyword-only parameter in the run method, with a default value of False.

python/rebop/__init__.py

@@ -9,19 +9,21 @@
 og_run = Gillespie.run
 
 
-def run_xarray(
+def run_xarray(  # noqa: PLR0913 too many parameters in function definition
     self: Gillespie,
     init: dict[str, int],
     tmax: float,
     nb_steps: int,
     seed: int | None = None,
+    *,
+    sparse: bool = False,
 ) -> 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, seed)
+    times, result = og_run(self, init, tmax, nb_steps, seed, sparse=sparse)
     ds = xr.Dataset(
         data_vars={
             name: xr.DataArray(values, dims="time", coords={"time": times})

python/rebop/rebop.pyi

@@ -30,6 +30,8 @@ class Gillespie:
         tmax: float,
         nb_steps: int,
         seed: int | None = None,
+        *,
+        sparse: bool = False,
     ) -> xarray.Dataset:
         """Run the system until `tmax` with `nb_steps` steps.
 

src/gillespie.rs

@@ -129,25 +129,28 @@ pub struct Gillespie {
     t: f64,
     reactions: Vec<(Rate, Jump)>,
     rng: SmallRng,
+    sparse: bool,
 }
 
 impl Gillespie {
     /// Creates a new problem instance, with `N` different species of
     /// specified initial conditions.
-    pub fn new<V: AsRef<[isize]>>(species: V) -> Self {
+    pub fn new<V: AsRef<[isize]>>(species: V, sparse: bool) -> Self {
         Gillespie {
             species: species.as_ref().to_vec(),
             t: 0.,
             reactions: Vec::new(),
             rng: SmallRng::from_entropy(),
+            sparse: sparse,
         }
     }
-    pub fn new_with_seed<V: AsRef<[isize]>>(species: V, seed: u64) -> Self {
+    pub fn new_with_seed<V: AsRef<[isize]>>(species: V, sparse: bool, seed: u64) -> Self {
         Gillespie {
             species: species.as_ref().to_vec(),
             t: 0.,
             reactions: Vec::new(),
             rng: SmallRng::seed_from_u64(seed),
+            sparse: sparse,
         }
     }
     /// Seeds the random number generator.
@@ -158,7 +161,7 @@ impl Gillespie {
     ///
     /// ```
     /// use rebop::gillespie::Gillespie;
-    /// let mut p: Gillespie = Gillespie::new([0, 1, 10, 100]);
+    /// let mut p: Gillespie = Gillespie::new([0, 1, 10, 100], false);
     /// assert_eq!(p.nb_species(), 4);
     /// ```
     pub fn nb_species(&self) -> usize {
@@ -168,7 +171,7 @@ impl Gillespie {
     ///
     /// ```
     /// use rebop::gillespie::Gillespie;
-    /// let mut p: Gillespie = Gillespie::new([0, 1, 10, 100]);
+    /// let mut p: Gillespie = Gillespie::new([0, 1, 10, 100], false);
     /// assert_eq!(p.nb_reactions(), 0);
     /// ```
     pub fn nb_reactions(&self) -> usize {
@@ -180,7 +183,7 @@ impl Gillespie {
     /// describing the state change as a result of the reaction.
     /// ```
     /// use rebop::gillespie::{Gillespie, Rate};
-    /// let mut sir = Gillespie::new([9999, 1, 0]);
+    /// let mut sir = Gillespie::new([9999, 1, 0], false);
     /// //                           [   S, I, R]
     /// // S + I -> I + I with rate 1e-5
     /// sir.add_reaction(Rate::lma(1e-5, [1, 1, 0]), [-1, 1, 0]);
@@ -191,7 +194,11 @@ impl Gillespie {
         // This assert ensures that the jump does not go out of bounds of the species
         assert_eq!(differences.as_ref().len(), self.species.len());
         let jump = Jump::new(differences);
-        self.reactions.push((rate.sparse(), jump));
+        if self.sparse {
+            self.reactions.push((rate.sparse(), jump.sparse()));
+        } else {
+            self.reactions.push((rate, jump));
+        }
     }
     /// Returns the current time in the model.
     pub fn get_time(&self) -> f64 {
@@ -205,7 +212,7 @@ impl Gillespie {
     ///
     /// ```
     /// use rebop::gillespie::Gillespie;
-    /// let p: Gillespie = Gillespie::new([0, 1, 10, 100]);
+    /// let p: Gillespie = Gillespie::new([0, 1, 10, 100], false);
     /// assert_eq!(p.get_species(2), 10);
     /// ```
     pub fn get_species(&self, s: usize) -> isize {
@@ -250,7 +257,7 @@ impl Gillespie {
     ///
     /// ```
     /// use rebop::gillespie::{Gillespie, Rate};
-    /// let mut dimers = Gillespie::new([1, 0, 0, 0]);
+    /// let mut dimers = Gillespie::new([1, 0, 0, 0], false);
     /// //                              [G, M, P, D]
     /// dimers.add_reaction(Rate::lma(25., [1, 0, 0, 0]), [0, 1, 0, 0]);
     /// dimers.add_reaction(Rate::lma(1000., [0, 1, 0, 0]), [0, 0, 1, 0]);
@@ -365,7 +372,7 @@ mod tests {
     use crate::gillespie::{Gillespie, Rate};
     #[test]
     fn sir() {
-        let mut sir = Gillespie::new([9999, 1, 0]);
+        let mut sir = Gillespie::new([9999, 1, 0], false);
         sir.add_reaction(Rate::lma(0.1 / 10000., [1, 1, 0]), [-1, 1, 0]);
         sir.add_reaction(Rate::lma(0.01, [0, 1, 0]), [0, -1, 1]);
         sir.advance_until(250.);
@@ -376,7 +383,7 @@ mod tests {
     }
     #[test]
     fn dimers() {
-        let mut dimers = Gillespie::new([1, 0, 0, 0]);
+        let mut dimers = Gillespie::new([1, 0, 0, 0], false);
         dimers.add_reaction(Rate::lma(25., [1, 0, 0, 0]), [0, 1, 0, 0]);
         dimers.add_reaction(Rate::lma(1000., [0, 1, 0, 0]), [0, 0, 1, 0]);
         dimers.add_reaction(Rate::lma(0.001, [0, 0, 2, 0]), [0, 0, -2, 1]);

src/lib.rs

@@ -119,7 +119,7 @@
 //! ```rust
 //! use rebop::gillespie::{Gillespie, Rate};
 //!
-//! let mut sir = Gillespie::new([999, 1, 0]);
+//! let mut sir = Gillespie::new([999, 1, 0], false);
 //! //                           [  S, I, R]
 //! // S + I => 2 I with rate 1e-4
 //! sir.add_reaction(Rate::lma(1e-4, [1, 1, 0]), [-1, 1, 0]);
@@ -303,13 +303,14 @@ impl Gillespie {
     /// values at the given time points.  One can specify a random `seed` for reproducibility.
     /// If `nb_steps` is `0`, then returns all reactions, ending with the first that happens at
     /// or after `tmax`.
-    #[pyo3(signature = (init, tmax, nb_steps, seed=None))]
+    #[pyo3(signature = (init, tmax, nb_steps, seed=None, sparse=false))]
     fn run(
         &self,
         init: HashMap<String, usize>,
         tmax: f64,
         nb_steps: usize,
         seed: Option<u64>,
+        sparse: bool,
     ) -> PyResult<(Vec<f64>, HashMap<String, Vec<isize>>)> {
         let mut x0 = vec![0; self.species.len()];
         for (name, &value) in &init {
@@ -318,8 +319,8 @@ impl Gillespie {
             }
         }
         let mut g = match seed {
-            Some(seed) => gillespie::Gillespie::new_with_seed(x0, seed),
-            None => gillespie::Gillespie::new(x0),
+            Some(seed) => gillespie::Gillespie::new_with_seed(x0, sparse, seed),
+            None => gillespie::Gillespie::new(x0, sparse),
         };
 
         for (rate, reactants, products) in self.reactions.iter() {

tests/test_rebop.py

@@ -50,3 +50,12 @@ def test_all_reactions(seed: int) -> None:
     assert set(dds.S.to_numpy()) <= {-1, 0}
     assert set(dds.I.to_numpy()) <= {-1, 1}
     assert set(dds.R.to_numpy()) <= {0, 1}
+
+
+def test_dense_vs_sparse() -> None:
+    sir = sir_model()
+    init = {"S": 999, "I": 1}
+    kwargs = {"tmax": 250, "nb_steps": 250, "seed": 42}
+    ds_dense = sir.run(init, **kwargs, sparse=False)
+    ds_sparse = sir.run(init, **kwargs, sparse=True)
+    assert (ds_dense == ds_sparse).all()