Audit cluster mates of flagged leaders together; expose max_rounds

rfoldsearch/experimental_calibrate.py

@@ -86,6 +86,34 @@ def _missing_high_score(scores, leaders, lengths, entropies, model, tm,
     return missing
 
 
+def _add_cluster_mates(missing, scores, tm, lengths, entropies, model,
+                       query_indices, ids, members, column, label_sets, ceiling):
+    """Also audit the unaudited eligible mates of every flagged leader.
+
+    The loop audits one leader per cluster and promotes the next member when
+    the leader is related, so a large related but unlabelled cluster would take
+    one round per member. Mates score no higher than their eligible leader, so
+    these audits are the ones later rounds would request; a mate with TM < .45
+    stays in the background as if it had not been audited.
+    """
+    for row, target in list(missing):
+        qi = query_indices[row]
+        query_families, query_clans = label_sets[qi]
+        mates = [j for j in members[column[target]]
+                 if j != target and not np.isfinite(tm[row, j])
+                 and np.isfinite(scores[row, j])
+                 and query_families.isdisjoint(label_sets[j][0])
+                 and query_clans.isdisjoint(label_sets[j][1])]
+        if not mates:
+            continue
+        e0 = np.atleast_1d(EV.initial_evalue(scores[row, mates], lengths[row],
+                                             entropies[row], model))
+        for mate, value in zip(mates, e0):
+            if value <= ceiling:
+                missing.setdefault((row, mate), (ids[qi], ids[mate], float(value)))
+    return missing
+
+
 def _audit_missing(missing, entries, query_indices, ids, tm, path, binary):
     pairs = [(query, target) for query, target, _ in missing.values()]
     by_id = dict(zip(ids, entries))
@@ -112,6 +140,10 @@ def fit_from_matrices(entries, tokens, ids, clusters, family, clan, query_indice
     for index, label in enumerate(clusters):
         members_by_cluster[str(label)].append(index)
     members = [members_by_cluster[key] for key in sorted(members_by_cluster)]
+    column = np.empty(len(clusters), dtype=np.int64)
+    for index, group in enumerate(members):
+        column[group] = index
+    label_sets = [(_labels(f), _labels(c)) for f, c in zip(family, clan)]
     n_clusters = len(members)
     lengths_all = np.asarray([len(token) for token in tokens], dtype=float)
     cluster_lengths = np.asarray([np.median(lengths_all[group]) for group in members])
@@ -136,7 +168,10 @@ def fit_from_matrices(entries, tokens, ids, clusters, family, clan, query_indice
         missing.update(_missing_high_score(scores, annotated, query_lengths,
                                            entropies, fitted, tm, query_indices,
                                            ids, audit_ceiling))
-        log(f"round {iteration}: {len(missing)} unaudited leaders at E0<={audit_ceiling:g}")
+        _add_cluster_mates(missing, scores, tm, query_lengths, entropies, fitted,
+                           query_indices, ids, members, column, label_sets,
+                           audit_ceiling)
+        log(f"round {iteration}: {len(missing)} unaudited pairs at E0<={audit_ceiling:g}")
         if not missing:
             break
         if audit_file is None:

rfoldsearch/predicted_calibrate.py

@@ -4,7 +4,8 @@
 import numpy as np
 
 from . import experimental_evalue as EV
-from .experimental_calibrate import (_audit_missing, _leader_tables,
+from .experimental_calibrate import (_add_cluster_mates, _audit_missing,
+                                     _labels, _leader_tables,
                                      _missing_high_score)
 from .usalign import find_usalign
 
@@ -13,7 +14,7 @@ def fit_from_matrices(entries, tokens, ids, clusters, family, clan,
                       training_indices, calibration_indices, scores, tm, *,
                       auxiliary_count=0,
                       audit_file=None, usalign=None, audit_ceiling=100.0,
-                      log=print):
+                      max_rounds=75, log=print):
     """Refit Gumbel on training queries and the empirical map on held-out queries.
 
     Rows of scores and tm are training rows followed by calibration rows.
@@ -34,6 +35,10 @@ def fit_from_matrices(entries, tokens, ids, clusters, family, clan,
     for index, label in enumerate(clusters):
         groups[str(label)].append(index)
     members = [groups[key] for key in sorted(groups)]
+    column = np.empty(len(clusters), dtype=np.int64)
+    for index, group in enumerate(members):
+        column[group] = index
+    label_sets = [(_labels(f), _labels(c)) for f, c in zip(family, clan)]
     lengths = np.asarray([len(token) for token in tokens], dtype=float)
     cluster_lengths = np.asarray([np.median(lengths[group]) for group in members])
     query_lengths = lengths[indices]
@@ -46,7 +51,7 @@ def fit_from_matrices(entries, tokens, ids, clusters, family, clan,
     training_weights = (np.r_[np.full(n_core, .8 / n_core),
                               np.full(auxiliary_count, .2 / auxiliary_count)]
                         if auxiliary_count else np.full(n_core, 1. / n_core))
-    for iteration in range(1, 76):
+    for iteration in range(1, max_rounds + 1):
         leaders, _, _ = _leader_tables(scores, indices, clusters, members,
                                        family, clan, tm)
         fitted = EV.fit_gumbel(scores[:len(training_indices)],
@@ -62,7 +67,10 @@ def fit_from_matrices(entries, tokens, ids, clusters, family, clan,
         missing = _missing_high_score(scores, leaders, query_lengths,
                                       entropies, fitted, tm, indices, ids,
                                       audit_ceiling)
-        log(f"round {iteration}: {len(missing)} unaudited leaders at E0<={audit_ceiling:g}")
+        _add_cluster_mates(missing, scores, tm, query_lengths, entropies, fitted,
+                           indices, ids, members, column, label_sets,
+                           audit_ceiling)
+        log(f"round {iteration}: {len(missing)} unaudited pairs at E0<={audit_ceiling:g}")
         if not missing:
             break
         if audit_file is None:
@@ -71,7 +79,8 @@ def fit_from_matrices(entries, tokens, ids, clusters, family, clan,
             binary = find_usalign(usalign)
         _audit_missing(missing, entries, indices, ids, tm, audit_file, binary)
     else:
-        raise RuntimeError("predicted structural audit did not converge")
+        raise RuntimeError(
+            f"predicted structural audit did not settle after {max_rounds} rounds")
     evalue_rows = []
     for row in range(len(training_indices), len(indices)):
         targets = leaders[row, leaders[row] >= 0]

tests/test_experimental_empirical_evalue.py

@@ -5,6 +5,7 @@
 import torch
 
 from helpers import add_coords
+from rfoldsearch import experimental_calibrate as XC
 from rfoldsearch import experimental_evalue as EV
 from rfoldsearch.cli import main
 from rfoldsearch.database import (calibration_context_fingerprint,
@@ -90,6 +91,48 @@ def test_fitter_uses_clean_background_for_both_stages():
     assert uncensored["censoring"] == "none"
 
 
+def _related_cluster_panel(seed, n_labelled, n_related, n_codes=20):
+    """Labelled singletons plus one unlabelled cluster related to every query."""
+    rng = np.random.default_rng(seed)
+    n = n_labelled + n_related
+    tokens = [rng.integers(0, n_codes, size=30 + i % 7).tolist() for i in range(n)]
+    entries = [{"id": f"t{i}", "tokens": token} for i, token in enumerate(tokens)]
+    ids = [entry["id"] for entry in entries]
+    clusters = np.asarray([f"c{i}" for i in range(n_labelled)] + ["big"] * n_related)
+    family = np.asarray([f"RF{i % 4}" for i in range(n_labelled)] + [""] * n_related)
+    clan = np.asarray([""] * n)
+    truth = np.full((n, n), .1)
+    truth[:, n_labelled:] = .8
+    return entries, tokens, ids, clusters, family, clan, truth
+
+
+def test_audit_loop_audits_related_cluster_mates_together(monkeypatch, tmp_path):
+    entries, tokens, ids, clusters, family, clan, truth = _related_cluster_panel(5, 24, 12)
+    queries = np.arange(24)
+    rng = np.random.default_rng(6)
+    scores = np.maximum(rng.gumbel(12, 3, (24, 36)), 0)
+    scores[:, 24:] += 30
+    rounds = []
+
+    def audit(missing, entries, query_indices, ids, tm, path, binary):
+        rounds.append(len(missing))
+        for row, target in missing:
+            tm[row, target] = truth[query_indices[row], target]
+    monkeypatch.setattr(XC, "_audit_missing", audit)
+    monkeypatch.setattr(XC, "find_usalign", lambda path: "USalign")
+    tm = np.full(scores.shape, np.nan)
+    model, calibration, _ = fit_from_matrices(
+        entries, tokens, ids, clusters, family, clan, queries, scores, tm,
+        audit_file=tmp_path / "audit.tsv", max_rounds=2, log=lambda _: None)
+    assert len(rounds) == 1
+    audited = truth[queries].copy()
+    expected, expected_calibration, _ = fit_from_matrices(
+        entries, tokens, ids, clusters, family, clan, queries, scores, audited,
+        log=lambda _: None)
+    assert model == expected
+    assert calibration == expected_calibration
+
+
 @pytest.mark.parametrize("ceiling", [100, 400])
 def test_experimental_cli_fit_attach_and_search(tmp_path, ceiling):
     rng = np.random.default_rng(41)

tests/test_predicted_empirical_evalue.py

@@ -6,10 +6,12 @@
 import torch
 
 from helpers import add_coords
+from test_experimental_empirical_evalue import _related_cluster_panel
 from rfoldsearch import compbias, experimental_evalue as EV
 from rfoldsearch.cli import main
 from rfoldsearch.database import (calibration_context_fingerprint,
                                   database_fingerprint)
+from rfoldsearch import predicted_calibrate as PC
 from rfoldsearch.predicted_calibrate import fit_from_matrices
 from rfoldsearch.experimental_calibrate import sha256
 from rfoldsearch.search import GAP_EXTEND, GAP_OPEN, Searcher, load_matrix
@@ -48,6 +50,39 @@ def test_clean_predicted_fit_uses_full_background_without_expansion():
     assert model["censoring"] == "left-censored-95th-percentile"
 
 
+def test_predicted_audit_settles_large_related_cluster_in_one_audit(monkeypatch, tmp_path):
+    entries, tokens, ids, clusters, family, clan, truth = _related_cluster_panel(
+        13, 40, 15, n_codes=10)
+    rng = np.random.default_rng(14)
+    scores = np.maximum(rng.gumbel(12, 2, (32, 55)), 0)
+    scores[:, 40:] += 25
+    rounds = []
+
+    def audit(missing, entries, query_indices, ids, tm, path, binary):
+        rounds.append(len(missing))
+        for row, target in missing:
+            tm[row, target] = truth[query_indices[row], target]
+    monkeypatch.setattr(PC, "_audit_missing", audit)
+    monkeypatch.setattr(PC, "find_usalign", lambda path: "USalign")
+    tm = np.full(scores.shape, np.nan)
+    model, empirical, _ = fit_from_matrices(
+        entries, tokens, ids, clusters, family, clan, np.arange(24),
+        np.arange(24, 32), scores, tm, audit_file=tmp_path / "audit.tsv",
+        max_rounds=2, log=lambda _: None)
+    assert len(rounds) == 1
+    expected, expected_empirical, _ = fit_from_matrices(
+        entries, tokens, ids, clusters, family, clan, np.arange(24),
+        np.arange(24, 32), scores, truth[:32].copy(), log=lambda _: None)
+    assert model == expected
+    assert empirical == expected_empirical
+    with pytest.raises(RuntimeError, match="after 1 rounds"):
+        fit_from_matrices(entries, tokens, ids, clusters, family, clan,
+                          np.arange(24), np.arange(24, 32), scores,
+                          np.full(scores.shape, np.nan),
+                          audit_file=tmp_path / "audit.tsv", max_rounds=1,
+                          log=lambda _: None)
+
+
 def test_compressed_gumbel_likelihood_matches_uncompressed_fit():
     rng = np.random.default_rng(73)
     scores = rng.gumbel(10, 2, (24, 40))