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

arviz/tests/base_tests/test_plots_matplotlib.py

@@ -2,21 +2,23 @@
 
 # pylint: disable=redefined-outer-name,too-many-lines
 import os
+import re
 from copy import deepcopy
 
 import matplotlib.pyplot as plt
 import numpy as np
 import pytest
+import xarray as xr
 from matplotlib import animation
 from pandas import DataFrame
 from scipy.stats import gaussian_kde, norm
-import xarray as xr
 
 from ...data import from_dict, load_arviz_data
+from ...labels import MapLabeller
 from ...plots import (
     plot_autocorr,
-    plot_bpv,
     plot_bf,
+    plot_bpv,
     plot_compare,
     plot_density,
     plot_dist,
@@ -43,20 +45,20 @@ from ...plots import (
     plot_ts,
     plot_violin,
 )
+from ...plots.dotplot import wilkinson_algorithm
+from ...plots.plot_utils import plot_point_interval
 from ...rcparams import rc_context, rcParams
 from ...stats import compare, hdi, loo, waic
 from ...stats.density_utils import kde as _kde
-from ...utils import _cov, BehaviourChangeWarning
-from ...plots.plot_utils import plot_point_interval
-from ...plots.dotplot import wilkinson_algorithm
+from ...utils import BehaviourChangeWarning, _cov
 from ..helpers import (  # pylint: disable=unused-import
+    RandomVariableTestClass,
     create_model,
     create_multidimensional_model,
     does_not_warn,
     eight_schools_params,
     models,
     multidim_models,
-    RandomVariableTestClass,
 )
 
 rcParams["data.load"] = "eager"
@@ -598,6 +600,21 @@ def test_plot_kde_inference_data(models):
             "reference_values": {"mu": 0, "tau": 0},
             "reference_values_kwargs": {"c": "C0", "marker": "*"},
         },
+        {
+            "var_names": ["mu", "tau"],
+            "reference_values": {"mu": 0, "tau": 0},
+            "labeller": MapLabeller({"mu": r"$\mu$", "theta": r"$\theta"}),
+        },
+        {
+            "var_names": ["theta"],
+            "reference_values": {"theta": [0.0] * 8},
+            "labeller": MapLabeller({"theta": r"$\theta$"}),
+        },
+        {
+            "var_names": ["theta"],
+            "reference_values": {"theta": np.zeros(8)},
+            "labeller": MapLabeller({"theta": r"$\theta$"}),
+        },
     ],
 )
 def test_plot_pair(models, kwargs):
@@ -1236,6 +1253,23 @@ def test_plot_hdi_dataset_error(models):
         plot_hdi(np.arange(8), hdi_data=hdi_data)
 
 
+def test_plot_hdi_string_error():
+    """Check x as type string raises an error."""
+    x_data = ["a", "b", "c", "d"]
+    y_data = np.random.normal(0, 5, (1, 200, len(x_data)))
+    hdi_data = hdi(y_data)
+    with pytest.raises(
+        NotImplementedError,
+        match=re.escape(
+            (
+                "The `arviz.plot_hdi()` function does not support categorical data. "
+                "Consider using `arviz.plot_forest()`."
+            )
+        ),
+    ):
+        plot_hdi(x=x_data, y=y_data, hdi_data=hdi_data)
+
+
 def test_plot_hdi_datetime_error():
     """Check x as datetime raises an error."""
     x_data = np.arange(start="2022-01-01", stop="2022-03-01", dtype=np.datetime64)
@@ -1896,7 +1930,7 @@ def test_wilkinson_algorithm(continuous_model):
         },
     ],
 )
-def test_plot_lm(models, kwargs):
+def test_plot_lm_1d(models, kwargs):
     """Test functionality for 1D data."""
     idata = models.model_1
     if "constant_data" not in idata.groups():
@@ -2082,7 +2116,65 @@ def test_plot_bf():
     idata = from_dict(
         posterior={"a": np.random.normal(1, 0.5, 5000)}, prior={"a": np.random.normal(0, 1, 5000)}
     )
