Harden shared empirical E-value fitting for both database types

README.md

@@ -232,6 +232,13 @@ Fitting audits high-scoring background hits with US-align and can take a
 while. An interrupted fit resumes from `--workdir`, and precomputed scores or
 TM-scores can be supplied to skip those steps (see `--help`).
 
+By default, background hits are audited up to an uncorrected E-value of 100.
+Above the highest audited value, reported E-values are extrapolated along a
+power law and capped at the number of clusters in the database. The parameter
+file records the highest E-value that the audit supports directly; if a
+database needs reliable E-values beyond it, refit with a larger
+`--audit-ceiling`.
+
 ### Python API
 
 #### Tokenize a structure into 3Dr

rfoldsearch/cli.py

@@ -232,6 +232,8 @@ def main(argv=None):
     p.add_argument("--usalign")
     p.add_argument("--uncensored", action="store_true",
                    help="fit all Gumbel scores exactly instead of left-censoring")
+    p.add_argument("--audit-ceiling", type=float, default=100.0,
+                   help="audit cluster leaders through this initial E-value (at least 100)")
     p.set_defaults(fn=cmd_fit_experimental_evalues)
 
     p = sub.add_parser("fit-predicted-evalues",
@@ -244,6 +246,8 @@ def main(argv=None):
         p.add_argument("--" + name.replace("_", "-"))
     p.add_argument("--matrix")
     p.add_argument("--usalign")
+    p.add_argument("--audit-ceiling", type=float, default=100.0,
+                   help="audit cluster leaders through this initial E-value (at least 100)")
     p.set_defaults(fn=cmd_fit_predicted_evalues)
 
     p = sub.add_parser("search", help="search a query against a token database")

rfoldsearch/experimental_calibrate.py

@@ -11,6 +11,7 @@
 import csv
 import hashlib
 import json
+import re
 from collections import defaultdict
 from pathlib import Path
 
@@ -33,29 +34,36 @@ def sha256(path):
     return digest.hexdigest()
 
 
+def _labels(value):
+    return {label for label in re.split(r"[,;|\s]+", str(value).strip())
+            if label.upper() not in {"", "NA", "N/A", "NONE", ".", "-"}}
+
+
 def _leader_tables(scores, query_indices, clusters, members, family, clan, tm):
     """Apply member exclusions before maxima; related leaders backfill in cluster."""
     n_queries, n_clusters = len(query_indices), len(members)
     clean = np.full((n_queries, n_clusters), -1, dtype=np.int32)
     annotated = np.full_like(clean, -1)
     before = np.zeros(n_queries, dtype=np.int32)
+    families = [_labels(label) for label in family]
+    clans = [_labels(label) for label in clan]
     for row, qi in enumerate(query_indices):
         for column, group in enumerate(members):
             if clusters[group[0]] == clusters[qi]:
                 continue
             eligible = [j for j in group if np.isfinite(scores[row, j])
-                        and (not family[qi] or family[j] != family[qi])
-                        and (not clan[qi] or clan[j] != clan[qi])]
+                        and families[qi].isdisjoint(families[j])
+                        and clans[qi].isdisjoint(clans[j])]
             if not eligible:
                 continue
-            if any(family[j] for j in eligible):
+            if any(families[j] for j in eligible):
                 before[row] += 1
             remaining = [j for j in eligible if not (np.isfinite(tm[row, j])
                                                    and tm[row, j] >= .45)]
             if remaining:
                 clean[row, column] = max(remaining,
                                          key=lambda j: (scores[row, j], -j))
-            annotated_members = [j for j in remaining if family[j]]
+            annotated_members = [j for j in remaining if families[j]]
             if annotated_members:
                 annotated[row, column] = max(annotated_members,
                                              key=lambda j: (scores[row, j], -j))
@@ -166,12 +174,16 @@ def main(argv=None):
     parser.add_argument("--matrix")
     parser.add_argument("--usalign")
     parser.add_argument("--uncensored", action="store_true")
