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

arviz/data/io_datatree.py

@@ -1,23 +0,0 @@
-"""Conversion between InferenceData and DataTree."""
-
-from .inference_data import InferenceData
-
-
-def to_datatree(data):
-    """Convert InferenceData object to a :class:`~xarray.DataTree`.
-
-    Parameters
-    ----------
-    data : InferenceData
-    """
-    return data.to_datatree()
-
-
-def from_datatree(datatree):
-    """Create an InferenceData object from a :class:`~xarray.DataTree`.
-
-    Parameters
-    ----------
-    datatree : DataTree
-    """
-    return InferenceData.from_datatree(datatree)

arviz/data/io_dict.py

@@ -1,462 +0,0 @@
-"""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
-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