Fix directional Foldseek lDDT in custom scoring

src/plinder/core/scores/custom.py

@@ -1373,11 +1373,34 @@ def map_custom_alignment_hits(
     alignment_chain_lookup: Path,
     output_path: Path,
     exclude_target_entries: Iterable[str] = (),
+    reverse_alignment: Path | None = None,
 ) -> Path:
     """Replace backend identifiers with custom and PLINDER label-asym IDs."""
     if backend not in SEARCH_BACKENDS:
         raise ValueError(f"unsupported search backend: {backend}")
     raw = pd.read_parquet(raw_alignment)
+    if reverse_alignment is not None and not raw.empty:
+        keys = ["query", "target", "qstart", "qend", "tstart", "tend", "qaln", "taln"]
+        reverse_keys = [
+            "target",
+            "query",
+            "tstart",
+            "tend",
+            "qstart",
+            "qend",
+            "taln",
+            "qaln",
+        ]
+        reverse = pd.read_parquet(reverse_alignment, columns=keys + ["lddt"])
+        reverse = reverse.rename(
+            columns=dict(zip(keys, reverse_keys, strict=True))
+            | {"lddt": "reverse_lddt"}
+        )
+        raw = raw.merge(
+            reverse.drop_duplicates(), on=keys, how="left", validate="many_to_one"
+        )
+        if raw["reverse_lddt"].isna().any():
+            raise ValueError("Foldseek hits are missing reverse lDDT scores")
     excluded = {str(entry).lower() for entry in exclude_target_entries}
     if excluded and not raw.empty:
         target_entries = {
@@ -1943,14 +1966,18 @@ def prepare_custom_protein_score_alignments(
         hits["fident_qcov"] = hits["fident"] * hits["qcov"]
         hits["seqsim_qcov"] = hits["seqsim"] * hits["qcov"]
         if backend == "foldseek":
-            if "lddt" not in hits:
-                raise ValueError("foldseek custom hit table is missing column lddt")
+            if "reverse_lddt" not in hits:
+                raise ValueError(
+                    "foldseek custom hit table is missing column reverse_lddt"
+                )
+            hits["lddt"] = hits["reverse_lddt"]
             hits["lddt_qcov"] = hits["lddt"] * hits["qcov"]
         hits = hits.drop(
             columns=[
                 "qaln",
                 "taln",
                 "cigar",
+                "reverse_lddt",
                 "target_selected_residue_indices",
             ],
             errors="ignore",
@@ -2030,6 +2057,11 @@ def run_custom_protein_searches(
             raw_prefix = raw_root / (
                 backend if label == "base" else f"{backend}_{label}"
             )
+            reverse_prefix = (
+                raw_root / f"{raw_prefix.name}_reverse"
+                if backend == "foldseek"
+                else None
+            )
             search_config = alignment_config
             if label == "base" and assets.shadowed_entries:
                 lookup = Path(f"{bundle.search_target}.lookup")
@@ -2058,6 +2090,7 @@ def run_custom_protein_searches(
                 tmp_dir=backend_scratch / f"tmp_{label}",
                 remove_tmp=True,
                 threads=threads,
+                reverse_aln_file=reverse_prefix,
             )
             mapped = output_dir / f"{backend}_{label}.parquet" if combine else output
             map_custom_alignment_hits(
@@ -2069,6 +2102,11 @@ def run_custom_protein_searches(
                 exclude_target_entries=(
                     assets.shadowed_entries if label == "base" else ()
                 ),
+                reverse_alignment=(
+                    reverse_prefix.with_suffix(".parquet")
+                    if reverse_prefix is not None
+                    else None
+                ),
             )
             if combine:
                 frame = pd.read_parquet(mapped)

src/plinder/data/annotations/get_similarity_scores.py

@@ -1267,6 +1267,7 @@ def run_alignment(
     include_target_pdb_id: bool = True,
     query_ids_only: bool = False,
     id_column: str = "query",
+    reverse_aln_file: Path | None = None,
 ) -> set[str] | None:
     if id_column not in {"query", "target"}:
         raise ValueError(f"unsupported alignment ID column: {id_column}")
@@ -1374,6 +1375,50 @@ def remove_search_results() -> None:
     ]
     subprocess.check_call(convert_commands, stdout=subprocess.DEVNULL)
 
+    if reverse_aln_file is not None:
+        # A clustered Foldseek bundle keeps all expanded chains in *_seq.
+        reverse_query_db = Path(f"{conversion_target_db}_seq")
+        if not Path(f"{reverse_query_db}.dbtype").is_file():
+            reverse_query_db = conversion_target_db
+        reverse_db = search_db.parent / f"{search_db.name}_reverse"
+        subprocess.check_call(
+            [
+                aln_type,
+                "swapresults",
+                str(query_db),
+                str(reverse_query_db),
+                str(search_db),
+                str(reverse_db),
+                "-e",
+                "inf",
+                "--threads",
+                str(threads),
+            ],
+            stdout=subprocess.DEVNULL,
+        )
+        # swapresults retains NEED_SRC (bit 2 of the extended type), although
+        # the new target is the unclustered custom query DB. Foldseek 10's
+        # setextendeddbtype only adds bits, so clear it on this scratch DB.
+        dbtype_path = Path(f"{reverse_db}.dbtype")
+        dbtype = int.from_bytes(dbtype_path.read_bytes(), "little")
+        dbtype_path.write_bytes((dbtype & ~(2 << 16)).to_bytes(4, "little"))
+        reverse_commands = convert_commands.copy()
+        reverse_commands[2:6] = [
+            str(reverse_query_db),
+            str(query_db),
+            str(reverse_db),
+            str(reverse_aln_file.with_suffix(".tsv")),
+        ]
+        subprocess.check_call(reverse_commands, stdout=subprocess.DEVNULL)
+        _stream_alignment_tsv_to_dataset(
+            reverse_aln_file.with_suffix(".tsv"),
+            reverse_aln_file.with_suffix(".parquet"),
+            aln_type=aln_type,
+            include_target_pdb_id=include_target_pdb_id,
+        )
+        if remove_tmp:
+            reverse_aln_file.with_suffix(".tsv").unlink()
+
     query_ids = None
     if query_ids_only:
         with aln_file.with_suffix(".tsv").open() as handle:

tests/core/test_custom_scores.py

@@ -787,7 +787,8 @@ def test_plinder_entry_subset_rejects_foldseek(tmp_path):
         )
 
 
-def test_map_custom_alignment_hits_maps_query_and_target_chains(tmp_path):
+@pytest.mark.parametrize("with_reverse", [False, True])
+def test_map_custom_alignment_hits_maps_query_and_target_chains(tmp_path, with_reverse):
     input_root = tmp_path / "query_inputs"
     input_root.mkdir()
     chain_manifest = input_root / "query_chains.parquet"
@@ -832,14 +833,44 @@ def test_map_custom_alignment_hits_maps_query_and_target_chains(tmp_path):
                 "pdb_00002def_xyz-enrich.cif.gz_OLD",
             ],
             "qstart": [1, 1],
+            "qend": [2, 2],
             "tstart": [2, 2],
+            "tend": [3, 3],
             "qcov": [0.8, 0.8],
             "fident": [0.5, 0.5],
             "qaln": ["AC", "AC"],
             "taln": ["AC", "AC"],
             "lddt": [0.7, 0.7],
         }
     ).to_parquet(raw, index=False)
