Drop the sw rerank mode; tidy CLI warnings, README, add CI

.github/workflows/tests.yml

@@ -0,0 +1,27 @@
+name: tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  pytest:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+      # CPU-only torch keeps the install small; the tests never touch a GPU.
+      - name: Install
+        run: |
+          python -m pip install --upgrade pip
+          pip install --extra-index-url https://download.pytorch.org/whl/cpu . pytest
+      - name: Test
+        run: python -m pytest -q tests

README.md

@@ -2,7 +2,7 @@
 
 <p>
   <a href="LICENSE"><img src="https://img.shields.io/badge/License-Apache_2.0-blue.svg" alt="License: Apache 2.0" /></a>
-  <img src="https://img.shields.io/badge/python-%3E%3D3.9-blue.svg" alt="Python >= 3.9" />
+  <img src="https://img.shields.io/badge/python-%3E%3D3.10-blue.svg" alt="Python >= 3.10" />
   <a href="#installation"><img src="https://img.shields.io/badge/install-conda-44A833.svg" alt="Install with conda" /></a>
 </p>
 
@@ -114,8 +114,6 @@ then a reranking of the candidates. What differs is how the candidates are reran
   **R<sub>str</sub> = S<sub>SW</sub> × TM<sub>a</sub>**. `TM_a` is the mean of
   query- and target-length-normalised TM terms for the fixed token alignment.
   The E-value is still calculated from `S_SW`, not from `R_str`.
-- **`sw`**: token-only Smith-Waterman ranking. It can search databases
-  without coordinates and reports the same Smith-Waterman E-value.
 - **`tm`** (RFoldSearch-TM): the candidate pool is handed to
   [US-align](https://github.com/pylelab/USalign) and ranked by the mean of its two length-normalised TM-scores. 
   This measures structural similarity directly instead of scoring a token alignment, at a large cost in speed.
@@ -138,7 +136,7 @@ A comment line names the query, database, and ranking mode, followed by one
 row per hit. Under the default mode, hits are ranked by `R_str`:
 
 ```
-# query: query_chain_id  (128 tokens)   database: target_db.pt   rerank: R_str = S_SW x TM_a
+# query: query_chain_id  (128 tokens)   database: target_db.pt  (2000 chains, 1500 clusters)   rerank: R_str = S_SW x TM_a
 rank  target                       R_str     TM_a      S_SW    E-value   aligned    len
    1  4V9F_A                       358.44   0.8700    412.00    2.1e-30       124    129
    2  6XU8_B                       221.38   0.7700    287.50    5.4e-19       111    121
@@ -148,7 +146,7 @@ rank  target                       R_str     TM_a      S_SW    E-value   aligned
 Under `--rerank tm` — RFoldSearch-TM — the comment line names the mode and the score columns change:
 
 ```
-# query: query_chain_id  (74 tokens)   database: target_db.pt   rerank: RFoldSearch-TM
+# query: query_chain_id  (74 tokens)   database: target_db.pt  (2000 chains)   rerank: RFoldSearch-TM
 rank  target                    TM_mean  TM_query  TM_target    len
    1  PDB_00008B2L_1_W2          0.7076    0.7143     0.7008     76
    2  PDB_00008KD9_1_M           0.6724    0.6674     0.6773     72
@@ -160,7 +158,7 @@ rank  target                    TM_mean  TM_query  TM_target    len
 | Option       | Category    | Description                                                                        |
 | ------------ | ----------- | ---------------------------------------------------------------------------------- |
 | `--top_k`    | Output      | Number of top hits to return (default: 10)                                         |
-| `--rerank`   | Scoring     | `tm_a` (default), `sw`, or `tm`                                                    |
+| `--rerank`   | Scoring     | `tm_a` (default) or `tm`                                                           |
 | `--top_n`    | Sensitivity | Candidate-pool depth passed to the reranker (default: 200) |
 | `--min_diag` | Sensitivity | k-mer hits required on a common diagonal to admit a target (default: 2) |
 | `--query_id` | Input       | Chain id to look up inside `--query` when it is a database file                    |
@@ -176,8 +174,8 @@ identity) expected to score at least as high as a hit by chance. Smaller
 values mean the hit is less likely to be a chance match. E-values are
 computed from the raw Smith--Waterman token-alignment score, with a
 composition adjustment for predicted-structure databases (below). They are
