"""Dictionary specific conversion code."""
import warnings
from typing import Optional
from ..rcparams import rcParams
from .base import dict_to_dataset, requires
from .inference_data import WARMUP_TAG, InferenceData
# pylint: disable=too-many-instance-attributes
class DictConverter:
"""Encapsulate Dictionary specific logic."""
def __init__(
self,
*,
posterior=None,
posterior_predictive=None,
predictions=None,
sample_stats=None,
log_likelihood=None,
prior=None,
prior_predictive=None,
sample_stats_prior=None,
observed_data=None,
constant_data=None,
predictions_constant_data=None,
warmup_posterior=None,
warmup_posterior_predictive=None,
warmup_predictions=None,
warmup_log_likelihood=None,
warmup_sample_stats=None,
save_warmup=None,
index_origin=None,
coords=None,
dims=None,
pred_dims=None,
pred_coords=None,
attrs=None,
**kwargs,
):
self.posterior = posterior
self.posterior_predictive = posterior_predictive
self.predictions = predictions
self.sample_stats = sample_stats
self.log_likelihood = log_likelihood
self.prior = prior
self.prior_predictive = prior_predictive
self.sample_stats_prior = sample_stats_prior
self.observed_data = observed_data
self.constant_data = constant_data
self.predictions_constant_data = predictions_constant_data
self.warmup_posterior = warmup_posterior
self.warmup_posterior_predictive = warmup_posterior_predictive
self.warmup_predictions = warmup_predictions
self.warmup_log_likelihood = warmup_log_likelihood
self.warmup_sample_stats = warmup_sample_stats
self.save_warmup = rcParams["data.save_warmup"] if save_warmup is None else save_warmup
self.coords = (
coords
if pred_coords is None
else pred_coords if coords is None else {**coords, **pred_coords}
)
self.index_origin = index_origin
self.coords = coords
self.dims = dims
self.pred_dims = dims if pred_dims is None else pred_dims
self.attrs = {} if attrs is None else attrs
self.attrs.pop("created_at", None)
self.attrs.pop("arviz_version", None)
self._kwargs = kwargs
def _init_dict(self, attr_name):
dict_or_none = getattr(self, attr_name, {})
return {} if dict_or_none is None else dict_or_none
@requires(["posterior", f"{WARMUP_TAG}posterior"])
def posterior_to_xarray(self):
"""Convert posterior samples to xarray."""
data = self._init_dict("posterior")
data_warmup = self._init_dict(f"{WARMUP_TAG}posterior")
if not isinstance(data, dict):
raise TypeError("DictConverter.posterior is not a dictionary")
if not isinstance(data_warmup, dict):
raise TypeError("DictConverter.warmup_posterior is not a dictionary")
if "log_likelihood" in data:
warnings.warn(
"log_likelihood variable found in posterior group."
" For stats functions log likelihood data needs to be in log_likelihood group.",
UserWarning,
)
posterior_attrs = self._kwargs.get("posterior_attrs")
posterior_warmup_attrs = self._kwargs.get("posterior_warmup_attrs")
return (
dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
attrs=posterior_attrs,
index_origin=self.index_origin,
),
dict_to_dataset(
data_warmup,
library=None,
coords=self.coords,
dims=self.dims,
attrs=posterior_warmup_attrs,
index_origin=self.index_origin,
),
)
@requires(["sample_stats", f"{WARMUP_TAG}sample_stats"])
def sample_stats_to_xarray(self):
"""Convert sample_stats samples to xarray."""
data = self._init_dict("sample_stats")
data_warmup = self._init_dict(f"{WARMUP_TAG}sample_stats")
if not isinstance(data, dict):
raise TypeError("DictConverter.sample_stats is not a dictionary")
if not isinstance(data_warmup, dict):
raise TypeError("DictConverter.warmup_sample_stats is not a dictionary")
if "log_likelihood" in data:
warnings.warn(
"log_likelihood variable found in sample_stats."
" Storing log_likelihood data in sample_stats group will be deprecated in "
"favour of storing them in the log_likelihood group.",
PendingDeprecationWarning,
)
sample_stats_attrs = self._kwargs.get("sample_stats_attrs")
sample_stats_warmup_attrs = self._kwargs.get("sample_stats_warmup_attrs")
return (
dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
attrs=sample_stats_attrs,
index_origin=self.index_origin,
),
dict_to_dataset(
data_warmup,
library=None,
coords=self.coords,
dims=self.dims,
attrs=sample_stats_warmup_attrs,
index_origin=self.index_origin,
),
)
@requires(["log_likelihood", f"{WARMUP_TAG}log_likelihood"])
def log_likelihood_to_xarray(self):
"""Convert log_likelihood samples to xarray."""
data = self._init_dict("log_likelihood")
data_warmup = self._init_dict(f"{WARMUP_TAG}log_likelihood")
if not isinstance(data, dict):
raise TypeError("DictConverter.log_likelihood is not a dictionary")
if not isinstance(data_warmup, dict):
raise TypeError("DictConverter.warmup_log_likelihood is not a dictionary")
log_likelihood_attrs = self._kwargs.get("log_likelihood_attrs")
log_likelihood_warmup_attrs = self._kwargs.get("log_likelihood_warmup_attrs")
return (
dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
attrs=log_likelihood_attrs,
index_origin=self.index_origin,
skip_event_dims=True,
),
dict_to_dataset(
data_warmup,
library=None,
coords=self.coords,
dims=self.dims,
attrs=log_likelihood_warmup_attrs,
index_origin=self.index_origin,
skip_event_dims=True,
),
)
@requires(["posterior_predictive", f"{WARMUP_TAG}posterior_predictive"])
def posterior_predictive_to_xarray(self):
"""Convert posterior_predictive samples to xarray."""
data = self._init_dict("posterior_predictive")
data_warmup = self._init_dict(f"{WARMUP_TAG}posterior_predictive")
if not isinstance(data, dict):
raise TypeError("DictConverter.posterior_predictive is not a dictionary")
if not isinstance(data_warmup, dict):
raise TypeError("DictConverter.warmup_posterior_predictive is not a dictionary")
posterior_predictive_attrs = self._kwargs.get("posterior_predictive_attrs")
posterior_predictive_warmup_attrs = self._kwargs.get("posterior_predictive_warmup_attrs")
return (
dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
attrs=posterior_predictive_attrs,
index_origin=self.index_origin,
),
dict_to_dataset(
data_warmup,
library=None,
coords=self.coords,
dims=self.dims,
attrs=posterior_predictive_warmup_attrs,
index_origin=self.index_origin,
),
)
@requires(["predictions", f"{WARMUP_TAG}predictions"])
def predictions_to_xarray(self):
"""Convert predictions to xarray."""
data = self._init_dict("predictions")
data_warmup = self._init_dict(f"{WARMUP_TAG}predictions")
if not isinstance(data, dict):
raise TypeError("DictConverter.predictions is not a dictionary")
if not isinstance(data_warmup, dict):
raise TypeError("DictConverter.warmup_predictions is not a dictionary")
predictions_attrs = self._kwargs.get("predictions_attrs")
predictions_warmup_attrs = self._kwargs.get("predictions_warmup_attrs")
return (
dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.pred_dims,
attrs=predictions_attrs,
index_origin=self.index_origin,
),
dict_to_dataset(
data_warmup,
library=None,
coords=self.coords,
dims=self.pred_dims,
attrs=predictions_warmup_attrs,
index_origin=self.index_origin,
),
)
@requires("prior")
def prior_to_xarray(self):
"""Convert prior samples to xarray."""
data = self.prior
if not isinstance(data, dict):
raise TypeError("DictConverter.prior is not a dictionary")
prior_attrs = self._kwargs.get("prior_attrs")
return dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
attrs=prior_attrs,
index_origin=self.index_origin,
)
@requires("sample_stats_prior")
def sample_stats_prior_to_xarray(self):
"""Convert sample_stats_prior samples to xarray."""
data = self.sample_stats_prior
if not isinstance(data, dict):
raise TypeError("DictConverter.sample_stats_prior is not a dictionary")
sample_stats_prior_attrs = self._kwargs.get("sample_stats_prior_attrs")
return dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
attrs=sample_stats_prior_attrs,
index_origin=self.index_origin,
)
@requires("prior_predictive")
def prior_predictive_to_xarray(self):
"""Convert prior_predictive samples to xarray."""
data = self.prior_predictive
if not isinstance(data, dict):
raise TypeError("DictConverter.prior_predictive is not a dictionary")
prior_predictive_attrs = self._kwargs.get("prior_predictive_attrs")
return dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
attrs=prior_predictive_attrs,
index_origin=self.index_origin,
)
def data_to_xarray(self, data, group, dims=None):
"""Convert data to xarray."""
if not isinstance(data, dict):
raise TypeError(f"DictConverter.{group} is not a dictionary")
if dims is None:
dims = {} if self.dims is None else self.dims
return dict_to_dataset(
data,
library=None,
coords=self.coords,
dims=self.dims,
default_dims=[],
attrs=self.attrs,
index_origin=self.index_origin,
)
@requires("observed_data")
def observed_data_to_xarray(self):
"""Convert observed_data to xarray."""
return self.data_to_xarray(self.observed_data, group="observed_data", dims=self.dims)
@requires("constant_data")
def constant_data_to_xarray(self):
"""Convert constant_data to xarray."""
return self.data_to_xarray(self.constant_data, group="constant_data")
@requires("predictions_constant_data")
def predictions_constant_data_to_xarray(self):
"""Convert predictions_constant_data to xarray."""
return self.data_to_xarray(
self.predictions_constant_data, group="predictions_constant_data", dims=self.pred_dims
)
def to_inference_data(self):
"""Convert all available data to an InferenceData object.
Note that if groups can not be created, then the InferenceData
will not have those groups.
"""
return InferenceData(
**{
"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(),
"observed_data": self.observed_data_to_xarray(),
"constant_data": self.constant_data_to_xarray(),
"predictions_constant_data": self.predictions_constant_data_to_xarray(),
"save_warmup": self.save_warmup,
"attrs": self.attrs,
}
)
# pylint: disable=too-many-instance-attributes
[docs]
def from_dict(
posterior=None,
*,
posterior_predictive=None,
predictions=None,
sample_stats=None,
log_likelihood=None,
prior=None,
prior_predictive=None,
sample_stats_prior=None,
observed_data=None,
constant_data=None,
predictions_constant_data=None,
warmup_posterior=None,
warmup_posterior_predictive=None,
warmup_predictions=None,
warmup_log_likelihood=None,
warmup_sample_stats=None,
save_warmup=None,
index_origin: Optional[int] = None,
coords=None,
dims=None,
pred_dims=None,
pred_coords=None,
attrs=None,
**kwargs,
):
"""Convert Dictionary data into an InferenceData object.
For a usage example read the
:ref:`Creating InferenceData section on from_dict <creating_InferenceData>`
Parameters
----------
posterior : dict, optional
posterior_predictive : dict, optional
predictions: dict, optional
sample_stats : dict, optional
log_likelihood : dict, optional
For stats functions, log likelihood data should be stored here.
prior : dict, optional
prior_predictive : dict, optional
observed_data : dict, optional
constant_data : dict, optional
predictions_constant_data: dict, optional
warmup_posterior : dict, optional
warmup_posterior_predictive : dict, optional
warmup_predictions : dict, optional
warmup_log_likelihood : dict, optional
warmup_sample_stats : dict, optional
save_warmup : bool, optional
Save warmup iterations InferenceData object. If not defined, use default
defined by the rcParams.
index_origin : int, optional
coords : dict of {str : list}, optional
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 of {str : list of str}, optional
A mapping from variables to a list of coordinate names for the variable.
pred_dims : dict of {str : list of str}, optional
A mapping from variables to a list of coordinate names for predictions.
pred_coords : dict of {str : list}, optional
A mapping from variables to a list of coordinate values for predictions.
attrs : dict, optional
A dictionary containing attributes for different groups.
kwargs : dict, optional
A dictionary containing group attrs. Accepted kwargs are:
- posterior_attrs, posterior_warmup_attrs : attrs for posterior group
- sample_stats_attrs, sample_stats_warmup_attrs : attrs for sample_stats group
- log_likelihood_attrs, log_likelihood_warmup_attrs : attrs for log_likelihood group
- posterior_predictive_attrs, posterior_predictive_warmup_attrs : attrs for
posterior_predictive group
- predictions_attrs, predictions_warmup_attrs : attrs for predictions group
- prior_attrs : attrs for prior group
- sample_stats_prior_attrs : attrs for sample_stats_prior group
- prior_predictive_attrs : attrs for prior_predictive group
Returns
-------
InferenceData
"""
return DictConverter(
posterior=posterior,
posterior_predictive=posterior_predictive,
predictions=predictions,
sample_stats=sample_stats,
log_likelihood=log_likelihood,
prior=prior,
prior_predictive=prior_predictive,
sample_stats_prior=sample_stats_prior,
observed_data=observed_data,
constant_data=constant_data,
predictions_constant_data=predictions_constant_data,
warmup_posterior=warmup_posterior,
warmup_posterior_predictive=warmup_posterior_predictive,
warmup_predictions=warmup_predictions,
warmup_log_likelihood=warmup_log_likelihood,
warmup_sample_stats=warmup_sample_stats,
save_warmup=save_warmup,
index_origin=index_origin,
coords=coords,
dims=dims,
pred_dims=pred_dims,
pred_coords=pred_coords,
attrs=attrs,
**kwargs,
).to_inference_data()
from_pytree = from_dict