-    bf_dict0, _ = plot_bf(idata, var_name="a", ref_val=0)
-    bf_dict1, _ = plot_bf(idata, prior=np.random.normal(0, 10, 5000), var_name="a", ref_val=0)
-    assert bf_dict0["BF10"] > bf_dict0["BF01"]
-    assert bf_dict1["BF10"] < bf_dict1["BF01"]
+    _, bf_plot = plot_bf(idata, var_name="a", ref_val=0)
+    assert bf_plot is not None
+
+
+def generate_lm_1d_data():
+    rng = np.random.default_rng()
+    return from_dict(
+        observed_data={"y": rng.normal(size=7)},
+        posterior_predictive={"y": rng.normal(size=(4, 1000, 7)) / 2},
+        posterior={"y_model": rng.normal(size=(4, 1000, 7))},
+        dims={"y": ["dim1"]},
+        coords={"dim1": range(7)},
+    )
+
+
+def generate_lm_2d_data():
+    rng = np.random.default_rng()
+    return from_dict(
+        observed_data={"y": rng.normal(size=(5, 7))},
+        posterior_predictive={"y": rng.normal(size=(4, 1000, 5, 7)) / 2},
+        posterior={"y_model": rng.normal(size=(4, 1000, 5, 7))},
+        dims={"y": ["dim1", "dim2"]},
+        coords={"dim1": range(5), "dim2": range(7)},
+    )
+
+
+@pytest.mark.parametrize("data", ("1d", "2d"))
+@pytest.mark.parametrize("kind", ("lines", "hdi"))
+@pytest.mark.parametrize("use_y_model", (True, False))
+def test_plot_lm(data, kind, use_y_model):
+    if data == "1d":
+        idata = generate_lm_1d_data()
+    else:
+        idata = generate_lm_2d_data()
+
+    kwargs = {"idata": idata, "y": "y", "kind_model": kind}
+    if data == "2d":
+        kwargs["plot_dim"] = "dim1"
+    if use_y_model:
+        kwargs["y_model"] = "y_model"
+    if kind == "lines":
+        kwargs["num_samples"] = 50
+
+    ax = plot_lm(**kwargs)
+    assert ax is not None
+
+
+@pytest.mark.parametrize(
+    "coords, expected_vars",
+    [
+        ({"school": ["Choate"]}, ["theta"]),
+        ({"school": ["Lawrenceville"]}, ["theta"]),
+        ({}, ["theta"]),
+    ],
+)
+def test_plot_autocorr_coords(coords, expected_vars):
+    """Test plot_autocorr with coords kwarg."""
+    idata = load_arviz_data("centered_eight")
+
+    axes = plot_autocorr(idata, var_names=expected_vars, coords=coords, show=False)
+
+    assert axes is not None

arviz/tests/base_tests/test_stats.py

@@ -14,10 +14,11 @@ from scipy.stats import linregress, norm, halfcauchy
 from xarray import DataArray, Dataset
 from xarray_einstats.stats import XrContinuousRV
 
-from ...data import concat, convert_to_inference_data, from_dict, load_arviz_data
+from ...data import concat, convert_to_inference_data, from_dict, load_arviz_data, InferenceData
 from ...rcparams import rcParams
 from ...stats import (
     apply_test_function,
+    bayes_factor,
     compare,
     ess,
     hdi,
@@ -871,3 +872,54 @@ def test_priorsens_coords(psens_data):
     assert "mu" in result
     assert "theta" in result
     assert "school" in result.theta_t.dims
+
+
+def test_bayes_factor():
+    idata = from_dict(
+        posterior={"a": np.random.normal(1, 0.5, 5000)}, prior={"a": np.random.normal(0, 1, 5000)}
+    )
+    bf_dict0 = bayes_factor(idata, var_name="a", ref_val=0)
+    bf_dict1 = bayes_factor(idata, prior=np.random.normal(0, 10, 5000), var_name="a", ref_val=0)
+    assert bf_dict0["BF10"] > bf_dict0["BF01"]
+    assert bf_dict1["BF10"] < bf_dict1["BF01"]
+
+
+def test_compare_sorting_consistency():
+    chains, draws = 4, 1000
+
+    # Model 1 - good fit
+    log_lik1 = np.random.normal(-2, 1, size=(chains, draws))
+    posterior1 = Dataset(
+        {"theta": (("chain", "draw"), np.random.normal(0, 1, size=(chains, draws)))},
+        coords={"chain": range(chains), "draw": range(draws)},
+    )
+    log_like1 = Dataset(
+        {"y": (("chain", "draw"), log_lik1)},
+        coords={"chain": range(chains), "draw": range(draws)},
+    )
+    data1 = InferenceData(posterior=posterior1, log_likelihood=log_like1)
+
+    # Model 2 - poor fit (higher variance)
+    log_lik2 = np.random.normal(-5, 2, size=(chains, draws))
+    posterior2 = Dataset(
+        {"theta": (("chain", "draw"), np.random.normal(0, 1, size=(chains, draws)))},
+        coords={"chain": range(chains), "draw": range(draws)},
+    )
+    log_like2 = Dataset(
+        {"y": (("chain", "draw"), log_lik2)},
+        coords={"chain": range(chains), "draw": range(draws)},
+    )
+    data2 = InferenceData(posterior=posterior2, log_likelihood=log_like2)
+
+    # Compare models in different orders
+    comp_dict1 = {"M1": data1, "M2": data2}
+    comp_dict2 = {"M2": data2, "M1": data1}
+
+    comparison1 = compare(comp_dict1, method="bb-pseudo-bma")
+    comparison2 = compare(comp_dict2, method="bb-pseudo-bma")
+
+    assert comparison1.index.tolist() == comparison2.index.tolist()
+
+    se1 = comparison1["se"].values
+    se2 = comparison2["se"].values
+    np.testing.assert_array_almost_equal(se1, se2)

arviz/tests/external_tests/test_data_numpyro.py

@@ -46,7 +46,9 @@ class TestDataNumPyro:
         )
         return predictions
 
