6. Save and resume results#

Download this notebook (.ipynb) · Read the saved outputs here, or run the cells in Jupyter.

A posterior export is useful for plots and summaries. Continuing MCMC also requires sampler state and the same analysis configuration. This chapter extends the Quickstart with both kinds of persistence.

This notebook is self-contained. Install JeansPy with the numpyro_cpu and plotting extras as described in the installation guide, then select that environment as your Jupyter kernel and run cells from top to bottom. Saved outputs are an example run; timings and short-chain results can vary.

Setup#

Run this setup in a fresh kernel. It defines all objects used below.

The generalized NFW halo uses ZhaoModel with fixed alpha=1, beta=3 and an untruncated cutoff. Its inner slope gamma is sampled in the inference examples; truth["gamma"]=1 generates standard-NFW mocks.

import os
os.environ.setdefault("JEANSPY_JAX_PLATFORM", "cpu")
os.environ.setdefault("JEANSPY_JAX_ENABLE_X64", "true")

# The example uses no progress widgets; ignore only their optional-import warning.
import warnings
warnings.filterwarnings("ignore", message="IProgress not found.*")

# Import JeansPy before JAX so its runtime settings take effect.
from jeanspy.model_jax import (
    DSphModel, PlummerModel, ZhaoModel, ConstantAnisotropyModel,
)
from jeanspy.sampler_numpyro import JeansLikelihoodModel, ParameterSpec
import jax
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, init_to_value
import numpy as np
import matplotlib.pyplot as plt

model = DSphModel(submodels={
    "StellarModel": PlummerModel(),
    "DMModel": ZhaoModel(),
    "AnisotropyModel": ConstantAnisotropyModel(),
})
fixed = dict(re_pc=200., alpha=1., beta=3., r_t_pc=np.inf)
truth = dict(**fixed, rs_pc=500., rhos_Msunpc3=.1, gamma=1., beta_ani=0.)
options = dict(solver="kernel", n_u=128, n_kernel=32,
               dm_mass_method="numeric", dm_mass_n_steps=128)

rng = np.random.default_rng(123)
u = rng.uniform(size=32)
R_pc = 200. * np.sqrt(u / (1. - u))  # Projected Plummer radii.
e_vlos_kms = np.full(32, 2.)
variance = np.asarray(model.sigmalos2(R_pc, params=truth, **options))
vlos_kms = rng.normal(0., np.sqrt(variance + e_vlos_kms**2))
data = dict(R_pc=R_pc, vlos_kms=vlos_kms, e_vlos_kms=e_vlos_kms)

def physical_parameters(sampled):
    return {**fixed, **sampled}

likelihood = JeansLikelihoodModel(
    model,
    [ParameterSpec.pow10("log10_rs_pc", dist.Uniform(1.5, 4.), param_name="rs_pc"),
     ParameterSpec.pow10("log10_rhos_Msunpc3", dist.Uniform(-3., 1.),
                         param_name="rhos_Msunpc3"),
     ParameterSpec("gamma", dist.Uniform(0., 2.)),
     ParameterSpec("beta_ani", dist.Uniform(-1., .75)),
     ParameterSpec("vmem_kms", dist.Uniform(-50., 50.))],
    parameter_postprocess=physical_parameters,
    sigmalos2_kwargs=options,
)

kernel = NUTS(likelihood, target_accept_prob=.85,
              init_strategy=init_to_value(values={
                  "log10_rs_pc": np.log10(500.), "log10_rhos_Msunpc3": -1.,
                  "gamma": 1., "beta_ani": 0., "vmem_kms": 0.}))
mcmc = MCMC(kernel, num_warmup=24, num_samples=32,
            num_chains=2, chain_method="sequential", progress_bar=False)
mcmc.run(jax.random.PRNGKey(42), **data)

Export a finished Quickstart run#

The setup below runs the same likelihood with two chains, 24 warmup steps and 32 posterior draws each. These reduced settings exercise storage only. Preserve the chain axes in a NumPy file:

from pathlib import Path

from uuid import uuid4
run_id = uuid4().hex[:8]
output = Path("jeanspy-result-" + run_id)
output.mkdir(exist_ok=False)  # Choose a new directory for a new analysis.
np.savez(output / "posterior.npz",
         **{name: np.asarray(value) for name, value in
            mcmc.get_samples(group_by_chain=True).items()})
