remove arviz dependency, move wrapping how-to-guide

- remove wrapping logic needed for arviz
- remove lengthy tutorial on using arviz plots
- instead, add a how-to-guide for how to wrap posterior samples into arviz object.

Author: janfb

docs/advanced_tutorials.rst

@@ -65,7 +65,6 @@ Diagnostics
    advanced_tutorials/11_diagnostics_simulation_based_calibration.ipynb
    advanced_tutorials/13_diagnostics_lc2st.ipynb
    advanced_tutorials/21_diagnostics_misspecification_checks.ipynb
-   advanced_tutorials/14_mcmc_diagnostics_with_arviz.ipynb
 
 
 Visualization

docs/advanced_tutorials/14_mcmc_diagnostics_with_arviz.ipynb

@@ -57,6 +57,7 @@ Sampling
    how_to_guide/09_sampler_interface.ipynb
    how_to_guide/10_refine_posterior_with_importance_sampling.ipynb
    how_to_guide/11_iid_sampling_with_nle_or_nre.ipynb
+   how_to_guide/12_mcmc_diagnostics_with_arviz.ipynb
 
 
 Diagnostics

docs/how_to_guide.rst

@@ -37,7 +37,6 @@ inference.
 - [Posterior predictive checks](10_diagnostics_posterior_predictive_checks.md)
 - [Simulation-based calibration](11_diagnostics_simulation_based_calibration.md)
 - [Local-C2ST coverage checks](13_diagnostics_lc2st.md)
-- [Density plots and MCMC diagnostics with ArviZ](14_mcmc_diagnostics_with_arviz.md)
 </div>
 
 ## Analysis

docs/how_to_guide/12_mcmc_diagnostics_with_arviz.ipynb