+    reverse_path = None
+    if with_reverse:
+        reverse_path = tmp_path / "reverse.parquet"
+        forward = pd.read_parquet(raw)
+        extra = forward.iloc[[0]].copy()
+        extra["qaln"] = "A-C"
+        extra["taln"] = "ACC"
+        extra["tend"] = 4
+        forward = pd.concat([forward, extra], ignore_index=True)
+        forward.to_parquet(raw, index=False)
+        reversed_hits = forward.rename(
+            columns={
+                "query": "target",
+                "target": "query",
+                "qstart": "tstart",
+                "tstart": "qstart",
+                "qend": "tend",
+                "tend": "qend",
+                "qaln": "taln",
+                "taln": "qaln",
+            }
+        )
+        reversed_hits["lddt"] = [0.6, 0.5, 0.4]
+        # Reversed results need not have the same order, and identical
+        # tracebacks may occur more than once.
+        pd.concat([reversed_hits.iloc[::-1], reversed_hits.iloc[[0]]]).to_parquet(
+            reverse_path, index=False
+        )
     lookup = tmp_path / "alignment_chain_lookup.parquet"
     pd.DataFrame(
         {
@@ -859,17 +890,21 @@ def test_map_custom_alignment_hits_maps_query_and_target_chains(tmp_path):
         alignment_chain_lookup=lookup,
         output_path=output,
         exclude_target_entries={"2def"},
+        reverse_alignment=reverse_path,
     )
     result = pd.read_parquet(output)
 
-    assert len(result) == 1
+    assert len(result) == (2 if with_reverse else 1)
     assert result.loc[0, "query_chain_id"] == "model__A"
     assert result.loc[0, "target_entry"] == "1abc"
     assert result.loc[0, "target_chain_asym_id"] == "B"
     assert result.loc[0, "source"] == "foldseek"
     assert result.loc[0, "query_sequence_source"] == "polymer"
     assert result.loc[0, "query_resolved_residue_numbers"].tolist() == [1, 2]
     assert result.loc[0, "target_selected_residue_numbers"].tolist() == [2, 4]
+    if with_reverse:
+        assert result["lddt"].tolist() == pytest.approx([0.7, 0.7])
+        assert result["reverse_lddt"].tolist() == pytest.approx([0.6, 0.4])
 
 
 def test_custom_search_defaults_keep_coverage_and_disable_identity_filter():
@@ -1079,6 +1114,7 @@ def test_prepare_custom_protein_score_alignments_reverses_direction(
     }
     if backend == "foldseek":
         row["lddt"] = 0.9
+        row["reverse_lddt"] = 0.8
     pd.DataFrame([row]).to_parquet(hit_path, index=False)
 
     outputs = custom.prepare_custom_protein_score_alignments(
@@ -1101,6 +1137,10 @@ def test_prepare_custom_protein_score_alignments_reverses_direction(
     assert result.loc[0, "selected_residue_identity_bits"] == b"\x01"
     assert result.loc[0, "fident_qcov"] == pytest.approx(0.75)
     assert "cigar" not in result
+    if backend == "foldseek":
+        assert result.loc[0, "lddt"] == pytest.approx(0.8)
+        assert result.loc[0, "lddt_qcov"] == pytest.approx(0.8 * 0.75)
+        assert "reverse_lddt" not in result
 
 
 def test_prepare_custom_score_alignments_maps_coordinate_fasta_positions(tmp_path):

tests/data/test_similarity_scores.py

@@ -3,16 +3,20 @@
 from __future__ import annotations
 
 import json
+import shutil
+import subprocess
 from collections import Counter
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import replace
 from pathlib import Path
 
+import biotite.structure as struc
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
 import pytest
+from biotite.structure.io.pdb import PDBFile
 from rdkit import Chem
 
 from plinder.core.scores.entries import (
@@ -793,6 +797,134 @@ def test_foldseek_cluster_target_is_also_used_for_conversion(
     assert stream_options == [{"aln_type": "foldseek", "include_target_pdb_id": False}]
 
 
+@pytest.mark.skipif(
+    shutil.which("foldseek") is None, reason="Foldseek is not installed"
+)
+@pytest.mark.parametrize("clustered", [False, True])
+def test_foldseek_recomputes_lddt_from_reversed_traceback(
+    tmp_path, monkeypatch, clustered
+):
+    coordinates = [
+        [[0, 0, 0], [1, 0, 0], [20, 0, 0], [21, 0, 0], [22, 0, 0]],
+        [[0, 0, 0], [1, 0, 0], [10, 0, 0], [11, 0, 0], [12, 0, 0]],
+    ]
+    databases = []
+    for name, coords in zip(["query", "target"], coordinates, strict=True):
+        atoms = struc.AtomArray(5)
+        atoms.chain_id[:] = "A"
+        atoms.res_id = np.arange(1, 6)
+        atoms.res_name[:] = "ALA"
+        atoms.atom_name[:] = "CA"
+        atoms.element[:] = "C"
+        atoms.coord = np.array(coords, dtype=np.float32)
+        pdb = PDBFile()
+        pdb.set_structure(atoms)
+        source = tmp_path / f"{name}.pdb"
+        pdb.write(source)
+        database = tmp_path / name
+        sources = [str(source)]
+        if name == "target" and clustered:
+            copy = tmp_path / "target_copy.pdb"
+            shutil.copyfile(source, copy)
+            sources.append(str(copy))
+        subprocess.check_call(
+            ["foldseek", "createdb", *sources, str(database), "--threads", "1"],
+            stdout=subprocess.DEVNULL,
+        )
+        databases.append(database)
+
+    if clustered:
+        cluster_tsv = tmp_path / "clusters.tsv"
+        cluster_tsv.write_text("0\t0\n0\t1\n")
+        subprocess.check_call(
+            [
+                "foldseek",
+                "tsv2db",
+                str(cluster_tsv),
+                str(tmp_path / "clusters"),
+                "--output-dbtype",
+                "6",
+            ],
+            stdout=subprocess.DEVNULL,
+        )
+        target = tmp_path / "clustered"
+        subprocess.check_call(
+            [
+                "foldseek",
+                "createclusearchdb",
+                str(databases[1]),
+                str(tmp_path / "clusters"),
+                str(target),
+                "--threads",
+                "1",
+            ],
+            stdout=subprocess.DEVNULL,
+        )
+        databases[1] = target
+        subprocess.check_call(
+            [
+                "foldseek",
+                "createindex",
+                str(target),
+                str(tmp_path / "index_tmp"),
+                "--index-subset",
+                "2",
+                "--threads",
+                "1",
+            ],
+            stdout=subprocess.DEVNULL,
+        )
+
+    # Supply one fixed alignment instead of running a search. The distance
+    # changes only affect pairs outside the forward query's 15-Angstrom cutoff.
+    traceback = tmp_path / "traceback.tsv"
+    traceback.write_text(f"0\t{int(clustered)}\t100\t1.0\t1e-8\t0\t4\t5\t0\t4\t5\t5M\n")
+    check_call = subprocess.check_call
+    commands = []
+
+    def run(command, **kwargs):
+        commands.append(command)
+        if command[1] == "search":
+            command = [
+                "foldseek",
+                "tsv2db",
+                str(traceback),
+                command[4],
+                "--output-dbtype",
+                str(5 | (2 << 16)) if clustered else "5",
+            ]
+        return check_call(command, **kwargs)
+
+    monkeypatch.setattr(scoring_module.subprocess, "check_call", run)
+    run_alignment(
+        aln_type="foldseek",
+        query_db=databases[0],
+        target_db=databases[1],
+        search_target_db=databases[1],
+        search_db=tmp_path / "result",
+        aln_file=tmp_path / "forward",
+        reverse_aln_file=tmp_path / "reverse",
+        alignment_config=FoldseekConfig(),
+        tmp_dir=tmp_path / "scratch",
+        threads=1,
+    )
+
+    forward = pd.read_parquet(tmp_path / "forward.parquet")
+    reverse = pd.read_parquet(tmp_path / "reverse.parquet")
+    assert forward.loc[0, "lddt"] == pytest.approx(1.0)
+    assert reverse.loc[0, "lddt"] == pytest.approx(0.4)
+    assert forward.loc[0, "query"] == reverse.loc[0, "target"]
+    assert forward.loc[0, "target"] == reverse.loc[0, "query"]
+    assert [command[1] for command in commands] == [
+        "search",
+        "convertalis",
+        "swapresults",
+        "convertalis",
+    ]
+    assert not list(tmp_path.glob("result*"))
+    assert not (tmp_path / "reverse.tsv").exists()
+
+
 def test_alignment_tsv_is_streamed_to_query_partitions(tmp_path) -> None:
     tsv_path = tmp_path / "alignment.tsv"
     tsv_path.write_text(