arviz 0.21.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

arviz/plots/backends/matplotlib/khatplot.py

@@ -7,6 +7,7 @@ from matplotlib import cm
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.colors import to_rgba_array
+from packaging import version
 
 from ....stats.density_utils import histogram
 from ...plot_utils import _scale_fig_size, color_from_dim, set_xticklabels, vectorized_to_hex
@@ -39,7 +40,13 @@ def plot_khat(
     show,
 ):
     """Matplotlib khat plot."""
-    if hover_label and mpl.get_backend() not in mpl.rcsetup.interactive_bk:
+    if version.parse(mpl.__version__) >= version.parse("3.9.0.dev0"):
+        interactive_backends = mpl.backends.backend_registry.list_builtin(
+            mpl.backends.BackendFilter.INTERACTIVE
+        )
+    else:
+        interactive_backends = mpl.rcsetup.interactive_bk
+    if hover_label and mpl.get_backend() not in interactive_backends:
         hover_label = False
         warnings.warn(
             "hover labels are only available with interactive backends. To switch to an "

arviz/plots/backends/matplotlib/lmplot.py

@@ -115,12 +115,18 @@ def plot_lm(
 
         if y_model is not None:
             _, _, _, y_model_plotters = y_model[i]
+
             if kind_model == "lines":
-                for j in range(num_samples):
-                    ax_i.plot(x_plotters, y_model_plotters[..., j], **y_model_plot_kwargs)
-                ax_i.plot([], **y_model_plot_kwargs, label="Uncertainty in mean")
+                # y_model_plotters should be (points, samples)
+                y_points = y_model_plotters.shape[0]
+                if x_plotters.shape[0] == y_points:
+                    for j in range(num_samples):
+                        ax_i.plot(x_plotters, y_model_plotters[:, j], **y_model_plot_kwargs)
+
+                    ax_i.plot([], **y_model_plot_kwargs, label="Uncertainty in mean")
+                    y_model_mean = np.mean(y_model_plotters, axis=1)
+                    ax_i.plot(x_plotters, y_model_mean, **y_model_mean_kwargs, label="Mean")
 
-                y_model_mean = np.mean(y_model_plotters, axis=1)
             else:
                 plot_hdi(
                     x_plotters,
@@ -128,10 +134,10 @@ def plot_lm(
                     fill_kwargs=y_model_fill_kwargs,
                     ax=ax_i,
                 )
-                ax_i.plot([], color=y_model_fill_kwargs["color"], label="Uncertainty in mean")
 
-                y_model_mean = np.mean(y_model_plotters, axis=(0, 1))
-            ax_i.plot(x_plotters, y_model_mean, **y_model_mean_kwargs, label="Mean")
+                ax_i.plot([], color=y_model_fill_kwargs["color"], label="Uncertainty in mean")
+                y_model_mean = np.mean(y_model_plotters, axis=0)
+                ax_i.plot(x_plotters, y_model_mean, **y_model_mean_kwargs, label="Mean")
 
         if legend:
             ax_i.legend(fontsize=xt_labelsize, loc="upper left")

arviz/plots/backends/matplotlib/pairplot.py

@@ -30,6 +30,8 @@ def plot_pair(
     diverging_mask,
     divergences_kwargs,
     flat_var_names,
+    flat_ref_slices,
+    flat_var_labels,
     backend_kwargs,
     marginal_kwargs,
     show,
@@ -77,24 +79,12 @@ def plot_pair(
         kde_kwargs["contour_kwargs"].setdefault("colors", "k")
 
     if reference_values:
-        reference_values_copy = {}
-        label = []
-        for variable in list(reference_values.keys()):
-            if " " in variable:
-                variable_copy = variable.replace(" ", "\n", 1)
-            else:
-                variable_copy = variable
-
-            label.append(variable_copy)
-            reference_values_copy[variable_copy] = reference_values[variable]
-
-        difference = set(flat_var_names).difference(set(label))
+        difference = set(flat_var_names).difference(set(reference_values.keys()))
 
         if difference:
-            warn = [diff.replace("\n", " ", 1) for diff in difference]
             warnings.warn(
                 "Argument reference_values does not include reference value for: {}".format(
-                    ", ".join(warn)
+                    ", ".join(difference)
                 ),
                 UserWarning,
             )
@@ -211,12 +201,12 @@ def plot_pair(
 
         if reference_values:
             ax.plot(
-                reference_values_copy[flat_var_names[0]],
-                reference_values_copy[flat_var_names[1]],
+                np.array(reference_values[flat_var_names[0]])[flat_ref_slices[0]],
+                np.array(reference_values[flat_var_names[1]])[flat_ref_slices[1]],
                 **reference_values_kwargs,
             )
-        ax.set_xlabel(f"{flat_var_names[0]}", fontsize=ax_labelsize, wrap=True)
-        ax.set_ylabel(f"{flat_var_names[1]}", fontsize=ax_labelsize, wrap=True)
+        ax.set_xlabel(f"{flat_var_labels[0]}", fontsize=ax_labelsize, wrap=True)
+        ax.set_ylabel(f"{flat_var_labels[1]}", fontsize=ax_labelsize, wrap=True)
         ax.tick_params(labelsize=xt_labelsize)
 
     else:
@@ -336,20 +326,22 @@ def plot_pair(
                         y_name = flat_var_names[j + not_marginals]
                         if (x_name not in difference) and (y_name not in difference):
                             ax[j, i].plot(
-                                reference_values_copy[x_name],
-                                reference_values_copy[y_name],
+                                np.array(reference_values[x_name])[flat_ref_slices[i]],
+                                np.array(reference_values[y_name])[
+                                    flat_ref_slices[j + not_marginals]
+                                ],
                                 **reference_values_kwargs,
                             )
 
                 if j != vars_to_plot - 1:
                     plt.setp(ax[j, i].get_xticklabels(), visible=False)
                 else:
-                    ax[j, i].set_xlabel(f"{flat_var_names[i]}", fontsize=ax_labelsize, wrap=True)
+                    ax[j, i].set_xlabel(f"{flat_var_labels[i]}", fontsize=ax_labelsize, wrap=True)
                 if i != 0:
                     plt.setp(ax[j, i].get_yticklabels(), visible=False)
                 else:
                     ax[j, i].set_ylabel(
-                        f"{flat_var_names[j + not_marginals]}",
+                        f"{flat_var_labels[j + not_marginals]}",
                         fontsize=ax_labelsize,
                         wrap=True,
                     )

arviz/plots/bpvplot.py

@@ -251,7 +251,7 @@ def plot_bpv(
     total_pp_samples = predictive_dataset.sizes["chain"] * predictive_dataset.sizes["draw"]
 
     for key in coords.keys():
-        coords[key] = np.where(np.in1d(observed[key], coords[key]))[0]
+        coords[key] = np.where(np.isin(observed[key], coords[key]))[0]
 
     obs_plotters = filter_plotters_list(
         list(

arviz/plots/dotplot.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 
+
 from ..rcparams import rcParams
 from .plot_utils import get_plotting_function
 
@@ -148,6 +149,7 @@ def plot_dot(
         raise ValueError("marker argument is valid only for matplotlib backend")
 
     values = np.ravel(values)
+    values = values[np.isfinite(values)]
     values.sort()
 
     if hdi_prob is None:

arviz/plots/forestplot.py

@@ -51,7 +51,7 @@ def plot_forest(
     data : InferenceData
         Any object that can be converted to an :class:`arviz.InferenceData` object
         Refer to documentation of :func:`arviz.convert_to_dataset` for details.
-    kind : {"foresplot", "ridgeplot"}, default "forestplot"
+    kind : {"forestplot", "ridgeplot"}, default "forestplot"
         Specify the kind of plot:
 
         * The ``kind="forestplot"`` generates credible intervals, where the central points are the
@@ -75,8 +75,8 @@ def plot_forest(
         interpret `var_names` as substrings of the real variables names. If "regex",
         interpret `var_names` as regular expressions on the real variables names. See
         :ref:`this section <common_filter_vars>` for usage examples.
-    transform : callable, optional
-        Function to transform data (defaults to None i.e.the identity function).
+    transform : callable or dict, optional
+        Function to transform the data. Defaults to None, i.e., the identity function.
     coords : dict, optional
         Coordinates of ``var_names`` to be plotted. Passed to :meth:`xarray.Dataset.sel`.
         See :ref:`this section <common_coords>` for usage examples.
@@ -228,7 +228,19 @@ def plot_forest(
 
     datasets = [convert_to_dataset(datum) for datum in reversed(data)]
     if transform is not None:
-        datasets = [transform(dataset) for dataset in datasets]
+        if callable(transform):
+            datasets = [transform(dataset) for dataset in datasets]
+        elif isinstance(transform, dict):
+            transformed_datasets = []
+            for dataset in datasets:
+                new_dataset = dataset.copy()
+                for var_name, func in transform.items():
+                    if var_name in new_dataset:
+                        new_dataset[var_name] = func(new_dataset[var_name])
+                transformed_datasets.append(new_dataset)
+            datasets = transformed_datasets
+        else:
+            raise ValueError("transform must be either a callable or a dict {var_name: callable}")
     datasets = get_coords(
         datasets, list(reversed(coords)) if isinstance(coords, (list, tuple)) else coords
     )

arviz/plots/kdeplot.py

@@ -255,6 +255,10 @@ def plot_kde(
             "or plot_pair instead of plot_kde"
         )
 
+    if backend is None:
+        backend = rcParams["plot.backend"]
+    backend = backend.lower()
+
     if values2 is None:
         if bw == "default":
             bw = "taylor" if is_circular else "experimental"
@@ -346,10 +350,6 @@ def plot_kde(
         **kwargs,
     )
 
-    if backend is None:
-        backend = rcParams["plot.backend"]
-    backend = backend.lower()
-
     # TODO: Add backend kwargs
     plot = get_plotting_function("plot_kde", "kdeplot", backend)
     ax = plot(**kde_plot_args)

arviz/plots/lmplot.py

@@ -300,20 +300,47 @@ def plot_lm(
     # Filter out the required values to generate plotters
     if y_model is not None:
         if kind_model == "lines":
-            y_model = y_model.stack(__sample__=("chain", "draw"))[..., pp_sample_ix]
-
-        y_model = [
-            tup
-            for _, tup in zip(
-                range(len_y),
-                xarray_var_iter(
-                    y_model,
-                    skip_dims=set(y_model.dims),
-                    combined=True,
-                ),
-            )
-        ]
-        y_model = _repeat_flatten_list(y_model, len_x)
+            var_name = y_model.name if y_model.name else "y_model"
+            data = y_model.values
+
+            total_samples = data.shape[0] * data.shape[1]
+            data = data.reshape(total_samples, *data.shape[2:])
+
+            if pp_sample_ix is not None:
+                data = data[pp_sample_ix]
+
+            if plot_dim is not None:
+                # For plot_dim case, transpose to get dimension first
+                data = data.transpose(1, 0, 2)[..., 0]
+
+            # Create plotter tuple(s)
+            if plot_dim is not None:
+                y_model = [(var_name, {}, {}, data) for _ in range(length_plotters)]
+            else:
+                y_model = [(var_name, {}, {}, data)]
+                y_model = _repeat_flatten_list(y_model, len_x)
+
+        elif kind_model == "hdi":
+            var_name = y_model.name if y_model.name else "y_model"
+            data = y_model.values
+
+            if plot_dim is not None:
+                # First transpose to get plot_dim first
+                data = data.transpose(2, 0, 1, 3)
+                # For plot_dim case, we just want HDI for first dimension
+                data = data[..., 0]
+
+                # Reshape to (samples, points)
+                data = data.transpose(1, 2, 0).reshape(-1, data.shape[0])
+                y_model = [(var_name, {}, {}, data) for _ in range(length_plotters)]
+
+            else:
+                data = data.reshape(-1, data.shape[-1])
+                y_model = [(var_name, {}, {}, data)]
+                y_model = _repeat_flatten_list(y_model, len_x)
+
+        if len(y_model) == 1:
+            y_model = _repeat_flatten_list(y_model, len_x)
 
     rows, cols = default_grid(length_plotters)
 

arviz/plots/pairplot.py

@@ -196,9 +196,14 @@ def plot_pair(
             get_coords(dataset, coords), var_names=var_names, skip_dims=combine_dims, combined=True
         )
     )
-    flat_var_names = [
-        labeller.make_label_vert(var_name, sel, isel) for var_name, sel, isel, _ in plotters
-    ]
+    flat_var_names = []
+    flat_ref_slices = []
+    flat_var_labels = []
+    for var_name, sel, isel, _ in plotters:
+        dims = [dim for dim in dataset[var_name].dims if dim not in ["chain", "draw"]]
+        flat_var_names.append(var_name)
+        flat_ref_slices.append(tuple(isel[dim] if dim in isel else slice(None) for dim in dims))
+        flat_var_labels.append(labeller.make_label_vert(var_name, sel, isel))
 
     divergent_data = None
     diverging_mask = None
@@ -253,6 +258,8 @@ def plot_pair(
         diverging_mask=diverging_mask,
         divergences_kwargs=divergences_kwargs,
         flat_var_names=flat_var_names,
+        flat_ref_slices=flat_ref_slices,
+        flat_var_labels=flat_var_labels,
         backend_kwargs=backend_kwargs,
         marginal_kwargs=marginal_kwargs,
         show=show,

arviz/plots/ppcplot.py

@@ -304,7 +304,7 @@ def plot_ppc(
     pp_sample_ix = np.random.choice(total_pp_samples, size=num_pp_samples, replace=False)
 
     for key in coords.keys():
-        coords[key] = np.where(np.in1d(observed_data[key], coords[key]))[0]
+        coords[key] = np.where(np.isin(observed_data[key], coords[key]))[0]
 
     obs_plotters = filter_plotters_list(
         list(

arviz/preview.py

@@ -8,41 +8,51 @@ info = ""
 
 try:
     from arviz_base import *
+    import arviz_base as base
 
-    status = "arviz_base available, exposing its functions as part of arviz.preview"
-    _log.info(status)
+    _status = "arviz_base available, exposing its functions as part of arviz.preview"
+    _log.info(_status)
 except ModuleNotFoundError:
-    status = "arviz_base not installed"
-    _log.info(status)
+    _status = "arviz_base not installed"
+    _log.info(_status)
 except ImportError:
-    status = "Unable to import arviz_base"
-    _log.info(status, exc_info=True)
+    _status = "Unable to import arviz_base"
+    _log.info(_status, exc_info=True)
 
-info += status + "\n"
+info += _status + "\n"
 
 try:
     from arviz_stats import *
 
-    status = "arviz_stats available, exposing its functions as part of arviz.preview"
-    _log.info(status)
+    # the base computational module fron arviz_stats will override the alias to arviz-base
+    # arviz.stats.base will still be available
+    import arviz_base as base
+    import arviz_stats as stats
+
+    _status = "arviz_stats available, exposing its functions as part of arviz.preview"
+    _log.info(_status)
 except ModuleNotFoundError:
-    status = "arviz_stats not installed"
-    _log.info(status)
+    _status = "arviz_stats not installed"
+    _log.info(_status)
 except ImportError:
-    status = "Unable to import arviz_stats"
-    _log.info(status, exc_info=True)
-info += status + "\n"
+    _status = "Unable to import arviz_stats"
+    _log.info(_status, exc_info=True)
+info += _status + "\n"
 
 try:
     from arviz_plots import *
+    import arviz_plots as plots
 
-    status = "arviz_plots available, exposing its functions as part of arviz.preview"
-    _log.info(status)
+    _status = "arviz_plots available, exposing its functions as part of arviz.preview"
+    _log.info(_status)
 except ModuleNotFoundError:
-    status = "arviz_plots not installed"
-    _log.info(status)
+    _status = "arviz_plots not installed"
+    _log.info(_status)
 except ImportError:
-    status = "Unable to import arviz_plots"
-    _log.info(status, exc_info=True)
+    _status = "Unable to import arviz_plots"
+    _log.info(_status, exc_info=True)
+
+info += _status + "\n"
 
-info += status + "\n"
+# clean namespace
+del logging, _status, _log

arviz/rcparams.py

@@ -12,11 +12,11 @@ from pathlib import Path
 from typing import Any, Dict
 from typing_extensions import Literal
 
-NO_GET_ARGS: bool = False
+NO_GET_ARGS: bool = False  # pylint: disable=invalid-name
 try:
     from typing_extensions import get_args
 except ImportError:
-    NO_GET_ARGS = True
+    NO_GET_ARGS = True  # pylint: disable=invalid-name
 
 
 import numpy as np

arviz/stats/density_utils.py

@@ -635,7 +635,7 @@ def _kde_circular(
     cumulative: bool, optional
         Whether return the PDF or the cumulative PDF. Defaults to False.
     grid_len: int, optional
-        The number of intervals used to bin the data pointa i.e. the length of the grid used in the
+        The number of intervals used to bin the data point i.e. the length of the grid used in the
         estimation. Defaults to 512.
     """
     # All values between -pi and pi

arviz/stats/stats.py

@@ -1,7 +1,6 @@
 # pylint: disable=too-many-lines
 """Statistical functions in ArviZ."""
 
-import itertools
 import warnings
 from copy import deepcopy
 from typing import List, Optional, Tuple, Union, Mapping, cast, Callable
@@ -11,14 +10,14 @@ import pandas as pd
 import scipy.stats as st
 from xarray_einstats import stats
 import xarray as xr
-from scipy.optimize import minimize
+from scipy.optimize import minimize, LinearConstraint, Bounds
 from typing_extensions import Literal
 
-NO_GET_ARGS: bool = False
+NO_GET_ARGS: bool = False  # pylint: disable=invalid-name
 try:
     from typing_extensions import get_args
 except ImportError:
-    NO_GET_ARGS = True
+    NO_GET_ARGS = True  # pylint: disable=invalid-name
 
 from .. import _log
 from ..data import InferenceData, convert_to_dataset, convert_to_inference_data, extract
@@ -225,37 +224,23 @@ def compare(
     if method.lower() == "stacking":
         rows, cols, ic_i_val = _ic_matrix(ics, ic_i)
         exp_ic_i = np.exp(ic_i_val / scale_value)
-        km1 = cols - 1
-
-        def w_fuller(weights):
-            return np.concatenate((weights, [max(1.0 - np.sum(weights), 0.0)]))
 
         def log_score(weights):
-            w_full = w_fuller(weights)
-            score = 0.0
-            for i in range(rows):
-                score += np.log(np.dot(exp_ic_i[i], w_full))
-            return -score
+            return -np.sum(np.log(exp_ic_i @ weights))
 
         def gradient(weights):
-            w_full = w_fuller(weights)
-            grad = np.zeros(km1)
-            for k, i in itertools.product(range(km1), range(rows)):
-                grad[k] += (exp_ic_i[i, k] - exp_ic_i[i, km1]) / np.dot(exp_ic_i[i], w_full)
-            return -grad
-
-        theta = np.full(km1, 1.0 / cols)
-        bounds = [(0.0, 1.0) for _ in range(km1)]
-        constraints = [
-            {"type": "ineq", "fun": lambda x: -np.sum(x) + 1.0},
-            {"type": "ineq", "fun": np.sum},
-        ]
+            denominator = exp_ic_i @ weights
+            return -np.sum(exp_ic_i / denominator[:, np.newaxis], axis=0)
 
-        weights = minimize(
+        theta = np.full(cols, 1.0 / cols)
+        bounds = Bounds(lb=np.zeros(cols), ub=np.ones(cols))
+        constraints = LinearConstraint(np.ones(cols), lb=1.0, ub=1.0)
+
+        minimize_result = minimize(
             fun=log_score, x0=theta, jac=gradient, bounds=bounds, constraints=constraints
         )
 
-        weights = w_fuller(weights["x"])
+        weights = minimize_result["x"]
         ses = ics["se"]
 
     elif method.lower() == "bb-pseudo-bma":
@@ -869,7 +854,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
     )
 
 
-def psislw(log_weights, reff=1.0):
+def psislw(log_weights, reff=1.0, normalize=True):
     """
     Pareto smoothed importance sampling (PSIS).
 
@@ -887,11 +872,13 @@ def psislw(log_weights, reff=1.0):
         Array of size (n_observations, n_samples)
     reff : float, default 1
         relative MCMC efficiency, ``ess / n``
+    normalize : bool, default True
+        return normalized log weights
 
     Returns
     -------
     lw_out : DataArray or (..., N) ndarray
-        Smoothed, truncated and normalized log weights.
+        Smoothed, truncated and possibly normalized log weights.
     kss : DataArray or (...) ndarray
         Estimates of the shape parameter *k* of the generalized Pareto
         distribution.
@@ -936,7 +923,12 @@ def psislw(log_weights, reff=1.0):
     out = np.empty_like(log_weights), np.empty(shape)
 
     # define kwargs
-    func_kwargs = {"cutoff_ind": cutoff_ind, "cutoffmin": cutoffmin, "out": out}
+    func_kwargs = {
+        "cutoff_ind": cutoff_ind,
+        "cutoffmin": cutoffmin,
+        "out": out,
+        "normalize": normalize,
+    }
     ufunc_kwargs = {"n_dims": 1, "n_output": 2, "ravel": False, "check_shape": False}
     kwargs = {"input_core_dims": [["__sample__"]], "output_core_dims": [["__sample__"], []]}
     log_weights, pareto_shape = _wrap_xarray_ufunc(
@@ -953,7 +945,7 @@ def psislw(log_weights, reff=1.0):
     return log_weights, pareto_shape
 
 
-def _psislw(log_weights, cutoff_ind, cutoffmin):
+def _psislw(log_weights, cutoff_ind, cutoffmin, normalize):
     """
     Pareto smoothed importance sampling (PSIS) for a 1D vector.
 
@@ -963,7 +955,7 @@ def _psislw(log_weights, cutoff_ind, cutoffmin):
         Array of length n_observations
     cutoff_ind: int
     cutoffmin: float
-    k_min: float
+    normalize: bool
 
     Returns
     -------
@@ -975,7 +967,8 @@ def _psislw(log_weights, cutoff_ind, cutoffmin):
     x = np.asarray(log_weights)
 
     # improve numerical accuracy
-    x -= np.max(x)
+    max_x = np.max(x)
+    x -= max_x
     # sort the array
     x_sort_ind = np.argsort(x)
     # divide log weights into body and right tail
@@ -1007,8 +1000,12 @@ def _psislw(log_weights, cutoff_ind, cutoffmin):
             x[tailinds[x_tail_si]] = smoothed_tail
             # truncate smoothed values to the largest raw weight 0
             x[x > 0] = 0
+
     # renormalize weights
-    x -= _logsumexp(x)
+    if normalize:
+        x -= _logsumexp(x)
+    else:
+        x += max_x
 
     return x, k
 

arviz/tests/base_tests/test_data.py

@@ -1501,10 +1501,6 @@ class TestJSON:
         assert not os.path.exists(filepath)
 
 
-@pytest.mark.skipif(
-    not (importlib.util.find_spec("datatree") or "ARVIZ_REQUIRE_ALL_DEPS" in os.environ),
-    reason="test requires xarray-datatree library",
-)
 class TestDataTree:
     def test_datatree(self):
         idata = load_arviz_data("centered_eight")
@@ -1514,6 +1510,15 @@ class TestDataTree:
             assert_identical(ds, idata_back[group])
         assert all(group in dt.children for group in idata.groups())
 
+    def test_datatree_attrs(self):
+        idata = load_arviz_data("centered_eight")
+        idata.attrs = {"not": "empty"}
+        assert idata.attrs
+        dt = idata.to_datatree()
+        idata_back = from_datatree(dt)
+        assert dt.attrs == idata.attrs
+        assert idata_back.attrs == idata.attrs
+
 
 class TestConversions:
     def test_id_conversion_idempotent(self):
@@ -1656,3 +1661,19 @@ class TestExtractDataset:
         post = extract(idata, num_samples=10)
         assert post.sizes["sample"] == 10
         assert post.attrs == idata.posterior.attrs
+
+
+def test_convert_to_inference_data_with_array_like():
+    class ArrayLike:
+        def __init__(self, data):
+            self._data = np.asarray(data)
+
+        def __array__(self):
+            return self._data
+
+    array_like = ArrayLike(np.random.randn(4, 100))
+    idata = convert_to_inference_data(array_like, group="posterior")
+
+    assert hasattr(idata, "posterior")
+    assert "x" in idata.posterior.data_vars
+    assert idata.posterior["x"].shape == (4, 100)