NumPyroSampler#

jeanspy.sampler_numpyro.NumPyroSampler

See methods and properties for individual lookup pages, or the alphabetical API dictionary to search all classes. The full class contract and existing member anchors are retained below.

class jeanspy.sampler_numpyro.NumPyroSampler(mcmc, *, output_dir, storage_backend='zarr', async_writes=True, arviz_converter=None)[source]#

Bases: object

Composition-based helper around numpyro.infer.MCMC.

It keeps the wrapped MCMC instance untouched while adding:

  • automatic reuse of post_warmup_state for repeated runs,

  • checkpoint save/load of last_state,

  • chunked ArviZ 1.0 persistence using backend-backed DataTree stores,

  • optional background writes for heavy output serialization.

Notes

Inputs and units. mcmc is numpyro.infer.MCMC; output_dir is a filesystem path; storage_backend chooses zarr/netcdf4/h5netcdf; async_writes toggles the writer; arviz_converter optionally returns xarray.DataTree. run(rng_key,*args,**kwargs) forwards data to the model, with explicit resume/save/write flags.

Returns and shape. run returns SamplerRunResult. Samples use ArviZ chain/draw dimensions; load_samples(combine=True) returns a combined DataTree, False a list. save_samples_chunk returns (index,path,submitted). save_checkpoint/load_checkpoint return paths; flush/close return None after awaiting writes.

Validity. Use a context manager or close/flush before relying on completed files. resume=’auto’ checks available state; resume=True requires valid state. Metadata binds data, prior, model/source/dependencies and effective runtime configuration. A checkpoint contains trusted Python serialization; open only your trusted analysis output.

Errors. Missing/inconsistent analysis metadata or checkpoint identity raises before restarting. Write failures propagate from flush/close. Use a new directory for a changed analysis.

Backend. NumPyro/JAX inference, host-side ArviZ storage.

Differentiation. The model may be differentiated; sampling control, disk I/O and checkpoint state are not differentiable.

Examples. examples/docs_numpyro_inference.py

Parameters:

  • mcmc (MCMC)

  • output_dir (str | Path)

  • storage_backend (StorageBackend | str)

  • async_writes (bool)

  • arviz_converter (Callable[..., xr.DataTree] | None)

property checkpoint_path: pathlib.Path#

Return the Path for the trusted, pickled NumPyro transition-state checkpoint.

property metadata_path: pathlib.Path#

Return the Path for persisted format and sampling-identity metadata.

property chunks_dir: pathlib.Path#

Return the Path containing persisted sample chunks.

property chunk_suffix: str#

Return the filename suffix for the selected sample-storage backend.

property uses_directory_stores: bool#

Whether the selected backend stores each sample chunk as a directory.

list_chunk_paths()[source]#

Find sample chunks in index order.

Notes

Inputs and units. No arguments.

Returns and shape. List of Paths; an empty list means no saved chunks were found.

Validity. Lists recognized chunk names in this directory; it does not inspect their scientific content or verify identity.

Errors. Filesystem access errors may propagate; an empty directory returns an empty list.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Return type:

list[pathlib.Path]

pending_write_count()[source]#

Count unfinished writes, retaining failures for flush/close.

Return type:

int

flush()[source]#

Wait for outstanding sample writes.

Notes

Inputs and units. No arguments.

Returns and shape. None; asynchronous exceptions propagate.

Validity. Waits for outstanding writes and propagates their failures; it does not change the MCMC target or validate a checkpoint.

Errors. A pending write’s exception is re-raised after waiting for the submitted futures.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Return type:

None

close()[source]#

Finish writes and shut down the background writer.

Notes

Inputs and units. No arguments; also called by context-manager exit.

Returns and shape. None; failures propagate.

Validity. Finish pending writes before releasing the executor. Prefer a context manager to ensure this happens.

Errors. Write failures propagate after the executor is shut down.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Return type:

None

clear_resume_state()[source]#

Clear the in-memory warmup-resume pointer.

Notes

Inputs and units. No arguments.

Returns and shape. None; sets mcmc.post_warmup_state = None. It does not delete or reset persisted analysis files.

Validity. Affects only the in-memory post_warmup_state pointer. A later auto-resume can still use last_state or an on-disk checkpoint.

Errors. No additional validation or documented domain exception.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Return type:

None

save_checkpoint()[source]#

Save trusted NumPyro transition state.

Notes