+    parser.add_argument("--audit-ceiling", type=float, default=100.0)
     return fit_files(parser.parse_args(argv))
 
 
 def fit_files(args, *, predicted=False):
     """Run the package fitting workflow from parsed CLI arguments."""
     db_path = Path(args.db)
+    audit_ceiling = float(args.audit_ceiling)
+    if not np.isfinite(audit_ceiling) or audit_ceiling < 100:
+        raise ValueError("--audit-ceiling must be finite and at least 100")
     workdir = Path(args.workdir)
     workdir.mkdir(parents=True, exist_ok=True)
     _, entries, tokens, ids, n_codes, clusters = load_database(str(db_path))
@@ -231,7 +243,7 @@ def fit_files(args, *, predicted=False):
     if (manifest.get("database_sha256") != sha256(db_path)
             or manifest.get("matrix_sha256") != sha256(matrix_path)
             or manifest.get("n_targets") != len(ids)
-            or (predicted and manifest.get("score_definition") != score_name)
+            or manifest.get("score_definition") != score_name
             or manifest.get("scores_sha256") != sha256(scores_path)):
         raise ValueError("score manifest does not match the target database and matrix")
     if len(manifest["query_ids"]) != len(set(manifest["query_ids"])):
@@ -254,14 +266,22 @@ def fit_files(args, *, predicted=False):
     if scores.shape != (len(query_indices), len(ids)) or tm.shape != scores.shape:
         raise ValueError("score and structural-audit matrices do not match the panel")
     audit_file = workdir / "new_usalign.tsv"
+    audit_manifest = workdir / "new_usalign.manifest.json"
+    audit_source = {"database_sha256": sha256(db_path), "query_ids": query_ids}
     if audit_file.exists():
+        if (not audit_manifest.exists() or
+                json.loads(audit_manifest.read_text(encoding="utf-8")) != audit_source):
+            raise ValueError("US-align cache does not match this database and query panel")
         with audit_file.open(newline="") as handle:
             index = {chain_id: row for row, chain_id in enumerate(query_ids)}
             target_index = {chain_id: column for column, chain_id in enumerate(ids)}
             for record in csv.DictReader(handle, delimiter="\t"):
                 if record["query_id"] in index and record["target_id"] in target_index:
                     tm[index[record["query_id"]], target_index[record["target_id"]]] = .5 * (
                         float(record["tm_query"]) + float(record["tm_target"]))
+    else:
+        audit_manifest.write_text(json.dumps(audit_source, indent=2) + "\n",
+                                  encoding="utf-8")
     if predicted:
         from . import compbias
         from . import predicted_calibrate
@@ -270,12 +290,13 @@ def fit_files(args, *, predicted=False):
             entries, tokens, ids, np.asarray(clusters), family, clan,
             training, calibration_indices=calibration_indices, scores=scores, tm=tm,
             auxiliary_count=len(auxiliary_indices),
-            audit_file=audit_file, usalign=args.usalign)
+            audit_file=audit_file, usalign=args.usalign,
+            audit_ceiling=audit_ceiling)
     else:
         model, calibration, audit = fit_from_matrices(
             entries, tokens, ids, np.asarray(clusters), family, clan,
             query_indices, scores, tm, audit_file=audit_file, usalign=args.usalign,
-            censored=not args.uncensored)
+            censored=not args.uncensored, audit_ceiling=audit_ceiling)
     target_digest = database_fingerprint(entries, clusters, n_codes)
     context_digest = calibration_context_fingerprint(
         target_digest, matrix, GAP_OPEN, GAP_EXTEND)
@@ -285,6 +306,7 @@ def fit_files(args, *, predicted=False):
               "target_database_fingerprint": target_digest,
               "calibration_context_fingerprint": context_digest,
               "validation_status": "audited-e0-le-100",
+              "audit_raw_evalue_ceiling": audit_ceiling,
               "target_aware_gumbel": model,
               "empirical_calibration": calibration,
               "audit": audit,

rfoldsearch/experimental_evalue.py

