Add optimized simultaneous ECDF bands

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## v0.x.x Unreleased
 
 ### New features
+- Add optimized simultaneous ECDF confidence bands ([2368](https://github.com/arviz-devs/arviz/pull/2368))
 
 ### Maintenance and fixes
 

arviz/plots/ecdfplot.py

@@ -73,6 +73,7 @@ def plot_ecdf(
         - False: No confidence bands are plotted (default).
         - True: Plot bands computed with the default algorithm (subject to change)
         - "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.
 
@@ -216,8 +217,7 @@ def plot_ecdf(
         >>> pit_vals = distribution.cdf(sample)
         >>> uniform_dist = uniform(0, 1)
         >>> az.plot_ecdf(
-        >>>     pit_vals, cdf=uniform_dist.cdf,
-        >>>     rvs=uniform_dist.rvs, confidence_bands=True
+        >>>     pit_vals, cdf=uniform_dist.cdf, confidence_bands=True,
         >>> )
 
     Plot an ECDF-difference plot of PIT values.
@@ -226,8 +226,8 @@ def plot_ecdf(
         :context: close-figs
 
         >>> az.plot_ecdf(
-        >>>     pit_vals, cdf = uniform_dist.cdf, rvs = uniform_dist.rvs,
-        >>>     confidence_bands = True, difference = True
+        >>>     pit_vals, cdf = uniform_dist.cdf, confidence_bands = True,
+        >>>     difference = True
         >>> )
     """
     if confidence_bands is True:
@@ -238,9 +238,12 @@ def plot_ecdf(
             )
             confidence_bands = "pointwise"
         else:
-            confidence_bands = "simulated"
-    elif confidence_bands == "simulated" and pointwise:
-        raise ValueError("Cannot specify both `confidence_bands='simulated'` and `pointwise=True`")
+            confidence_bands = "optimized"
+        # if pointwise especified, confidence_bands must be a bool or 'pointwise'
+    elif confidence_bands not in [False, "pointwise"] and pointwise:
+        raise ValueError(
+            f"Cannot specify both `confidence_bands='{confidence_bands}'` and `pointwise=True`"
+        )
 
     if fpr is not None:
         warnings.warn(
@@ -298,7 +301,7 @@ def plot_ecdf(
             "`eval_points` explicitly.",
             BehaviourChangeWarning,
         )
-        if confidence_bands == "simulated":
+        if confidence_bands in ["optimized", "simulated"]:
             warnings.warn(
                 "For simultaneous bands to be correctly calibrated, specify `eval_points` "
                 "independent of the `values`"

arviz/stats/ecdf_utils.py

@@ -5,6 +5,7 @@
 
 import numpy as np
 from scipy.stats import uniform, binom
+from scipy.optimize import minimize_scalar
 
 
 def compute_ecdf(sample: np.ndarray, eval_points: np.ndarray) -> np.ndarray:
@@ -73,7 +74,7 @@
     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.
@@ -92,6 +93,7 @@
     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
@@ -115,12 +117,18 @@
 
     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
@@ -129,6 +137,90 @@
     return prob_lower, prob_upper
 
 
+def _update_ecdf_band_interior_probabilities(
+    prob_left: np.ndarray,
+    interval_left: np.ndarray,
+    interval_right: np.ndarray,
+    cdf_left: float,
+    cdf_right: 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.
+    cdf_left : float
+        The CDF at the previous point.
+    cdf_right : float
+        The CDF at the current point.
+    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.
+    """
+    z_tilde = (cdf_right - cdf_left) / (1 - cdf_left)
+    interval_left = interval_left[:, np.newaxis]
+    prob_conditional = binom.pmf(interval_right, ndraws - interval_left, z_tilde, loc=interval_left)
+    prob_right = prob_left.dot(prob_conditional)
+    return prob_right
+
+
+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
+
+    interval_left = np.zeros(1)
+    cdf_left = 0
+    prob_interior = np.ones(1)
+    for i in range(cdf_at_eval_points.shape[0]):
+        interval_right = np.arange(lower_count[i], upper_count[i])
+        cdf_right = cdf_at_eval_points[i]
+        prob_interior = _update_ecdf_band_interior_probabilities(
+            prob_interior, interval_left, interval_right, cdf_left, cdf_right, ndraws
+        )
+        interval_left = interval_right
+        cdf_left = cdf_right
+    prob_interior = prob_interior.sum()
+    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,

arviz/tests/base_tests/test_plots_matplotlib.py

@@ -1285,7 +1285,7 @@ def test_plot_ecdf_eval_points():
     assert axes is not None
 
 
-@pytest.mark.parametrize("confidence_bands", [True, "pointwise", "simulated"])
+@pytest.mark.parametrize("confidence_bands", [True, "pointwise", "optimized", "simulated"])
 def test_plot_ecdf_confidence_bands(confidence_bands):
     """Check that all confidence_bands values correctly accepted"""
     data = np.random.randn(4, 1000)
@@ -1326,6 +1326,8 @@ def test_plot_ecdf_error():
     # contradictory confidence band types
     with pytest.raises(ValueError):
         plot_ecdf(data, cdf=dist.cdf, confidence_bands="simulated", pointwise=True)
+    with pytest.raises(ValueError):
+        plot_ecdf(data, cdf=dist.cdf, confidence_bands="optimized", pointwise=True)
     plot_ecdf(data, cdf=dist.cdf, confidence_bands=True, pointwise=True)
     plot_ecdf(data, cdf=dist.cdf, confidence_bands="pointwise")
 

arviz/tests/base_tests/test_stats_ecdf_utils.py

@@ -109,7 +109,7 @@ def test_get_pointwise_confidence_band(dist, prob, ndraws, num_trials=1_000, see
     ids=["continuous", "continuous default rvs", "discrete"],
 )
 @pytest.mark.parametrize("ndraws", [10_000])
-@pytest.mark.parametrize("method", ["pointwise", "simulated"])
+@pytest.mark.parametrize("method", ["pointwise", "optimized", "simulated"])
 def test_ecdf_confidence_band(dist, rvs, prob, ndraws, method, num_trials=1_000, seed=57):
     """Test test_ecdf_confidence_band."""
     eval_points = np.linspace(*dist.interval(0.99), 10)