Pair search over every candidate, up to a pair budget

docs/skeleton.rst

@@ -10,7 +10,9 @@ multivariate transfer-entropy search of IDTxl
 2. add source lags conditioned on everything selected so far;
 3. revisit the target's past conditioned on the selected sources;
 4. search pairs of candidates that are only jointly informative, as in an XOR
-   (``synergy_search``);
+   (``synergy_search``). The pool holds every remaining candidate, up to ``max_pairs``
+   pairs (5000, i.e. 100 candidates) taken in order of lag, an order that does not
+   look at the target and so keeps the maximum statistic valid;
 5. prune with the minimum statistic, run an omnibus test of all sources jointly,
    and a sequential maximum-statistic test per source;
 6. control the false discovery rate across targets on the omnibus p-values.

docs/validation.rst

@@ -59,6 +59,35 @@ somewhat anti-conservative for strongly autocorrelated processes such as these V
 series, while weakly autocorrelated ones (the logistic maps) keep their precision.
 The options are therefore not the default.
 
+Pair search
+-----------
+
+The pair search (step 4 of :doc:`skeleton`) finds parents that inform the target only
+jointly, as in an XOR. Its pool takes candidates in order of lag and process, an order that
+does not look at the target, and it used to stop at the first 20 candidates. In larger
+networks that hid XOR parents at longer lags.
+
+``validation/pair_pool.py`` puts an XOR of two random processes, at random lags 1--4,
+into networks of independent fair bits. Each run uses 4000 samples, the default permutation
+counts, and 10 XOR and 10 null targets per size. The old pool is the first 20 candidates
+without curtailment; the new one holds up to ``max_pairs=5000`` pairs:
+
+==========  =========  ===============  ==============  ============
+Processes   Pool       XOR recovered    False parents   s per target
+==========  =========  ===============  ==============  ============
+8           first 20   5 / 10           0.00            2.3
+8           5000       10 / 10          0.10            3.8
+12          first 20   2 / 10           0.05            2.7
+12          5000       10 / 10          0.05            7.8
+20          first 20   0 / 10           0.05            3.7
+20          5000       10 / 10          0.25            25.6
+==========  =========  ===============  ==============  ============
+
+False parents are counted per target, over the XOR and null targets. The pair test keeps
+its level with the larger pool. On null targets it rejected in 5 of 100 with 12 processes
+(1540 pairs) and in 6 of 100 with 20 processes (4560 pairs). Most of the extra time is
+the full permutation null of a pair that is found.
+
 Large networks
 --------------
 

infoflow/selection.py

@@ -11,7 +11,8 @@
    with a source, e.g. ``y_t = x_{t-1} xor y_{t-1}``). If single-variable additions
    stop, a *synergy-aware pair search* tests pairs of remaining variables jointly,
    which recovers parents that act only together (e.g. ``Y = X1 xor X2``, where
-   each alone carries no information).
+   each alone carries no information). Its pool holds every remaining candidate, up
+   to a budget of pairs.
 3. Prune source variables with a *minimum statistic* test.
 4. Test all selected sources jointly with an *omnibus* test, and give each source
    variable a p-value with the *sequential maximum statistic*.
@@ -27,9 +28,10 @@
 
 import os
 import time
+import warnings
 from dataclasses import dataclass, field
 from itertools import combinations
-from math import factorial
+from math import factorial, isqrt
 
 import numpy as np
 from dit.inference._symbols import as_generator
@@ -106,9 +108,19 @@ class SkeletonSettings:
     admitted each source. ``target_estimators`` maps targets to their own estimator
     (or ensemble), overriding ``estimator`` for them.
 
