# pylint: disable=too-many-instance-attributes,too-many-lines
"""PyStan-specific conversion code."""
import re
from collections import OrderedDict
from copy import deepcopy
from math import ceil
import numpy as np
import xarray as xr
from .. import _log
from ..rcparams import rcParams
from .base import dict_to_dataset, generate_dims_coords, infer_stan_dtypes, make_attrs, requires
from .inference_data import InferenceData
try:
import ujson as json
except ImportError:
# Can't find ujson using json
# mypy struggles with conditional imports expressed as catching ImportError:
# https://github.com/python/mypy/issues/1153
import json # type: ignore
class PyStanConverter:
"""Encapsulate PyStan specific logic."""
def __init__(
self,
*,
posterior=None,
posterior_predictive=None,
predictions=None,
prior=None,
prior_predictive=None,
observed_data=None,
constant_data=None,
predictions_constant_data=None,
log_likelihood=None,
coords=None,
dims=None,
save_warmup=None,
dtypes=None,
):
self.posterior = posterior
self.posterior_predictive = posterior_predictive
self.predictions = predictions
self.prior = prior
self.prior_predictive = prior_predictive
self.observed_data = observed_data
self.constant_data = constant_data
self.predictions_constant_data = predictions_constant_data
self.log_likelihood = (
rcParams["data.log_likelihood"] if log_likelihood is None else log_likelihood
)
self.coords = coords
self.dims = dims
self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
self.dtypes = dtypes
if (
self.log_likelihood is True
and self.posterior is not None
and "log_lik" in self.posterior.sim["pars_oi"]
):
self.log_likelihood = ["log_lik"]
elif isinstance(self.log_likelihood, bool):
self.log_likelihood = None
import pystan # pylint: disable=import-error
self.pystan = pystan
@requires("posterior")
def posterior_to_xarray(self):
"""Extract posterior samples from fit."""
posterior = self.posterior
# filter posterior_predictive and log_likelihood
posterior_predictive = self.posterior_predictive
if posterior_predictive is None:
posterior_predictive = []
elif isinstance(posterior_predictive, str):
posterior_predictive = [posterior_predictive]
predictions = self.predictions
if predictions is None:
predictions = []
elif isinstance(predictions, str):
predictions = [predictions]
log_likelihood = self.log_likelihood
if log_likelihood is None:
log_likelihood = []
elif isinstance(log_likelihood, str):
log_likelihood = [log_likelihood]
elif isinstance(log_likelihood, dict):
log_likelihood = list(log_likelihood.values())
ignore = posterior_predictive + predictions + log_likelihood + ["lp__"]
data, data_warmup = get_draws(
posterior, ignore=ignore, warmup=self.save_warmup, dtypes=self.dtypes
)
attrs = get_attrs(posterior)
return (
dict_to_dataset(
data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
),
dict_to_dataset(
data_warmup, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
),
)
@requires("posterior")
def sample_stats_to_xarray(self):
"""Extract sample_stats from posterior."""
posterior = self.posterior
data, data_warmup = get_sample_stats(posterior, warmup=self.save_warmup)
# lp__
stat_lp, stat_lp_warmup = get_draws(
posterior, variables="lp__", warmup=self.save_warmup, dtypes=self.dtypes
)
data["lp"] = stat_lp["lp__"]
if stat_lp_warmup:
data_warmup["lp"] = stat_lp_warmup["lp__"]
attrs = get_attrs(posterior)
return (
dict_to_dataset(
data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
),
dict_to_dataset(
data_warmup, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
),
)
@requires("posterior")
@requires("log_likelihood")
def log_likelihood_to_xarray(self):
"""Store log_likelihood data in log_likelihood group."""
fit = self.posterior
# log_likelihood values
log_likelihood = self.log_likelihood
if isinstance(log_likelihood, str):
log_likelihood = [log_likelihood]
if isinstance(log_likelihood, (list, tuple)):
log_likelihood = {name: name for name in log_likelihood}
log_likelihood_draws, log_likelihood_draws_warmup = get_draws(
fit,
variables=list(log_likelihood.values()),
warmup=self.save_warmup,
dtypes=self.dtypes,
)
data = {
obs_var_name: log_likelihood_draws[log_like_name]
for obs_var_name, log_like_name in log_likelihood.items()
if log_like_name in log_likelihood_draws
}
data_warmup = {
obs_var_name: log_likelihood_draws_warmup[log_like_name]
for obs_var_name, log_like_name in log_likelihood.items()
if log_like_name in log_likelihood_draws_warmup
}
return (
dict_to_dataset(
data, library=self.pystan, coords=self.coords, dims=self.dims, skip_event_dims=True
),
dict_to_dataset(
data_warmup,
library=self.pystan,
coords=self.coords,
dims=self.dims,
skip_event_dims=True,
),
)
@requires("posterior")
@requires("posterior_predictive")
def posterior_predictive_to_xarray(self):
"""Convert posterior_predictive samples to xarray."""
posterior = self.posterior
posterior_predictive = self.posterior_predictive
data, data_warmup = get_draws(
posterior, variables=posterior_predictive, warmup=self.save_warmup, dtypes=self.dtypes
)
return (
dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims),
dict_to_dataset(data_warmup, library=self.pystan, coords=self.coords, dims=self.dims),
)
@requires("posterior")
@requires("predictions")
def predictions_to_xarray(self):
"""Convert predictions samples to xarray."""
posterior = self.posterior
predictions = self.predictions
data, data_warmup = get_draws(
posterior, variables=predictions, warmup=self.save_warmup, dtypes=self.dtypes
)
return (
dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims),
dict_to_dataset(data_warmup, library=self.pystan, coords=self.coords, dims=self.dims),
)
@requires("prior")
def prior_to_xarray(self):
"""Convert prior samples to xarray."""
prior = self.prior
# filter posterior_predictive and log_likelihood
prior_predictive = self.prior_predictive
if prior_predictive is None:
prior_predictive = []
elif isinstance(prior_predictive, str):
prior_predictive = [prior_predictive]
ignore = prior_predictive + ["lp__"]
data, _ = get_draws(prior, ignore=ignore, warmup=False, dtypes=self.dtypes)
attrs = get_attrs(prior)
return dict_to_dataset(
data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
)
@requires("prior")
def sample_stats_prior_to_xarray(self):
"""Extract sample_stats_prior from prior."""
prior = self.prior
data, _ = get_sample_stats(prior, warmup=False)
# lp__
stat_lp, _ = get_draws(prior, variables="lp__", warmup=False, dtypes=self.dtypes)
data["lp"] = stat_lp["lp__"]
attrs = get_attrs(prior)
return dict_to_dataset(
data, library=self.pystan, attrs=attrs, coords=self.coords, dims=self.dims
)
@requires("prior")
@requires("prior_predictive")
def prior_predictive_to_xarray(self):
"""Convert prior_predictive samples to xarray."""
prior = self.prior
prior_predictive = self.prior_predictive
data, _ = get_draws(prior, variables=prior_predictive, warmup=False, dtypes=self.dtypes)
return dict_to_dataset(data, library=self.pystan, coords=self.coords, dims=self.dims)
@requires("posterior")
@requires(["observed_data", "constant_data", "predictions_constant_data"])
def data_to_xarray(self):
"""Convert observed, constant data and predictions constant data to xarray."""
posterior = self.posterior
dims = {} if self.dims is None else self.dims
obs_const_dict = {}
for group_name in ("observed_data", "constant_data", "predictions_constant_data"):
names = getattr(self, group_name)
if names is None:
continue
names = [names] if isinstance(names, str) else names
data = OrderedDict()
for key in names:
vals = np.atleast_1d(posterior.data[key])
val_dims = dims.get(key)
val_dims, coords = generate_dims_coords(
vals.shape, key, dims=val_dims, coords=self.coords
)
data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
obs_const_dict[group_name] = xr.Dataset(
data_vars=data, attrs=make_attrs(library=self.pystan)
)
return obs_const_dict
def to_inference_data(self):
"""Convert all available data to an InferenceData object.
Note that if groups can not be created (i.e., there is no `fit`, so
the `posterior` and `sample_stats` can not be extracted), then the InferenceData
will not have those groups.
"""
data_dict = self.data_to_xarray()
return InferenceData(
save_warmup=self.save_warmup,
**{
"posterior": self.posterior_to_xarray(),
"sample_stats": self.sample_stats_to_xarray(),
"log_likelihood": self.log_likelihood_to_xarray(),
"posterior_predictive": self.posterior_predictive_to_xarray(),
"predictions": self.predictions_to_xarray(),
"prior": self.prior_to_xarray(),
"sample_stats_prior": self.sample_stats_prior_to_xarray(),
"prior_predictive": self.prior_predictive_to_xarray(),
**({} if data_dict is None else data_dict),
},
)
class PyStan3Converter:
"""Encapsulate PyStan3 specific logic."""
# pylint: disable=too-many-instance-attributes
def __init__(
self,
*,
posterior=None,
posterior_model=None,
posterior_predictive=None,
predictions=None,
prior=None,
prior_model=None,
prior_predictive=None,
observed_data=None,
constant_data=None,
predictions_constant_data=None,
log_likelihood=None,
coords=None,
dims=None,
save_warmup=None,
dtypes=None,
):
self.posterior = posterior
self.posterior_model = posterior_model
self.posterior_predictive = posterior_predictive
self.predictions = predictions
self.prior = prior
self.prior_model = prior_model
self.prior_predictive = prior_predictive
self.observed_data = observed_data
self.constant_data = constant_data
self.predictions_constant_data = predictions_constant_data
self.log_likelihood = (
rcParams["data.log_likelihood"] if log_likelihood is None else log_likelihood
)
self.coords = coords
self.dims = dims
self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
self.dtypes = dtypes
if (
self.log_likelihood is True
and self.posterior is not None
and "log_lik" in self.posterior.param_names
):
self.log_likelihood = ["log_lik"]
elif isinstance(self.log_likelihood, bool):
self.log_likelihood = None
import stan # pylint: disable=import-error
self.stan = stan
@requires("posterior")
def posterior_to_xarray(self):
"""Extract posterior samples from fit."""
posterior = self.posterior
posterior_model = self.posterior_model
# filter posterior_predictive and log_likelihood
posterior_predictive = self.posterior_predictive
if posterior_predictive is None:
posterior_predictive = []
elif isinstance(posterior_predictive, str):
posterior_predictive = [posterior_predictive]
predictions = self.predictions
if predictions is None:
predictions = []
elif isinstance(predictions, str):
predictions = [predictions]
log_likelihood = self.log_likelihood
if log_likelihood is None:
log_likelihood = []
elif isinstance(log_likelihood, str):
log_likelihood = [log_likelihood]
elif isinstance(log_likelihood, dict):
log_likelihood = list(log_likelihood.values())
ignore = posterior_predictive + predictions + log_likelihood
data, data_warmup = get_draws_stan3(
posterior,
model=posterior_model,
ignore=ignore,
warmup=self.save_warmup,
dtypes=self.dtypes,
)
attrs = get_attrs_stan3(posterior, model=posterior_model)
return (
dict_to_dataset(
data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
dict_to_dataset(
data_warmup, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
)
@requires("posterior")
def sample_stats_to_xarray(self):
"""Extract sample_stats from posterior."""
posterior = self.posterior
posterior_model = self.posterior_model
data, data_warmup = get_sample_stats_stan3(
posterior, ignore="lp__", warmup=self.save_warmup, dtypes=self.dtypes
)
data_lp, data_warmup_lp = get_sample_stats_stan3(
posterior, variables="lp__", warmup=self.save_warmup
)
data["lp"] = data_lp["lp"]
if data_warmup_lp:
data_warmup["lp"] = data_warmup_lp["lp"]
attrs = get_attrs_stan3(posterior, model=posterior_model)
return (
dict_to_dataset(
data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
dict_to_dataset(
data_warmup, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
)
@requires("posterior")
@requires("log_likelihood")
def log_likelihood_to_xarray(self):
"""Store log_likelihood data in log_likelihood group."""
fit = self.posterior
log_likelihood = self.log_likelihood
model = self.posterior_model
if isinstance(log_likelihood, str):
log_likelihood = [log_likelihood]
if isinstance(log_likelihood, (list, tuple)):
log_likelihood = {name: name for name in log_likelihood}
log_likelihood_draws, log_likelihood_draws_warmup = get_draws_stan3(
fit,
model=model,
variables=list(log_likelihood.values()),
warmup=self.save_warmup,
dtypes=self.dtypes,
)
data = {
obs_var_name: log_likelihood_draws[log_like_name]
for obs_var_name, log_like_name in log_likelihood.items()
if log_like_name in log_likelihood_draws
}
data_warmup = {
obs_var_name: log_likelihood_draws_warmup[log_like_name]
for obs_var_name, log_like_name in log_likelihood.items()
if log_like_name in log_likelihood_draws_warmup
}
return (
dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims),
dict_to_dataset(data_warmup, library=self.stan, coords=self.coords, dims=self.dims),
)
@requires("posterior")
@requires("posterior_predictive")
def posterior_predictive_to_xarray(self):
"""Convert posterior_predictive samples to xarray."""
posterior = self.posterior
posterior_model = self.posterior_model
posterior_predictive = self.posterior_predictive
data, data_warmup = get_draws_stan3(
posterior,
model=posterior_model,
variables=posterior_predictive,
warmup=self.save_warmup,
dtypes=self.dtypes,
)
return (
dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims),
dict_to_dataset(data_warmup, library=self.stan, coords=self.coords, dims=self.dims),
)
@requires("posterior")
@requires("predictions")
def predictions_to_xarray(self):
"""Convert predictions samples to xarray."""
posterior = self.posterior
posterior_model = self.posterior_model
predictions = self.predictions
data, data_warmup = get_draws_stan3(
posterior,
model=posterior_model,
variables=predictions,
warmup=self.save_warmup,
dtypes=self.dtypes,
)
return (
dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims),
dict_to_dataset(data_warmup, library=self.stan, coords=self.coords, dims=self.dims),
)
@requires("prior")
def prior_to_xarray(self):
"""Convert prior samples to xarray."""
prior = self.prior
prior_model = self.prior_model
# filter posterior_predictive and log_likelihood
prior_predictive = self.prior_predictive
if prior_predictive is None:
prior_predictive = []
elif isinstance(prior_predictive, str):
prior_predictive = [prior_predictive]
ignore = prior_predictive
data, data_warmup = get_draws_stan3(
prior, model=prior_model, ignore=ignore, warmup=self.save_warmup, dtypes=self.dtypes
)
attrs = get_attrs_stan3(prior, model=prior_model)
return (
dict_to_dataset(
data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
dict_to_dataset(
data_warmup, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
)
@requires("prior")
def sample_stats_prior_to_xarray(self):
"""Extract sample_stats_prior from prior."""
prior = self.prior
prior_model = self.prior_model
data, data_warmup = get_sample_stats_stan3(
prior, warmup=self.save_warmup, dtypes=self.dtypes
)
attrs = get_attrs_stan3(prior, model=prior_model)
return (
dict_to_dataset(
data, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
dict_to_dataset(
data_warmup, library=self.stan, attrs=attrs, coords=self.coords, dims=self.dims
),
)
@requires("prior")
@requires("prior_predictive")
def prior_predictive_to_xarray(self):
"""Convert prior_predictive samples to xarray."""
prior = self.prior
prior_model = self.prior_model
prior_predictive = self.prior_predictive
data, data_warmup = get_draws_stan3(
prior,
model=prior_model,
variables=prior_predictive,
warmup=self.save_warmup,
dtypes=self.dtypes,
)
return (
dict_to_dataset(data, library=self.stan, coords=self.coords, dims=self.dims),
dict_to_dataset(data_warmup, library=self.stan, coords=self.coords, dims=self.dims),
)
@requires("posterior_model")
@requires(["observed_data", "constant_data"])
def observed_and_constant_data_to_xarray(self):
"""Convert observed data to xarray."""
posterior_model = self.posterior_model
dims = {} if self.dims is None else self.dims
obs_const_dict = {}
for group_name in ("observed_data", "constant_data"):
names = getattr(self, group_name)
if names is None:
continue
names = [names] if isinstance(names, str) else names
data = OrderedDict()
for key in names:
vals = np.atleast_1d(posterior_model.data[key])
val_dims = dims.get(key)
val_dims, coords = generate_dims_coords(
vals.shape, key, dims=val_dims, coords=self.coords
)
data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
obs_const_dict[group_name] = xr.Dataset(
data_vars=data, attrs=make_attrs(library=self.stan)
)
return obs_const_dict
@requires("posterior_model")
@requires("predictions_constant_data")
def predictions_constant_data_to_xarray(self):
"""Convert observed data to xarray."""
posterior_model = self.posterior_model
dims = {} if self.dims is None else self.dims
names = self.predictions_constant_data
names = [names] if isinstance(names, str) else names
data = OrderedDict()
for key in names:
vals = np.atleast_1d(posterior_model.data[key])
val_dims = dims.get(key)
val_dims, coords = generate_dims_coords(
vals.shape, key, dims=val_dims, coords=self.coords
)
data[key] = xr.DataArray(vals, dims=val_dims, coords=coords)
return xr.Dataset(data_vars=data, attrs=make_attrs(library=self.stan))
def to_inference_data(self):
"""Convert all available data to an InferenceData object.
Note that if groups can not be created (i.e., there is no `fit`, so
the `posterior` and `sample_stats` can not be extracted), then the InferenceData
will not have those groups.
"""
obs_const_dict = self.observed_and_constant_data_to_xarray()
predictions_const_data = self.predictions_constant_data_to_xarray()
return InferenceData(
save_warmup=self.save_warmup,
**{
"posterior": self.posterior_to_xarray(),
"sample_stats": self.sample_stats_to_xarray(),
"log_likelihood": self.log_likelihood_to_xarray(),
"posterior_predictive": self.posterior_predictive_to_xarray(),
"predictions": self.predictions_to_xarray(),
"prior": self.prior_to_xarray(),
"sample_stats_prior": self.sample_stats_prior_to_xarray(),
"prior_predictive": self.prior_predictive_to_xarray(),
**({} if obs_const_dict is None else obs_const_dict),
**(
{}
if predictions_const_data is None
else {"predictions_constant_data": predictions_const_data}
),
},
)
def get_draws(fit, variables=None, ignore=None, warmup=False, dtypes=None):
"""Extract draws from PyStan fit."""
if ignore is None:
ignore = []
if fit.mode == 1:
msg = "Model in mode 'test_grad'. Sampling is not conducted."
raise AttributeError(msg)
if fit.mode == 2 or fit.sim.get("samples") is None:
msg = "Fit doesn't contain samples."
raise AttributeError(msg)
if dtypes is None:
dtypes = {}
dtypes = {**infer_dtypes(fit), **dtypes}
if variables is None:
variables = fit.sim["pars_oi"]
elif isinstance(variables, str):
variables = [variables]
variables = list(variables)
for var, dim in zip(fit.sim["pars_oi"], fit.sim["dims_oi"]):
if var in variables and np.prod(dim) == 0:
del variables[variables.index(var)]
ndraws_warmup = fit.sim["warmup2"]
if max(ndraws_warmup) == 0:
warmup = False
ndraws = [s - w for s, w in zip(fit.sim["n_save"], ndraws_warmup)]
nchain = len(fit.sim["samples"])
# check if the values are in 0-based (<=2.17) or 1-based indexing (>=2.18)
shift = 1
if any(dim and np.prod(dim) != 0 for dim in fit.sim["dims_oi"]):
# choose variable with lowest number of dims > 1
par_idx = min(
(dim, i) for i, dim in enumerate(fit.sim["dims_oi"]) if (dim and np.prod(dim) != 0)
)[1]
offset = int(sum(map(np.prod, fit.sim["dims_oi"][:par_idx])))
par_offset = int(np.prod(fit.sim["dims_oi"][par_idx]))
par_keys = fit.sim["fnames_oi"][offset : offset + par_offset]
shift = len(par_keys)
for item in par_keys:
_, shape = item.replace("]", "").split("[")
shape_idx_min = min(int(shape_value) for shape_value in shape.split(","))
shift = min(shift, shape_idx_min)
# If shift is higher than 1, this will probably mean that Stan
# has implemented sparse structure (saves only non-zero parts),
# but let's hope that dims are still corresponding to the full shape
shift = int(min(shift, 1))
var_keys = OrderedDict((var, []) for var in fit.sim["pars_oi"])
for key in fit.sim["fnames_oi"]:
var, *tails = key.split("[")
loc = [Ellipsis]
for tail in tails:
loc = []
for i in tail[:-1].split(","):
loc.append(int(i) - shift)
var_keys[var].append((key, loc))
shapes = dict(zip(fit.sim["pars_oi"], fit.sim["dims_oi"]))
variables = [var for var in variables if var not in ignore]
data = OrderedDict()
data_warmup = OrderedDict()
for var in variables:
if var in data:
continue
keys_locs = var_keys.get(var, [(var, [Ellipsis])])
shape = shapes.get(var, [])
dtype = dtypes.get(var)
ndraw = max(ndraws)
ary_shape = [nchain, ndraw] + shape
ary = np.empty(ary_shape, dtype=dtype, order="F")
if warmup:
nwarmup = max(ndraws_warmup)
ary_warmup_shape = [nchain, nwarmup] + shape
ary_warmup = np.empty(ary_warmup_shape, dtype=dtype, order="F")
for chain, (pyholder, ndraw, ndraw_warmup) in enumerate(
zip(fit.sim["samples"], ndraws, ndraws_warmup)
):
axes = [chain, slice(None)]
for key, loc in keys_locs:
ary_slice = tuple(axes + loc)
ary[ary_slice] = pyholder.chains[key][-ndraw:]
if warmup:
ary_warmup[ary_slice] = pyholder.chains[key][:ndraw_warmup]
data[var] = ary
if warmup:
data_warmup[var] = ary_warmup
return data, data_warmup
def get_sample_stats(fit, warmup=False, dtypes=None):
"""Extract sample stats from PyStan fit."""
if dtypes is None:
dtypes = {}
dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64, **dtypes}
rename_dict = {
"divergent": "diverging",
"n_leapfrog": "n_steps",
"treedepth": "tree_depth",
"stepsize": "step_size",
"accept_stat": "acceptance_rate",
}
ndraws_warmup = fit.sim["warmup2"]
if max(ndraws_warmup) == 0:
warmup = False
ndraws = [s - w for s, w in zip(fit.sim["n_save"], ndraws_warmup)]
extraction = OrderedDict()
extraction_warmup = OrderedDict()
for chain, (pyholder, ndraw, ndraw_warmup) in enumerate(
zip(fit.sim["samples"], ndraws, ndraws_warmup)
):
if chain == 0:
for key in pyholder["sampler_param_names"]:
extraction[key] = []
if warmup:
extraction_warmup[key] = []
for key, values in zip(pyholder["sampler_param_names"], pyholder["sampler_params"]):
extraction[key].append(values[-ndraw:])
if warmup:
extraction_warmup[key].append(values[:ndraw_warmup])
data = OrderedDict()
for key, values in extraction.items():
values = np.stack(values, axis=0)
dtype = dtypes.get(key)
values = values.astype(dtype)
name = re.sub("__$", "", key)
name = rename_dict.get(name, name)
data[name] = values
data_warmup = OrderedDict()
if warmup:
for key, values in extraction_warmup.items():
values = np.stack(values, axis=0)
values = values.astype(dtypes.get(key))
name = re.sub("__$", "", key)
name = rename_dict.get(name, name)
data_warmup[name] = values
return data, data_warmup
def get_attrs(fit):
"""Get attributes from PyStan fit object."""
attrs = {}
try:
attrs["args"] = [deepcopy(holder.args) for holder in fit.sim["samples"]]
except Exception as exp: # pylint: disable=broad-except
_log.warning("Failed to fetch args from fit: %s", exp)
if "args" in attrs:
for arg in attrs["args"]:
if isinstance(arg["init"], bytes):
arg["init"] = arg["init"].decode("utf-8")
attrs["args"] = json.dumps(attrs["args"])
try:
attrs["inits"] = [holder.inits for holder in fit.sim["samples"]]
except Exception as exp: # pylint: disable=broad-except
_log.warning("Failed to fetch `args` from fit: %s", exp)
else:
attrs["inits"] = json.dumps(attrs["inits"])
attrs["step_size"] = []
attrs["metric"] = []
attrs["inv_metric"] = []
for holder in fit.sim["samples"]:
try:
step_size = float(
re.search(
r"step\s*size\s*=\s*([0-9]+.?[0-9]+)\s*",
holder.adaptation_info,
flags=re.IGNORECASE,
).group(1)
)
except AttributeError:
step_size = np.nan
attrs["step_size"].append(step_size)
inv_metric_match = re.search(
r"mass matrix:\s*(.*)\s*$", holder.adaptation_info, flags=re.DOTALL
)
if inv_metric_match:
inv_metric_str = inv_metric_match.group(1)
if "Diagonal elements of inverse mass matrix" in holder.adaptation_info:
metric = "diag_e"
inv_metric = [float(item) for item in inv_metric_str.strip(" #\n").split(",")]
else:
metric = "dense_e"
inv_metric = [
list(map(float, item.split(",")))
for item in re.sub(r"#\s", "", inv_metric_str).splitlines()
]
else:
metric = "unit_e"
inv_metric = None
attrs["metric"].append(metric)
attrs["inv_metric"].append(inv_metric)
attrs["inv_metric"] = json.dumps(attrs["inv_metric"])
if not attrs["step_size"]:
del attrs["step_size"]
attrs["adaptation_info"] = fit.get_adaptation_info()
attrs["stan_code"] = fit.get_stancode()
return attrs
def get_draws_stan3(fit, model=None, variables=None, ignore=None, warmup=False, dtypes=None):
"""Extract draws from PyStan3 fit."""
if ignore is None:
ignore = []
if dtypes is None:
dtypes = {}
if model is not None:
dtypes = {**infer_dtypes(fit, model), **dtypes}
if not fit.save_warmup:
warmup = False
num_warmup = ceil((fit.num_warmup * fit.save_warmup) / fit.num_thin)
if variables is None:
variables = fit.param_names
elif isinstance(variables, str):
variables = [variables]
variables = list(variables)
data = OrderedDict()
data_warmup = OrderedDict()
for var in variables:
if var in ignore:
continue
if var in data:
continue
dtype = dtypes.get(var)
new_shape = (*fit.dims[fit.param_names.index(var)], -1, fit.num_chains)
if 0 in new_shape:
continue
values = fit._draws[fit._parameter_indexes(var), :] # pylint: disable=protected-access
values = values.reshape(new_shape, order="F")
values = np.moveaxis(values, [-2, -1], [1, 0])
values = values.astype(dtype)
if warmup:
data_warmup[var] = values[:, num_warmup:]
data[var] = values[:, num_warmup:]
return data, data_warmup
def get_sample_stats_stan3(fit, variables=None, ignore=None, warmup=False, dtypes=None):
"""Extract sample stats from PyStan3 fit."""
if dtypes is None:
dtypes = {}
dtypes = {"divergent__": bool, "n_leapfrog__": np.int64, "treedepth__": np.int64, **dtypes}
rename_dict = {
"divergent": "diverging",
"n_leapfrog": "n_steps",
"treedepth": "tree_depth",
"stepsize": "step_size",
"accept_stat": "acceptance_rate",
}
if isinstance(variables, str):
variables = [variables]
if isinstance(ignore, str):
ignore = [ignore]
if not fit.save_warmup:
warmup = False
num_warmup = ceil((fit.num_warmup * fit.save_warmup) / fit.num_thin)
data = OrderedDict()
data_warmup = OrderedDict()
for key in fit.sample_and_sampler_param_names:
if (variables and key not in variables) or (ignore and key in ignore):
continue
new_shape = -1, fit.num_chains
values = fit._draws[fit._parameter_indexes(key)] # pylint: disable=protected-access
values = values.reshape(new_shape, order="F")
values = np.moveaxis(values, [-2, -1], [1, 0])
dtype = dtypes.get(key)
values = values.astype(dtype)
name = re.sub("__$", "", key)
name = rename_dict.get(name, name)
if warmup:
data_warmup[name] = values[:, :num_warmup]
data[name] = values[:, num_warmup:]
return data, data_warmup
def get_attrs_stan3(fit, model=None):
"""Get attributes from PyStan3 fit and model object."""
attrs = {}
for key in ["num_chains", "num_samples", "num_thin", "num_warmup", "save_warmup"]:
try:
attrs[key] = getattr(fit, key)
except AttributeError as exp:
_log.warning("Failed to access attribute %s in fit object %s", key, exp)
if model is not None:
for key in ["model_name", "program_code", "random_seed"]:
try:
attrs[key] = getattr(model, key)
except AttributeError as exp:
_log.warning("Failed to access attribute %s in model object %s", key, exp)
return attrs
def infer_dtypes(fit, model=None):
"""Infer dtypes from Stan model code.
Function strips out generated quantities block and searches for `int`
dtypes after stripping out comments inside the block.
"""
if model is None:
stan_code = fit.get_stancode()
model_pars = fit.model_pars
else:
stan_code = model.program_code
model_pars = fit.param_names
dtypes = {key: item for key, item in infer_stan_dtypes(stan_code).items() if key in model_pars}
return dtypes
# pylint disable=too-many-instance-attributes
[docs]
def from_pystan(
posterior=None,
*,
posterior_predictive=None,
predictions=None,
prior=None,
prior_predictive=None,
observed_data=None,
constant_data=None,
predictions_constant_data=None,
log_likelihood=None,
coords=None,
dims=None,
posterior_model=None,
prior_model=None,
save_warmup=None,
dtypes=None,
):
"""Convert PyStan data into an InferenceData object.
For a usage example read the
:ref:`Creating InferenceData section on from_pystan <creating_InferenceData>`
Parameters
----------
posterior : StanFit4Model or stan.fit.Fit
PyStan fit object for posterior.
posterior_predictive : str, a list of str
Posterior predictive samples for the posterior.
predictions : str, a list of str
Out-of-sample predictions for the posterior.
prior : StanFit4Model or stan.fit.Fit
PyStan fit object for prior.
prior_predictive : str, a list of str
Posterior predictive samples for the prior.
observed_data : str or a list of str
observed data used in the sampling.
Observed data is extracted from the `posterior.data`.
PyStan3 needs model object for the extraction.
See `posterior_model`.
constant_data : str or list of str
Constants relevant to the model (i.e. x values in a linear
regression).
predictions_constant_data : str or list of str
Constants relevant to the model predictions (i.e. new x values in a linear
regression).
log_likelihood : dict of {str: str}, list of str or str, optional
Pointwise log_likelihood for the data. log_likelihood is extracted from the
posterior. It is recommended to use this argument as a dictionary whose keys
are observed variable names and its values are the variables storing log
likelihood arrays in the Stan code. In other cases, a dictionary with keys
equal to its values is used. By default, if a variable ``log_lik`` is
present in the Stan model, it will be retrieved as pointwise log
likelihood values. Use ``False`` or set ``data.log_likelihood`` to
false to avoid this behaviour.
coords : dict[str, iterable]
A dictionary containing the values that are used as index. The key
is the name of the dimension, the values are the index values.
dims : dict[str, List(str)]
A mapping from variables to a list of coordinate names for the variable.
posterior_model : stan.model.Model
PyStan3 specific model object. Needed for automatic dtype parsing
and for the extraction of observed data.
prior_model : stan.model.Model
PyStan3 specific model object. Needed for automatic dtype parsing.
save_warmup : bool
Save warmup iterations into InferenceData object. If not defined, use default
defined by the rcParams.
dtypes: dict
A dictionary containing dtype information (int, float) for parameters.
By default dtype information is extracted from the model code.
Model code is extracted from fit object in PyStan 2 and from model object
in PyStan 3.
Returns
-------
InferenceData object
"""
check_posterior = (posterior is not None) and (type(posterior).__module__ == "stan.fit")
check_prior = (prior is not None) and (type(prior).__module__ == "stan.fit")
if check_posterior or check_prior:
return PyStan3Converter(
posterior=posterior,
posterior_model=posterior_model,
posterior_predictive=posterior_predictive,
predictions=predictions,
prior=prior,
prior_model=prior_model,
prior_predictive=prior_predictive,
observed_data=observed_data,
constant_data=constant_data,
predictions_constant_data=predictions_constant_data,
log_likelihood=log_likelihood,
coords=coords,
dims=dims,
save_warmup=save_warmup,
dtypes=dtypes,
).to_inference_data()
else:
return PyStanConverter(
posterior=posterior,
posterior_predictive=posterior_predictive,
predictions=predictions,
prior=prior,
prior_predictive=prior_predictive,
observed_data=observed_data,
constant_data=constant_data,
predictions_constant_data=predictions_constant_data,
log_likelihood=log_likelihood,
coords=coords,
dims=dims,
save_warmup=save_warmup,
dtypes=dtypes,
).to_inference_data()
