adding HDP based coverage test

README.md

@@ -137,25 +137,44 @@ greater than the current one.
 
 ### Coverage test
 
-In the coverage test we have some number of simulations `nsim` where there is a
-true value `g` and some posterior samples `s`. The procedure goes like this,
-first you sample from your prior: `g ~ Prior(G)`. Then you sample from your
-likelihood: `x ~ Likelihood(X | g)`. Then you sample from your posterior: 
-`s ~ Posterior(S | x)`, you will want many samples `s`. You repeat this 
-procedure `nsim` times. The `g` and `s` samples are what you need for the test.
+In the coverage test we have some number of simulations ``nsim`` where there is
+a true value ``g`` and some posterior samples ``s``. The procedure goes like
+this, first you sample from your prior: ``g ~ Prior(G)``. Then you sample from
+your likelihood: ``x ~ Likelihood(X | g)``. Then you sample from your posterior:
+``s ~ Posterior(S | x)``, you will want many samples ``s``. You repeat this
+procedure ``nsim`` times. The ``g`` and ``s`` samples are what you need for the
+test.
 
 Internally, for each simulation separately we use PTED to compute a p-value,
-essentially asking the question "was `g` drawn from the distribution that
-generated `s`?". Individually, these tests are possibly not especially
+essentially asking the question "was ``g`` drawn from the distribution that
+generated ``s``?". Individually, these tests are possibly not especially
 informative (unless the sampler is really bad), however their p-values must have
-been drawn from `U(0,1)` under the null-hypothesis[^2]. Thus we just need a way
-to combine their statistical power. It turns out that for some `p ~ U(0,1)`, we
-have that `- 2 ln(p)` is chi2 distributed with `dof = 2`. This means that we can
-sum the chi2 values for the PTED test on each simulation and compare with a chi2
-distribution with `dof = 2 * nsim`. We use a density based two tailed p-value
-test on this chi2 distribution meaning that if your posterior is underconfident
-or overconfident, you will get a small p-value that can be used to reject the
-null.
+been drawn from ``U(0,1)`` under the null-hypothesis[^2]. Thus we just need a
+way to combine their statistical power. It turns out that for some ``p ~
+U(0,1)``, we have that ``- 2 ln(p)`` is chi2 distributed with ``dof = 2``. This
+means that we can sum the chi2 values for the PTED test on each simulation and
+compare with a chi2 distribution with ``dof = 2 * nsim``. We use a simple
+doubling procedure (``2 * min(p_right, p_left)``) to get the p-value meaning
+that if your posterior is underconfident or overconfident, you will get a small
+p-value that can be used to reject the null.
+
+#### If you have posterior densities
+
+If you have posterior densities (and you trust them), then chances are the HDP
+region coverage test is a more powerful test. Essentially, instead of using a
+permutation test to determine the p-value for a single simulation, you just
+determine the fraction of posterior samples with higher posterior density than
+the ground truth[^3]. The package has a quick tool to let you do that:
+
+```python
+from pted import hdp_coverage_test
+
+ground_truth = np.random.uniform(100) # Nsim
+posterior_samples = np.random.uniform((200, 100)) # Nsamp, Nsim
+
+p_value = hdp_coverage_test(ground_truth, posterior_samples, two_tailed = True)
+print(f"p-value: {p_value:.3f}") # expect uniform random from 0-1
+```
 
 ## Example: Sensitivity comparison with KS-test
 
@@ -259,7 +278,7 @@ could not find any significant discrepancies. The samples could have been drawn
 from the same distribution, or PTED could be insensitive to the deviation, or
 maybe the test needs more samples. In some sense PTED (like all null hypothesis
 tests) is "necessary but not sufficient" in that failing the test is bad news
-for the null, but passing the test is possibly inconclusive[^3]. Use your judgement,
+for the null, but passing the test is possibly inconclusive[^4]. Use your judgement,
 and contact me or some smarter stat-oriented person if you are unsure about the
 results you are getting!
 
