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

arviz/tests/base_tests/test_plots_bokeh.py

@@ -8,6 +8,7 @@ from pandas import DataFrame  # pylint: disable=wrong-import-position
 from scipy.stats import norm  # pylint: disable=wrong-import-position
 
 from ...data import from_dict, load_arviz_data  # pylint: disable=wrong-import-position
+from ...labels import MapLabeller  # pylint: disable=wrong-import-position
 from ...plots import (  # pylint: disable=wrong-import-position
     plot_autocorr,
     plot_bpv,
@@ -773,7 +774,6 @@ def test_plot_mcse_no_divergences(models):
         {"divergences": True, "var_names": ["theta", "mu"]},
         {"kind": "kde", "var_names": ["theta"]},
         {"kind": "hexbin", "var_names": ["theta"]},
-        {"kind": "hexbin", "var_names": ["theta"]},
         {
             "kind": "hexbin",
             "var_names": ["theta"],
@@ -785,6 +785,21 @@ def test_plot_mcse_no_divergences(models):
             "reference_values": {"mu": 0, "tau": 0},
             "reference_values_kwargs": {"line_color": "blue"},
         },
+        {
+            "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):
@@ -1201,7 +1216,7 @@ def test_plot_dot_rotated(continuous_model, kwargs):
         },
     ],
 )
-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():
@@ -1228,3 +1243,46 @@ def test_plot_lm_list():
     """Test the plots when input data is list or ndarray."""
     y = [1, 2, 3, 4, 5]
     assert plot_lm(y=y, x=np.arange(len(y)), show=False, backend="bokeh")
+
+
+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, "backend": "bokeh", "show": False}
+    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

arviz/tests/base_tests/test_plots_matplotlib.py

@@ -14,6 +14,7 @@ from pandas import DataFrame
 from scipy.stats import gaussian_kde, norm
 
 from ...data import from_dict, load_arviz_data
+from ...labels import MapLabeller
 from ...plots import (
     plot_autocorr,
     plot_bf,
@@ -599,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):
@@ -1914,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():
@@ -2102,3 +2118,80 @@ def test_plot_bf():
     )
     _, 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
+
+
+def test_plot_forest_with_transform():
+    """Test if plot_forest runs successfully with a transform dictionary."""
+    data = xr.Dataset(
+        {
+            "var1": (["chain", "draw"], np.array([[1, 2, 3], [4, 5, 6]])),
+            "var2": (["chain", "draw"], np.array([[7, 8, 9], [10, 11, 12]])),
+        },
+        coords={"chain": [0, 1], "draw": [0, 1, 2]},
+    )
+    transform_dict = {
+        "var1": lambda x: x + 1,
+        "var2": lambda x: x * 2,
+    }
+
+    axes = plot_forest(data, transform=transform_dict, show=False)
+    assert axes is not None

arviz/tests/base_tests/test_stats.py

@@ -14,7 +14,7 @@ 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,
@@ -882,3 +882,44 @@ def test_bayes_factor():
     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/base_tests/test_stats_ecdf_utils.py

@@ -13,9 +13,9 @@ from ...stats.ecdf_utils import (
 try:
     import numba  # pylint: disable=unused-import
 
-    numba_options = [True, False]
+    numba_options = [True, False]  # pylint: disable=invalid-name
 except ImportError:
-    numba_options = [False]
+    numba_options = [False]  # pylint: disable=invalid-name
 
 
 def test_compute_ecdf():

arviz/tests/external_tests/test_data_numpyro.py

@@ -1,4 +1,4 @@
-# pylint: disable=no-member, invalid-name, redefined-outer-name
+# pylint: disable=no-member, invalid-name, redefined-outer-name, too-many-public-methods
 from collections import namedtuple
 import numpy as np
 import pytest
@@ -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,144 @@ 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
+
+    def test_potential_energy_sign_conversion(self):
+        """Test that potential energy is converted to log probability (lp) with correct sign."""
+        import numpyro
+        import numpyro.distributions as dist
+        from numpyro.infer import MCMC, NUTS
+
+        num_samples = 10
+
+        def simple_model():
+            numpyro.sample("x", dist.Normal(0, 1))
+
+        nuts_kernel = NUTS(simple_model)
+        mcmc = MCMC(nuts_kernel, num_samples=num_samples, num_warmup=5)
+        mcmc.run(PRNGKey(0), extra_fields=["potential_energy"])
+
+        # Get the raw extra fields from NumPyro
+        extra_fields = mcmc.get_extra_fields(group_by_chain=True)
+        # Convert to ArviZ InferenceData
+        inference_data = from_numpyro(mcmc)
+        arviz_lp = inference_data["sample_stats"]["lp"].values
+
+        np.testing.assert_array_equal(arviz_lp, -extra_fields["potential_energy"])

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.21.0.dist-info → arviz-0.23.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: arviz
-Version: 0.21.0
+Version: 0.23.0
 Summary: Exploratory analysis of Bayesian models
 Home-page: http://github.com/arviz-devs/arviz
 Author: ArviZ Developers
@@ -22,14 +22,14 @@ 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: typing_extensions>=4.1.0
 Requires-Dist: xarray-einstats>=0.3
 Provides-Extra: all
 Requires-Dist: numba; extra == "all"
@@ -39,12 +39,23 @@ 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"
 Requires-Dist: arviz-stats[xarray]; extra == "preview"
 Requires-Dist: arviz-plots; extra == "preview"
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 <img src="https://raw.githubusercontent.com/arviz-devs/arviz-project/main/arviz_logos/ArviZ.png#gh-light-mode-only" width=200></img>
 <img src="https://raw.githubusercontent.com/arviz-devs/arviz-project/main/arviz_logos/ArviZ_white.png#gh-dark-mode-only" width=200></img>