diff --git a/README.md b/README.md index bb76589..0fe5a8b 100644 --- a/README.md +++ b/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 diff --git a/rfoldsearch/cli.py b/rfoldsearch/cli.py index c975f1e..8030a28 100644 --- a/rfoldsearch/cli.py +++ b/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") 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 01cad25..23a1a7c 100644 --- a/tests/test_experimental_empirical_evalue.py +++ b/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,6 +30,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.])], @@ -34,11 +50,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(): @@ -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): diff --git a/tests/test_predicted_empirical_evalue.py b/tests/test_predicted_empirical_evalue.py index 9ca6136..b587ce4 100644 --- a/tests/test_predicted_empirical_evalue.py +++ b/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)