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

arviz/preview.py

@@ -0,0 +1,17 @@
+# pylint: disable=unused-import,unused-wildcard-import,wildcard-import
+"""Expose features from arviz-xyz refactored packages inside ``arviz.preview`` namespace."""
+
+try:
+    from arviz_base import *
+except ModuleNotFoundError:
+    pass
+
+try:
+    import arviz_stats
+except ModuleNotFoundError:
+    pass
+
+try:
+    from arviz_plots import *
+except ModuleNotFoundError:
+    pass

arviz/rcparams.py

@@ -26,6 +26,8 @@ _log = logging.getLogger(__name__)
 ScaleKeyword = Literal["log", "negative_log", "deviance"]
 ICKeyword = Literal["loo", "waic"]
 
+_identity = lambda x: x
+
 
 def _make_validate_choice(accepted_values, allow_none=False, typeof=str):
     """Validate value is in accepted_values.
@@ -300,7 +302,7 @@ defaultParams = {  # pylint: disable=invalid-name
         lambda x: x,
     ),
     "plot.matplotlib.show": (False, _validate_boolean),
-    "stats.hdi_prob": (0.94, _validate_probability),
+    "stats.ci_prob": (0.94, _validate_probability),
     "stats.information_criterion": (
         "loo",
         _make_validate_choice({"loo", "waic"} if NO_GET_ARGS else set(get_args(ICKeyword))),
@@ -318,6 +320,9 @@ defaultParams = {  # pylint: disable=invalid-name
     ),
 }
 
+# map from deprecated params to (version, new_param, fold2new, fnew2old)
+deprecated_map = {"stats.hdi_prob": ("0.18.0", "stats.ci_prob", _identity, _identity)}
+
 
 class RcParams(MutableMapping):
     """Class to contain ArviZ default parameters.
@@ -335,6 +340,15 @@ class RcParams(MutableMapping):
 
     def __setitem__(self, key, val):
         """Add validation to __setitem__ function."""
+        if key in deprecated_map:
+            version, key_new, fold2new, _ = deprecated_map[key]
+            warnings.warn(
+                f"{key} is deprecated since {version}, use {key_new} instead",
+                FutureWarning,
+            )
+            key = key_new
+            val = fold2new(val)
+
         try:
             try:
                 cval = self.validate[key](val)
@@ -349,7 +363,18 @@ class RcParams(MutableMapping):
 
     def __getitem__(self, key):
         """Use underlying dict's getitem method."""
-        return self._underlying_storage[key]
+        if key in deprecated_map:
+            version, key_new, _, fnew2old = deprecated_map[key]
+            warnings.warn(
+                f"{key} is deprecated since {version}, use {key_new} instead",
+                FutureWarning,
+            )
+            if key not in self._underlying_storage:
+                key = key_new
+        else:
+            fnew2old = _identity
+
+        return fnew2old(self._underlying_storage[key])
 
     def __delitem__(self, key):
         """Raise TypeError if someone ever tries to delete a key from RcParams."""

arviz/stats/diagnostics.py