-    def get_inference_data(self, data, eight_schools_params, predictions_data, predictions_params):
+    def get_inference_data(
+        self, data, eight_schools_params, predictions_data, predictions_params, infer_dims=False
+    ):
         posterior_samples = data.obj.get_samples()
         model = data.obj.sampler.model
         posterior_predictive = Predictive(model, posterior_samples)(
@@ -55,6 +57,12 @@ class TestDataNumPyro:
         prior = Predictive(model, num_samples=500)(
             PRNGKey(2), eight_schools_params["J"], eight_schools_params["sigma"]
         )
+        dims = {"theta": ["school"], "eta": ["school"], "obs": ["school"]}
+        pred_dims = {"theta": ["school_pred"], "eta": ["school_pred"], "obs": ["school_pred"]}
+        if infer_dims:
+            dims = None
+            pred_dims = None
+
         predictions = predictions_data
         return from_numpyro(
             posterior=data.obj,
@@ -65,8 +73,8 @@ class TestDataNumPyro:
                 "school": np.arange(eight_schools_params["J"]),
                 "school_pred": np.arange(predictions_params["J"]),
             },
-            dims={"theta": ["school"], "eta": ["school"], "obs": ["school"]},
-            pred_dims={"theta": ["school_pred"], "eta": ["school_pred"], "obs": ["school_pred"]},
+            dims=dims,
+            pred_dims=pred_dims,
         )
 
     def test_inference_data_namedtuple(self, data):
@@ -77,6 +85,7 @@ class TestDataNumPyro:
         data.obj.get_samples = lambda *args, **kwargs: data_namedtuple
         inference_data = from_numpyro(
             posterior=data.obj,
+            dims={},  # This mock test needs to turn off autodims like so or mock group_by_chain
         )
         assert isinstance(data.obj.get_samples(), Samples)
         data.obj.get_samples = _old_fn
@@ -282,3 +291,121 @@ class TestDataNumPyro:
         mcmc.run(PRNGKey(0))
         inference_data = from_numpyro(mcmc)
         assert inference_data.observed_data