@@ -491,4 +510,5 @@ If you think those are neat, then you'll probably also like this paper, which us
 
 [^1]: See the Simulation-Based Calibration paper by Talts et al. 2018 for what "SBC" is.
 [^2]: Since PTED works by a permutation test, we only get the p-value from a discrete uniform distribution. By default we use 1000 permutations, if you are running an especially sensitive test you may need more permutations, but for most purposes this is sufficient.
-[^3]: actual "necessary but not sufficient" conditions are a different thing than null hypothesis tests, but they have a similar intuitive meaning.
+[^3]: Actually, we take (q + 1) / (Nsamp + 1) rather than q/Nsamp where q is the number of posterior samples with posterior density greater than the ground truth. This just turns out to be a better estimator for finite Nsamp.
+[^4]: actual "necessary but not sufficient" conditions are a different thing than null hypothesis tests, but they have a similar intuitive meaning.

src/pted/__init__.py

@@ -1,5 +1,6 @@
 from .pted import pted, pted_coverage_test
 from .tests import test
+from .utils import hdp_coverage_test
 from ._version import version as __version__  # noqa
 
 __author__ = "Connor Stone"
@@ -8,6 +9,7 @@
 __all__ = [
     "pted",
     "pted_coverage_test",
+    "hdp_coverage_test",
     "test",
     "__version__",
     "__author__",

src/pted/utils.py

@@ -306,30 +306,38 @@ def pted_chunk_jax(
 
 def two_tailed_p(chi2, df):
     assert df > 2, "Degrees of freedom must be greater than 2 for two-tailed p-value calculation."
-    alpha = chi2_dist.pdf(chi2, df)
-    mode = df - 2
+    p_left = chi2_dist.cdf(chi2, df)
+    p_right = chi2_dist.sf(chi2, df)
+    return 2 * min(p_left, p_right)
 
-    if np.isclose(chi2, mode):
-        return 1.0
 
-    def root_eq(x):
-        return chi2_dist.pdf(x, df) - alpha
+###### This is a density based two tailed p-value, it is kept for reference but not used #######
+# def two_tailed_p(chi2, df):
+#     assert df > 2, "Degrees of freedom must be greater than 2 for two-tailed p-value calculation."
+#     alpha = chi2_dist.pdf(chi2, df)
+#     mode = df - 2
 
-    # Find left root
-    if chi2 < mode:
-        left = chi2_dist.cdf(chi2, df)
-    else:
-        res_left = root_scalar(root_eq, bracket=[0, mode], method="brentq")
-        left = chi2_dist.cdf(res_left.root, df)
+#     if np.isclose(chi2, mode):
+#         return 1.0
 
-    # Find right root
-    if chi2 > mode:
-        right = chi2_dist.sf(chi2, df)
-    else:
-        res_right = root_scalar(root_eq, bracket=[mode, 10000 * df], method="brentq")
-        right = chi2_dist.sf(res_right.root, df)
+#     def root_eq(x):
+#         return chi2_dist.pdf(x, df) - alpha
 
-    return left + right
+#     # Find left root
+#     if chi2 < mode:
+#         left = chi2_dist.cdf(chi2, df)
+#     else:
+#         res_left = root_scalar(root_eq, bracket=[0, mode], method="brentq")
+#         left = chi2_dist.cdf(res_left.root, df)
+
+#     # Find right root
+#     if chi2 > mode:
+#         right = chi2_dist.sf(chi2, df)
+#     else:
+#         res_right = root_scalar(root_eq, bracket=[mode, 10000 * df], method="brentq")
+#         right = chi2_dist.sf(res_right.root, df)
+
+#     return left + right
 
 
 class OverconfidenceWarning(UserWarning):
@@ -462,3 +470,34 @@ def pit_plot(pvals, saveto, confidence=0.95):
     ax.legend()
     fig.savefig(saveto, bbox_inches="tight")
     plt.close(fig)
+
+
+def hdp_coverage_test(
+    ground_truth: np.ndarray, posterior_samples: np.ndarray, two_tailed: bool = True
+) -> float:
+    """
+    Perform a Highest Density Posterior (HDP) coverage test. Essentially this
+    rank orders the posterior samples by their posterior density and also places
+    the ground truth in that ranking. The fraction of posterior samples with
+    higher density than the ground truth forms a p-value under the null
+    hypothesis. For many repeated experiments, we check that the p-values are
+    uniformly distributed.
+
+    Args:
+        ground_truth: Posterior density (or log density) at the ground-truth parameters, shape (Nsim,)
+        posterior_samples: Posterior density (or log density) for each posterior draw, shape (Nsamp, Nsim)
+        two_tailed: Whether to compute a two-tailed p-value (default: True)
+
+    Returns:
+        pvalue: The p-value for the coverage test
+    """
+    from scipy.stats import chi2 as chi2_dist
+
+    Nsamp, Nsim = posterior_samples.shape
+    q = np.sum(posterior_samples >= ground_truth[None], axis=0)
+    chi2_hdp = -2 * np.sum(np.log((q + 1) / (Nsamp + 1)))
+    pvalue_right = chi2_dist.sf(chi2_hdp, 2 * Nsim)
+    pvalue_left = chi2_dist.cdf(chi2_hdp, 2 * Nsim)
+    if two_tailed:
+        return 2 * min(pvalue_left, pvalue_right)
+    return pvalue_right

tests/test_utils.py

@@ -2,14 +2,19 @@
 
 import numpy as np
 
-from pted.utils import two_tailed_p, simulation_based_calibration_histogram, pit_plot
+from pted.utils import (
+    two_tailed_p,
+    simulation_based_calibration_histogram,
+    pit_plot,
+    hdp_coverage_test,
+)
 
 import pytest
 
 
 def test_two_tailed_p():
 
-    assert np.isclose(two_tailed_p(4, 6), 1.0), "p-value at mode should be 1.0"
+    # assert np.isclose(two_tailed_p(4, 6), 1.0), "p-value at mode should be 1.0"
 
     assert two_tailed_p(0.01, 10) < 0.01, "p-value should be less than 0.01 for small chi2"
     assert two_tailed_p(100, 10) < 0.01, "p-value should be less than 0.01 for large chi2"
@@ -41,3 +46,24 @@ def test_pit_plot_no_matplotlib(monkeypatch):
 
     with pytest.warns(UserWarning, match="matplotlib"):
         pit_plot(pvals, "pit_no_mpl.pdf")
+
+
+def test_hdp_coverage_test():
+    np.random.seed(42)
+    # Null is true
+    ground_truth = np.random.normal(loc=0, scale=1, size=128)
+    posterior_samples = np.random.normal(loc=0, scale=1, size=(1024, 128))
+    pvalue = hdp_coverage_test(ground_truth, posterior_samples)
+    assert 1e-3 <= pvalue <= 0.999, "p-value should be between 0 and 1"
+
+    # Posterior is biased
+    posterior_samples = np.random.normal(loc=5, scale=1, size=(1024, 128))
+    pvalue = hdp_coverage_test(ground_truth, posterior_samples)
+    assert pvalue < 0.01, "p-value should be small for poorly calibrated posterior samples"
+
+    # Posterior is underconfident
+    posterior_samples = np.random.normal(loc=0, scale=2, size=(1024, 128))
+    pvalue = hdp_coverage_test(ground_truth, posterior_samples)
+    assert pvalue < 0.01, "p-value should be small for poorly calibrated posterior samples"
+    pvalue = hdp_coverage_test(ground_truth, posterior_samples, two_tailed=False)
+    assert pvalue > 0.01, "p-value should not be small for underconfident and one_tailed test"