@@ -135,10 +135,11 @@ def ess(
 
     References
     ----------
-    * Vehtari et al. (2019) see https://arxiv.org/abs/1903.08008
-    * https://mc-stan.org/docs/2_18/reference-manual/effective-sample-size-section.html
-      Section 15.4.2
-    * Gelman et al. BDA (2014) Formula 11.8
+    * Vehtari et al. (2021). Rank-normalization, folding, and
+        localization: An improved Rhat for assessing convergence of
+        MCMC. Bayesian analysis, 16(2):667-718.
+    * https://mc-stan.org/docs/reference-manual/analysis.html#effective-sample-size.section
+    * Gelman et al. BDA3 (2013) Formula 11.8
 
     See Also
     --------
@@ -246,7 +247,7 @@ def rhat(data, *, var_names=None, method="rank", dask_kwargs=None):
         Names of variables to include in the rhat report
     method : str
         Select R-hat method. Valid methods are:
-        - "rank"        # recommended by Vehtari et al. (2019)
+        - "rank"        # recommended by Vehtari et al. (2021)
         - "split"
         - "folded"
         - "z_scale"
@@ -269,7 +270,7 @@ def rhat(data, *, var_names=None, method="rank", dask_kwargs=None):
     -----
     The diagnostic is computed by:
 
-      .. math:: \hat{R} = \frac{\hat{V}}{W}
+      .. math:: \hat{R} = \sqrt{\frac{\hat{V}}{W}}
 
     where :math:`W` is the within-chain variance and :math:`\hat{V}` is the posterior variance
     estimate for the pooled rank-traces. This is the potential scale reduction factor, which
@@ -277,12 +278,15 @@ def rhat(data, *, var_names=None, method="rank", dask_kwargs=None):
     greater than one indicate that one or more chains have not yet converged.
 
     Rank values are calculated over all the chains with ``scipy.stats.rankdata``.
-    Each chain is split in two and normalized with the z-transform following Vehtari et al. (2019).
+    Each chain is split in two and normalized with the z-transform following
+    Vehtari et al. (2021).
 
     References
     ----------
-    * Vehtari et al. (2019) see https://arxiv.org/abs/1903.08008
-    * Gelman et al. BDA (2014)
+    * Vehtari et al. (2021).  Rank-normalization, folding, and
+      localization: An improved Rhat for assessing convergence of
+      MCMC. Bayesian analysis, 16(2):667-718.
+    * Gelman et al. BDA3 (2013)
     * Brooks and Gelman (1998)
     * Gelman and Rubin (1992)
 

arviz/stats/ecdf_utils.py

@@ -1,10 +1,25 @@
 """Functions for evaluating ECDFs and their confidence bands."""
 
+import math
 from typing import Any, Callable, Optional, Tuple
 import warnings
 
 import numpy as np
 from scipy.stats import uniform, binom
+from scipy.optimize import minimize_scalar
+
+try:
+    from numba import jit, vectorize
+except ImportError:
+
+    def jit(*args, **kwargs):  # pylint: disable=unused-argument
+        return lambda f: f
+
+    def vectorize(*args, **kwargs):  # pylint: disable=unused-argument
+        return lambda f: f
+
+
+from ..utils import Numba
 
 
 def compute_ecdf(sample: np.ndarray, eval_points: np.ndarray) -> np.ndarray:
@@ -25,6 +40,13 @@ def _get_ecdf_points(
     return x, y
 
 
+def _call_rvs(rvs, ndraws, random_state):
+    if random_state is None:
+        return rvs(ndraws)
+    else:
+        return rvs(ndraws, random_state=random_state)
+
+
 def _simulate_ecdf(
     ndraws: int,
     eval_points: np.ndarray,
@@ -32,7 +54,7 @@ def _simulate_ecdf(
     random_state: Optional[Any] = None,
 ) -> np.ndarray:
     """Simulate ECDF at the `eval_points` using the given random variable sampler"""
-    sample = rvs(ndraws, random_state=random_state)
+    sample = _call_rvs(rvs, ndraws, random_state)
     sample.sort()
     return compute_ecdf(sample, eval_points)
 
@@ -66,7 +88,7 @@ def ecdf_confidence_band(
     eval_points: np.ndarray,
     cdf_at_eval_points: np.ndarray,
     prob: float = 0.95,
-    method="simulated",
+    method="optimized",
     **kwargs,
 ) -> Tuple[np.ndarray, np.ndarray]:
     """Compute the `prob`-level confidence band for the ECDF.
@@ -85,20 +107,17 @@ def ecdf_confidence_band(
     method : string, default "simulated"
         The method used to compute the confidence band. Valid options are:
         - "pointwise": Compute the pointwise (i.e. marginal) confidence band.
+        - "optimized": Use optimization to estimate a simultaneous confidence band.
         - "simulated": Use Monte Carlo simulation to estimate a simultaneous confidence band.
           `rvs` must be provided.
     rvs: callable, optional
         A function that takes an integer `ndraws` and optionally the object passed to
         `random_state` and returns an array of `ndraws` samples from the same distribution
         as the original dataset. Required if `method` is "simulated" and variable is discrete.
-    num_trials : int, default 1000
+    num_trials : int, default 500
         The number of random ECDFs to generate for constructing simultaneous confidence bands
         (if `method` is "simulated").
-    random_state : {None, int, `numpy.random.Generator`,
-                    `numpy.random.RandomState`}, optional
-        If `None`, the `numpy.random.RandomState` singleton is used. If an `int`, a new
-        ``numpy.random.RandomState`` instance is used, seeded with seed. If a `RandomState` or
-        `Generator` instance, the instance is used.
+    random_state : int, numpy.random.Generator or numpy.random.RandomState, optional
 
     Returns
     -------
@@ -112,12 +131,18 @@ def ecdf_confidence_band(
 
     if method == "pointwise":
         prob_pointwise = prob
+    elif method == "optimized":
+        prob_pointwise = _optimize_simultaneous_ecdf_band_probability(
+            ndraws, eval_points, cdf_at_eval_points, prob=prob, **kwargs
+        )
     elif method == "simulated":
         prob_pointwise = _simulate_simultaneous_ecdf_band_probability(
             ndraws, eval_points, cdf_at_eval_points, prob=prob, **kwargs
         )
     else:
-        raise ValueError(f"Unknown method {method}. Valid options are 'pointwise' or 'simulated'.")
+        raise ValueError(
+            f"Unknown method {method}. Valid options are 'pointwise', 'optimized', or 'simulated'."
+        )
 
     prob_lower, prob_upper = _get_pointwise_confidence_band(
         prob_pointwise, ndraws, cdf_at_eval_points
@@ -126,13 +151,146 @@ def ecdf_confidence_band(
     return prob_lower, prob_upper
 
 
+def _update_ecdf_band_interior_probabilities(
+    prob_left: np.ndarray,
+    interval_left: np.ndarray,
+    interval_right: np.ndarray,
+    p: float,
+    ndraws: int,
+) -> np.ndarray:
+    """Update the probability that an ECDF has been within the envelope including at the current
+    point.
+
+    Arguments
+    ---------
+    prob_left : np.ndarray
+        For each point in the interior at the previous point, the joint probability that it and all
+        points before are in the interior.
+    interval_left : np.ndarray
+        The set of points in the interior at the previous point.
+    interval_right : np.ndarray
+        The set of points in the interior at the current point.
+    p : float
+        The probability of any given point found between the previous point and the current one.
+    ndraws : int
+        Number of draws in the original dataset.
+
+    Returns
+    -------
+    prob_right : np.ndarray
+        For each point in the interior at the current point, the joint probability that it and all
+        previous points are in the interior.
+    """
+    interval_left = interval_left[:, np.newaxis]
+    prob_conditional = binom.pmf(interval_right, ndraws - interval_left, p, loc=interval_left)
+    prob_right = prob_left.dot(prob_conditional)
+    return prob_right
+
+
+@vectorize(["float64(int64, int64, float64, int64)"])
+def _binom_pmf(k, n, p, loc):
+    k -= loc
+    if k < 0 or k > n:
+        return 0.0
+    if p == 0:
+        return 1.0 if k == 0 else 0.0
+    if p == 1:
+        return 1.0 if k == n else 0.0
+    if k == 0:
+        return (1 - p) ** n
+    if k == n:
+        return p**n
+    lbinom = math.lgamma(n + 1) - math.lgamma(k + 1) - math.lgamma(n - k + 1)
+    return np.exp(lbinom + k * np.log(p) + (n - k) * np.log1p(-p))
+
+
+@jit(nopython=True)
+def _update_ecdf_band_interior_probabilities_numba(
+    prob_left: np.ndarray,
+    interval_left: np.ndarray,
+    interval_right: np.ndarray,
+    p: float,
+    ndraws: int,
+) -> np.ndarray:
+    interval_left = interval_left[:, np.newaxis]
+    prob_conditional = _binom_pmf(interval_right, ndraws - interval_left, p, interval_left)
+    prob_right = prob_left.dot(prob_conditional)
+    return prob_right
+
+
+def _ecdf_band_interior_probability(prob_between_points, ndraws, lower_count, upper_count):
+    interval_left = np.arange(1)
+    prob_interior = np.ones(1)
+    for i in range(prob_between_points.shape[0]):
+        interval_right = np.arange(lower_count[i], upper_count[i])
+        prob_interior = _update_ecdf_band_interior_probabilities(
+            prob_interior, interval_left, interval_right, prob_between_points[i], ndraws
+        )
+        interval_left = interval_right
+    return prob_interior.sum()
+
+
+@jit(nopython=True)
+def _ecdf_band_interior_probability_numba(prob_between_points, ndraws, lower_count, upper_count):
+    interval_left = np.arange(1)
+    prob_interior = np.ones(1)
+    for i in range(prob_between_points.shape[0]):
+        interval_right = np.arange(lower_count[i], upper_count[i])
+        prob_interior = _update_ecdf_band_interior_probabilities_numba(
+            prob_interior, interval_left, interval_right, prob_between_points[i], ndraws
+        )
+        interval_left = interval_right
+    return prob_interior.sum()
+
+
+def _ecdf_band_optimization_objective(
+    prob_pointwise: float,
+    cdf_at_eval_points: np.ndarray,
+    ndraws: int,
+    prob_target: float,
+) -> float:
+    """Objective function for optimizing the simultaneous confidence band probability."""
+    lower, upper = _get_pointwise_confidence_band(prob_pointwise, ndraws, cdf_at_eval_points)
+    lower_count = (lower * ndraws).astype(int)
+    upper_count = (upper * ndraws).astype(int) + 1
+    cdf_with_zero = np.insert(cdf_at_eval_points[:-1], 0, 0)
+    prob_between_points = (cdf_at_eval_points - cdf_with_zero) / (1 - cdf_with_zero)
+    if Numba.numba_flag:
+        prob_interior = _ecdf_band_interior_probability_numba(
+            prob_between_points, ndraws, lower_count, upper_count
+        )
+    else:
+        prob_interior = _ecdf_band_interior_probability(
+            prob_between_points, ndraws, lower_count, upper_count
+        )
+    return abs(prob_interior - prob_target)
+
+
+def _optimize_simultaneous_ecdf_band_probability(
+    ndraws: int,
+    eval_points: np.ndarray,  # pylint: disable=unused-argument
+    cdf_at_eval_points: np.ndarray,
+    prob: float = 0.95,
+    **kwargs,  # pylint: disable=unused-argument
+):
+    """Estimate probability for simultaneous confidence band using optimization.
+
+    This function simulates the pointwise probability needed to construct pointwise confidence bands
+    that form a `prob`-level confidence envelope for the ECDF of a sample.
+    """
+    cdf_at_eval_points = np.unique(cdf_at_eval_points)
+    objective = lambda p: _ecdf_band_optimization_objective(p, cdf_at_eval_points, ndraws, prob)
+    prob_pointwise = minimize_scalar(objective, bounds=(prob, 1), method="bounded").x
+    return prob_pointwise
+
+
 def _simulate_simultaneous_ecdf_band_probability(
     ndraws: int,
     eval_points: np.ndarray,
     cdf_at_eval_points: np.ndarray,
     prob: float = 0.95,
     rvs: Optional[Callable[[int, Optional[Any]], np.ndarray]] = None,
-    num_trials: int = 1000,
+    num_trials: int = 500,
     random_state: Optional[Any] = None,
 ) -> float:
     """Estimate probability for simultaneous confidence band using simulation.

arviz/stats/stats.py

@@ -270,12 +270,12 @@ def compare(
             weights[i] = u_weights / np.sum(u_weights)
 
         weights = weights.mean(axis=0)
-        ses = pd.Series(z_bs.std(axis=0), index=names)  # pylint: disable=no-member
+        ses = pd.Series(z_bs.std(axis=0), index=ics.index)  # pylint: disable=no-member
 
     elif method.lower() == "pseudo-bma":
         min_ic = ics.iloc[0][f"elpd_{ic}"]
         z_rv = np.exp((ics[f"elpd_{ic}"] - min_ic) / scale_value)
-        weights = z_rv / np.sum(z_rv)
+        weights = (z_rv / np.sum(z_rv)).to_numpy()
         ses = ics["se"]
 
     if np.any(weights):
@@ -471,7 +471,7 @@ def hdi(
         Refer to documentation of :func:`arviz.convert_to_dataset` for details.
     hdi_prob: float, optional
         Prob for which the highest density interval will be computed. Defaults to
-        ``stats.hdi_prob`` rcParam.
+        ``stats.ci_prob`` rcParam.
     circular: bool, optional
         Whether to compute the hdi taking into account `x` is a circular variable
         (in the range [-np.pi, np.pi]) or not. Defaults to False (i.e non-circular variables).
@@ -553,7 +553,7 @@ def hdi(
 
     """
     if hdi_prob is None:
-        hdi_prob = rcParams["stats.hdi_prob"]
+        hdi_prob = rcParams["stats.ci_prob"]
     elif not 1 >= hdi_prob > 0:
         raise ValueError("The value of hdi_prob should be in the interval (0, 1]")
 
@@ -711,15 +711,19 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
     Returns
     -------
     ELPDData object (inherits from :class:`pandas.Series`) with the following row/attributes:
-    elpd: approximated expected log pointwise predictive density (elpd)
+    elpd_loo: approximated expected log pointwise predictive density (elpd)
     se: standard error of the elpd
     p_loo: effective number of parameters
-    shape_warn: bool
-        True if the estimated shape parameter of
-        Pareto distribution is greater than 0.7 for one or more samples
-    loo_i: array of pointwise predictive accuracy, only if pointwise True
+    n_samples: number of samples
+    n_data_points: number of data points
+    warning: bool
+        True if the estimated shape parameter of Pareto distribution is greater than
+        ``good_k``.
+    loo_i: :class:`~xarray.DataArray` with the pointwise predictive accuracy,
+            only if pointwise=True
     pareto_k: array of Pareto shape values, only if pointwise True
     scale: scale of the elpd
+    good_k: For a sample size S, the thresold is compute as min(1 - 1/log10(S), 0.7)
 
         The returned object has a custom print method that overrides pd.Series method.
 
@@ -785,13 +789,15 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
     log_weights += log_likelihood
 
     warn_mg = False
-    if np.any(pareto_shape > 0.7):
+    good_k = min(1 - 1 / np.log10(n_samples), 0.7)
+
+    if np.any(pareto_shape > good_k):
         warnings.warn(
-            "Estimated shape parameter of Pareto distribution is greater than 0.7 for "
-            "one or more samples. You should consider using a more robust model, this is because "
-            "importance sampling is less likely to work well if the marginal posterior and "
-            "LOO posterior are very different. This is more likely to happen with a non-robust "
-            "model and highly influential observations."
+            f"Estimated shape parameter of Pareto distribution is greater than {good_k:.2f} "
+            "for one or more samples. You should consider using a more robust model, this is "
+            "because importance sampling is less likely to work well if the marginal posterior "
+            "and LOO posterior are very different. This is more likely to happen with a "
+            "non-robust model and highly influential observations."
         )
         warn_mg = True
 
@@ -816,8 +822,17 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
 
     if not pointwise:
         return ELPDData(
-            data=[loo_lppd, loo_lppd_se, p_loo, n_samples, n_data_points, warn_mg, scale],
-            index=["elpd_loo", "se", "p_loo", "n_samples", "n_data_points", "warning", "scale"],
+            data=[loo_lppd, loo_lppd_se, p_loo, n_samples, n_data_points, warn_mg, scale, good_k],
+            index=[
+                "elpd_loo",
+                "se",
+                "p_loo",
+                "n_samples",
+                "n_data_points",
+                "warning",
+                "scale",
+                "good_k",
+            ],
         )
     if np.equal(loo_lppd, loo_lppd_i).all():  # pylint: disable=no-member
         warnings.warn(
@@ -835,6 +850,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
             loo_lppd_i.rename("loo_i"),
             pareto_shape,
             scale,
+            good_k,
         ],
         index=[
             "elpd_loo",
@@ -846,6 +862,7 @@ def loo(data, pointwise=None, var_name=None, reff=None, scale=None):
             "loo_i",
             "pareto_k",
             "scale",
+            "good_k",
         ],
     )
 
@@ -879,7 +896,8 @@ def psislw(log_weights, reff=1.0):
 
     References
     ----------
-    * Vehtari et al. (2015) see https://arxiv.org/abs/1507.02646
+    * Vehtari et al. (2024). Pareto smoothed importance sampling. Journal of Machine
+      Learning Research, 25(72):1-58.
 
     See Also
     --------
@@ -899,6 +917,7 @@ def psislw(log_weights, reff=1.0):
            ...: az.psislw(-log_likelihood, reff=0.8)
 
     """
+    log_weights = deepcopy(log_weights)
     if hasattr(log_weights, "__sample__"):
         n_samples = len(log_weights.__sample__)
         shape = [
@@ -1322,7 +1341,7 @@ def summary(
     if labeller is None:
         labeller = BaseLabeller()
     if hdi_prob is None:
-        hdi_prob = rcParams["stats.hdi_prob"]
+        hdi_prob = rcParams["stats.ci_prob"]
     elif not 1 >= hdi_prob > 0:
         raise ValueError("The value of hdi_prob should be in the interval (0, 1]")
 
@@ -1565,7 +1584,9 @@ def waic(data, pointwise=None, var_name=None, scale=None, dask_kwargs=None):
     elpd_waic: approximated expected log pointwise predictive density (elpd)
     se: standard error of the elpd
     p_waic: effective number parameters
-    var_warn: bool
+    n_samples: number of samples
+    n_data_points: number of data points
+    warning: bool
         True if posterior variance of the log predictive densities exceeds 0.4
     waic_i: :class:`~xarray.DataArray` with the pointwise predictive accuracy,
             only if pointwise=True

arviz/stats/stats_utils.py

@@ -454,10 +454,9 @@ POINTWISE_LOO_FMT = """------
 
 Pareto k diagnostic values:
                          {{0:>{0}}} {{1:>6}}
-(-Inf, 0.5]   (good)     {{2:{0}d}} {{6:6.1f}}%
- (0.5, 0.7]   (ok)       {{3:{0}d}} {{7:6.1f}}%
-   (0.7, 1]   (bad)      {{4:{0}d}} {{8:6.1f}}%
-   (1, Inf)   (very bad) {{5:{0}d}} {{9:6.1f}}%
+(-Inf, {{8:.2f}}]   (good)     {{2:{0}d}} {{5:6.1f}}%
+   ({{8:.2f}}, 1]   (bad)      {{3:{0}d}} {{6:6.1f}}%
+   (1, Inf)   (very bad) {{4:{0}d}} {{7:6.1f}}%
 """
 SCALE_DICT = {"deviance": "deviance", "log": "elpd", "negative_log": "-elpd"}
 
@@ -488,11 +487,14 @@ class ELPDData(pd.Series):  # pylint: disable=too-many-ancestors
             base += "\n\nThere has been a warning during the calculation. Please check the results."
 
         if kind == "loo" and "pareto_k" in self:
-            bins = np.asarray([-np.inf, 0.5, 0.7, 1, np.inf])
+            bins = np.asarray([-np.inf, self.good_k, 1, np.inf])
             counts, *_ = _histogram(self.pareto_k.values, bins)
             extended = POINTWISE_LOO_FMT.format(max(4, len(str(np.max(counts)))))
             extended = extended.format(
-                "Count", "Pct.", *[*counts, *(counts / np.sum(counts) * 100)]
+                "Count",
+                "Pct.",
+                *[*counts, *(counts / np.sum(counts) * 100)],
+                self.good_k,
             )
             base = "\n".join([base, extended])
         return base

arviz/tests/base_tests/test_data.py

@@ -42,9 +42,12 @@ from ..helpers import (  # pylint: disable=unused-import
     draws,
     eight_schools_params,
     models,
-    running_on_ci,
 )
 
+# Check if dm-tree is installed
+dm_tree_installed = importlib.util.find_spec("tree") is not None  # pylint: disable=invalid-name
+skip_tests = (not dm_tree_installed) and ("ARVIZ_REQUIRE_ALL_DEPS" not in os.environ)
+
 
 @pytest.fixture(autouse=True)
 def no_remote_data(monkeypatch, tmpdir):
@@ -896,6 +899,11 @@ class TestInferenceData:  # pylint: disable=too-many-public-methods
         assert escape(repr(idata)) in html
         xr.set_options(display_style=display_style)
 
+    def test_setitem(self, data_random):
+        data_random["new_group"] = data_random.posterior
+        assert "new_group" in data_random.groups()
+        assert hasattr(data_random, "new_group")
+
     def test_add_groups(self, data_random):
         data = np.random.normal(size=(4, 500, 8))
         idata = data_random
@@ -1077,6 +1085,7 @@ def test_dict_to_dataset():
     assert set(dataset.b.coords) == {"chain", "draw", "c"}
 
 
+@pytest.mark.skipif(skip_tests, reason="test requires dm-tree which is not installed")
 def test_nested_dict_to_dataset():
     datadict = {
         "top": {"a": np.random.randn(100), "b": np.random.randn(1, 100, 10)},
@@ -1469,7 +1478,7 @@ class TestJSON:
 
 
 @pytest.mark.skipif(
-    not (importlib.util.find_spec("datatree") or running_on_ci()),
+    not (importlib.util.find_spec("datatree") or "ARVIZ_REQUIRE_ALL_DEPS" in os.environ),
     reason="test requires xarray-datatree library",
 )
 class TestDataTree:

arviz/tests/base_tests/test_data_zarr.py

@@ -16,7 +16,6 @@ from ..helpers import (  # pylint: disable=unused-import
     draws,
     eight_schools_params,
     importorskip,
-    running_on_ci,
 )
 
 zarr = importorskip("zarr")  # pylint: disable=invalid-name

arviz/tests/base_tests/test_diagnostics_numba.py

@@ -1,7 +1,5 @@
 """Test Diagnostic methods"""
 
-import importlib
-
 # pylint: disable=redefined-outer-name, no-member, too-many-public-methods
 import numpy as np
 import pytest
@@ -11,13 +9,10 @@ from ...rcparams import rcParams
 from ...stats import bfmi, mcse, rhat
 from ...stats.diagnostics import _mc_error, ks_summary
 from ...utils import Numba
-from ..helpers import running_on_ci
+from ..helpers import importorskip
 from .test_diagnostics import data  # pylint: disable=unused-import
 
-pytestmark = pytest.mark.skipif(  # pylint: disable=invalid-name
-    (importlib.util.find_spec("numba") is None) and not running_on_ci(),
-    reason="test requires numba which is not installed",
-)
+importorskip("numba")
 
 rcParams["data.load"] = "eager"
 

arviz/tests/base_tests/test_helpers.py

@@ -6,13 +6,13 @@ from ..helpers import importorskip
 
 def test_importorskip_local(monkeypatch):
     """Test ``importorskip`` run on local machine with non-existent module, which should skip."""
-    monkeypatch.delenv("ARVIZ_CI_MACHINE", raising=False)
+    monkeypatch.delenv("ARVIZ_REQUIRE_ALL_DEPS", raising=False)
     with pytest.raises(Skipped):
         importorskip("non-existent-function")
 
 
 def test_importorskip_ci(monkeypatch):
     """Test ``importorskip`` run on CI machine with non-existent module, which should fail."""
-    monkeypatch.setenv("ARVIZ_CI_MACHINE", 1)
+    monkeypatch.setenv("ARVIZ_REQUIRE_ALL_DEPS", 1)
     with pytest.raises(ModuleNotFoundError):
         importorskip("non-existent-function")