diff --git a/rfoldsearch/experimental_calibrate.py b/rfoldsearch/experimental_calibrate.py index d9f613c..527b204 100644 --- a/rfoldsearch/experimental_calibrate.py +++ b/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: diff --git a/rfoldsearch/predicted_calibrate.py b/rfoldsearch/predicted_calibrate.py index bd8af8f..771ad42 100644 --- a/rfoldsearch/predicted_calibrate.py +++ b/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] diff --git a/tests/test_experimental_empirical_evalue.py b/tests/test_experimental_empirical_evalue.py index 22bd411..1cde4da 100644 --- a/tests/test_experimental_empirical_evalue.py +++ b/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) diff --git a/tests/test_predicted_empirical_evalue.py b/tests/test_predicted_empirical_evalue.py index a12f929..acd7c71 100644 --- a/tests/test_predicted_empirical_evalue.py +++ b/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))