feat: add rng to Model.math()

Author: aseyboldt

examples/csv_trace.rs

@@ -179,7 +179,7 @@ impl Model for MvnModel {
     where
         Self: 'model;
 
-    fn math(&self) -> Result<Self::Math<'_>> {
+    fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
         Ok(self.math.clone())
     }
 

examples/hashmap_storage.rs

@@ -27,6 +27,7 @@ impl MultivariateNormal {
 }
 
 // Custom LogpError implementation
+#[allow(dead_code)]
 #[derive(Debug, Error)]
 enum MyLogpError {
     #[error("Recoverable error in logp calculation: {0}")]
@@ -114,7 +115,7 @@ impl Model for MvnModel {
     where
         Self: 'model;
 
-    fn math(&self) -> Result<Self::Math<'_>> {
+    fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
         Ok(self.math.clone())
     }
 
@@ -228,6 +229,9 @@ fn main() -> Result<()> {
                             nuts_rs::HashMapValue::U64(vec) => {
                                 println!("    {}: {} u64 draws", name, vec.len());
                             }
+                            nuts_rs::HashMapValue::String(vec) => {
+                                println!("    {}: {} string draws", name, vec.len());
+                            }
                         }
                     }
 

examples/ndarray_storage.rs

@@ -116,7 +116,7 @@ impl Model for MvnModel {
     where
         Self: 'model;
 
-    fn math(&self) -> Result<Self::Math<'_>> {
+    fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
         Ok(self.math.clone())
     }
 

examples/zarr_async_trace.rs

@@ -186,7 +186,7 @@ impl Model for MvnModel {
     where
         Self: 'model;
 
-    fn math(&self) -> Result<Self::Math<'_>> {
+    fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
         Ok(self.math.clone())
     }
 

examples/zarr_trace.rs

@@ -185,7 +185,7 @@ impl Model for MvnModel {
     where
         Self: 'model;
 
-    fn math(&self) -> Result<Self::Math<'_>> {
+    fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
         Ok(self.math.clone())
     }
 

src/csv_storage.rs

@@ -712,7 +712,7 @@ mod tests {
         where
             Self: 'model;
 
-        fn math(&self) -> Result<Self::Math<'_>> {
+        fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
             Ok(self.math.clone())
         }
 
@@ -790,7 +790,7 @@ mod tests {
         where
             Self: 'model;
 
-        fn math(&self) -> Result<Self::Math<'_>> {
+        fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
             Ok(self.math.clone())
         }
 

src/model.rs

@@ -27,7 +27,7 @@ pub trait Model: Send + Sync + 'static {
         Self: 'model;
 
     /// Returns the math backend for this model.
-    fn math(&self) -> Result<Self::Math<'_>>;
+    fn math<R: Rng + ?Sized>(&self, rng: &mut R) -> Result<Self::Math<'_>>;
 
     /// Initializes the starting position for MCMC sampling.
     ///

src/sampler.rs

@@ -563,7 +563,7 @@ impl<T: TraceStorage> ChainProcess<T> {
         let (stop_marker_tx, stop_marker_rx) = channel();
 
         let mut rng = ChaCha8Rng::seed_from_u64(seed);
-        rng.set_stream(chain_id);
+        rng.set_stream(chain_id + 1);
 
         let chain_trace = Arc::new(Mutex::new(Some(chain_trace)));
         let progress = Arc::new(Mutex::new(ChainProgress::new(
@@ -578,7 +578,9 @@ impl<T: TraceStorage> ChainProcess<T> {
             let progress = progress_inner;
 
             let mut sample = move || {
-                let logp = model.math().context("Failed to create model density")?;
+                let logp = model
+                    .math(&mut rng)
+                    .context("Failed to create model density")?;
                 let dim = logp.dim();
 
                 let mut sampler = settings.new_chain(chain_id, logp, &mut rng);
@@ -660,7 +662,7 @@ impl<T: TraceStorage> ChainProcess<T> {
 
             let result = sample();
 
-            // We intentially ignore errors here, because this means some other
+            // We intentionally ignore errors here, because this means some other
             // chain already failed, and should have reported the error.
             let _ = results.send(result);
             drop(results);
@@ -749,7 +751,12 @@ impl<F: Send + 'static> Sampler<F> {
                 let results = results_tx;
                 let mut chains = Vec::with_capacity(settings.num_chains());
 
-                let math = model_ref.math().context("Could not create model density")?;
+                let mut rng = ChaCha8Rng::seed_from_u64(settings.seed());
+                rng.set_stream(0);
+
+                let math = model_ref
+                    .math(&mut rng)
+                    .context("Could not create model density")?;
                 let trace = trace_config
                     .new_trace(settings_ref, &math)
                     .context("Could not create trace object")?;
@@ -962,6 +969,7 @@ pub mod test_logps {
     };
     use anyhow::Result;
     use nuts_storable::HasDims;
+    use rand::Rng;
     use thiserror::Error;
 
     #[derive(Clone, Debug)]
@@ -1103,7 +1111,7 @@ pub mod test_logps {
     {
         type Math<'model> = CpuMath<&'model F>;
 
-        fn math(&self) -> Result<Self::Math<'_>> {
+        fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> Result<Self::Math<'_>> {
             Ok(CpuMath::new(&self.logp))
         }
 

tests/sample_normal.rs

@@ -95,7 +95,7 @@ impl Model for NormalModel {
     where
         Self: 'model;
 
-    fn math(&self) -> anyhow::Result<Self::Math<'_>> {
+    fn math<R: Rng + ?Sized>(&self, _rng: &mut R) -> anyhow::Result<Self::Math<'_>> {
         Ok(CpuMath::new(NormalLogp {
             dim: self.mu.len(),
             mu: &self.mu,