np.savez(output / "sample_stats.npz",
         **{name: np.asarray(value) for name, value in
            mcmc.get_extra_fields(group_by_chain=True).items()})
np.savez(output / "observations.npz", **data)
with np.load(output / "posterior.npz") as saved:
    print(saved["beta_ani"].shape)  # (2, 32) for this storage example.

(2, 32)

These files allow read-back of values; they do not contain a restartable NUTS checkpoint. Keep the script defining the model, transforms and priors, the random seed, numerical settings, source revision and environment with the exported data. A collection of marginal intervals alone loses the joint posterior and cannot be used for chain diagnostics.

Enable checkpointing before sampling#

For automatic storage and restart, construct a fresh MCMC object with the same likelihood and wrap it in jeanspy.sampler_numpyro.NumPyroSampler:

from jeanspy.sampler_numpyro import NumPyroSampler

def make_mcmc():
    return MCMC(kernel, num_warmup=24, num_samples=32,
                num_chains=2, chain_method="sequential", progress_bar=False)

with NumPyroSampler(make_mcmc(), output_dir="jeanspy-checkpoint-" + run_id,
                     storage_backend="h5netcdf", async_writes=False) as stored:
    result = stored.run(jax.random.PRNGKey(42), **data, resume=False)
    print(result.checkpoint_path.name)
    tree = stored.load_samples(combine=True)
    posterior = tree["posterior"].ds
    print(posterior["beta_ani"].shape)

last_state.pkl
(2, 32)

Use a new directory for the first run. This launches sampling with storage enabled; it does not import the already completed Quickstart draws. kernel and data are defined in this notebook’s setup. The context manager flushes writes before closing. Setting async_writes=False makes the disk write complete before run returns.

Reconstruct and resume the same analysis#

Recreate the same model, likelihood, kernel and observations, then open the existing output directory:

with NumPyroSampler(make_mcmc(), output_dir="jeanspy-checkpoint-" + run_id,
                     storage_backend="h5netcdf", async_writes=False) as stored:
    result = stored.run(jax.random.PRNGKey(43), **data, resume=True)
    print("Resumed:", result.resumed)
    tree = stored.load_samples(combine=True)
    print(tree["posterior"].ds["beta_ani"].shape)  # (2, 64) after two chunks.

Resumed: True
(2, 64)

The checkpoint restores the continuation state, including its random state, and skips completed warmup. A new seed argument is not a request for an independent chain when resuming. load_samples(combine=True) concatenates chunks along draw, preserving chain. In the locked environment it returns an Xarray DataTree; use tree["posterior"].ds to obtain the dataset.

The wrapper checks data, model/prior configuration, computational source/dependencies and effective JAX backend/precision against the stored analysis identity. A changed analysis needs a new output directory. If an existing analysis lacks metadata.json, restore the original metadata or start separately; do not assign current inputs as a replacement identity for old samples.

Identity format 2 allows documentation-only edits while retaining checks on calculation code, data and numerical settings. Full source/data byte hashes are kept separately in metadata.json under source_provenance. Format-1 chains require their original code/environment to resume; see the API migration guide.

Understand the stored files#

File	Purpose
`metadata.json`	Analysis identity, source provenance and storage configuration
`last_state.pkl`	Sampler continuation state; load only trusted local checkpoints
`chunks/*.nc` with `h5netcdf`	Posterior draws and available sampler statistics
`observations.csv` in the worked example	Input data saved by the example, not by the wrapper automatically

The NumPyroSampler API describes the other storage backends and asynchronous-write options. To run the complete save/reconstruct/resume example:

python examples/docs_quickstart_numpyro.py --output-dir /tmp/jeanspy-numpyro-example

That script deliberately requires a fresh directory, then performs both runs itself. Its recorded results include a check that the first chunk is unchanged after resuming.

Save emcee results#

emcee.backends.HDFBackend stores the chain and random state. Reopen the backend and pass None as the initial state to append. The direct emcee route requires you to preserve the model, priors and data yourself; JeansPy’s jeanspy.sampler.Sampler adds identity checks for its estimation models. See the NumPy/SciPy wrapper example and complete spherical analysis.

Continue with a worked axisymmetric analysis, or use the API dictionary for an individual method’s arguments.