Skip to content

Export to ArviZ

Install the optional dependency. The example below uses the result and observations from the quickstart.

pip install "smcx[arviz]"

import jax.random as jr
import smcx

result = smcx.to_arviz(
    posterior,
    key=jr.key(7),
    num_draws=1_000,
    emissions=observations,
)

posterior above can be a ParticleFilterPosterior or TemperedPosterior. Pass a sequence of independent results to represent multiple chains. num_draws controls how many equal-weight draws are resampled from each particle cloud; the key makes that resampling reproducible.

Adaptive tempered runs may require different numbers of stages. Their stage-wise diagnostics are padded with NaN to the longest run, and the Boolean particle_diagnostics.stage_valid mask identifies each run's recorded stages.

Particle-filter export requires the default store_history=True. A final-only result keeps the full ESS and evidence traces but only one particle cloud, so those arrays do not share an ArviZ time dimension. Rerun with full history before exporting.

Groups

Group or attribute Contents
posterior Resampled particle values
particle_diagnostics Source log weights and algorithm diagnostics
observed_data Emissions, when supplied
unconstrained_posterior Aligned u-space values, when supplied
posterior.attrs["marginal_loglik"] Evidence estimate for each run

Particle-filter output has dimensions (chain, draw, time, ...). Each time slice is a filtering marginal \(p(x_t \mid y_{0:t})\); draws with the same index across time do not form a joint trajectory. Use reconstruct_trajectories when ancestry is needed.

The particle_diagnostics group describes the source particle clouds, not the resampled posterior draws. Its leading dimension is run, followed by time for particle filters or stage for tempered SMC; source log weights also have a particle dimension. These values are stored once per run rather than repeated across draw.

Structured particle states use their PyTree paths as variable names. Supply var_names to rename them and dims to label event dimensions. Values passed through unconstrained= follow the same resampling indices as the constrained particles. Raw dotted paths must be unambiguous before aliases are applied, and aliases must resolve to unique names within each group. Variable names cannot shadow sample or event dimensions. Event labels must be unique within a variable and one size when shared in a group. Posterior and unconstrained schemas may differ in rank or extent; particle labels do not affect observed_data.

ArviZ 0.x returns InferenceData; ArviZ 1.x returns DataTree. smcx dispatches to the constructor available in the installed generation: arviz.from_dict for 0.23.4 or arviz_base.from_dict for 1.x. ArviZ is distributed under the Apache License 2.0.