feat: add zarr_store argument to write trace while sampling

Author: aseyboldt

Cargo.lock

@@ -21,7 +21,7 @@ name = "_lib"
 crate-type = ["cdylib"]
 
 [dependencies]
-nuts-rs = { version = "0.16.1", features = ["zarr", "arrow"] }
+nuts-rs = { version = "0.17.0", features = ["zarr", "arrow"] }
 numpy = "0.26.0"
 rand = "0.9.0"
 thiserror = "2.0.3"

Cargo.toml

@@ -25,7 +25,7 @@ trace = nutpie.sample(
     tune=500,            # Number of warmup draws for adaptation
     chains=6,            # Number of independent chains
     cores=None,          # Number chains that are allowed to run simultainiously
-    seed=12345          # Random seed for reproducibility
+    seed=12345           # Random seed for reproducibility
 )
 ```
 
@@ -143,6 +143,60 @@ trace = nutpie.sample(
 )
 ```
 
+## Zarr Storage (Experimental)
+
+Nutpie includes experimental support for writing traces directly to zarr storage, which can be useful for large traces that don't fit in memory or for distributed storage scenarios. The zarr format provides efficient, chunked, compressed storage for multi-dimensional arrays.
+
+### Basic Usage
+
+You can write traces directly to zarr storage by providing a `zarr_store` parameter:
+
+```python
+import nutpie
+import pymc as pm
+
+with pm.Model() as model:
+    pm.HalfNormal("a")
+
+compiled = nutpie.compile_pymc_model(model, backend="numba")
+
+# Create a local zarr store
+path = "trace.zarr"
+store = nutpie.zarr_store.LocalStore(path)
+
+trace = nutpie.sample(
+    compiled,
+    chains=2,
+    seed=123,
+    draws=100,
+    tune=100,
+    zarr_store=store
+)
+```
+
+### Memory Considerations
+
+When using zarr storage, the trace object supports lazy loading:
+
+```python
+# The trace is not loaded into memory by default
+posterior_data = trace.posterior.a  # Lazy access
+
+# Explicitly load the entire trace into memory (optional)
+loaded_trace = trace.load()
+posterior_data = loaded_trace.posterior.a  # In-memory access
+```
+
+### Available Store Types
+
+Nutpie supports several zarr store backends:
+
+- `nutpie.zarr_store.LocalStore(path)` - Local filesystem storage
+- `nutpie.zarr_store.S3Store(...)` - Amazon S3 storage
+- `nutpie.zarr_store.GCSStore(...)` - Google Cloud Storage
+- `nutpie.zarr_store.AzureStore(...)` - Azure Blob Storage
+- `nutpie.zarr_store.HTTPStore(...)` - HTTP-based storage
+
 ## Progress Monitoring
 
 Customize the sampling progress display:

docs/sampling-options.qmd

@@ -21,6 +21,8 @@ dependencies = [
   "pandas >= 2.0",
   "xarray >= 2025.01.2",
   "arviz >= 0.20.0",
+  "obstore >= 0.8.0",
+  "zarr >= 3.1.0",
 ]
 dynamic = ["version"]
 

pyproject.toml

@@ -2,6 +2,13 @@
 from nutpie.compile_pymc import compile_pymc_model
 from nutpie.compile_stan import compile_stan_model
 from nutpie.sample import sample
+from nutpie._lib import store as zarr_store
 
 __version__: str = _lib.__version__
-__all__ = ["__version__", "compile_pymc_model", "compile_stan_model", "sample"]
+__all__ = [
+    "__version__",
+    "compile_pymc_model",
+    "compile_stan_model",
+    "sample",
+    "zarr_store",
+]

python/nutpie/__init__.py

@@ -4,7 +4,6 @@
 from pathlib import Path
 from typing import Any, Optional
 
-import pandas as pd
 from numpy.typing import NDArray
 
 from nutpie import _lib

python/nutpie/compile_stan.py

@@ -778,10 +778,12 @@ def sample(
     transform_adapt: bool, default=False
         Use the experimental transform adaptation algorithm
         during tuning.
-    zarr_store: nutpie.store.Store
-        A store created using nutpie.store to store the samples
+    zarr_store: nutpie.zarr_store.*
+        A store created using nutpie.zarr_store to store the samples
         in. If None (default), the samples will be stored in
-        memory using an arrow table.
+        memory using an arrow table. This can be used to write
+        the trace directly into a zarr store, for instance
+        on disk or to S3 or GCS.
     **kwargs
         Pass additional arguments to nutpie._lib.PySamplerArgs
 

python/nutpie/sample.py

@@ -445,3 +445,28 @@ def test_deterministic_sampling_jax():
     compiled = nutpie.compile_pymc_model(model, backend="jax", gradient_backend="jax")
     trace = nutpie.sample(compiled, chains=2, seed=123, draws=100, tune=100)
     return trace.posterior.a.values.ravel()
+
+
+@pytest.mark.pymc
+def test_zarr_store(tmp_path):
+    with pm.Model() as model:
+        pm.HalfNormal("a")
+
+    compiled = nutpie.compile_pymc_model(model, backend="numba")
+
+    path = tmp_path / "trace.zarr"
+    path.mkdir()
+    store = nutpie.zarr_store.LocalStore(str(path))
+    trace = nutpie.sample(
+        compiled, chains=2, seed=123, draws=100, tune=100, zarr_store=store
+    )
+    trace.load().posterior.a  # noqa: B018
+
+
+@pytest.fixture
+def tmp_path():
+    import tempfile
+    from pathlib import Path
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        yield Path(tmpdirname)