From 8918f36a7878641b70db6467ebd7a8bbad2c959a Mon Sep 17 00:00:00 2001 From: Hongxiang Li Date: Mon, 26 Oct 2026 22:02:39 +0800 Subject: [PATCH 1/2] Unify audited empirical E-value fitting and continuous tails --- README.md | 21 ++++++++---- rfoldsearch/cli.py | 4 +++ rfoldsearch/experimental_calibrate.py | 36 +++++++++++++++++---- rfoldsearch/experimental_evalue.py | 20 ++++++++---- tests/test_experimental_empirical_evalue.py | 26 +++++++++++++-- tests/test_predicted_empirical_evalue.py | 13 ++++++-- 6 files changed, 96 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index e78a536..dbbc85b 100644 --- a/README.md +++ b/README.md @@ -195,7 +195,8 @@ for every query--target pair and audits high-scoring cluster leaders with US-align. To reuse caches, supply `--scores` with `--score-manifest`, and `--tm-matrix` with `--tm-manifest`; the manifests must record the database SHA-256, target count and query IDs (the score manifest also records the scoring-matrix -SHA-256). `scripts/export_experimental_rfam_labels.py` can export the Rfam +SHA-256, score-file SHA-256 and `"score_definition": "raw_sw"`). +`scripts/export_experimental_rfam_labels.py` can export the Rfam PDB-chain labels. For example: ``` @@ -213,13 +214,15 @@ The experimental fit excludes the query's own cluster and same Rfam family/clan, then removes US-align-related members (mean of the two normalised TM-scores >= 0.45) and uses the next eligible member in each cluster. There is no length-ratio filter. It refits the left-censored Gumbel after each round of -missing leader audits with `E0 <= 100`; `--uncensored` is an optional alternative +missing leader audits within the specified initial E-value ceiling (100 by +default); `--uncensored` is an optional alternative for a separately refitted database. The annotated background is expanded by `N_D/B_q`, where `B_q` is the eligible annotated-cluster count before structural exclusion. The empirical map interpolates cumulative counts inside the observed -range, uses the fitted global power slope below it, and assigns `N_D` above -the upper endpoint. These tails are extrapolations, not -independently validated estimates. +range and uses the fitted global power slope below it. New fits continue the +endpoint-anchored power slope above the last knot and cap it at `N_D`; older +parameter files retain their original upper-endpoint rule. These tails are +extrapolations, not independently validated estimates. To fit the selected predicted-database model, supply distinct `.npy` index panels for core Gumbel training, auxiliary short/low-entropy Gumbel training, @@ -252,7 +255,13 @@ ordered query IDs, score-file SHA-256, and and TM-file hashes, target count, and ordered query IDs. The generated JSON is tied to the target-token database, matrix, gap costs, cluster count and composition window. No predicted-database parameters are bundled with this -PR: the deployment NuFold release must be fitted and audited separately. +package: the deployment NuFold release must be fitted and audited separately. +Both fitting commands audit leaders through initial `E0 <= 100` by default; +`--audit-ceiling` can extend this range. The new empirical parameter file +records its highest directly supported reported E-value. Increase the audit +ceiling and refit if that support does not reach the reporting range needed +for a particular database. Cached US-align results in `--workdir` are reused +only when their database SHA-256 and query panel match. The existing `calibrate` command remains the legacy power-law workflow: diff --git a/rfoldsearch/cli.py b/rfoldsearch/cli.py index 086b767..9bb0887 100644 --- a/rfoldsearch/cli.py +++ b/rfoldsearch/cli.py @@ -238,6 +238,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", @@ -250,6 +252,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") diff --git a/rfoldsearch/experimental_calibrate.py b/rfoldsearch/experimental_calibrate.py index 1562d35..6fba9e8 100644 --- a/rfoldsearch/experimental_calibrate.py +++ b/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,7 +266,12 @@ 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)} @@ -262,6 +279,9 @@ def fit_files(args, *, predicted=False): 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, diff --git a/rfoldsearch/experimental_evalue.py b/rfoldsearch/experimental_evalue.py index 918e267..4866039 100644 --- a/rfoldsearch/experimental_evalue.py +++ b/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 diff --git a/tests/test_experimental_empirical_evalue.py b/tests/test_experimental_empirical_evalue.py index baee79a..8b76ea6 100644 --- a/tests/test_experimental_empirical_evalue.py +++ b/tests/test_experimental_empirical_evalue.py @@ -26,6 +26,19 @@ 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.])], @@ -33,11 +46,20 @@ def test_empirical_map_is_monotone_and_uses_declared_tails(): 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(): diff --git a/tests/test_predicted_empirical_evalue.py b/tests/test_predicted_empirical_evalue.py index d95b15d..11c5fcb 100644 --- a/tests/test_predicted_empirical_evalue.py +++ b/tests/test_predicted_empirical_evalue.py @@ -158,15 +158,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)) @@ -177,3 +181,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) From cb3cc28147c9c035c963606b3de1e2e3ddb618e3 Mon Sep 17 00:00:00 2001 From: Hongxiang Li Date: Mon, 26 Oct 2026 22:43:37 +0800 Subject: [PATCH 2/2] Test experimental E-value fit through database attachment and search --- tests/test_experimental_empirical_evalue.py | 60 ++++++++++++++++++++- 1 file changed, 59 insertions(+), 1 deletion(-) diff --git a/tests/test_experimental_empirical_evalue.py b/tests/test_experimental_empirical_evalue.py index 8b76ea6..37d9e5d 100644 --- a/tests/test_experimental_empirical_evalue.py +++ b/tests/test_experimental_empirical_evalue.py @@ -5,10 +5,13 @@ import torch 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(): @@ -80,6 +83,61 @@ 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 = [{"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(np.asarray(entries[0]["tokens"]), + rerank="sw", 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):