diff --git a/src/plinder/core/scores/custom.py b/src/plinder/core/scores/custom.py index 7492e8ee..be242a4f 100644 --- a/src/plinder/core/scores/custom.py +++ b/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) diff --git a/src/plinder/data/annotations/get_similarity_scores.py b/src/plinder/data/annotations/get_similarity_scores.py index e0487d8e..f213a78f 100644 --- a/src/plinder/data/annotations/get_similarity_scores.py +++ b/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: diff --git a/tests/core/test_custom_scores.py b/tests/core/test_custom_scores.py index fd51d531..780ddbb4 100644 --- a/tests/core/test_custom_scores.py +++ b/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,7 +833,9 @@ 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"], @@ -840,6 +843,34 @@ def test_map_custom_alignment_hits_maps_query_and_target_chains(tmp_path): "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,10 +890,11 @@ 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" @@ -870,6 +902,9 @@ def test_map_custom_alignment_hits_maps_query_and_target_chains(tmp_path): 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): diff --git a/tests/data/test_similarity_scores.py b/tests/data/test_similarity_scores.py index 57da19f6..65f3e58a 100644 --- a/tests/data/test_similarity_scores.py +++ b/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(