@@ -1,8 +1,4 @@
-"""Target-aware cluster E-values for the experimental RNA database.
-
-The fitted score is the *raw* token Smith--Waterman score. The null counts
-cluster maxima, while the empirical map counts unrelated clusters per query.
-"""
+"""Shared target-aware cluster E-value model and empirical calibration."""
 from __future__ import annotations
 
 import math
@@ -246,15 +242,17 @@ def fit_empirical(evalue_rows, families, expansions, n_clusters, ceiling=100.0,
         raise ValueError("low-tail power slope must be positive")
     return {"scheme": "family-balanced-monotone", "raw_evalues": unique.tolist(),
             "observed_counts": cumulative.tolist(), "low_tail_power_b": float(slope),
+            "upper_tail_rule": "anchored-power-cap",
             "fit_raw_evalue_min": float(unique[0]),
             "fit_raw_evalue_max": float(unique[-1]), "fit_raw_evalue_ceiling": float(ceiling),
+            "fit_reported_evalue_max": float(cumulative[-1]),
             "n_clusters": int(n_clusters), "n_queries": int(len(families)),
             "n_families": int(len(np.unique(families))),
             "background_scaling": background_scaling}
 
 
 def apply_empirical(evalues, calibration):
-    """Linear monotone interpolation; power below support, cluster cap above."""
+    """Interpolate observed counts and apply endpoint-anchored tails."""
     x = np.asarray(evalues, dtype=float)
     xp = np.asarray(calibration["raw_evalues"], dtype=float)
     yp = np.asarray(calibration["observed_counts"], dtype=float)
@@ -276,7 +274,15 @@ def apply_empirical(evalues, calibration):
             np.log(yp[0]) + b * (np.log(low) - np.log(xp[0])),
             np.log(np.nextafter(0., 1.))))
         y = np.where(x < xp[0], np.where(x == 0, 0., extrapolated), y)
-    y = np.where(x > xp[-1], float(n_clusters), y)
+    if calibration.get("upper_tail_rule") == "anchored-power-cap":
+        upper = np.exp(np.minimum(
+            np.log(yp[-1]) + b * (np.log(np.maximum(x, xp[-1])) - np.log(xp[-1])),
+            np.log(n_clusters)))
+        y = np.where(x > xp[-1], upper, y)
+    elif calibration.get("upper_tail_rule") not in (None, "cluster-cap"):
+        raise ValueError("invalid empirical upper-tail rule")
+    else:  # Preserve the behaviour of parameters fitted before this change.
+        y = np.where(x > xp[-1], float(n_clusters), y)
     y = np.minimum(y, n_clusters)
     return float(y) if np.ndim(y) == 0 else y
 

tests/test_experimental_empirical_evalue.py

@@ -6,10 +6,13 @@
 
 from helpers import add_coords
 from rfoldsearch import experimental_evalue as EV
+from rfoldsearch.cli import main
 from rfoldsearch.database import (calibration_context_fingerprint,
                                   database_fingerprint)
 from rfoldsearch.search import GAP_EXTEND, GAP_OPEN, Searcher, load_matrix
-from rfoldsearch.experimental_calibrate import _leader_tables, fit_from_matrices
+from rfoldsearch.experimental_calibrate import (_leader_tables, fit_from_matrices,
+                                                sha256)
+from rfoldsearch.smith_waterman import smith_waterman_numba
 
 
 def test_structural_exclusion_backfills_within_cluster():
@@ -27,18 +30,40 @@ def test_structural_exclusion_backfills_within_cluster():
     assert before.tolist() == [2]
 
 