@@ -29,7 +29,6 @@ dynamic = ["version"]
 readme = "README.md"
 keywords = ["Bayesian inference", "simulation-based inference", "PyTorch"]
 dependencies = [
-    "arviz",
     "joblib>=1.0.0",
     "matplotlib",
     "numpy",

mkdocs/docs/tutorials/index.md

@@ -8,10 +8,8 @@
 from typing import Any, Callable, Dict, Literal, Optional, Union
 from warnings import warn
 
-import arviz as az
 import torch
 import torch.distributions.transforms as torch_tf
-from arviz.data import InferenceData
 from joblib import Parallel, delayed
 from numpy import ndarray
 from pyro.infer.mcmc import HMC, NUTS
@@ -1040,64 +1038,6 @@ def map(
             force_update=force_update,
         )
 
-    def get_arviz_inference_data(self) -> InferenceData:
-        """Returns arviz InferenceData object constructed most recent samples.
-
-        Note: the InferenceData is constructed using the posterior samples generated in
-        most recent call to `.sample(...)`.
-
-        For Pyro and PyMC samplers, InferenceData will contain diagnostics, but for
-        sbi slice samplers, only the samples are added.
-
-        Returns:
-            inference_data: Arviz InferenceData object.
-        """
-        assert self._posterior_sampler is not None, (
-            """No samples have been generated, call .sample() first."""
-        )
-
-        sampler: Union[
-            MCMC, SliceSamplerSerial, SliceSamplerVectorized, PyMCSampler
-        ] = self._posterior_sampler
-
-        # If Pyro sampler and samples not transformed, use arviz' from_pyro.
-        if isinstance(sampler, (HMC, NUTS)) and isinstance(
-            self.theta_transform, torch_tf.IndependentTransform
-        ):
-            inference_data = az.from_pyro(sampler)
-        # If PyMC sampler and samples not transformed, get cached InferenceData.
-        elif isinstance(sampler, PyMCSampler) and isinstance(
-            self.theta_transform, torch_tf.IndependentTransform
-        ):
-            inference_data = sampler.get_inference_data()
-
-        # otherwise get samples from sampler and transform to original space.
-        else:
-            transformed_samples = sampler.get_samples(group_by_chain=True)
-            # Pyro samplers returns dicts, get values.
-            if isinstance(transformed_samples, Dict):
-                # popitem gets last items, [1] get the values as tensor.
-                transformed_samples = transformed_samples.popitem()[1]
-            # Our slice samplers return numpy arrays.
-            elif isinstance(transformed_samples, ndarray):
-                transformed_samples = torch.from_numpy(transformed_samples).type(
-                    torch.float32
-                )
-            # For MultipleIndependent priors transforms first dim must be batch dim.
-            # thus, reshape back and forth to have batch dim in front.
-            samples_shape = transformed_samples.shape
-            samples = self.theta_transform.inv(  # type: ignore
-                transformed_samples.reshape(-1, samples_shape[-1])
-            ).reshape(  # type: ignore
-                *samples_shape
-            )
-
-            inference_data = az.convert_to_inference_data({
-                f"{self.param_name}": samples
-            })
-
-        return inference_data
-
     def __getstate__(self) -> Dict:
         """Get state of MCMCPosterior.
 
@@ -1124,12 +1064,6 @@ def _process_thin_default(thin: int) -> int:
     """
     if thin == -1:
         thin = 1
-        warn(
-            "The default value for thinning in MCMC sampling has been changed from "
-            "10 to 1. This might cause the results differ from the last benchmark.",
-            UserWarning,
-            stacklevel=2,
-        )
 
     return thin
 

pyproject.toml

@@ -7,7 +7,6 @@
 import pymc
 import pytensor.tensor as pt
 import torch
-from arviz.data import InferenceData
 
 from sbi.utils.torchutils import tensor2numpy
 
@@ -206,7 +205,7 @@ def get_samples(
         else:
             return samples[-num_samples:, :]
 
-    def get_inference_data(self) -> InferenceData:
+    def get_inference_data(self) -> Any:
         """Returns InferenceData from last call to self.run,
         which contains diagnostic information in addition to samples
 

sbi/inference/posteriors/mcmc_posterior.py

@@ -3,35 +3,22 @@
 
 from __future__ import annotations
 
-from dataclasses import asdict
-
 import numpy as np
 import pytest
 import torch
 from torch import eye, ones, zeros
-from torch.distributions import Uniform
 
-from sbi.inference import (
-    NLE,
-    MCMCPosterior,
-    likelihood_estimator_based_potential,
-)
 from sbi.inference.posteriors.mcmc_posterior import build_from_potential
 from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters
-from sbi.neural_nets import likelihood_nn
 from sbi.samplers.mcmc.pymc_wrapper import PyMCSampler
 from sbi.samplers.mcmc.slice_numpy import (
     SliceSampler,
     SliceSamplerSerial,
     SliceSamplerVectorized,
 )
-from sbi.simulators.linear_gaussian import (
-    diagonal_linear_gaussian,
-    true_posterior_linear_gaussian_mvn_prior,
-)
+from sbi.simulators.linear_gaussian import true_posterior_linear_gaussian_mvn_prior
 from sbi.utils import BoxUniform
 from sbi.utils.metrics import check_c2st
-from sbi.utils.user_input_checks import process_prior
 
 
 @pytest.mark.mcmc
@@ -198,63 +185,6 @@ def lp_f(x, track_gradients=True):
         "slice_np_vectorized",
     ),
 )
-def test_getting_inference_diagnostics(
-    method, mcmc_params_fast: MCMCPosteriorParameters
-):
-    num_simulations = 100
-    num_samples = 10
-    num_dim = 2
-
-    # Use composed prior to test MultipleIndependent case.
-    prior = [
-        Uniform(low=-ones(1), high=ones(1)),
-        Uniform(low=-ones(1), high=ones(1)),
-    ]
-
-    simulator = diagonal_linear_gaussian
-    density_estimator = likelihood_nn("maf", num_transforms=3)
-    inference = NLE(density_estimator=density_estimator, show_progress_bars=False)
-    prior, *_ = process_prior(prior)
-    theta = prior.sample((num_simulations,))
-    x = simulator(theta)
-    likelihood_estimator = inference.append_simulations(theta, x).train(
-        training_batch_size=num_simulations, max_num_epochs=2
-    )
-
-    x_o = zeros((1, num_dim))
-    potential_fn, theta_transform = likelihood_estimator_based_potential(
-        prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o
-    )
-    posterior = MCMCPosterior(
-        proposal=prior,
-        potential_fn=potential_fn,
-        theta_transform=theta_transform,
-        **asdict(mcmc_params_fast),
-    )
-    posterior.sample(
-        sample_shape=(num_samples,),
-        method=method,
-    )
-    idata = posterior.get_arviz_inference_data()
-
-    assert hasattr(idata, "posterior"), (
-        f"`MCMCPosterior.get_arviz_inference_data()` for method {method} "
-        f"returned invalid InferenceData. Must contain key 'posterior', "
-        f"but found only {list(idata.keys())}"
-    )
-    samples = getattr(idata.posterior, posterior.param_name).data
-    samples = samples.reshape(-1, samples.shape[-1])[:: mcmc_params_fast.thin][
-        :num_samples
-    ]
-    assert samples.shape == (
-        num_samples,
-        num_dim,
-    ), (
-        f"MCMC samples for method {method} have incorrect shape (n_samples, n_dims). "
-        f"Expected {(num_samples, num_dim)}, got {samples.shape}"
-    )
-
-
 @pytest.mark.mcmc
 def test_direct_mcmc_unconditional():
     "Test MCMCPosterior from user defined potential (unconditional)"