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

arviz/__init__.py

@@ -1,6 +1,6 @@
 # pylint: disable=wildcard-import,invalid-name,wrong-import-position
 """ArviZ is a library for exploratory analysis of Bayesian models."""
-__version__ = "0.21.0"
+__version__ = "0.22.0"
 
 import logging
 import os
@@ -8,6 +8,7 @@ import os
 from matplotlib.colors import LinearSegmentedColormap
 from matplotlib.pyplot import style
 import matplotlib as mpl
+from packaging import version
 
 
 class Logger(logging.Logger):
@@ -41,8 +42,12 @@ from . import preview
 
 # add ArviZ's styles to matplotlib's styles
 _arviz_style_path = os.path.join(os.path.dirname(__file__), "plots", "styles")
-style.core.USER_LIBRARY_PATHS.append(_arviz_style_path)
-style.core.reload_library()
+if version.parse(mpl.__version__) >= version.parse("3.11.0.dev0"):
+    style.USER_LIBRARY_PATHS.append(_arviz_style_path)
+    style.reload_library()
+else:
+    style.core.USER_LIBRARY_PATHS.append(_arviz_style_path)
+    style.core.reload_library()
 
 
 if not logging.root.handlers:

arviz/data/inference_data.py

@@ -532,24 +532,27 @@ class InferenceData(Mapping[str, xr.Dataset]):
         return filename
 
     def to_datatree(self):
-        """Convert InferenceData object to a :class:`~datatree.DataTree`."""
+        """Convert InferenceData object to a :class:`~xarray.DataTree`."""
         try:
-            from datatree import DataTree
-        except ModuleNotFoundError as err:
-            raise ModuleNotFoundError(
-                "datatree must be installed in order to use InferenceData.to_datatree"
+            from xarray import DataTree
+        except ImportError as err:
+            raise ImportError(
+                "xarray must be have DataTree in order to use InferenceData.to_datatree. "
+                "Update to xarray>=2024.11.0"
             ) from err
         return DataTree.from_dict({group: ds for group, ds in self.items()})
 
     @staticmethod
     def from_datatree(datatree):
-        """Create an InferenceData object from a :class:`~datatree.DataTree`.
+        """Create an InferenceData object from a :class:`~xarray.DataTree`.
 
         Parameters
         ----------
         datatree : DataTree
         """
-        return InferenceData(**{group: sub_dt.to_dataset() for group, sub_dt in datatree.items()})
+        return InferenceData(
+            **{group: child.to_dataset() for group, child in datatree.children.items()}
+        )
 
     def to_dict(self, groups=None, filter_groups=None):
         """Convert InferenceData to a dictionary following xarray naming conventions.
@@ -797,12 +800,20 @@ class InferenceData(Mapping[str, xr.Dataset]):
         ----------
         https://zarr.readthedocs.io/
         """
-        try:  # Check zarr
+        try:
             import zarr
-
-            assert version.parse(zarr.__version__) >= version.parse("2.5.0")
-        except (ImportError, AssertionError) as err:
-            raise ImportError("'to_zarr' method needs Zarr (2.5.0+) installed.") from err
+        except ImportError as err:
+            raise ImportError("'to_zarr' method needs Zarr (>=2.5.0,<3) installed.") from err
+        if version.parse(zarr.__version__) < version.parse("2.5.0"):
+            raise ImportError(
+                "Found zarr<2.5.0, please upgrade to a zarr (>=2.5.0,<3) to use 'to_zarr'"
+            )
+        if version.parse(zarr.__version__) >= version.parse("3.0.0.dev0"):
+            raise ImportError(
+                "Found zarr>=3, which is not supported by ArviZ. Instead, you can use "
+                "'dt = InfereceData.to_datatree' followed by 'dt.to_zarr()' "
+                "(needs xarray>=2024.11.0)"
+            )
 
         # Check store type and create store if necessary
         if store is None:
@@ -851,10 +862,18 @@ class InferenceData(Mapping[str, xr.Dataset]):
         """
         try:
             import zarr
-
-            assert version.parse(zarr.__version__) >= version.parse("2.5.0")
-        except (ImportError, AssertionError) as err:
-            raise ImportError("'to_zarr' method needs Zarr (2.5.0+) installed.") from err
+        except ImportError as err:
+            raise ImportError("'from_zarr' method needs Zarr (>=2.5.0,<3) installed.") from err
+        if version.parse(zarr.__version__) < version.parse("2.5.0"):
+            raise ImportError(
+                "Found zarr<2.5.0, please upgrade to a zarr (>=2.5.0,<3) to use 'from_zarr'"
+            )
+        if version.parse(zarr.__version__) >= version.parse("3.0.0.dev0"):
+            raise ImportError(
+                "Found zarr>=3, which is not supported by ArviZ. Instead, you can use "
+                "'xarray.open_datatree' followed by 'arviz.InferenceData.from_datatree' "
+                "(needs xarray>=2024.11.0)"
+            )
 
         # Check store type and create store if necessary
         if isinstance(store, str):
@@ -1531,9 +1550,8 @@ class InferenceData(Mapping[str, xr.Dataset]):
             import xarray as xr
             from xarray_einstats.stats import XrDiscreteRV
             from scipy.stats import poisson
-            dist = XrDiscreteRV(poisson)
-            log_lik = xr.Dataset()
-            log_lik["home_points"] = dist.logpmf(obs["home_points"], np.exp(post["atts"]))
+            dist = XrDiscreteRV(poisson, np.exp(post["atts"]))
+            log_lik = dist.logpmf(obs["home_points"]).to_dataset(name="home_points")
             idata2.add_groups({"log_likelihood": log_lik})
             idata2
 

arviz/data/io_datatree.py

@@ -4,7 +4,7 @@ from .inference_data import InferenceData
 
 
 def to_datatree(data):
-    """Convert InferenceData object to a :class:`~datatree.DataTree`.
+    """Convert InferenceData object to a :class:`~xarray.DataTree`.
 
     Parameters
     ----------
@@ -14,7 +14,7 @@ def to_datatree(data):
 
 
 def from_datatree(datatree):
-    """Create an InferenceData object from a :class:`~datatree.DataTree`.
+    """Create an InferenceData object from a :class:`~xarray.DataTree`.
 
     Parameters
     ----------

arviz/data/io_numpyro.py

@@ -1,7 +1,8 @@
 """NumPyro-specific conversion code."""
 
+from collections import defaultdict
 import logging
-from typing import Callable, Optional
+from typing import Any, Callable, Optional, Dict, List, Tuple
 
 import numpy as np
 
@@ -13,6 +14,70 @@ from .inference_data import InferenceData
 _log = logging.getLogger(__name__)
 
 
+def _add_dims(dims_a: Dict[str, List[str]], dims_b: Dict[str, List[str]]) -> Dict[str, List[str]]:
+    merged = defaultdict(list)
+
+    for k, v in dims_a.items():
+        merged[k].extend(v)
+
+    for k, v in dims_b.items():
+        merged[k].extend(v)
+
+    # Convert back to a regular dict
+    return dict(merged)
+
+
+def infer_dims(
+    model: Callable,
+    model_args: Optional[Tuple[Any, ...]] = None,
+    model_kwargs: Optional[Dict[str, Any]] = None,
+) -> Dict[str, List[str]]:
+
+    from numpyro import handlers, distributions as dist
+    from numpyro.ops.pytree import PytreeTrace
+    from numpyro.infer.initialization import init_to_sample
+    import jax
+
+    model_args = tuple() if model_args is None else model_args
+    model_kwargs = dict() if model_args is None else model_kwargs
+
+    def _get_dist_name(fn):
+        if isinstance(fn, (dist.Independent, dist.ExpandedDistribution, dist.MaskedDistribution)):
+            return _get_dist_name(fn.base_dist)
+        return type(fn).__name__
+
+    def get_trace():
+        # We use `init_to_sample` to get around ImproperUniform distribution,
+        # which does not have `sample` method.
+        subs_model = handlers.substitute(
+            handlers.seed(model, 0),
+            substitute_fn=init_to_sample,
+        )
+        trace = handlers.trace(subs_model).get_trace(*model_args, **model_kwargs)
+        # Work around an issue where jax.eval_shape does not work
+        # for distribution output (e.g. the function `lambda: dist.Normal(0, 1)`)
+        # Here we will remove `fn` and store its name in the trace.
+        for _, site in trace.items():
+            if site["type"] == "sample":
+                site["fn_name"] = _get_dist_name(site.pop("fn"))
+            elif site["type"] == "deterministic":
+                site["fn_name"] = "Deterministic"
+        return PytreeTrace(trace)
+
+    # We use eval_shape to avoid any array computation.
+    trace = jax.eval_shape(get_trace).trace
+
+    named_dims = {}
+
+    for name, site in trace.items():
+        batch_dims = [frame.name for frame in sorted(site["cond_indep_stack"], key=lambda x: x.dim)]
+        event_dims = list(site.get("infer", {}).get("event_dims", []))
+        if site["type"] in ["sample", "deterministic"] and (batch_dims or event_dims):
+            named_dims[name] = batch_dims + event_dims
+
+    return named_dims
+
+
 class NumPyroConverter:
     """Encapsulate NumPyro specific logic."""
 
@@ -36,6 +101,7 @@ class NumPyroConverter:
         coords=None,
         dims=None,
         pred_dims=None,
+        extra_event_dims=None,
         num_chains=1,
     ):
         """Convert NumPyro data into an InferenceData object.
@@ -58,9 +124,12 @@ class NumPyroConverter:
         coords : dict[str] -> list[str]
             Map of dimensions to coordinates
         dims : dict[str] -> list[str]
-            Map variable names to their coordinates
+            Map variable names to their coordinates. Will be inferred if they are not provided.
         pred_dims: dict
             Dims for predictions data. Map variable names to their coordinates.
+        extra_event_dims: dict
+            Extra event dims for deterministic sites. Maps event dims that couldnt be inferred to
+            their coordinates.
         num_chains: int
             Number of chains used for sampling. Ignored if posterior is present.
         """
