Compute exact p-values for small null configurations

src/copairs/compute.py

@@ -3,6 +3,7 @@
 import os
 import warnings
 import itertools
+from math import comb
 from typing import Tuple, Union, Callable, Optional
 from pathlib import Path
 from multiprocessing.pool import ThreadPool
@@ -460,6 +461,72 @@ def random_ap(num_perm: int, num_pos: int, total: int, seed: int):
     return null_dist
 
 
+def exact_ap(num_pos: int, total: int) -> np.ndarray:
+    """Compute exact null distribution by enumerating all possible rankings.
+
+    When the number of possible rankings is small, this computes the exact
+    null distribution instead of sampling.
+
+    Parameters
+    ----------
+    num_pos : int
+        Number of positive samples in each relevance list.
+    total : int
+        Total number of samples in each relevance list.
+
+    Returns
+    -------
+    np.ndarray
+        A 1D array containing the Average Precision scores for all possible
+        rankings of `num_pos` positives among `total` items.
+    """
+    # Generate all combinations of positions for positives
+    # Each combination is a tuple of sorted indices where positives could appear
+    all_combinations = itertools.combinations(range(total), num_pos)
+
+    # Convert to array of shape (n_combinations, num_pos)
+    rel_k = np.array(list(all_combinations), dtype=np.uint16)
+
+    # Compute AP for each combination
+    return average_precision(rel_k)
+
+
+def _compute_null_dist(
+    num_pos: int, total: int, seed: int, null_size: int
+) -> np.ndarray:
+    """Compute null distribution, using exact method when feasible.
+
+    Parameters
+    ----------
+    num_pos : int
+        Number of positive pairs in the configuration.
+    total : int
+        Total number of pairs (positive + negative) in the configuration.
+    seed : int
+        Random seed for reproducibility.
+    null_size : int
+        Number of samples to generate in the null distribution.
+
+    Returns
+    -------
+    np.ndarray
+        Null distribution of shape (null_size,).
+    """
+    n_combinations = comb(total, num_pos)
+
+    if n_combinations <= null_size:
+        # Use exact computation - enumerate all possible rankings
+        exact_dist = exact_ap(num_pos, total)
+        # Tile to fill null_size for consistent interface
+        n_repeats = (null_size + n_combinations - 1) // n_combinations
+        null_dist = np.tile(exact_dist, n_repeats)[:null_size]
+    else:
+        # Use random sampling
+        null_dist = random_ap(null_size, num_pos, total, seed)
+
+    return null_dist
+
+
 def null_dist_cached(
     num_pos: int, total: int, seed: int, null_size: int, cache_dir: Path
 ) -> np.ndarray:
@@ -469,6 +536,10 @@ def null_dist_cached(
     pairs (`num_pos`) and total pairs (`total`). It uses caching to store and
     retrieve precomputed distributions, saving time and computational resources.
 
+    When the total number of possible rankings (C(total, num_pos)) is less than
+    or equal to null_size, the exact null distribution is computed by enumerating
+    all possible rankings instead of random sampling.
+
     Parameters
     ----------
     num_pos : int
@@ -504,19 +575,19 @@ def null_dist_cached(
                 cache_file.unlink(missing_ok=True)
 
                 # Compute the null distribution
-                null_dist = random_ap(null_size, num_pos, total, seed)
+                null_dist = _compute_null_dist(num_pos, total, seed, null_size)
 
                 # Save the new distribution to the cache
                 np.save(cache_file, null_dist)
         else:
             # If the cache file doesn't exist, compute the null distribution
-            null_dist = random_ap(null_size, num_pos, total, seed)
+            null_dist = _compute_null_dist(num_pos, total, seed, null_size)
 
             # Save the computed distribution to the cache
             np.save(cache_file, null_dist)
     else:
         # If no seed is provided, compute the null distribution without caching
-        null_dist = random_ap(null_size, num_pos, total, seed)
+        null_dist = _compute_null_dist(num_pos, total, seed, null_size)
 
     # Return the null distribution (loaded or computed)
     return null_dist

tests/test_compute.py

@@ -1,6 +1,7 @@
 """Test pairwise distance calculation functions."""
 
 import tempfile
+from math import comb
 from pathlib import Path
 
 import numpy as np
@@ -237,3 +238,79 @@ def test_null_dist_cached_corrupt():
         assert len(null_dist) == 100
         assert np.all(null_dist >= 0)
         assert np.all(null_dist <= 1)
+
+
+def test_exact_ap():
+    """Test exact_ap computes AP for all possible rankings."""
+    num_pos = 3
+    total = 6
+    n_combinations = comb(total, num_pos)  # 20 combinations
+
+    exact_dist = compute.exact_ap(num_pos, total)
+
+    # Should have exactly C(total, num_pos) AP scores
+    assert len(exact_dist) == n_combinations
+    assert np.all(exact_dist >= 0)
+    assert np.all(exact_dist <= 1)
+
+    # Best case: all positives at the start (positions 0, 1, 2)
+    # AP = (1/1 + 2/2 + 3/3) / 3 = 1.0
+    assert np.max(exact_dist) == pytest.approx(1.0)
+
+    # Worst case: all positives at the end (positions 3, 4, 5)
+    # AP = (1/4 + 2/5 + 3/6) / 3 = (0.25 + 0.4 + 0.5) / 3 ≈ 0.383
+    assert np.min(exact_dist) == pytest.approx((1 / 4 + 2 / 5 + 3 / 6) / 3)
+
+
+def test_exact_null_dist_used_when_small():
+    """Test that exact computation is used when combinations < null_size."""
+    num_pos = 3
+    total = 6
+    null_size = 100
+    n_combinations = comb(total, num_pos)  # 20 < 100
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        cache_dir = Path(tmpdir)
+
+        null_dist = compute.null_dist_cached(
+            num_pos=num_pos,
+            total=total,
+            seed=42,
+            null_size=null_size,
+            cache_dir=cache_dir,
+        )
+
+        # Should be padded to null_size
+        assert len(null_dist) == null_size
+
+        # Should contain exactly n_combinations unique values (the exact distribution)
+        unique_values = np.unique(null_dist)
+        assert len(unique_values) == n_combinations
+
+
+def test_random_null_dist_used_when_large():
+    """Test that random sampling is used when combinations > null_size."""
+    num_pos = 10
+    total = 100
+    null_size = 1000
+    # comb(100, 10) = 17,310,309,456,440 >> 1000
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        cache_dir = Path(tmpdir)
+
+        null_dist = compute.null_dist_cached(
+            num_pos=num_pos,
+            total=total,
+            seed=42,
+            null_size=null_size,
+            cache_dir=cache_dir,
+        )
+
+        # Should have null_size samples
+        assert len(null_dist) == null_size
+
+        # With random sampling, we expect many unique values (not exactly n_combinations)
+        unique_values = np.unique(null_dist)
+        # Random sampling won't produce exactly 1000 unique values due to collisions
+        # but should be reasonably close
+        assert len(unique_values) > 500