convert.to_datatree()

Convert an already-drawn posterior into an ArviZ-schema xarray.DataTree.

Usage

convert.to_datatree(
    rng_key,
    model,
    posterior,
    data,
    covariates,
    *,
    num_chains=1,
    batch_size=None,
    device="host",
    coords=None,
    time_coord=None,
    posterior_dims=None,
    covariate_dims=None,
    obs_dims=None
)

Posterior-first: callers draw their own posterior (mcmc.get_samples() for MCMC, draw_posterior() for a variational fit) and pass it in; to_datatree never draws a posterior of its own. rng_key is consumed only by the in-sample posterior-predictive draws and, when a forecast horizon is present, the forecast draws.

Parameters

rng_key: Array: PRNG key for the in-sample predictive draws and, when a horizon is present, the forecast draws.
model: ForecastModel: The forecasting model that produced posterior.
posterior: Mapping[str, ArrayLike]: Posterior samples of the latent sites of the training run, with a single flattened sample axis leading (NumPyro’s mcmc.get_samples() order, or the output of draw_posterior()). Host-committed leaves (the output of draw_posterior(..., device="host")) and NumPy leaves are accepted directly. Forecast-horizon keys follow the rule of forecast(): one with values (a key that ends in "_future" or contains "_future_", or "forecast") raises ValueError. One with an empty leaf (a site that the model registers on every call, recorded by the training run with a size-0 time axis) is kept in the posterior group as a size-0 variable, and the predictive drivers drop it. An "obs" key is dropped, as in predict_in_sample(): it is not in the posterior group, and the in-sample draws come from the latent sites. Every other key becomes a posterior variable, and xarray.DataTree rejects a variable name that contains "/", so a key that contains "/" raises ValueError before any draw. numpyro.handlers.scope joins its prefix and the site name with "/" by default (trend/drift): pass divider="_" to every scope call of a model you export (trend_drift). The divider changes only the names: a fit drawn under "/" exports without a refit once the model uses divider="_" and the posterior keys are renamed the same way. Both must change: the drivers substitute posterior sites by name, and forecast() and predict_in_sample() ignore a key that matches no site and draw that site from the prior with no error. to_datatree traces the in-sample model call of predict_in_sample() abstractly, with the posterior substituted as the driver does (no draw, no compile), and raises ValueError before any draw when such a key names a latent sample site that has no key of its own: the key equals the site name with "/" replaced by "_" (a half-applied rename, trend_drift with a model that still has trend/drift), or does so without a trailing "_haar" or "_dct" (the posterior of a time_reparam() fit that holds drift_haar but no drift, passed with the plain model) or without an inner "_decentered" (a fit of a site reparameterized with LocScaleReparam, such as drift_decentered or drift_decentered_haar, passed with a model without that reparam). Any other key that matches no site, such as _auto_latent from Predictive over an AutoContinuous guide or a derived quantity you add, is stored in the posterior group and changes no draw.
data: Array: In-sample data shaped (*batch, time, obs): time at axis -2, the observation dim at axis -1 and any batch axes to the left (obs_dims names them).
covariates: Array: Covariates with time at axis -2. When covariates extends beyond data along the time axis (the package-wide shape convention for a forecast horizon), the trailing rows are treated as future covariates: the returned tree additionally carries predictions (forecast obs draws from forecast()) and predictions_constant_data groups.
num_chains: int = 1: Number of chains to split posterior’s flattened sample axis into (and, identically, the in-sample/forecast predictive draws, which are drawn with the same sample count). Defaults to 1 (a single pseudo-chain, correct for a posterior with no chain structure, e.g. SVI or Pathfinder draws). For an MCMC posterior, pass the num_chains the sampler was run with; see _reshape_chains() for the reshape contract and its divisibility requirement. Must be a positive integer.
batch_size: int | None = None: Optional chunk size that bounds how many draws touch the accelerator at once, across both the in-sample and forecast predictive sampling. When set, sampling runs in chunks of this many draws, each chunk moved to device before the next is drawn. The per-chunk accelerator footprint is a handful of (batch_size, time, series) buffers, so it scales linearly with this value times the panel width: on wide panels lower it until a chunk fits. The batch size must be strictly below the draw count for that bound to hold: at or above it, sampling falls back to the single-shot path and the full array is materialized on the default device before the single transfer. Chunking changes the PRNG stream layout of the predictive draws, so results are reproducible per (rng_key, batch_size). None (default) samples everything in one shot (the results are still moved to device).
device: jax.Device | str | None = "host": Where the predictive draws are moved as they are sampled, forwarded to the device argument of predict_in_sample() and forecast() (the placement contract of draw_posterior()). It is resolved once and the same placement is handed to both, so an unmet "cpu" warns once per export. The default "host" keeps the predictive draws in pageable host memory (jax Arrays the tree views as NumPy without a copy, or NumPy arrays when no CPU backend is initialized), which is what bounds accelerator memory when batch_size is set; pass None to keep the draws on the default device (chunked compute without per-chunk host transfers, for when the draws fit on the accelerator and transfers would dominate runtime).
coords: Mapping[str, Collection[Any]] | None = None: Optional extra coordinates; these take precedence over the generated time coordinate. They also propagate to the forecast groups, where the generated forecast time takes precedence instead (a user time entry covers the in-sample window; use time_coord for explicit forecast time values). Each value is a one-dimensional collection of labels (a list, a NumPy array, a pandas index) and is stored as np.asarray(value), so the name of a pandas index is dropped.
time_coord: Collection[Any] | None = None: Optional explicit time coordinate values. Without a forecast horizon it covers the in-sample window (defaults to range(n_time)); with a horizon it must cover the full covariates length and is split into the in-sample and forecast time coordinates (the default is the integer continuation). Any one-dimensional collection is accepted (a list, a NumPy array, a pandas index) and stored as np.asarray(time_coord). A timezone-aware pd.DatetimeIndex keeps its zone and clock time, as an object array of aware pd.Timestamp values (arviz_base.dict_to_dataset() cannot store a zoned datetime dtype). xarray does not treat that coordinate as a datetime: .dt, resample() and selection by a date string fail (select with a pd.Timestamp), and to_netcdf() cannot write it. For a datetime coordinate pass index.tz_convert("UTC").tz_localize(None) (UTC times), or index.tz_localize(None) (local clock times, which repeat an hour when daylight saving time ends).
posterior_dims: Mapping[str, Collection[str]] | None = None: Optional mapping from a posterior site name to its non-sample dimension names, e.g. {"drift": ["time"]}. Sites listed here share the tree-wide time coordinate; unlisted sites keep ArviZ’s auto-named dims. This is an explicit opt-in on purpose: inferring time-indexed sites from trace shapes is fragile (a coincidental n_params == n_time would misattribute the axis). Each value, like covariate_dims and obs_dims, is any ordered one-dimensional collection of str (a list, a tuple, a NumPy array, a pandas index).
covariate_dims: Collection[str] | None = None: Optional dimension names for the stored covariates, one per axis, as any ordered one-dimensional collection of str; defaults to the 2-D ("time", "covariate_dim") layout. Use this when covariates carries extra batch axes, e.g. a panel tensor shaped (channel, time, series) with covariate_dims=["channel", "time", "series"]. The time axis is always -2 (the package-wide convention), so its entry should be named "time" to share the tree-wide time coordinate.
obs_dims: Collection[str] | None = None: Optional dimension names for the obs variables of observed_data, posterior_predictive and predictions, one per data axis, as any ordered one-dimensional collection of str. Defaults to ("time", "obs_dim") for 2-D data and names the leading batch axes of a panel obs_batch_0, obs_batch_1, …, e.g. ("obs_batch_0", "time", "obs_dim") for (batch, time, obs) data. Give a batch axis the name of the matching covariate_dims entry so the groups share it, e.g. obs_dims=["origin", "time", "destin"] with covariate_dims=["origin", "time", "covariate_dim"]; coords labels and posterior_dims entries refer to the same names. Axis -2 must be named "time" (it carries the tree-wide time coordinate), the names must be distinct, and "chain" and "draw" are reserved for the sample dims.

