feat: generalize sample and stats storage

Author: aseyboldt

Cargo.toml

@@ -5,7 +5,7 @@ authors = [
   "Adrian Seyboldt <[email protected]>",
   "PyMC Developers <[email protected]>",
 ]
-edition = "2021"
+edition = "2024"
 license = "MIT"
 repository = "https://github.com/pymc-devs/nuts-rs"
 keywords = ["statistics", "bayes"]
@@ -22,22 +22,39 @@ rand = { version = "0.9.0", features = ["small_rng"] }
 rand_distr = "0.5.0"
 itertools = "0.14.0"
 thiserror = "2.0.3"
-arrow = { version = "56.1.0", default-features = false, features = ["ffi"] }
 rand_chacha = "0.9.0"
 anyhow = "1.0.72"
 faer = { version = "0.22.6", default-features = false, features = ["linalg"] }
 pulp = "0.21.4"
 rayon = "1.10.0"
+zarrs = { version = "0.21.0", features = [
+  "filesystem",
+  "gzip",
+  "sharding",
+], optional = true }
+ndarray = { version = "0.16.1", optional = true }
+nuts-derive = { path = "./nuts-derive" }
+nuts-storable = { path = "./nuts-storable" }
+serde = { version = "1.0.219", features = ["derive"] }
+serde_json = "1.0"
 
 [dev-dependencies]
 proptest = "1.6.0"
 pretty_assertions = "1.4.0"
 criterion = "0.7.0"
 nix = { version = "0.30.0", features = ["sched"] }
 approx = "0.5.1"
-ndarray = "0.16.1"
 equator = "0.4.2"
+serde_json = "1.0"
+ndarray = "0.16.1"
+tempfile = "3.0"
+
+[features]
+zarr = ["dep:zarrs"]
+ndarray = ["dep:ndarray"]
 
 [[bench]]
 name = "sample"
 harness = false
+
+[workspace]

benches/sample.rs

@@ -1,9 +1,10 @@
 use std::hint::black_box;
 
-use criterion::{criterion_group, criterion_main, Criterion};
-use nix::sched::{sched_setaffinity, CpuSet};
+use criterion::{Criterion, criterion_group, criterion_main};
+use nix::sched::{CpuSet, sched_setaffinity};
 use nix::unistd::Pid;
 use nuts_rs::{Chain, CpuLogpFunc, CpuMath, LogpError, Math, Settings};
+use nuts_storable::HasDims;
 use rand::SeedableRng;
 use rayon::ThreadPoolBuilder;
 use thiserror::Error;
@@ -22,11 +23,20 @@ impl LogpError for PosteriorLogpError {
     }
 }
 
+impl HasDims for PosteriorDensity {
+    fn dim_sizes(&self) -> std::collections::HashMap<String, u64> {
+        vec![("unconstrained_parameter".to_string(), self.dim() as u64)]
+            .into_iter()
+            .collect()
+    }
+}
+
 impl CpuLogpFunc for PosteriorDensity {
     type LogpError = PosteriorLogpError;
+    type ExpandedVector = Vec<f64>;
 
     // Only used for transforming adaptation.
-    type TransformParams = ();
+    type FlowParameters = ();
 
     // We define a 10 dimensional normal distribution
     fn dim(&self) -> usize {
@@ -48,6 +58,17 @@ impl CpuLogpFunc for PosteriorDensity {
             .sum();
         return Ok(logp);
     }
+
+    fn expand_vector<R>(
+        &mut self,
+        _rng: &mut R,
+        array: &[f64],
+    ) -> Result<Self::ExpandedVector, nuts_rs::CpuMathError>
+    where
+        R: rand::Rng + ?Sized,
+    {
+        Ok(array.to_vec())
+    }
 }
 
 fn make_sampler(dim: usize) -> impl Chain<CpuMath<PosteriorDensity>> {

examples/adam_adaptation.rs

@@ -7,6 +7,7 @@ use nuts_rs::{
     AdamOptions, Chain, CpuLogpFunc, CpuMath, DiagGradNutsSettings, LogpError, Settings,
     StepSizeAdaptMethod,
 };
+use nuts_storable::HasDims;
 use thiserror::Error;
 
 // Define a function that computes the unnormalized posterior density
@@ -23,11 +24,20 @@ impl LogpError for PosteriorLogpError {
     }
 }
 
+impl HasDims for PosteriorDensity {
+    fn dim_sizes(&self) -> std::collections::HashMap<String, u64> {
+        vec![("unconstrained_parameter".to_string(), self.dim() as u64)]
+            .into_iter()
+            .collect()
+    }
+}
+
 impl CpuLogpFunc for PosteriorDensity {
     type LogpError = PosteriorLogpError;
+    type ExpandedVector = Vec<f64>;
 
     // Only used for transforming adaptation.
-    type TransformParams = ();
+    type FlowParameters = ();
 
     // We define a 10 dimensional normal distribution
     fn dim(&self) -> usize {
@@ -49,6 +59,17 @@ impl CpuLogpFunc for PosteriorDensity {
             .sum();
         return Ok(logp);
     }
+
+    fn expand_vector<R>(
+        &mut self,
+        _rng: &mut R,
+        array: &[f64],
+    ) -> Result<Self::ExpandedVector, nuts_rs::CpuMathError>
+    where
+        R: rand::Rng + ?Sized,
+    {
+        Ok(array.to_vec())
+    }
 }
 
 fn main() {