From 552edcbd2fd2ed9ac75e8e0661afd2e76be4f218 Mon Sep 17 00:00:00 2001 From: Ryan James Date: Sat, 14 Nov 2026 10:58:18 -0700 Subject: [PATCH] Pair search over every candidate, up to a pair budget The synergy-aware pair search took the remaining candidates in (lag, process) order and stopped at the first 20, so jointly informative (XOR) parents at longer lags or in larger networks were never searched. With 12 processes the old pool recovered 2 of 10 XOR pairs; with 20 processes it recovered none. - SkeletonSettings.max_pairs (default 5000 pairs, i.e. 100 candidates) replaces the 20-candidate cap. max_pair_candidates now defaults to None and still caps the pool if set. The order is unchanged and does not look at the target, so the maximum statistic stays valid: on null targets the pair test rejected 5 of 100 with 1540 pairs and 6 of 100 with 4560. - Candidates left out are listed in TargetSkeleton.pair_pool_dropped, with one warning. - The pair test is curtailed for every estimator unless curtail=False: it reports no p-value, and curtailment never changes a decision. - Joint codes are built pair by pair; materialising them all would mean up to 4950 columns as long as the data. - validation/pair_pool.py and a table in the validation docs: 10 of 10 XOR pairs recovered with 8, 12 and 20 processes, at 1.7 to 7 times the time per target. --- docs/skeleton.rst | 4 +- docs/validation.rst | 29 ++++++++++++++ infoflow/selection.py | 81 ++++++++++++++++++++++++++++++++----- tests/test_selection.py | 55 ++++++++++++++++++++++++++ validation/pair_pool.py | 88 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 246 insertions(+), 11 deletions(-) create mode 100644 validation/pair_pool.py diff --git a/docs/skeleton.rst b/docs/skeleton.rst index 30386db..0adf01b 100644 --- a/docs/skeleton.rst +++ b/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. diff --git a/docs/validation.rst b/docs/validation.rst index d6da2f0..8991eee 100644 --- a/docs/validation.rst +++ b/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 -------------- diff --git a/infoflow/selection.py b/infoflow/selection.py index 0707441..7de906d 100644 --- a/infoflow/selection.py +++ b/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: diff --git a/tests/test_selection.py b/tests/test_selection.py index b0b9ce0..e7b81bb 100644 --- a/tests/test_selection.py +++ b/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 diff --git a/validation/pair_pool.py b/validation/pair_pool.py new file mode 100644 index 0000000..9c9a5e4 --- /dev/null +++ b/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()