+def test_multilabel_family_and_clan_overlaps_are_excluded():
+    scores = np.array([[50., 40., 30., 20.]])
+    clusters = np.array(["a", "b", "c", "q"])
+    family = np.array(["RF1;RF2", "RF3", "RF4", "RF2"])
+    clan = np.array(["", "CL2|CL3", "", "CL3"])
+    leaders, annotated, before = _leader_tables(
+        scores, np.array([3]), clusters, [[0], [1], [2], [3]],
+        family, clan, np.full((1, 4), .1))
+    assert leaders.tolist() == [[-1, -1, 2, -1]]
+    assert annotated.tolist() == [[-1, -1, 2, -1]]
+    assert before.tolist() == [1]
+
+
 def test_empirical_map_is_monotone_and_uses_declared_tails():
     fitted = EV.fit_empirical([np.array([.1, .3, 1., 3., 10.]),
                                np.array([.2, .4, 2., 4., 20.])],
                               np.array(["RF1", "RF2"]), np.ones(2), 100)
     values = EV.apply_empirical(np.array([0., .01, .1, 1., 20., 21., 10000.]), fitted)
     assert np.all(np.diff(values) >= 0)
     assert values[0] == 0
-    assert values[-1] == 100
-    assert values[-2] == 100
+    assert values[-1] == pytest.approx(min(100, fitted["observed_counts"][-1]
+        * (10000 / fitted["raw_evalues"][-1]) ** fitted["low_tail_power_b"]))
+    assert EV.apply_empirical(1e100, fitted) == 100
+    assert fitted["upper_tail_rule"] == "anchored-power-cap"
+    assert values[-2] == pytest.approx(min(100, fitted["observed_counts"][-1]
+        * (21 / fitted["raw_evalues"][-1]) ** fitted["low_tail_power_b"]))
+    assert EV.apply_empirical(np.nextafter(20., np.inf), fitted) == pytest.approx(
+        fitted["observed_counts"][-1])
     assert values[1] == pytest.approx(
         fitted["observed_counts"][0] * (.01 / .1) ** fitted["low_tail_power_b"])
     assert EV.apply_empirical(1e-300, fitted) > 0
+    legacy = dict(fitted)
+    legacy.pop("upper_tail_rule")
+    assert EV.apply_empirical(21., legacy) == 100
 
 
 def test_fitter_uses_clean_background_for_both_stages():
@@ -59,6 +84,60 @@ def test_fitter_uses_clean_background_for_both_stages():
     assert model["n_observations"] == 24 * 18 - 1
     assert calibration["n_queries"] == 24
     assert audit["n_retained_background"] == 24 * 18 - 1
+    uncensored, _, _ = fit_from_matrices(
+        entries, tokens, ids, clusters, family, clan, np.arange(24),
+        scores, tm, censored=False, log=lambda _: None)
+    assert uncensored["censoring"] == "none"
+
+
+def test_experimental_cli_fit_attach_and_search(tmp_path):
+    rng = np.random.default_rng(41)
+    entries = add_coords([{"id": f"t{i}", "tokens": rng.integers(0, 20, 28 + i % 7).tolist(),
+                           "cluster": f"c{i}"} for i in range(32)])
+    db = tmp_path / "db.pt"
+    torch.save({"n_codes": 80, "entries": entries, "n_clusters": 32}, db)
+    labels = tmp_path / "labels.tsv"
+    labels.write_text("chain_id\tfamily\tclan\n" + "".join(
+        f"t{i}\tRF{i % 6}\t\n" for i in range(32)))
+    matrix = load_matrix()
+    scores = np.array([[smith_waterman_numba(
+        np.asarray(q["tokens"]), np.asarray(t["tokens"]), matrix,
+        GAP_OPEN, GAP_EXTEND) for t in entries] for q in entries], dtype=np.float32)
+    score_path = tmp_path / "scores.npy"
+    np.save(score_path, scores)
+    from rfoldsearch.search import asset
+    matrix_path = asset("sub_matrix.pt")
+    score_manifest = tmp_path / "scores.json"
+    score_manifest.write_text(json.dumps({
+        "database_sha256": sha256(db), "matrix_sha256": sha256(matrix_path),
+        "n_targets": 32, "query_ids": [f"t{i}" for i in range(32)],
+        "score_definition": "raw_sw", "scores_sha256": sha256(score_path)}))
+    tm_path = tmp_path / "tm.npy"
+    np.save(tm_path, np.full((32, 32), .1))
+    tm_manifest = tmp_path / "tm.json"
+    tm_manifest.write_text(json.dumps({
+        "database_sha256": sha256(db), "tm_sha256": sha256(tm_path),
+        "n_targets": 32, "query_ids": [f"t{i}" for i in range(32)]}))
+    params_path = tmp_path / "params.json"
+    main(["fit-experimental-evalues", "--db", str(db), "--labels", str(labels),
+          "--workdir", str(tmp_path / "fit"), "--output", str(params_path),
+          "--scores", str(score_path), "--score-manifest", str(score_manifest),
+          "--tm-matrix", str(tm_path), "--tm-manifest", str(tm_manifest)])
+    params = json.loads(params_path.read_text())
+    assert params["target_aware_gumbel"]["censoring"] == "left-censored-95th-percentile"
+    assert params["empirical_calibration"]["upper_tail_rule"] == "anchored-power-cap"
+    cluster_file = tmp_path / "clusters.tsv"
+    cluster_file.write_text("".join(f"t{i}\tc{i}\n" for i in range(32)))
+    attached = tmp_path / "attached.pt"
+    main(["attach-calibration", "--db", str(db), "--clusters", str(cluster_file),
+          "--evalue_params", str(params_path), "--output", str(attached)])
+    searcher = Searcher(str(attached))
+    assert searcher.experimental_ev is not None
+    hit = searcher.search_id("t0", top_k=1)[0]
+    assert np.isfinite(hit["evalue"])
+    assert hit["evalue"] >= 0
+    assert hit["evalue"] == pytest.approx(EV.apply_empirical(
+        hit["raw_evalue"], params["empirical_calibration"]))
 
 
 def test_searcher_uses_matching_experimental_empirical_parameters(tmp_path):

