arviz 0.23.3__py3-none-any.whl → 1.0.0rc0__py3-none-any.whl

arviz/data/io_pystan.py

@@ -1,1095 +0,0 @@
-#  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
-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()