+    ``synergy_search`` tests pairs of candidates jointly once single additions stop, since
+    an XOR target needs both of its parents at once. The pair pool takes the remaining
+    candidates in a fixed order (lag, then process) that does not look at the target, so
+    the maximum statistic over its pairs stays a valid test. It holds as many candidates
+    as keep the number of pairs within ``max_pairs`` (5000 pairs, i.e. 100 candidates),
+    and at most ``max_pair_candidates`` if that is set. When candidates are left out, a
+    warning is issued and :attr:`TargetSkeleton.pair_pool_dropped` lists them.
+
     ``curtail`` (default: on for KSG and Gaussian) evaluates permutations ``curtail_batch`` at a time
     and stops a test as soon as its p-value must exceed the stage's level, which leaves
     every decision unchanged and saves most permutations in the many tests that fail.
+    The pair search is curtailed for every estimator unless ``curtail=False``: it
+    reports no p-value, and its pool can hold thousands of pairs.
     ``prescreen_alpha`` first tests each source process on its own (its best lag's
     transfer entropy given the selected target past, ``n_perm_prescreen`` permutations)
     and keeps only those passing for the multivariate search, at the risk of missing a
@@ -151,7 +163,8 @@ class SkeletonSettings:
     forced_conditionals: tuple = ()
     faes: bool = False
     synergy_search: bool = True
-    max_pair_candidates: int = 20
+    max_pair_candidates: int | None = None
+    max_pairs: int | None = 5000
     max_source_lag: int | None = None
     n_perm_pairs: int = 100
     alpha_pairs: float = 0.05
@@ -196,6 +209,10 @@ def check(self):
             check_n_perm(getattr(self, f"n_perm_{name}"), getattr(self, f"alpha_{name}") / k)
         if self.synergy_search:
             check_n_perm(self.n_perm_pairs, self.alpha_pairs / k)
+        for name in ("max_pairs", "max_pair_candidates"):
+            value = getattr(self, name)
+            if value is not None and value < 0:
+                raise ValueError(f"{name} must be None or non-negative")
         if self.tdmi_screen:
             check_n_perm(self.n_perm_tdmi, self.alpha_tdmi / k)
         if self.ksg_null not in ("permutation", "local"):
@@ -226,6 +243,7 @@ class TargetSkeleton:
     significant: bool = False
     prescreened: list | None = None
     pairs: list = field(default_factory=list)
+    pair_pool_dropped: list = field(default_factory=list)
     tdmi_candidates: dict = field(default_factory=dict)
     n_samples: int = 0
     admitted_by: dict = field(default_factory=dict)
@@ -1093,23 +1111,52 @@ def _pair_search(cols, candidates, cond, settings, perm, batch=None):
     return pair if pair is not None and p <= settings.alpha_pairs else None
 
 
+def _largest_pool(max_pairs):
+    """
+    The largest number of candidates :math:`m` with :math:`m (m - 1) / 2` pairs at most `max_pairs`.
+    """
+    m = (1 + isqrt(1 + 8 * max_pairs)) // 2
+    while m * (m - 1) // 2 > max_pairs:
+        m -= 1
+    return m
+
+
+def _pair_pool(candidates, settings):
+    """
+    The candidates the pair search considers, and those it leaves out.
+
+    Candidates are taken in an order that does not look at the target (lag, then process),
+    so the maximum statistic over the pool's pairs stays a valid test: at most
+    ``max_pair_candidates`` of them, and no more than keep the pairs within ``max_pairs``.
+    """
+    ordered = sorted(candidates, key=lambda v: (v[1], v[0]))
+    size = len(ordered)
+    if settings.max_pair_candidates is not None:
+        size = min(size, settings.max_pair_candidates)
+    if settings.max_pairs is not None:
+        size = min(size, _largest_pool(settings.max_pairs))
+    return ordered[:size], ordered[size:]
+
+
 def _pair_test(cols, candidates, cond, settings, perm, batch=None, alpha=None):
     """
     The best pair of candidates and its max-statistic p-value (curtailed at `alpha`,
     default ``settings.alpha_pairs``), or ``(None, 1.0)`` without pairs.
+
+    Joint codes are built pair by pair rather than kept: a pool of 100 candidates has
+    4950 pairs, each a column as long as the data.
     """
     alpha = settings.alpha_pairs if alpha is None else alpha
     debiased = settings.statistic == "debiased"