+
+    def test_mcmc_infer_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            # note: group2 gets assigned dim=-1 and group1 is assigned dim=-2
+            with numpyro.plate("group2", 5), numpyro.plate("group1", 10):
+                _ = numpyro.sample("param", dist.Normal(0, 1))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(
+            mcmc, coords={"group1": np.arange(10), "group2": np.arange(5)}
+        )
+        assert inference_data.posterior.param.dims == ("chain", "draw", "group1", "group2")
+        assert all(dim in inference_data.posterior.param.coords for dim in ("group1", "group2"))
+
+    def test_mcmc_infer_unsorted_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            group1_plate = numpyro.plate("group1", 10, dim=-1)
+            group2_plate = numpyro.plate("group2", 5, dim=-2)
+
+            # the plate contexts are entered in a different order than the pre-defined dims
+            # we should make sure this still works because the trace has all of the info it needs
+            with group2_plate, group1_plate:
+                _ = numpyro.sample("param", dist.Normal(0, 1))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(
+            mcmc, coords={"group1": np.arange(10), "group2": np.arange(5)}
+        )
+        assert inference_data.posterior.param.dims == ("chain", "draw", "group2", "group1")
+        assert all(dim in inference_data.posterior.param.coords for dim in ("group1", "group2"))
+
+    def test_mcmc_infer_dims_no_coords(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            with numpyro.plate("group", 5):
+                _ = numpyro.sample("param", dist.Normal(0, 1))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(mcmc)
+        assert inference_data.posterior.param.dims == ("chain", "draw", "group")
+
+    def test_mcmc_event_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            _ = numpyro.sample(
+                "gamma", dist.ZeroSumNormal(1, event_shape=(10,)), infer={"event_dims": ["groups"]}
+            )
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(mcmc, coords={"groups": np.arange(10)})
+        assert inference_data.posterior.gamma.dims == ("chain", "draw", "groups")
+        assert "groups" in inference_data.posterior.gamma.coords
+
+    @pytest.mark.xfail
+    def test_mcmc_inferred_dims_univariate(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+        import jax.numpy as jnp
+
+        def model():
+            alpha = numpyro.sample("alpha", dist.Normal(0, 1))
+            sigma = numpyro.sample("sigma", dist.HalfNormal(1))
+            with numpyro.plate("obs_idx", 3):
+                # mu is plated by obs_idx, but isnt broadcasted to the plate shape
+                # the expected behavior is that this should cause a failure
+                mu = numpyro.deterministic("mu", alpha)
+                return numpyro.sample("y", dist.Normal(mu, sigma), obs=jnp.array([-1, 0, 1]))
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(mcmc, coords={"obs_idx": np.arange(3)})
+        assert inference_data.posterior.mu.dims == ("chain", "draw", "obs_idx")
+        assert "obs_idx" in inference_data.posterior.mu.coords
+
+    def test_mcmc_extra_event_dims(self):
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        def model():
+            gamma = numpyro.sample("gamma", dist.ZeroSumNormal(1, event_shape=(10,)))
+            _ = numpyro.deterministic("gamma_plus1", gamma + 1)
+
+        mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+        mcmc.run(PRNGKey(0))
+        inference_data = from_numpyro(
+            mcmc, coords={"groups": np.arange(10)}, extra_event_dims={"gamma_plus1": ["groups"]}
+        )
+        assert inference_data.posterior.gamma_plus1.dims == ("chain", "draw", "groups")
+        assert "groups" in inference_data.posterior.gamma_plus1.coords
+
+    def test_mcmc_predictions_infer_dims(
+        self, data, eight_schools_params, predictions_data, predictions_params
+    ):
+        inference_data = self.get_inference_data(
+            data, eight_schools_params, predictions_data, predictions_params, infer_dims=True
+        )
+        assert inference_data.predictions.obs.dims == ("chain", "draw", "J")
+        assert "J" in inference_data.predictions.obs.coords

arviz/utils.py

@@ -330,6 +330,7 @@ class Numba:
     """A class to toggle numba states."""
 
     numba_flag = numba_check()
+    """bool: Indicates whether Numba optimizations are enabled. Defaults to False."""
 
     @classmethod
     def disable_numba(cls):
@@ -732,7 +733,10 @@ class Dask:
     """
 
     dask_flag = False
+    """bool: Enables Dask parallelization when set to True. Defaults to False."""
     dask_kwargs = None
+    """dict: Additional keyword arguments for Dask configuration.
+    Defaults to an empty dictionary."""
 
     @classmethod
     def enable_dask(cls, dask_kwargs=None):

arviz/wrappers/base.py

@@ -197,7 +197,7 @@ class SamplingWrapper:
         """Check that all methods listed are implemented.
 
         Not all functions that require refitting need to have all the methods implemented in
-        order to work properly. This function shoulg be used before using the SamplingWrapper and
+        order to work properly. This function should be used before using the SamplingWrapper and
         its subclasses to get informative error messages.
 
         Parameters

arviz/wrappers/wrap_stan.py

@@ -44,7 +44,7 @@ class StanSamplingWrapper(SamplingWrapper):
         excluded_observed_data : str
             Variable name containing the pointwise log likelihood data of the excluded
             data. As PyStan cannot call C++ functions and log_likelihood__i is already
-            calculated *during* the simultion, instead of the value on which to evaluate
+            calculated *during* the simulation, instead of the value on which to evaluate
             the likelihood, ``log_likelihood__i`` expects a string so it can extract the
             corresponding data from the InferenceData object.
         """

{arviz-0.20.0.dist-info → arviz-0.22.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arviz
-Version: 0.20.0
+Version: 0.22.0
 Summary: Exploratory analysis of Bayesian models
 Home-page: http://github.com/arviz-devs/arviz
 Author: ArviZ Developers
@@ -22,12 +22,12 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: setuptools>=60.0.0
-Requires-Dist: matplotlib>=3.5
-Requires-Dist: numpy>=1.23.0
-Requires-Dist: scipy>=1.9.0
+Requires-Dist: matplotlib>=3.8
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: scipy>=1.11.0
 Requires-Dist: packaging
-Requires-Dist: pandas>=1.5.0
-Requires-Dist: xarray>=2022.6.0
+Requires-Dist: pandas>=2.1.0
+Requires-Dist: xarray>=2023.7.0
 Requires-Dist: h5netcdf>=1.0.2
 Requires-Dist: typing-extensions>=4.1.0
 Requires-Dist: xarray-einstats>=0.3
@@ -39,7 +39,7 @@ Requires-Dist: contourpy; extra == "all"
 Requires-Dist: ujson; extra == "all"
 Requires-Dist: dask[distributed]; extra == "all"
 Requires-Dist: zarr<3,>=2.5.0; extra == "all"
-Requires-Dist: xarray-datatree; extra == "all"
+Requires-Dist: xarray>=2024.11.0; extra == "all"
 Requires-Dist: dm-tree>=0.1.8; extra == "all"
 Provides-Extra: preview
 Requires-Dist: arviz-base[h5netcdf]; extra == "preview"