-reported in the default mode and with `--rerank sw`, and they do not measure
-the structure score R<sub>str</sub> used for ranking in the default mode.
+reported in the default mode, and they do not measure the structure score
+R<sub>str</sub> used for ranking.
 RFoldSearch-TM (`--rerank tm`) aligns with US-align instead and reports no
 E-values.
 
@@ -272,8 +270,8 @@ from rfoldsearch import Searcher
 
 s = Searcher("target_db.pt")
 hits = s.search_id("target_chain_id", top_k=10)
-# [{'rank': 1, 'target': …, 'structure_score': …, 'tm_a': …,
-#   'aligned': …, 'sw_score': …, 'bits': …, 'evalue': …}, …]
+# [{'rank': 1, 'target': …, 'sw_score': …, 'length': …,
+#   'structure_score': …, 'tm_a': …, 'aligned': …, 'bits': …, 'evalue': …}, …]
 ```
 
 Anything else is searched by handing `search` the tokens from above:
@@ -284,8 +282,8 @@ import numpy as np
 hits = s.search(np.asarray(enc["tokens"], np.int64), query_entry=enc, top_k=10)
 ```
 
-`query_entry` is required by the default `rerank="tm_a"` and by `rerank="tm"`,
-because both use query coordinates. Plain `rerank="sw"` can leave it out.
+`query_entry` is required by both `rerank="tm_a"` (default) and `rerank="tm"`,
+because both use query coordinates.
 `search_id` forwards its keyword arguments to `search`, so `top_n`, `min_diag`, `rerank`
 and `tie_seed` are available from both. `bits` and `evalue` are present when
 E-value parameters matching the database are loaded (by default, only for the

pyproject.toml

@@ -8,11 +8,9 @@ version = "0.1.0"
 description = "Fast RNA tertiary-structure search with the learned 3Dr structural alphabet"
 readme = "README.md"
 license = {file = "LICENSE"}
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 dependencies = [
     "numpy>=1.21",
-    # The last biotite for Python 3.9 (0.40) is built against NumPy 1.
-    "numpy<2; python_version < '3.10'",
     "torch>=2.0",
     "numba>=0.55",
     # SciPy 1.15.x wheels fail to load on current macOS (dyld rejects their

rfoldsearch/cli.py

@@ -3,6 +3,7 @@
 import glob
 import json
 import sys
+import warnings
 
 import numpy as np
 import torch
@@ -111,7 +112,7 @@ def cmd_search(a):
         head += "   E-values: power-calibrated"
     if a.rerank == "tm":
         head += "   rerank: RFoldSearch-TM"
-    elif a.rerank == "tm_a":
+    else:
         head += "   rerank: R_str = S_SW x TM_a"
     print(head)
     if not hits:
@@ -129,29 +130,22 @@ def cmd_search(a):
                   f"  {h['length']:>5}")
         return
     has_ev = "evalue" in hits[0]
-    if "structure_score" in hits[0]:
-        header = (f"{'rank':>4}  {'target':<24}{'R_str':>10}{'TM_a':>9}"
-                  f"{'S_SW':>10}")
-        if has_ev:
-            header += f"{'E-value':>11}"
-        print(header + f"{'aligned':>10}  {'len':>5}")
-        for h in hits:
-            line = (f"{h['rank']:>4}  {h['target']:<24}"
-                    f"{h['structure_score']:>10.2f}{h['tm_a']:>9.4f}"
-                    f"{h['sw_score']:>10.2f}")
-            if has_ev:
-                line += f"{h['evalue']:>11.2g}"
-            print(line + f"{h['aligned']:>10}  {h['length']:>5}")
-        return
-    header = f"{'rank':>4}  {'target':<24}{'SW_score':>10}"
+    header = (f"{'rank':>4}  {'target':<24}{'R_str':>10}{'TM_a':>9}"
+              f"{'S_SW':>10}")
     if has_ev:
-        header += f"{'bits':>8}{'E-value':>11}"
-    print(header + f"  {'len':>5}")
+        header += f"{'E-value':>11}"
+    print(header + f"{'aligned':>10}  {'len':>5}")
     for h in hits:
-        line = f"{h['rank']:>4}  {h['target']:<24}{h['sw_score']:>10.2f}"
+        line = (f"{h['rank']:>4}  {h['target']:<24}"
+                f"{h['structure_score']:>10.2f}{h['tm_a']:>9.4f}"
+                f"{h['sw_score']:>10.2f}")
         if has_ev:
-            line += f"{h['bits']:>8.1f}{h['evalue']:>11.2g}"
-        print(line + f"  {h['length']:>5}")
+            line += f"{h['evalue']:>11.2g}"
+        print(line + f"{h['aligned']:>10}  {h['length']:>5}")
+
+
+def _show_warning(message, category, filename, lineno, file=None, line=None):
+    print(f"rfoldsearch: warning: {message}", file=sys.stderr)
 
 
 def main(argv=None):
@@ -265,8 +259,8 @@ def main(argv=None):
     p.add_argument("--top_n", type=int, default=DEFAULT_TOP_N)
     p.add_argument("--min_diag", type=int, default=DEFAULT_MIN_DIAG)
     p.add_argument("--rerank", choices=RERANKERS, default=DEFAULT_RERANK,
-                   help="rerank by R_str = S_SW x TM_a (tm_a, default), "
-                        "token Smith-Waterman (sw), or US-align TM-score (tm)")
+                   help="rerank by R_str = S_SW x TM_a (tm_a, default) "
+                        "or US-align TM-score (tm)")
     p.add_argument("--usalign", default=None,
                    help="path to the USalign binary (default: $USALIGN, then PATH)")
     p.add_argument("--tie_seed", type=int, default=None,
@@ -276,10 +270,14 @@ def main(argv=None):
     p.set_defaults(fn=cmd_search)
 
     a = ap.parse_args(argv)
-    try:
-        a.fn(a)
-    except (USalignNotFound, USalignFailed) as e:
-        raise SystemExit(str(e))
+    with warnings.catch_warnings():
+        # One line per warning on stderr, without the source line Python
+        # would otherwise echo after it.
+        warnings.showwarning = _show_warning
+        try:
+            a.fn(a)
+        except (USalignNotFound, USalignFailed) as e:
+            raise SystemExit(str(e))
 
 
 if __name__ == "__main__":

rfoldsearch/evalue.py

@@ -73,7 +73,9 @@ def calibrate_clustermax(
         raise ValueError("need at least two clusters to calibrate")
 
     lens = np.array([len(t) for t in toks])
-    order = np.argsort(np.abs(lens - float(length)))
+    # Stable, so chains tied on length are picked by index: numpy's default
+    # sort orders ties differently across array sizes and CPUs.
+    order = np.argsort(np.abs(lens - float(length)), kind="stable")
     pool = [int(i) for i in order[: max(2, min(len(order), n_queries))]]
 
     maxima = []
@@ -135,7 +137,7 @@ def calibrate_clustermax_length_dependent(
     table, skipped = [], []
     n_clusters = len(set(np.asarray(clusters, dtype=object).tolist()))
     for L in lengths:
-        order = np.argsort(np.abs(lens - float(L)))
+        order = np.argsort(np.abs(lens - float(L)), kind="stable")
         pool = lens[order[: max(2, min(len(order), n_queries))]]
         med = float(np.median(pool))
         if med > max_length_gap * L or L > max_length_gap * med:

rfoldsearch/search.py

@@ -1,9 +1,8 @@
 """The search pipeline: encode -> prefilter -> rerank.
 
-Three reranking backends share the token prefilter. ``tm_a`` (default) aligns
+Two reranking backends share the token prefilter. ``tm_a`` (default) aligns
 candidates by affine-gap token Smith-Waterman, superposes the aligned C4'
-atoms and ranks by R_str = S_SW * TM_a. ``sw`` keeps the token-only
-Smith-Waterman ranking. Both report the E-value of S_SW. ``tm``
+atoms, ranks by R_str = S_SW * TM_a and reports the E-value of S_SW. ``tm``
 (RFoldSearch-TM) hands the candidate pool to US-align and ranks by the mean of
 its two length-normalised TM-scores; that backend reports no E-value.
 """
@@ -31,7 +30,7 @@
 DEFAULT_TOP_N = 200      # candidate-pool size after the prefilter (validation-selected)
 DEFAULT_MIN_DIAG = 2     # Foldseek double-match rule
 DEFAULT_RERANK = "tm_a"
-RERANKERS = ("tm_a", "sw", "tm")
+RERANKERS = ("tm_a", "tm")
 _NO_CUT = 1 << 62        # ask the prefilter for every candidate, uncut
 
 _ASSETS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets")
@@ -232,8 +231,7 @@ def search(self, qtok, top_k=10, top_n=DEFAULT_TOP_N,
         ``rerank="tm_a"`` (default) ranks by ``structure_score``
         (R_str = S_SW * TM_a) and reports ``tm_a``, ``aligned`` (aligned
         nucleotide pairs), and the S_SW-derived fields ``sw_score``, ``bits``,
-        and E-value. ``rerank="sw"`` retains the token-only Smith-Waterman
-        ranking. ``rerank="tm"`` ranks by US-align TM-score and reports
+        and E-value. ``rerank="tm"`` ranks by US-align TM-score and reports
         ``tm_mean``/``tm_query``/``tm_target`` instead.
 
         ``tie_seed`` makes equal scores break by a hash of the chain id
@@ -265,21 +263,12 @@ def search(self, qtok, top_k=10, top_n=DEFAULT_TOP_N,
                     )
                 query_entry = self.ents[self_idx]
             return self._rerank_tm(query_entry, cand, top_k, usalign, tie_seed)
-        tm_a = aligned = None
-        if rerank == "tm_a":
-            (toks, offsets, xyz), qxyz = self._tm_a_inputs(
-                qtok, query_entry, self_idx)
-            sw, tm_a, aligned = _ss.score_candidates(
-                qtok, qxyz, np.asarray(cand, np.int64), toks, offsets, xyz,
-                self.sub, np.float32(GAP_OPEN), np.float32(GAP_EXTEND))
-            score = sw * tm_a
-        else:
-            sw = np.array([
-                smith_waterman_numba(qtok, self.toks[t], self.sub,
-                                     GAP_OPEN, GAP_EXTEND)
-                for t in cand
-            ])
-            score = sw
+        (toks, offsets, xyz), qxyz = self._tm_a_inputs(
+            qtok, query_entry, self_idx)
+        sw, tm_a, aligned = _ss.score_candidates(
+            qtok, qxyz, np.asarray(cand, np.int64), toks, offsets, xyz,
+            self.sub, np.float32(GAP_OPEN), np.float32(GAP_EXTEND))
+        score = sw * tm_a
         if tie_seed is None:
             order = np.argsort(-score, kind="stable")
         else:
@@ -296,11 +285,9 @@ def search(self, qtok, top_k=10, top_n=DEFAULT_TOP_N,
         for rank, oi in enumerate(order[:top_k], 1):
             t, sc = int(cand[oi]), float(sw[oi])
             hit = {"rank": rank, "target": self.ids[t], "sw_score": sc,
-                   "length": len(self.toks[t])}
-            if tm_a is not None:
-                hit["structure_score"] = float(score[oi])
-                hit["tm_a"] = float(tm_a[oi])
-                hit["aligned"] = int(aligned[oi])
+                   "length": len(self.toks[t]),
+                   "structure_score": float(score[oi]),
+                   "tm_a": float(tm_a[oi]), "aligned": int(aligned[oi])}
             if self.ev_par is not None:
                 if self.experimental_ev is not None or self.predicted_ev is not None:
                     if self.predicted_ev is not None:

tests/test_experimental_empirical_evalue.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 
+from helpers import add_coords
 from rfoldsearch import experimental_evalue as EV
 from rfoldsearch.database import (calibration_context_fingerprint,
                                   database_fingerprint)
@@ -61,8 +62,8 @@ def test_fitter_uses_clean_background_for_both_stages():
 
 
 def test_searcher_uses_matching_experimental_empirical_parameters(tmp_path):
-    entries = [{"id": f"t{i}", "tokens": [i % 5, (i + 1) % 5] * 15,
-                "cluster": f"c{i}"} for i in range(5)]
+    entries = add_coords([{"id": f"t{i}", "tokens": [i % 5, (i + 1) % 5] * 15,
+                           "cluster": f"c{i}"} for i in range(5)])
     target = database_fingerprint(entries, [e["cluster"] for e in entries], 80)
     context = calibration_context_fingerprint(
         target, load_matrix(), GAP_OPEN, GAP_EXTEND)
@@ -88,7 +89,7 @@ def test_searcher_uses_matching_experimental_empirical_parameters(tmp_path):
                 "evalue_params": params}, db)
     searcher = Searcher(str(db))
     hit = searcher.search(np.asarray(entries[0]["tokens"]),
-                          rerank="sw", top_k=1)[0]
+                          query_entry=entries[0], top_k=1)[0]
     assert hit["raw_evalue"] == pytest.approx(EV.initial_evalue(
         hit["sw_score"], 30, EV.token_entropy(entries[0]["tokens"]),
         params["target_aware_gumbel"]))

tests/test_paper_db.py

@@ -15,8 +15,8 @@
 @pytest.mark.skipif(not DB, reason="RFOLDSEARCH_PAPER_DB not set")
 def test_readme_example():
     s = Searcher(DB)
-    hits = s.search_id("PDB_00001FFZ_1_A", top_k=3, rerank="sw")
-    got = [(h["target"], round(h["sw_score"], 2)) for h in hits]
-    assert got == [("PDB_00001FFK_1_0", 203.54),
-                   ("PDB_00005AN9_1_N", 180.02),
-                   ("PDB_00004V9F_1_0", 168.13)]
+    hits = s.search_id("PDB_00001FFZ_1_A", top_k=200)
+    sw = {h["target"]: round(h["sw_score"], 2) for h in hits}
+    assert sw["PDB_00001FFK_1_0"] == 203.54
+    assert sw["PDB_00005AN9_1_N"] == 180.02
+    assert sw["PDB_00004V9F_1_0"] == 168.13

tests/test_predicted_empirical_evalue.py

@@ -5,6 +5,7 @@
 import pytest
 import torch
 
+from helpers import add_coords
 from rfoldsearch import compbias, experimental_evalue as EV
 from rfoldsearch.cli import main
 from rfoldsearch.database import (calibration_context_fingerprint,
@@ -65,8 +66,8 @@ def test_compressed_gumbel_likelihood_matches_uncompressed_fit():
 
 
 def test_search_uses_compbias_only_for_evalue_and_rejects_wrong_context(tmp_path):
-    entries = [{"id": f"t{i}", "tokens": [i % 5, (i + 1) % 5] * 15,
-                "cluster": f"c{i}"} for i in range(5)]
+    entries = add_coords([{"id": f"t{i}", "tokens": [i % 5, (i + 1) % 5] * 15,
+                           "cluster": f"c{i}"} for i in range(5)])
     tokens = [entry["tokens"] for entry in entries]
     matrix = load_matrix()
     background = compbias.background_frequencies(tokens, matrix.shape[0])
@@ -97,7 +98,7 @@ def test_search_uses_compbias_only_for_evalue_and_rejects_wrong_context(tmp_path
                 "evalue_params": params}, db)
     searcher = Searcher(str(db))
     query = np.asarray(tokens[0])
-    hit = searcher.search(query, rerank="sw", top_k=1)[0]
+    hit = searcher.search(query, query_entry=entries[0], top_k=1)[0]
     positions, profile = compbias.query_profile(query, matrix, background)
     target_index = searcher.id2i[hit["target"]]
     adjusted = smith_waterman_numba(positions, np.asarray(tokens[target_index]),
@@ -119,8 +120,8 @@ def test_search_uses_compbias_only_for_evalue_and_rejects_wrong_context(tmp_path
 
 def test_fitting_cli_writes_attachable_predicted_parameters(tmp_path):
     rng = np.random.default_rng(30)
-    entries = [{"id": f"t{i}", "tokens": rng.integers(0, 15, 25 + i % 8).tolist(),
-                "cluster": f"c{i}"} for i in range(40)]
+    entries = add_coords([{"id": f"t{i}", "tokens": rng.integers(0, 15, 25 + i % 8).tolist(),
+                           "cluster": f"c{i}"} for i in range(40)])
     tokens = [entry["tokens"] for entry in entries]
     db = tmp_path / "db.pt"
     torch.save({"n_codes": 80, "entries": entries, "n_clusters": 40}, db)
@@ -170,7 +171,8 @@ def test_fitting_cli_writes_attachable_predicted_parameters(tmp_path):
     assert params["empirical_calibration"]["background_scaling"] == "complete-eligible-clusters"
     searcher = Searcher(str(db), matrix_path=str(matrix_path),
                         evalue_params_path=str(output))
-    assert searcher.search(np.asarray(tokens[0]), rerank="sw", top_k=1)[0]["evalue"] >= 0
+    assert searcher.search(np.asarray(tokens[0]), query_entry=entries[0],
+                           top_k=1)[0]["evalue"] >= 0
     clusters = tmp_path / "clusters.tsv"
     clusters.write_text("".join(f"t{i}\tc{i}\n" for i in range(40)))
     attached = tmp_path / "attached.pt"