@@ -80,6 +149,7 @@ class NumPyroConverter:
         self.coords = coords
         self.dims = dims
         self.pred_dims = pred_dims
+        self.extra_event_dims = extra_event_dims
         self.numpyro = numpyro
 
         def arbitrary_element(dct):
@@ -107,6 +177,10 @@ class NumPyroConverter:
             # model arguments and keyword arguments
             self._args = self.posterior._args  # pylint: disable=protected-access
             self._kwargs = self.posterior._kwargs  # pylint: disable=protected-access
+            self.dims = self.dims if self.dims is not None else self.infer_dims()
+            self.pred_dims = (
+                self.pred_dims if self.pred_dims is not None else self.infer_pred_dims()
+            )
         else:
             self.nchains = num_chains
             get_from = None
@@ -325,6 +399,23 @@ class NumPyroConverter:
             }
         )
 
+    @requires("posterior")
+    @requires("model")
+    def infer_dims(self) -> Dict[str, List[str]]:
+        dims = infer_dims(self.model, self._args, self._kwargs)
+        if self.extra_event_dims:
+            dims = _add_dims(dims, self.extra_event_dims)
+        return dims
+
+    @requires("posterior")
+    @requires("model")
+    @requires("predictions")
+    def infer_pred_dims(self) -> Dict[str, List[str]]:
+        dims = infer_dims(self.model, self._args, self._kwargs)
+        if self.extra_event_dims:
+            dims = _add_dims(dims, self.extra_event_dims)
+        return dims
+
 
 def from_numpyro(
     posterior=None,
@@ -339,10 +430,25 @@ def from_numpyro(
     coords=None,
     dims=None,
     pred_dims=None,
+    extra_event_dims=None,
     num_chains=1,
 ):
     """Convert NumPyro data into an InferenceData object.
 
+    If no dims are provided, this will infer batch dim names from NumPyro model plates.
+    For event dim names, such as with the ZeroSumNormal, `infer={"event_dims":dim_names}`
+    can be provided in numpyro.sample, i.e.::
+
+        # equivalent to dims entry, {"gamma": ["groups"]}
+        gamma = numpyro.sample(
+            "gamma",
+            dist.ZeroSumNormal(1, event_shape=(n_groups,)),
+            infer={"event_dims":["groups"]}
+        )
+
+    There is also an additional `extra_event_dims` input to cover any edge cases, for instance
+    deterministic sites with event dims (which dont have an `infer` argument to provide metadata).
+
     For a usage example read the
     :ref:`Creating InferenceData section on from_numpyro <creating_InferenceData>`
 
@@ -364,9 +470,10 @@ def from_numpyro(
     coords : dict[str] -> list[str]
         Map of dimensions to coordinates
     dims : dict[str] -> list[str]
-        Map variable names to their coordinates
+        Map variable names to their coordinates. Will be inferred if they are not provided.
     pred_dims: dict
-        Dims for predictions data. Map variable names to their coordinates.
+        Dims for predictions data. Map variable names to their coordinates. Default behavior is to
+        infer dims if this is not provided
     num_chains: int
         Number of chains used for sampling. Ignored if posterior is present.
     """
@@ -382,5 +489,6 @@ def from_numpyro(
         coords=coords,
         dims=dims,
         pred_dims=pred_dims,
+        extra_event_dims=extra_event_dims,
         num_chains=num_chains,
     ).to_inference_data()

arviz/plots/autocorrplot.py

@@ -4,7 +4,7 @@ from ..data import convert_to_dataset
 from ..labels import BaseLabeller
 from ..sel_utils import xarray_var_iter
 from ..rcparams import rcParams
-from ..utils import _var_names
+from ..utils import _var_names, get_coords
 from .plot_utils import default_grid, filter_plotters_list, get_plotting_function
 
 
@@ -14,6 +14,7 @@ def plot_autocorr(
     filter_vars=None,
     max_lag=None,
     combined=False,
+    coords=None,
     grid=None,
     figsize=None,
     textsize=None,
@@ -42,6 +43,8 @@ def plot_autocorr(
         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.
+    coords: mapping, optional
+        Coordinates of var_names to be plotted. Passed to :meth:`xarray.Dataset.sel`
     max_lag : int, optional
         Maximum lag to calculate autocorrelation. By Default, the plot displays the
         first 100 lag or the total number of draws, whichever is smaller.
@@ -124,11 +127,18 @@ def plot_autocorr(
     if max_lag is None:
         max_lag = min(100, data["draw"].shape[0])
 
+    if coords is None:
+        coords = {}
+
     if labeller is None:
         labeller = BaseLabeller()
 
     plotters = filter_plotters_list(
-        list(xarray_var_iter(data, var_names, combined, dim_order=["chain", "draw"])),
+        list(
+            xarray_var_iter(
+                get_coords(data, coords), var_names, combined, dim_order=["chain", "draw"]
+            )
+        ),
         "plot_autocorr",
     )
     rows, cols = default_grid(len(plotters), grid=grid)

arviz/plots/backends/bokeh/hdiplot.py

@@ -21,9 +21,13 @@ def plot_hdi(ax, x_data, y_data, color, figsize, plot_kwargs, fill_kwargs, backe
     plot_kwargs["color"] = vectorized_to_hex(plot_kwargs.get("color", color))
     plot_kwargs.setdefault("alpha", 0)
 
-    fill_kwargs = {} if fill_kwargs is None else fill_kwargs
-    fill_kwargs["color"] = vectorized_to_hex(fill_kwargs.get("color", color))
-    fill_kwargs.setdefault("alpha", 0.5)
+    fill_kwargs = {} if fill_kwargs is None else fill_kwargs.copy()
+    # Convert matplotlib color to bokeh fill_color if needed
+    if "color" in fill_kwargs and "fill_color" not in fill_kwargs:
+        fill_kwargs["fill_color"] = vectorized_to_hex(fill_kwargs.pop("color"))
+    else:
+        fill_kwargs["fill_color"] = vectorized_to_hex(fill_kwargs.get("fill_color", color))
+    fill_kwargs.setdefault("fill_alpha", fill_kwargs.pop("alpha", 0.5))
 
     figsize, *_ = _scale_fig_size(figsize, None)
 
@@ -38,9 +42,6 @@ def plot_hdi(ax, x_data, y_data, color, figsize, plot_kwargs, fill_kwargs, backe
     plot_kwargs.setdefault("line_color", plot_kwargs.pop("color"))
     plot_kwargs.setdefault("line_alpha", plot_kwargs.pop("alpha", 0))
 
-    fill_kwargs.setdefault("fill_color", fill_kwargs.pop("color"))
-    fill_kwargs.setdefault("fill_alpha", fill_kwargs.pop("alpha", 0))
-
     ax.patch(
         np.concatenate((x_data, x_data[::-1])),
         np.concatenate((y_data[:, 0], y_data[:, 1][::-1])),

arviz/plots/backends/bokeh/lmplot.py

@@ -68,7 +68,13 @@ def plot_lm(
 
     if y_hat_fill_kwargs is None:
         y_hat_fill_kwargs = {}
-    y_hat_fill_kwargs.setdefault("color", "orange")
+    else:
+        y_hat_fill_kwargs = y_hat_fill_kwargs.copy()
+    # Convert matplotlib color to bokeh fill_color if needed
+    if "color" in y_hat_fill_kwargs and "fill_color" not in y_hat_fill_kwargs:
+        y_hat_fill_kwargs["fill_color"] = y_hat_fill_kwargs.pop("color")
+    y_hat_fill_kwargs.setdefault("fill_color", "orange")
+    y_hat_fill_kwargs.setdefault("fill_alpha", 0.5)
 
     if y_model_plot_kwargs is None:
         y_model_plot_kwargs = {}
@@ -78,8 +84,13 @@ def plot_lm(
 
     if y_model_fill_kwargs is None:
         y_model_fill_kwargs = {}
-    y_model_fill_kwargs.setdefault("color", "black")
-    y_model_fill_kwargs.setdefault("alpha", 0.5)
+    else:
+        y_model_fill_kwargs = y_model_fill_kwargs.copy()
+    # Convert matplotlib color to bokeh fill_color if needed
+    if "color" in y_model_fill_kwargs and "fill_color" not in y_model_fill_kwargs:
+        y_model_fill_kwargs["fill_color"] = y_model_fill_kwargs.pop("color")
+    y_model_fill_kwargs.setdefault("fill_color", "black")
+    y_model_fill_kwargs.setdefault("fill_alpha", 0.5)
 
     if y_model_mean_kwargs is None:
         y_model_mean_kwargs = {}
@@ -149,6 +160,11 @@ def plot_lm(
                 )
 
                 y_model_mean = np.mean(y_model_plotters, axis=(0, 1))
+                # Plot mean line across all x values instead of just edges
+                mean_legend = ax_i.line(x_plotters, y_model_mean, **y_model_mean_kwargs)
+                legend_it.append(("Mean", [mean_legend]))
+                continue  # Skip the edge plotting since we plotted full line
+
             x_plotters_edge = [min(x_plotters), max(x_plotters)]
             y_model_mean_edge = [min(y_model_mean), max(y_model_mean)]
             mean_legend = ax_i.line(x_plotters_edge, y_model_mean_edge, **y_model_mean_kwargs)

arviz/plots/backends/bokeh/pairplot.py

@@ -37,6 +37,8 @@ def plot_pair(
     diverging_mask,
     divergences_kwargs,
     flat_var_names,
+    flat_ref_slices,
+    flat_var_labels,
     backend_kwargs,
     marginal_kwargs,
     show,
@@ -72,50 +74,12 @@ def plot_pair(
         kde_kwargs["contour_kwargs"].setdefault("line_alpha", 1)
 
     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))
-
-        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)
-                ),
-                UserWarning,
-            )
-
-    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))
-
-        for dif in difference:
-            reference_values_copy[dif] = None
+        difference = set(flat_var_names).difference(set(reference_values.keys()))
 
         if difference:
-            warn = [dif.replace("\n", " ", 1) for dif in difference]
             warnings.warn(
                 "Argument reference_values does not include reference value for: {}".format(
-                    ", ".join(warn)
+                    ", ".join(difference)
                 ),
                 UserWarning,
             )
@@ -262,8 +226,8 @@ def plot_pair(
                     **marginal_kwargs,
                 )
 
-                ax[j, i].xaxis.axis_label = flat_var_names[i]
-                ax[j, i].yaxis.axis_label = flat_var_names[j + marginals_offset]
+                ax[j, i].xaxis.axis_label = flat_var_labels[i]
+                ax[j, i].yaxis.axis_label = flat_var_labels[j + marginals_offset]
 
             elif j + marginals_offset > i:
                 if "scatter" in kind:
@@ -346,12 +310,18 @@ def plot_pair(
                             ax[-1, -1].add_layout(ax_pe_hline)
 
                 if reference_values:
-                    x = reference_values_copy[flat_var_names[j + marginals_offset]]
-                    y = reference_values_copy[flat_var_names[i]]
-                    if x and y:
-                        ax[j, i].scatter(y, x, **reference_values_kwargs)
-                ax[j, i].xaxis.axis_label = flat_var_names[i]
-                ax[j, i].yaxis.axis_label = flat_var_names[j + marginals_offset]
+                    x_name = flat_var_names[j + marginals_offset]
+                    y_name = flat_var_names[i]
+                    if (x_name not in difference) and (y_name not in difference):
+                        ax[j, i].scatter(
+                            np.array(reference_values[y_name])[flat_ref_slices[i]],
+                            np.array(reference_values[x_name])[
+                                flat_ref_slices[j + marginals_offset]
+                            ],
+                            **reference_values_kwargs,
+                        )
+                ax[j, i].xaxis.axis_label = flat_var_labels[i]
+                ax[j, i].yaxis.axis_label = flat_var_labels[j + marginals_offset]
 
     show_layout(ax, show)
 

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")