Inputs and units. No arguments; requires a last_state and bound analysis identity.

Returns and shape. Path of the written checkpoint.

Validity. A completed MCMC last_state and an unchanged bound analysis are required. The file contains Python pickle data.

Errors. RuntimeError before last_state exists; ValueError for unverified/changed analysis. Filesystem errors propagate.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Return type:

pathlib.Path

load_checkpoint()[source]#

Load a trusted matching transition state.

Notes

Inputs and units. No arguments; reads this output directory’s checkpoint and metadata.

Returns and shape. Checkpoint Path; sets post_warmup_state only after identity checks.

Validity. Read only a trusted checkpoint. Stored target/metadata identity is checked here; supplied observation arguments are checked later by run.

Errors. FileNotFoundError for a missing checkpoint; ValueError for format/identity mismatch. Pickle and filesystem errors propagate.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Return type:

pathlib.Path

to_datatree(**arviz_kwargs)[source]#

Convert the wrapped MCMC result to ArviZ storage.

Notes

Inputs and units. Keyword arguments forwarded to arviz_converter.

Returns and shape. xarray.DataTree with chain/draw dimensions; the converter must return this format.

Validity. The wrapped MCMC must contain a result supported by the chosen converter.

Errors. TypeError if the converter returns anything other than xarray.DataTree; converter exceptions propagate.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Parameters:

arviz_kwargs (Any)

Return type:

xarray.core.datatree.DataTree

save_samples_chunk(*, datatree=None, wait=False, **arviz_kwargs)[source]#

Write an ArviZ sample chunk.

Notes

Inputs and units. Optional datatree; wait=True writes synchronously; other keywords go to the converter.

Returns and shape. Tuple (chunk_index, path, write_submitted). False in the third entry means the write was synchronous.

Validity. The sampler assigns a fresh chunk index. Wait for flush/close before relying on an asynchronous write.

Errors. FileExistsError prevents replacing a reserved chunk. Conversion, array loading and write failures propagate; asynchronous failures are raised by flush/close.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Parameters:

  • datatree (xarray.core.datatree.DataTree | None)

  • wait (bool)

  • arviz_kwargs (Any)

Return type:

tuple[int, pathlib.Path, bool]

load_samples(*, combine=True)[source]#

Read persisted sample chunks.

Notes

Inputs and units. combine=True concatenates matching draw dimensions; False retains separate trees.

Returns and shape. Combined xarray.DataTree or ordered list of DataTrees. Pending writes are flushed first.

Validity. Reads this directory’s chunks. Completed arrays are loaded into memory before closing backing stores.

Errors. FileNotFoundError when no chunks exist. Pending-write, reader and concatenation errors propagate.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Parameters:

combine (bool)

Return type:

xarray.core.datatree.DataTree | list[xarray.core.datatree.DataTree]

static combine_trees(trees)[source]#

Combine compatible ArviZ chunks.

Notes

Inputs and units. Nonempty sequence of DataTrees from the same analysis, with consistent groups, variables and nondraw coordinates.

Returns and shape. DataTree concatenated over draws, with draw coordinates renumbered from zero. No independent sampling-identity verification occurs here.

Validity. Supply chunks from the same analysis. Static groups must be equal. The draw-group concatenation uses xarray compat=override and is not a general identity validator.

Errors. ValueError for an empty list, missing groups, inconsistent draw-axis presence or changed static groups. Other xarray errors propagate.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Parameters:

trees (Sequence[xarray.core.datatree.DataTree])

Return type:

xarray.core.datatree.DataTree

run(rng_key, *args, resume='auto', save_checkpoint=True, save_samples=True, wait_for_write=False, arviz_kwargs=None, **kwargs)[source]#

Run one chunk and optionally persist it.

Notes

Inputs and units. rng_key plus model data arguments; resume is auto/True/False. save_checkpoint/save_samples/wait_for_write control persistence; arviz_kwargs are conversion options.

Returns and shape. SamplerRunResult; wait for flush/close before relying on asynchronous files.

Validity. Use an unchanged model, prior, observations and numerical configuration when resuming. Identity is checked before sampling.

Errors. ValueError for identity mismatch or invalid resume mode; FileNotFoundError for required missing state. Model, NumPyro and storage exceptions propagate.

Differentiation. Host-side sampling/storage control; this method has no physical-parameter derivative.

Parameters:
Return type:

jeanspy.sampler_numpyro.SamplerRunResult