Returns

xarray.DataTree: A tree with posterior ((chain, draw, ...), split per num_chains), posterior_predictive (in-sample obs), observed_data, and constant_data groups. When covariates extends beyond data, also predictions and predictions_constant_data groups (sharing the same num_chains split). The obs variable has dims obs_dims in observed_data and (chain, draw, *obs_dims) in posterior_predictive and predictions.

Raises

ValueError: If data or covariates has fewer than two axes (time must be at axis -2), if data has no time steps, if covariates is shorter than data along the time axis, if a posterior leaf has no sample axis, if a forecast-horizon key of posterior (one that ends in "_future" or contains "_future_", or "forecast") has a non-empty leaf, if "obs" is the only key of posterior, if a posterior key contains "/" (a site scoped with the default divider of numpyro.handlers.scope; an empty forecast-horizon key counts, because the export keeps it as a variable; see posterior above), if a posterior key that matches no site of the model names a latent site that has no key (a half-applied divider rename, or a time_reparam() or LocScaleReparam key with a model without that reparam; see posterior above), if time_coord or a coords value is not a one-dimensional collection of labels (a set, a dict, a scalar, a 2-D array, ragged input) or is a pd.MultiIndex, if time_coord mixes time zones (or aware and naive times), if obs_dims, covariate_dims or a posterior_dims value is not an ordered one-dimensional collection (a set, a dict, a single str), if time_coord is given but its length does not match the in-sample window plus the forecast horizon, if num_chains is not a positive integer (bool included), if posterior’s sample count is not evenly divisible by num_chains, if a posterior value of an innovations() site does not have one step per data row (a posterior drawn for a different training length), or if the in-sample draws do not have one time step per data row or the forecast draws one per covariates row past data (see predict_in_sample() and forecast()). Also if obs_dims does not name every data axis, does not name axis -2 "time", repeats a name or uses "chain" or "draw", if the in-sample or forecast draws do not have one axis per obs_dims name (the model’s obs site must have the layout of data), or if an in-sample draw does not have the shape of data (a predictor shared by the series of data without a series plate draws one column).
TypeError: If a name in obs_dims, covariate_dims or a posterior_dims value is not a str.
CovariateDimsError: If covariate_dims does not name every covariates axis.
HostMemoryKindError: If device="pinned_host" is requested on a device that exposes no host memory kind (see _host_memory_kind()).
DevicePlatformError: If device names a platform whose backend is not initialized (see resolve_device()).

Warns

UserWarning: If device="cpu" is requested and the JAX CPU backend is not initialized, so the predictive draws take the NumPy path of "host" instead (once per call).

Notes

When a forecast horizon is present, rng_key is split internally into a predictive subkey and a forecast subkey, so passing the same key twice never correlates the two sample sets. When there is no horizon, rng_key is used unsplit for the in-sample predictive draw. batch_size is the built-in route to memory-bounded predictive sampling; for fully manual control over the forecast draws, build the in-sample tree with matching-length covariates and attach the horizon with add_forecast_groups().