-    pool = sorted(candidates, key=lambda v: (v[1], v[0]))[: settings.max_pair_candidates]
+    pool, _ = _pair_pool(candidates, settings)
     pairs = [(a, b) for a, b in combinations(pool, 2) if a[0] != b[0] or a[1] != b[1]]
     if not pairs:
         return None, 1.0
     z, Kz = cols.context(cond)
-    codes = [cols.joint([a, b]) for a, b in pairs]
-    obs = np.array(list(cols.map(lambda code: cols.cmi(code[0], code[1], z, Kz), codes)))
+    obs = np.array(list(cols.map(lambda pair: cols.cmi(*cols.joint(list(pair)), z, Kz), pairs)))
     if debiased:
         idx = [perm(z) for _ in range(settings.n_perm_pairs)]
-        nulls = np.array([cols.cmi_batch(c, K, z, Kz, idx) for c, K in codes])
+        nulls = np.array([cols.cmi_batch(*cols.joint(list(pair)), z, Kz, idx) for pair in pairs])
         shift = _centers(nulls, True)
         best = int(np.argmax(obs - shift))
         null = (nulls - shift[:, None]).max(axis=0)
@@ -1119,7 +1166,10 @@ def _pair_test(cols, candidates, cond, settings, perm, batch=None, alpha=None):
         observed = obs[best]
 
         def evaluate(idx):
-            return np.max([cols.cmi_batch(c, K, z, Kz, idx) for c, K in codes], axis=0)
+            top = np.full(len(idx), -np.inf)
+            for pair in pairs:
+                top = np.maximum(top, cols.cmi_batch(*cols.joint(list(pair)), z, Kz, idx))
+            return top
 
         null = _null_statistics(settings.n_perm_pairs, lambda: perm(z), evaluate, observed, alpha, batch)
     return pairs[best], _pvalue(null, observed)
@@ -1327,6 +1377,9 @@ def __init__(self, name, cols, perm, settings):
         self.deb = settings.statistic == "debiased"
         curtail = settings.curtail if settings.curtail is not None else name in ("ksg", "mixed", "gaussian")
         self.bt = settings.curtail_batch if curtail else None
+        # The pair search reports no p-value and can test thousands of pairs, so it is
+        # curtailed for every estimator unless curtailment is switched off explicitly.
+        self.pair_bt = settings.curtail_batch if settings.curtail is not False else None
         self.admitted = {}
 
     @property
@@ -1348,7 +1401,7 @@ def prescreen_pvalue(self, lags, cond):
         )[2]
 
     def pair(self, pool, cond):
-        pair = _pair_search(self.cols, pool, cond, self.settings, self.perm, self.bt)
+        pair = _pair_search(self.cols, pool, cond, self.settings, self.perm, self.pair_bt)
         if pair is not None:
             self.admitted.update(dict.fromkeys(pair, self.name))
         return pair
@@ -1448,7 +1501,7 @@ def pair(self, pool, cond):
         s = self.settings
         best = None
         for m in self.members:
-            pair, p = self._timed(m, _pair_test, m.cols, pool, cond, s, m.perm, m.bt, s.alpha_pairs / self.k)
+            pair, p = self._timed(m, _pair_test, m.cols, pool, cond, s, m.perm, m.pair_bt, s.alpha_pairs / self.k)
             if pair is not None and (best is None or p < best[1]):
                 best = (pair, p, m.name)
         if best is None or self._combine(best[1]) > s.alpha_pairs:
@@ -1588,6 +1641,8 @@ def select_parents(data, target, embeddings=None, sources=None, settings=None, p
         pool = remaining + screened + [v for v in target_cands if v not in past]
         if len(pool) < 2:
             break
+        pool, dropped = _pair_pool(pool, settings)
+        result.pair_pool_dropped += [v for v in dropped if v not in result.pair_pool_dropped]
         pair = stages.pair(pool, conditionals + past + selected)
         if pair is None:
             break
@@ -1596,6 +1651,12 @@ def select_parents(data, target, embeddings=None, sources=None, settings=None, p
         result.pairs.append(pair)
         remaining = [v for v in remaining if v not in pair]
         screened = [v for v in screened if v not in pair]
+    if result.pair_pool_dropped:
+        warnings.warn(
+            "the pair search left some candidates out (max_pairs, max_pair_candidates); "
+            "see TargetSkeleton.pair_pool_dropped, and raise max_pairs or set it to None to search them",
+            stacklevel=2,
+        )
     base = conditionals + past
 
     if selected:

tests/test_selection.py

@@ -69,6 +69,61 @@ def test_xor_needs_pair_search():
     assert not {(0, 1), (1, 1)} <= set(without.sources)
 
 
+def _xor_target(n, n_processes, parents, flip, seed):
+    """
+    Independent fair bits, with the last process the noisy XOR of the given (process, lag) parents.
+    """
+    rng = np.random.default_rng(seed)
+    x = rng.integers(0, 2, (n, n_processes))
+    y = np.zeros(n, dtype=int)
+    for p, lag in parents:
+        y ^= np.roll(x[:, p], lag)
+    x[:, -1] = y ^ (rng.random(n) < flip)
+    return DiscreteData.from_discrete(x)
+
+
+def test_pair_pool_order_and_budget():
+    from infoflow.selection import _largest_pool, _pair_pool
+
+    assert [_largest_pool(m) for m in (0, 1, 2, 3, 190, 4950, 5000)] == [1, 2, 2, 3, 20, 100, 100]
+    candidates = [(p, lag) for p in range(4) for lag in (1, 2, 3)]
+    pool, dropped = _pair_pool(candidates, SkeletonSettings(max_pairs=None))
+    assert pool == sorted(candidates, key=lambda v: (v[1], v[0])) and dropped == []
+    pool, dropped = _pair_pool(candidates, SkeletonSettings(max_pairs=10))
+    assert pool == [(0, 1), (1, 1), (2, 1), (3, 1), (0, 2)] and len(dropped) == 7
+    pool, _ = _pair_pool(candidates, SkeletonSettings(max_pairs=None, max_pair_candidates=3))
+    assert pool == [(0, 1), (1, 1), (2, 1)]
+    with pytest.raises(ValueError, match="max_pairs"):
+        SkeletonSettings(max_pairs=-1).check()
+
+
+def test_pair_search_reaches_candidates_beyond_twenty():
+    """
+    XOR parents at lag 3 of process 5 and lag 2 of process 6 (of 8): in lag order (5, 3) is
+    the 21st candidate, beyond the 20 the pool used to stop at.
+    """
+    data = _xor_target(3000, 8, [(5, 3), (6, 2)], 0.1, 0)
+    found = select_parents(data, 7, Embedding(max_lag=1), settings=SkeletonSettings(**FAST), prng=0)
+    assert sorted(found.sources) == [(5, 3), (6, 2)] and found.significant
+    assert found.pair_pool_dropped == []
+    with pytest.warns(UserWarning, match="pair search left some candidates out"):
+        capped = select_parents(
+            data, 7, Embedding(max_lag=1), settings=SkeletonSettings(**FAST, max_pair_candidates=20), prng=0
+        )
+    assert not {(5, 3), (6, 2)} <= set(capped.sources)
+    assert (5, 3) in capped.pair_pool_dropped and (6, 2) not in capped.pair_pool_dropped
+
+
+def test_pair_search_curtailment_keeps_the_decision():
+    data = _xor_target(3000, 3, [(0, 1), (1, 1)], 0.1, 3)
+    curtailed, full = (
+        select_parents(data, 2, Embedding(max_lag=1), settings=SkeletonSettings(**FAST, curtail=c), prng=0)
+        for c in (None, False)
+    )
+    assert sorted(curtailed.sources) == sorted(full.sources) == [(0, 1), (1, 1)]
+    assert curtailed.pairs == full.pairs
+
+
 def test_target_past_selected():
     rng = np.random.default_rng(2)
     n = 3000

validation/pair_pool.py

@@ -0,0 +1,88 @@
+"""
+Pair-search coverage: XOR parents at random processes and lags, old pool (the first 20
+candidates in lag order, no curtailment) against the default pool (up to max_pairs pairs,
+curtailed), with null targets for false parents and time.
+
+Example::
+
+    python validation/pair_pool.py --processes 8 12 20 --reps 10
+"""
+
+import argparse
+import time
+
+import numpy as np
+
+from infoflow.data import DiscreteData
+from infoflow.embedding import Embedding
+from infoflow.selection import SkeletonSettings, select_parents
+
+SETTINGS = {
+    "first 20": SkeletonSettings(max_pair_candidates=20, max_pairs=None, curtail=False),
+    "max_pairs": SkeletonSettings(),
+}
+
+
+def simulate(n, n_processes, parents, flip, rng):
+    """
+    Independent fair bits; the last process is the noisy XOR of the (process, lag) parents.
+    """
+    x = rng.integers(0, 2, (n, n_processes))
+    y = np.zeros(n, dtype=int)
+    for p, lag in parents:
+        y ^= np.roll(x[:, p], lag)
+    x[:, -1] = y ^ (rng.random(n) < flip) if parents else x[:, -1]
+    return DiscreteData.from_discrete(x)
+
+
+def run(n_processes, reps, n, flip, seed):
+    rows = []
+    for kind in ("xor", "null"):
+        for r in range(reps):
+            rng = np.random.default_rng([seed, n_processes, r, kind == "xor"])
+            if kind == "xor":
+                a, b = rng.choice(n_processes - 1, 2, replace=False)
+                parents = [(int(a), int(rng.integers(1, 5))), (int(b), int(rng.integers(1, 5)))]
+            else:
+                parents = []
+            data = simulate(n, n_processes, parents, flip, rng)
+            for name, settings in SETTINGS.items():
+                start = time.perf_counter()
+                sk = select_parents(data, n_processes - 1, Embedding(max_lag=1), settings=settings, prng=r)
+                elapsed = time.perf_counter() - start
+                found = set(sk.sources) if sk.significant else set()
+                rows.append(
+                    {
+                        "processes": n_processes,
+                        "kind": kind,
+                        "settings": name,
+                        "recovered": set(parents) <= found if parents else None,
+                        "false": len(found - set(parents)),
+                        "seconds": elapsed,
+                    }
+                )
+    return rows
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[1])
+    parser.add_argument("--processes", type=int, nargs="+", default=[8, 12, 20])
+    parser.add_argument("--reps", type=int, default=10)
+    parser.add_argument("--samples", type=int, default=4000)
+    parser.add_argument("--flip", type=float, default=0.1)
+    parser.add_argument("--seed", type=int, default=0)
+    args = parser.parse_args()
+    print(f"{'processes':>9}  {'settings':>9}  {'XOR recovered':>13}  {'false/target':>12}  {'s/target':>8}")
+    for P in args.processes:
+        rows = run(P, args.reps, args.samples, args.flip, args.seed)
+        for name in SETTINGS:
+            xor = [r for r in rows if r["settings"] == name and r["kind"] == "xor"]
+            null = [r for r in rows if r["settings"] == name and r["kind"] == "null"]
+            recovered = sum(r["recovered"] for r in xor)
+            false = np.mean([r["false"] for r in xor + null])
+            seconds = np.mean([r["seconds"] for r in xor + null])
+            print(f"{P:>9}  {name:>9}  {recovered:>6}/{len(xor):<6}  {false:>12.2f}  {seconds:>8.1f}")
+
+
+if __name__ == "__main__":
+    main()