Merge branch 'main' into rank_gap

CHANGELOG.md

@@ -8,12 +8,15 @@
 
 ### Maintenance and fixes
 - Fix `reloo` outdated usage of `ELPDData` ([2158](https://github.com/arviz-devs/arviz/pull/2158))
+- plot_bpv smooth discrete data only when computing u_values ([2179](https://github.com/arviz-devs/arviz/pull/2179))
+- Fix bug when beanmachine objects lack some fields ([2154](https://github.com/arviz-devs/arviz/pull/2154))
 - Fix gap for `plot_trace` with option `kind="rank_bars"` ([2180](https://github.com/arviz-devs/arviz/pull/2180))
 
 ### Deprecation
 
 ### Documentation
 - Add PyMC and CmdStanPy sampling wrapper examples ([2158](https://github.com/arviz-devs/arviz/pull/2158))
+- Fix docstring for plot_trace chain_prop and compact_prop parameters ([2176](https://github.com/arviz-devs/arviz/pull/2176))
 
 ## v0.14.0 (2022 Nov 15)
 

arviz/data/io_beanmachine.py

@@ -24,15 +24,23 @@ def __init__(
 
         if "posterior" in self.sampler.namespaces:
             self.posterior = self.sampler.namespaces["posterior"].samples
+        else:
+            self.posterior = None
 
         if "posterior_predictive" in self.sampler.namespaces:
             self.posterior_predictive = self.sampler.namespaces["posterior_predictive"].samples
+        else:
+            self.posterior_predictive = None
 
         if self.sampler.log_likelihoods is not None:
             self.log_likelihoods = self.sampler.log_likelihoods
+        else:
+            self.log_likelihoods = None
 
         if self.sampler.observations is not None:
             self.observations = self.sampler.observations
+        else:
+            self.observations = None
 
     @requires("posterior")
     def posterior_to_xarray(self):

arviz/plots/backends/bokeh/bpvplot.py

@@ -89,9 +89,6 @@ def plot_bpv(
         obs_vals = obs_vals.flatten()
         pp_vals = pp_vals.reshape(total_pp_samples, -1)
 
-        if obs_vals.dtype.kind == "i" or pp_vals.dtype.kind == "i":
-            obs_vals, pp_vals = smooth_data(obs_vals, pp_vals)
-
         if kind == "p_value":
             tstat_pit = np.mean(pp_vals <= obs_vals, axis=-1)
             x_s, tstat_pit_dens = kde(tstat_pit)
@@ -117,6 +114,9 @@ def plot_bpv(
                     )
 
         elif kind == "u_value":
+            if obs_vals.dtype.kind == "i" or pp_vals.dtype.kind == "i":
+                obs_vals, pp_vals = smooth_data(obs_vals, pp_vals)
+
             tstat_pit = np.mean(pp_vals <= obs_vals, axis=0)
             x_s, tstat_pit_dens = kde(tstat_pit)
             ax_i.line(x_s, tstat_pit_dens, color=color)

arviz/plots/backends/matplotlib/bpvplot.py

@@ -86,9 +86,6 @@ def plot_bpv(
         obs_vals = obs_vals.flatten()
         pp_vals = pp_vals.reshape(total_pp_samples, -1)
 
-        if obs_vals.dtype.kind == "i" or pp_vals.dtype.kind == "i":
-            obs_vals, pp_vals = smooth_data(obs_vals, pp_vals)
-
         if kind == "p_value":
             tstat_pit = np.mean(pp_vals <= obs_vals, axis=-1)
             x_s, tstat_pit_dens = kde(tstat_pit)
@@ -113,6 +110,9 @@ def plot_bpv(
                     ax_i.plot(x_ss, u_dens, linewidth=linewidth, **plot_ref_kwargs)
 
         elif kind == "u_value":
+            if obs_vals.dtype.kind == "i" or pp_vals.dtype.kind == "i":
+                obs_vals, pp_vals = smooth_data(obs_vals, pp_vals)
+
             tstat_pit = np.mean(pp_vals <= obs_vals, axis=0)
             x_s, tstat_pit_dens = kde(tstat_pit)
             ax_i.plot(x_s, tstat_pit_dens, color=color)

arviz/plots/traceplot.py

@@ -82,13 +82,18 @@ def plot_trace(
     compact: bool, optional
         Plot multidimensional variables in a single plot.
     compact_prop: str or dict {str: array_like}, optional
-        Tuple containing the property name and the property values to distinguish different
-        dimensions with compact=True
+         Defines the property name and the property values to distinguish different
+        dimensions with compact=True.
+        When compact=True it defaults to color, it is
+        ignored otherwise.
     combined: bool, optional
         Flag for combining multiple chains into a single line. If False (default), chains will be
         plotted separately.
     chain_prop: str or dict {str: array_like}, optional
-        Tuple containing the property name and the property values to distinguish different chains
+        Defines the property name and the property values to distinguish different chains.
+        If compact=True it defaults to linestyle,
+        otherwise it uses the color to distinguish
+        different chains.
     legend: bool, optional
         Add a legend to the figure with the chain color code.
     plot_kwargs, fill_kwargs, rug_kwargs, hist_kwargs: dict, optional

arviz/tests/external_tests/test_data_beanmachine.py

@@ -0,0 +1,76 @@
+# pylint: disable=no-member, invalid-name, redefined-outer-name
+import numpy as np
+import pytest
+
+from ...data.io_beanmachine import from_beanmachine  # pylint: disable=wrong-import-position
+from ..helpers import (  # pylint: disable=unused-import, wrong-import-position
+    chains,
+    draws,
+    eight_schools_params,
+    importorskip,
+    load_cached_models,
+)
+
+# Skip all tests if beanmachine or pytorch not installed
+torch = importorskip("torch")
+bm = importorskip("beanmachine.ppl")
+dist = torch.distributions
+
+
+class TestDataBeanMachine:
+    @pytest.fixture(scope="class")
+    def data(self, eight_schools_params, draws, chains):
+        class Data:
+            model, prior, obj = load_cached_models(
+                eight_schools_params,
+                draws,
+                chains,
+                "beanmachine",
+            )["beanmachine"]
+
+        return Data
+
+    @pytest.fixture(scope="class")
+    def predictions_data(self, data):
+        """Generate predictions for predictions_params"""
+        posterior_samples = data.obj
+        model = data.model
+        predictions = bm.inference.predictive.simulate([model.obs()], posterior_samples)
+        return predictions
+
+    def get_inference_data(self, eight_schools_params, predictions_data):
+        predictions = predictions_data
+        return from_beanmachine(
+            sampler=predictions,
+            coords={
+                "school": np.arange(eight_schools_params["J"]),
+                "school_pred": np.arange(eight_schools_params["J"]),
+            },
+        )
+
+    def test_inference_data(self, data, eight_schools_params, predictions_data):
+        inference_data = self.get_inference_data(eight_schools_params, predictions_data)
+        model = data.model
+        mu = model.mu()
+        tau = model.tau()
+        eta = model.eta()
+        obs = model.obs()
+
+        assert mu in inference_data.posterior
+        assert tau in inference_data.posterior
+        assert eta in inference_data.posterior
+        assert obs in inference_data.posterior_predictive
+
+    def test_inference_data_has_log_likelihood_and_observed_data(self, data):
+        idata = from_beanmachine(data.obj)
+        obs = data.model.obs()
+
+        assert obs in idata.log_likelihood
+        assert obs in idata.observed_data
+
+    def test_inference_data_no_posterior(self, data):
+        model = data.model
+        # only prior
+        inference_data = from_beanmachine(data.prior)
+        assert not model.obs() in inference_data.posterior
+        assert "observed_data" not in inference_data

arviz/tests/helpers.py

@@ -486,6 +486,52 @@ def pystan_noncentered_schools(data, draws, chains):
     return stan_model, fit
 
 
+def bm_schools_model(data, draws, chains):
+    import beanmachine.ppl as bm
+    import torch
+    import torch.distributions as dist
+
+    class EightSchools:
+        @bm.random_variable
+        def mu(self):
+            return dist.Normal(0, 5)
+
+        @bm.random_variable
+        def tau(self):
+            return dist.HalfCauchy(5)
+
+        @bm.random_variable
+        def eta(self):
+            return dist.Normal(0, 1).expand((data["J"],))
+
+        @bm.functional
+        def theta(self):
+            return self.mu() + self.tau() * self.eta()
+
+        @bm.random_variable
+        def obs(self):
+            return dist.Normal(self.theta(), torch.from_numpy(data["sigma"]).float())
+
+    model = EightSchools()
+
+    prior = bm.GlobalNoUTurnSampler().infer(
+        queries=[model.mu(), model.tau(), model.eta()],
+        observations={},
+        num_samples=draws,
+        num_adaptive_samples=500,
+        num_chains=chains,
+    )
+
+    posterior = bm.GlobalNoUTurnSampler().infer(
+        queries=[model.mu(), model.tau(), model.eta()],
+        observations={model.obs(): torch.from_numpy(data["y"]).float()},
+        num_samples=draws,
+        num_adaptive_samples=500,
+        num_chains=chains,
+    )
+    return model, prior, posterior
+
+
 def library_handle(library):
     """Import a library and return the handle."""
     if library == "pystan":
@@ -506,6 +552,7 @@ def load_cached_models(eight_schools_data, draws, chains, libs=None):
         ("emcee", emcee_schools_model),
         ("pyro", pyro_noncentered_schools),
         ("numpyro", numpyro_schools_model),
+        ("beanmachine", bm_schools_model),
     )
     data_directory = os.path.join(here, "saved_models")
     models = {}

azure-pipelines.yml

@@ -20,5 +20,4 @@ jobs:
   - template: .azure-pipelines/azure-pipelines-base.yml
   - template: .azure-pipelines/azure-pipelines-external.yml
   - template: .azure-pipelines/azure-pipelines-benchmarks.yml
-  - template: .azure-pipelines/azure-pipelines-docs.yml
   - template: .azure-pipelines/azure-pipelines-wheel.yml

doc/source/conf.py

@@ -187,7 +187,7 @@
             "name": "Mastodon",
             "url": "https://bayes.club/@ArviZ",
             "icon": "fa-brands fa-mastodon",
-        }
+        },
     ],
     "navbar_start": ["navbar-logo", "navbar-version"],
     "header_links_before_dropdown": 7,