tests/test_predicted_empirical_evalue.py

@@ -159,15 +159,19 @@ def test_fitting_cli_writes_attachable_predicted_parameters(tmp_path):
         "database_sha256": sha256(db), "tm_sha256": sha256(tm_path),
         "n_targets": 40, "query_ids": [f"t{i}" for i in range(32)]}))
     output = tmp_path / "predicted.json"
-    main(["fit-predicted-evalues", "--db", str(db), "--labels", str(labels),
+    command = ["fit-predicted-evalues", "--db", str(db), "--labels", str(labels),
           "--training-indices", str(training), "--calibration-indices", str(calibration),
           "--auxiliary-indices", str(auxiliary),
           "--workdir", str(tmp_path / "fit"), "--output", str(output),
+          "--audit-ceiling", "120",
           "--matrix", str(matrix_path), "--scores", str(scores_path),
           "--score-manifest", str(manifest), "--tm-matrix", str(tm_path),
-          "--tm-manifest", str(tm_manifest)])
+          "--tm-manifest", str(tm_manifest)]
+    main(command)
     params = json.loads(output.read_text())
     assert params["scheme"] == "target-aware-compbias20-raw-sw-empirical-v1"
+    assert params["audit_raw_evalue_ceiling"] == 120
+    assert params["empirical_calibration"]["upper_tail_rule"] == "anchored-power-cap"
     assert params["empirical_calibration"]["background_scaling"] == "complete-eligible-clusters"
     searcher = Searcher(str(db), matrix_path=str(matrix_path),
                         evalue_params_path=str(output))
@@ -179,3 +183,8 @@ def test_fitting_cli_writes_attachable_predicted_parameters(tmp_path):
     main(["attach-calibration", "--db", str(db), "--clusters", str(clusters),
           "--evalue_params", str(output), "--output", str(attached)])
     assert Searcher(str(attached), matrix_path=str(matrix_path)).predicted_ev is not None
+    audit_file = tmp_path / "fit" / "new_usalign.tsv"
+    audit_file.write_text("query_id\ttarget_id\ttm_query\ttm_target\n")
+    (tmp_path / "fit" / "new_usalign.manifest.json").write_text("{}")
+    with pytest.raises(ValueError, match="US-align cache does not match"):
+        main(command)