Use ci_prob for the tail ESS in pareto_khat, power_scale_lw and thin_factor

src/arviz_stats/base/array.py

@@ -320,7 +320,7 @@ def mcse(self, ary, chain_axis=-2, draw_axis=-1, method="mean", prob=None, circu
         mcse_array = make_ufunc(mcse_func, n_output=1, n_input=1, n_dims=2, ravel=False)
         return mcse_array(ary, **func_kwargs)
 
-    def pareto_min_ss(self, ary, chain_axis=-2, draw_axis=-1):
+    def pareto_min_ss(self, ary, chain_axis=-2, draw_axis=-1, prob=0.89):
         """Compute minimum effective sample size.
 
         See docstring of :func:`arviz_stats.pareto_min_ss` for full description of computation
@@ -331,11 +331,13 @@ def pareto_min_ss(self, ary, chain_axis=-2, draw_axis=-1):
         ary : array-like
         chain_axis : int, default -2
         draw_axis : int, default -1
+        prob : float, default 0.89
+            Probability for the tail ESS used to compute the relative efficiency.
         """
         ary, chain_axis, draw_axis = process_chain_none(ary, chain_axis, draw_axis)
         ary, _ = process_ary_axes(ary, [chain_axis, draw_axis])
         pms_array = make_ufunc(self._pareto_min_ss, n_output=1, n_input=1, n_dims=2, ravel=False)
-        return pms_array(ary)
+        return pms_array(ary, prob=prob)
 
     def psislw(self, ary, r_eff=1, axis=-1):
         """Compute log weights for Pareto-smoothed importance sampling (PSIS) method.
@@ -377,7 +379,14 @@ def bfmi(self, ary, chain_axis=-2, draw_axis=-1):
         return bfmi_ufunc(ary, out_shape=(ary.shape[chain_axis],))
 
     def pareto_khat(
-        self, ary, chain_axis=-2, draw_axis=-1, r_eff=None, tail="both", log_weights=False
+        self,
+        ary,
+        chain_axis=-2,
+        draw_axis=-1,
+        r_eff=None,
+        tail="both",
+        log_weights=False,
+        prob=0.89,
     ):
         """Compute Pareto k-hat diagnostic.
 
@@ -389,26 +398,30 @@ def pareto_khat(
         r_eff : float
         tail : str
         log_weights : bool
+        prob : float, default 0.89
+            Probability for the tail ESS used to compute `r_eff` when it is not given.
         """
         ary, chain_axis, draw_axis = process_chain_none(ary, chain_axis, draw_axis)
         ary, _ = process_ary_axes(ary, [chain_axis, draw_axis])
 
         def _khat_wrapper(current_slice):
             return self._pareto_khat(
-                current_slice, r_eff=r_eff, tail=tail, log_weights=log_weights
+                current_slice, prob, r_eff=r_eff, tail=tail, log_weights=log_weights
             )[1]
 
         khat_ufunc = make_ufunc(_khat_wrapper, n_output=1, n_input=1, n_dims=2, ravel=False)
         return khat_ufunc(ary, out_shape=())
 
-    def power_scale_lw(self, ary, alpha=0, axis=-1):
+    def power_scale_lw(self, ary, alpha=0, axis=-1, prob=0.89):
         """Compute log weights for power-scaling component by alpha.
 
         Parameters
         ----------
         ary : array-like
         alpha : float, default 0
         axis : int, sequence of int or None, default -1
+        prob : float, default 0.89
+            Probability for the tail ESS used to compute the relative efficiency.
         """
         ary, axes = process_ary_axes(ary, axis)
         sample_shape = tuple(ary.shape[i] for i in axes)
@@ -419,7 +432,7 @@ def power_scale_lw(self, ary, alpha=0, axis=-1):
             n_dims=len(axes),
             ravel=False,
         )
-        return psl_ufunc(ary, out_shape=(sample_shape, ()), alpha=alpha)
+        return psl_ufunc(ary, out_shape=(sample_shape, ()), alpha=alpha, prob=prob)
 
     def power_scale_sense(
         self, ary, lower_w, upper_w, lower_alpha, upper_alpha, chain_axis=-2, draw_axis=-1
@@ -1101,7 +1114,9 @@ def metrics(self, observed, predicted, kind):
         )
         return metrics_ufunc(observed, predicted)
 
-    def thin_factor(self, ary, target_ess=None, reduce_func="mean", chain_axis=-2, draw_axis=-1):
+    def thin_factor(
+        self, ary, target_ess=None, reduce_func="mean", chain_axis=-2, draw_axis=-1, prob=0.89
+    ):
         """Get thinning factor over draw dimension to preserve ESS in samples or target a given ESS.
 
         Parameters
@@ -1114,6 +1129,8 @@ def thin_factor(self, ary, target_ess=None, reduce_func="mean", chain_axis=-2, d
         reduce_func : {"mean", "min"}, default "mean"
         chain_axis : int, default -2
         draw_axis : int, default -1
+        prob : float, default 0.89
+            Probability for the tail ESS.
 
         Returns
         -------
@@ -1125,7 +1142,7 @@ def thin_factor(self, ary, target_ess=None, reduce_func="mean", chain_axis=-2, d
 
         bulk_ess = self.ess(ary, chain_axis=chain_axis, draw_axis=draw_axis, method="bulk")
         tail_ess = self.ess(
-            ary, chain_axis=chain_axis, draw_axis=draw_axis, method="tail", prob=0.95
+            ary, chain_axis=chain_axis, draw_axis=draw_axis, method="tail", prob=prob
         )
         ess = np.minimum(bulk_ess, tail_ess)
 

src/arviz_stats/base/dataarray.py

@@ -637,15 +637,16 @@ def thin(self, da, factor="auto", sample_dims=None):
 
         return da.sel({sample_dims: slice(None, None, factor)})
 
-    def pareto_min_ss(self, da, sample_dims=None):
+    def pareto_min_ss(self, da, sample_dims=None, prob=None):
         """Compute the minimum effective sample size."""
         dims, chain_axis, draw_axis = validate_dims_chain_draw_axis(sample_dims)
+        prob = validate_ci_prob(prob)
         return apply_ufunc(
             self.array_class.pareto_min_ss,
             da,
             input_core_dims=[dims],
             output_core_dims=[[]],
-            kwargs={"chain_axis": chain_axis, "draw_axis": draw_axis},
+            kwargs={"chain_axis": chain_axis, "draw_axis": draw_axis, "prob": prob},
         )
 
     def psislw(self, da, r_eff=1, dim=None):
@@ -674,9 +675,12 @@ def bfmi(self, da, sample_dims=None):
             },
         )
 
-    def pareto_khat(self, da, sample_dims=None, r_eff=None, tail="both", log_weights=False):
+    def pareto_khat(
+        self, da, sample_dims=None, r_eff=None, tail="both", log_weights=False, prob=None
+    ):
         """Compute Pareto k-hat diagnostic on DataArray input."""
         dims, chain_axis, draw_axis = validate_dims_chain_draw_axis(sample_dims)
+        prob = validate_ci_prob(prob)
         return apply_ufunc(
             self.array_class.pareto_khat,
             da,
@@ -688,6 +692,7 @@ def pareto_khat(self, da, sample_dims=None, r_eff=None, tail="both", log_weights
                 "r_eff": r_eff,
                 "tail": tail,
                 "log_weights": log_weights,
+                "prob": prob,
             },
             dask="parallelized",
             output_dtypes=[float],
@@ -1184,16 +1189,17 @@ def loo_r2(
             random_state=random_state,
         )
 
-    def power_scale_lw(self, da, alpha=0, dim=None):
+    def power_scale_lw(self, da, alpha=0, dim=None, prob=None):
         """Compute log weights for power-scaling component by alpha."""
         dims = validate_dims(dim)
+        prob = validate_ci_prob(prob)
         return apply_ufunc(
             self.array_class.power_scale_lw,
             da,
             alpha,
             input_core_dims=[dims, []],
             output_core_dims=[dims, []],
-            kwargs={"axis": np.arange(-len(dims), 0, 1)},
+            kwargs={"axis": np.arange(-len(dims), 0, 1), "prob": prob},
         )
 
     def power_scale_sense(self, da, lower_w, upper_w, lower_alpha, upper_alpha, sample_dims=None):

src/arviz_stats/base/diagnostics.py

@@ -340,7 +340,7 @@ def _mcse_quantile(self, ary, prob, circular=False):
         th2 = sorted_ary[int(np.ceil(np.nanmin((ppf_size[1], size - 1))))]
         return (th2 - th1) / 2
 
-    def _pareto_min_ss(self, ary):
+    def _pareto_min_ss(self, ary, prob):
         """Compute the minimum effective sample size."""
         ary = np.asarray(ary)
         if _not_valid(ary, check_nan=False, shape_kwargs={"min_draws": 4, "min_chains": 1}):
@@ -355,7 +355,7 @@ def _pareto_min_ss(self, ary):
             )
             return np.nan
 
-        _, kappa = self._pareto_khat(ary, tail="both", log_weights=False)
+        _, kappa = self._pareto_khat(ary, prob, tail="both", log_weights=False)
 
         if np.isnan(kappa):
             return np.nan
@@ -1077,7 +1077,7 @@ def _pareto_pit_single(self, draws, y_val, lw, rng, gpd_ok, tail_ids, ndraws_tai
 
         return float(np.clip(raw_pit, min_tail_prob, 1.0 - min_tail_prob))
 
-    def _pareto_khat(self, ary, r_eff=None, tail="both", log_weights=False):
+    def _pareto_khat(self, ary, prob, r_eff=None, tail="both", log_weights=False):
         """
         Compute Pareto k-hat diagnostic.
 
@@ -1086,6 +1086,8 @@ def _pareto_khat(self, ary, r_eff=None, tail="both", log_weights=False):
         Parameters
         ----------
         ary : Array
+        prob : float
+            Probability for the tail ESS used to compute `r_eff` when it is not given.
         r_eff : float, optional
             Relative efficiency. Effective sample size divided the number of samples.
         tail : srt, optional
@@ -1103,7 +1105,7 @@ def _pareto_khat(self, ary, r_eff=None, tail="both", log_weights=False):
             tail = "right"
 
         if r_eff is None:
-            r_eff = self._ess_tail(ary, prob=0.05, relative=True)
+            r_eff = self._ess_tail(ary, prob=prob, relative=True)
 
         ary = ary.flatten()
         n_draws = len(ary)
@@ -1371,13 +1373,13 @@ def _power_scale_sense(self, ary, lower_w, upper_w, lower_alpha, upper_alpha):
         upper_grad = upper_cjs / np.log2(upper_alpha)
         return (lower_grad + upper_grad) / 2
 
-    def _power_scale_lw(self, ary, alpha):
+    def _power_scale_lw(self, ary, alpha, prob):
         """Compute log weights for power-scaling component by alpha."""
         shape = ary.shape
         ary = np.ravel(ary)
         log_weights = (alpha - 1) * ary
         n_draws = len(log_weights)
-        r_eff = self._ess_tail(ary, prob=0.05, relative=True)
+        r_eff = self._ess_tail(ary, prob=prob, relative=True)
         n_draws_tail = self._get_ps_tails(n_draws, r_eff, tail="both")
         log_weights, pareto_k = self._ps_tail(
             log_weights,

tests/base/test_stats.py

@@ -12,6 +12,8 @@
 xr = importorskip("xarray")
 
 import arviz_stats.accessors  # noqa: F401  # pylint: disable=unused-import
+from arviz_stats.base import array_stats
+from arviz_stats.base.dataarray import dataarray_stats
 
 
 @pytest.fixture(scope="module")
@@ -358,3 +360,25 @@ def test_ecdf_boundary_values(centered_eight):
     y_vals = result.sel(plot_axis="y")
     assert np.all(y_vals >= 0)
     assert np.all(y_vals <= 1)
+
+
+@pytest.mark.parametrize("func", ["pareto_khat", "pareto_min_ss"])
+def test_pareto_ci_prob(centered_eight, func):
+    da = centered_eight.posterior["theta"]
+    with azb.rc_context(rc={"stats.ci_prob": 0.8}):
+        result = getattr(da.azstats, func)()
+    expected = getattr(array_stats, func)(da.values, chain_axis=0, draw_axis=1, prob=0.8)
+    assert_array_almost_equal(result, expected)
+
+
+def test_power_scale_lw_ci_prob(centered_eight):
+    da = centered_eight.posterior["mu"]
+    with azb.rc_context(rc={"stats.ci_prob": 0.8}):
+        _, khat = da.azstats.power_scale_lw(alpha=0.5, dim=["chain", "draw"])
+    _, expected = array_stats.power_scale_lw(da.values, alpha=0.5, axis=(0, 1), prob=0.8)
+    assert_array_almost_equal(khat, expected)
+
+
+def test_thin_factor_array_dataarray(centered_eight):
+    da = centered_eight.posterior["mu"]
+    assert dataarray_stats.thin_factor(da) == array_stats.thin_factor(da.values)