diff --git a/pyproject.toml b/pyproject.toml index 9461214..7bbdfdd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,10 +4,10 @@ version = "0.1.0" requires-python = ">=3.12,<3.13" dependencies = [ "torch==2.8.0", - "torch-geometric", + "torch-geometric>=2.8.0", "torch-scatter", "torch-cluster", - "pyg-lib", + "pyg-lib>=0.6.0", "e3nn", "esm", "biotite", diff --git a/scripts/generate_esm_embeddings.py b/scripts/generate_esm_embeddings.py index 5c00020..cbb92ac 100644 --- a/scripts/generate_esm_embeddings.py +++ b/scripts/generate_esm_embeddings.py @@ -41,11 +41,12 @@ from loguru import logger from tqdm import tqdm -from src.constants import ONE_TO_THREE, THREE_TO_ONE +from src.constants import THREE_TO_ONE from src.dataset import parse_asu_with_biotite from src.utils import ( normalize_ins_code, parse_split_file, + sanitize_res_names_for_esm, setup_logging_for_tqdm, ) @@ -99,13 +100,11 @@ def compute_esm_embeddings( ] num_residues = len(biotite_seq) - # Sanitize the AtomArray so ESM accepts all residues + # Sanitize the AtomArray so ESM accepts all residues. Uses the shared + # helper so residue-name canonicalization stays identical to the residue + # indexing in src/dataset.py. + protein_atoms = sanitize_res_names_for_esm(protein_atoms) protein_atoms.hetero[:] = False - for i in range(len(protein_atoms)): - orig_res = protein_atoms.res_name[i] - # Map to 1-letter code, then convert back to 3-letter - aa1 = THREE_TO_ONE.get(orig_res, "X") - protein_atoms.res_name[i] = ONE_TO_THREE.get(aa1, "UNK") # Write sanitized array to an in-memory buffer sanitized_pdb = PDBFile() diff --git a/src/dataset.py b/src/dataset.py index e376961..4a6ba0b 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -31,10 +31,16 @@ from torch_geometric.data import Batch, HeteroData from tqdm import tqdm -from src.constants import EDGE_PP, ELEM_IDX, ELEMENT_VOCAB, NUM_RBF +from src.constants import ( + EDGE_PP, + ELEM_IDX, + ELEMENT_VOCAB, + NUM_RBF, +) from src.utils import ( compute_edge_features, normalize_ins_code, + sanitize_res_names_for_esm, ) @@ -157,25 +163,39 @@ def match_atoms_to_coords( atoms: bts.AtomArray, target_coords: np.ndarray, tolerance: float = 0.01 ) -> list[int]: """ - Match biotite atoms to target coordinates by nearest neighbor. (needed for mates when parsing with PyMOL) + Match biotite atoms to coordinates from a second parse of the same structure. + + Reconciles biotite against PyMOL. The two disagree on count by design: PyMOL + keeps every altloc conformer while biotite takes the highest-occupancy one, + so PyMOL's atom set is a superset. Every biotite atom should still be found + in it; the caller drops any that are not. Args: atoms: Biotite AtomArray with coord attribute - target_coords: (N, 3) array of target coordinates to match + target_coords: (N, 3) coordinates to match against tolerance: Maximum distance in Angstroms for a valid match Returns: - List of indices into atoms array for matched atoms + Index into atoms for each target coordinate whose nearest atom lies + within tolerance. May repeat an index if two targets share an atom. """ - if target_coords.shape[0] == 0: + if target_coords.shape[0] == 0 or len(atoms) == 0: return [] - matched = [] - for i, coord in enumerate(target_coords): - dists = np.linalg.norm(atoms.coord - coord, axis=1) - min_idx = np.argmin(dists) - if dists[min_idx] < tolerance: - matched.append(min_idx) + from scipy.spatial import cKDTree + + tree = cKDTree(atoms.coord) + dists, nearest = tree.query(target_coords, k=1, distance_upper_bound=tolerance) + within = np.isfinite(dists) & (nearest < len(atoms)) + matched = nearest[within].tolist() + + # A wholesale miss means the parses disagree (frame, cell), not that the + # atoms are bad. Warn, or the caller's drop looks like clean data. + if len(set(matched)) < len(atoms) / 2: + logger.warning( + f"Only {len(set(matched))}/{len(atoms)} atoms matched within " + f"{tolerance}A; parses may disagree. Unmatched atoms are dropped." + ) return matched @@ -696,7 +716,7 @@ def filter_waters_by_quality( class ProteinWaterDataset(Dataset): """ - Dataset for protein crystal contact prediction. + Dataset for predicting water positions in protein crystal structures. Returns HeteroData with: - 'protein' node type: ASU protein atoms + optionally symmetry mates @@ -968,9 +988,9 @@ def _preprocess_one(self, entry: dict, cache_path: Path): crystal_data = get_crystal_contacts_pymol(struc_path, self.cutoff) - # Ensure consistency between biotite and PyMOL parsing. - # Both parse the same ASU, but may differ in altloc selection, hydrogen - # handling, or edge cases. Keep only waters present in both representations. + # Keep only the waters PyMOL also saw. PyMOL's ASU is a superset of + # biotite's (it keeps every altloc conformer), so a water missing from it + # means the two parses disagree rather than that the water is unwanted. asu_water_indices = match_atoms_to_coords( water_atoms, crystal_data["asu_coords"] ) @@ -1067,20 +1087,19 @@ def _preprocess_one(self, entry: dict, cache_path: Path): protein_elements = [str(e).upper() for e in protein_atoms.element] protein_x = element_onehot(protein_elements) - # compute residue indices (including ins_code to match ESM/SLAE residue counting) - res_id = protein_atoms.res_id - chain_id_arr = protein_atoms.chain_id - ins_code_arr = np.array( - [normalize_ins_code(x) for x in protein_atoms.ins_code], dtype=object - ) - residue_keys = list(zip(chain_id_arr, res_id, ins_code_arr)) - unique_res = {k: i for i, k in enumerate(dict.fromkeys(residue_keys))} - protein_res_idx = torch.tensor( - [unique_res[k] for k in residue_keys], dtype=torch.long - ) - - # check water/residue ratio - num_residues = len(unique_res) + # protein_res_idx indexes cached ESM embedding rows, so it uses biotite's + # residue segmentation, not res_id (not 0-based, not contiguous, repeats + # across chains). Sanitize names and normalize ins_codes first so residues + # split exactly where the ESM script splits them. + sanitized_for_idx = sanitize_res_names_for_esm(protein_atoms) + for i in range(len(sanitized_for_idx)): + sanitized_for_idx.ins_code[i] = normalize_ins_code( + sanitized_for_idx.ins_code[i] + ) + num_residues = bts.get_residue_count(sanitized_for_idx) + protein_res_idx = torch.from_numpy( + bts.spread_residue_wise(sanitized_for_idx, np.arange(num_residues)) + ).long() num_waters = len(water_atoms) ratio_valid, ratio_reason = check_water_residue_ratio( num_waters, diff --git a/src/flow.py b/src/flow.py index 42f2953..6a4f151 100644 --- a/src/flow.py +++ b/src/flow.py @@ -176,12 +176,37 @@ def build_knn_edges( batch_dst: torch.Tensor | None = None, ) -> torch.Tensor: """ - KNN edges from src -> dst (source indices in row 0, dest in row 1). + Build KNN edges from src -> dst (source indices in row 0, dest in row 1). + + The KNN query is performed *per destination*: for each point in ``dst_pos`` + we look up its ``k`` nearest neighbors in ``src_pos`` (``knn(x=src_pos, + y=dst_pos, ...)``) and emit them as incoming edges. As a consequence every + destination node is guaranteed to have up to ``k`` incoming edges (and so + appears in row 1), whereas a source node that is no destination's nearest + neighbor may not appear in row 0 at all. Coverage checks ("every node has an + edge") must therefore be made against the destination row (row 1). + + For a homogeneous graph (``src_pos is dst_pos``) self-edges are dropped. + + Args: + src_pos: (N_src, 3) source node positions. + dst_pos: (N_dst, 3) destination node positions. + k: Number of nearest source neighbors to find per destination node. + batch_src: (N_src,) batch assignment for source nodes, or None. + batch_dst: (N_dst,) batch assignment for destination nodes, or None. + + Returns: + (2, E) edge index tensor with source indices in row 0, destination in + row 1. """ if src_pos.numel() == 0 or dst_pos.numel() == 0: return torch.empty(2, 0, dtype=torch.long, device=src_pos.device) - idx = knn(x=dst_pos, y=src_pos, k=k, batch_x=batch_dst, batch_y=batch_src) + # knn(x, y) returns row 0 = y (query), row 1 = x (neighbor); swap for this + # repo's src(row 0)->dst(row 1) convention. That row order is undocumented, + # so it is pinned by tests/test_flow.py::TestBuildKnnEdgesDirection. + idx = knn(x=src_pos, y=dst_pos, k=k, batch_x=batch_src, batch_y=batch_dst) + idx = torch.stack((idx[1], idx[0]), dim=0) # remove self-edges if homogeneous if src_pos.data_ptr() == dst_pos.data_ptr(): diff --git a/src/utils.py b/src/utils.py index 12f59e4..8e3b660 100644 --- a/src/utils.py +++ b/src/utils.py @@ -4,7 +4,7 @@ """ Utility functions organized by category: -1. Feature encoding (rbf, atom37_to_atoms, normalize_ins_code) +1. Feature encoding (rbf, normalize_ins_code) 2. Optimal transport (ot_coupling) 3. Metrics (recall_precision, compute_rmsd, compute_placement_metrics) 4. Visualization (plot_3d_frame, create_trajectory_gif, save_protein_plot) @@ -14,6 +14,7 @@ from collections.abc import Sequence from pathlib import Path +from typing import TYPE_CHECKING import matplotlib.pyplot as plt import numpy as np @@ -24,44 +25,13 @@ from PIL import Image from scipy.optimize import linear_sum_assignment from torch import Tensor -from torch_geometric.nn import knn from tqdm import tqdm -from src.constants import NUM_RBF, RBF_CUTOFF +from src.constants import NUM_RBF, ONE_TO_THREE, RBF_CUTOFF, THREE_TO_ONE -def build_knn_edges( - src_pos: torch.Tensor, - dst_pos: torch.Tensor, - k: int, - batch_src: torch.Tensor | None = None, - batch_dst: torch.Tensor | None = None, -) -> torch.Tensor: - """ - Build KNN edges from source to destination nodes. - - Args: - src_pos: (N_src, 3) source node positions - dst_pos: (N_dst, 3) destination node positions - k: Number of nearest neighbors per source node - batch_src: (N_src,) batch indices for source nodes, or None if single graph - batch_dst: (N_dst,) batch indices for destination nodes, or None if single graph - - Returns: - (2, E) edge index tensor with source indices in row 0, destination in row 1. - Self-edges are removed for homogeneous graphs (src_pos is dst_pos). - """ - if src_pos.numel() == 0 or dst_pos.numel() == 0: - return torch.empty(2, 0, dtype=torch.long, device=src_pos.device) - - idx = knn(x=dst_pos, y=src_pos, k=k, batch_x=batch_dst, batch_y=batch_src) - - # remove self-edges if homogeneous - if src_pos.data_ptr() == dst_pos.data_ptr(): - mask = idx[0] != idx[1] - idx = idx[:, mask] - - return idx.unique(dim=1) +if TYPE_CHECKING: + import biotite.structure as bts def setup_logging_for_tqdm( @@ -117,6 +87,36 @@ def normalize_ins_code(value) -> str: return ins +def sanitize_res_names_for_esm(atoms: bts.AtomArray) -> bts.AtomArray: + """ + Return a copy of an AtomArray with residue names canonicalized to match the + ESM embedding pipeline. + + Each residue name is mapped to its one-letter code and back + (``THREE_TO_ONE`` -> ``ONE_TO_THREE``), with anything unrecognized collapsed + to ``"UNK"``. This merges non-canonical names that share a residue position + (e.g. modified residues -> their canonical parent, unknowns -> ``UNK``) so + that biotite's ``get_residue_starts`` does not split them apart. + + This is the single source of truth for residue-name sanitization shared by + ``scripts/generate_esm_embeddings.py`` (which feeds the sanitized structure + to ESM3) and ``src/dataset.py`` (which derives residue indices that must line + up with the stored ESM embeddings). Insertion codes are normalized + separately via :func:`normalize_ins_code`. + + Args: + atoms: A biotite ``AtomArray`` with a ``res_name`` annotation. + + Returns: + A copy of ``atoms`` with ``res_name`` canonicalized. + """ + sanitized = atoms.copy() + for i in range(len(sanitized)): + aa1 = THREE_TO_ONE.get(sanitized.res_name[i], "X") + sanitized.res_name[i] = ONE_TO_THREE.get(aa1, "UNK") + return sanitized + + def parse_split_file(split_file: Path, base_pdb_dir: Path) -> list[dict]: """ Parse split file and construct entries with resolved structure paths. @@ -172,9 +172,6 @@ def parse_split_file(split_file: Path, base_pdb_dir: Path) -> list[dict]: return entries -ATOM37_FILL = 1e-5 - - def rbf(r: Tensor, num_gaussians: int = NUM_RBF, cutoff: float = RBF_CUTOFF) -> Tensor: """ Compute radial basis function encoding of distances. @@ -272,32 +269,6 @@ def compute_edge_features( return unit_vectors, rbf_features -def atom37_to_atoms( - atom_tensor: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """ - Convert atom37 representation to flat atom list. - - Args: - atom_tensor: (N_res, 37, 3) atom37 coordinates - - Returns: - coords: (N_atoms, 3) coordinates of present atoms - residue_index: (N_atoms,) which residue each atom belongs to - atom_type: (N_atoms,) atom type index (0-36) - """ - present = (atom_tensor != ATOM37_FILL).any(dim=-1) # (N_res, 37) - nz = present.nonzero(as_tuple=False) # (N_atoms, 2) - residue_index = nz[:, 0] - atom_type = nz[:, 1].long() - - flat = atom_tensor.reshape(-1, 3) - flat_mask = present.reshape(-1) - coords = flat[flat_mask] - - return coords, residue_index, atom_type - - @torch.no_grad() def ot_coupling( x1: torch.Tensor, diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 88373ef..b58f2f2 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -22,6 +22,7 @@ import json from pathlib import Path +import biotite.structure as bts import numpy as np import pytest import torch @@ -52,6 +53,7 @@ parse_asu_with_biotite, ProteinWaterDataset, ) +from src.utils import normalize_ins_code, sanitize_res_names_for_esm # PDB fixtures (pdb_base_dir, pdb_6eey, pdb_2b5w, pdb_8dzt, pdb_1deu) are @@ -73,6 +75,17 @@ def single_pdb_list_file(tmp_path, pdb_6eey): return str(list_file) +@pytest.fixture +def warning_log(): + """Collect loguru warning messages; loguru does not reach pytest's caplog.""" + from loguru import logger + + messages = [] + sink_id = logger.add(messages.append, level="WARNING", format="{message}") + yield messages + logger.remove(sink_id) + + @pytest.mark.unit class TestElementOnehot: """Tests for element one-hot encoding.""" @@ -249,6 +262,48 @@ def test_tolerance_parameter(self): matched_loose = match_atoms_to_coords(atoms, target_coords, tolerance=0.1) assert len(matched_loose) == 1 + def test_warns_when_most_atoms_unmatched(self, warning_log): + """A mostly-failed match must warn: the caller drops what doesn't match.""" + atoms = bts.AtomArray(4) + atoms.coord = np.array([[float(i), 0.0, 0.0] for i in range(4)]) + + # coordinates nowhere near the atoms + matched = match_atoms_to_coords(atoms, np.array([[99.0, 0.0, 0.0]])) + + assert matched == [] + assert "0/4 atoms matched" in warning_log[0] + + def test_no_warning_when_all_match(self, warning_log): + """A clean match must stay silent.""" + atoms = bts.AtomArray(4) + atoms.coord = np.array([[float(i), 0.0, 0.0] for i in range(4)]) + + matched = match_atoms_to_coords(atoms, atoms.coord.copy()) + + assert len(matched) == 4 + assert warning_log == [] + + @pytest.mark.parametrize( + "n_atoms,n_matched", + [ + (1, 0), # 0% + (3, 1), # 33% + (5, 2), # 40% + ], + ) + def test_warns_on_odd_counts_below_half(self, warning_log, n_atoms, n_matched): + """Fewer than half matched must warn even when half is fractional; a + threshold of len(atoms) // 2 rounds the cutoff down and stays silent.""" + atoms = bts.AtomArray(n_atoms) + atoms.coord = np.array([[float(i), 0.0, 0.0] for i in range(n_atoms)]) + + # hit exactly n_matched atoms, plus one coordinate far from every atom + targets = np.vstack([atoms.coord[:n_matched], [[99.0, 0.0, 0.0]]]) + matched = match_atoms_to_coords(atoms, targets) + + assert len(set(matched)) == n_matched + assert f"{n_matched}/{n_atoms} atoms matched" in warning_log[0] + @pytest.mark.unit class TestCheckComDistance: @@ -2611,6 +2666,81 @@ def test_num_asu_protein_metadata_correct( assert data.num_asu_protein_atoms > 0 +# ============== Tests for residue index assignment ============== + + +@pytest.mark.integration +class TestResidueIndexAssignment: + """protein_res_idx must be usable as an index into cached ESM embedding rows.""" + + @staticmethod + def _asu_res_idx(dataset): + """ASU protein residue indices for the first entry (excludes mates).""" + data = dataset[0] + return data["protein"].residue_index[: data.num_asu_protein_atoms] + + def _build(self, pdb_id, tmp_path, pdb_base_dir, **kwargs): + """Build a dataset with water filters off (irrelevant to indexing).""" + list_file = tmp_path / "list.txt" + list_file.write_text(f"{pdb_id}_final\n") + return ProteinWaterDataset( + pdb_list_file=str(list_file), + processed_dir=str(tmp_path / "processed"), + base_pdb_dir=str(pdb_base_dir), + include_mates=False, + preprocess=True, + filter_by_distance=False, + filter_by_edia=False, + filter_by_bfactor=False, + **kwargs, + ) + + def test_index_is_zero_based_and_contiguous(self, tmp_path, pdb_base_dir): + """Indices must cover 0..num_residues-1 with no gaps.""" + res_idx = self._asu_res_idx(self._build("6eey", tmp_path, pdb_base_dir)) + + assert res_idx.min().item() == 0 + assert torch.equal( + torch.unique(res_idx), torch.arange(res_idx.max().item() + 1) + ) + + def test_index_matches_num_residues(self, tmp_path, pdb_base_dir): + """Distinct index count must equal the ESM embedding table's row count.""" + dataset = self._build("6eey", tmp_path, pdb_base_dir) + data = dataset[0] + res_idx = data["protein"].residue_index[: data.num_asu_protein_atoms] + + assert len(torch.unique(res_idx)) == data["protein"].num_protein_residues + + def test_atoms_of_same_residue_share_index(self, tmp_path, pdb_base_dir): + """Each residue is one contiguous atom block owning >=1 atom.""" + res_idx = self._asu_res_idx(self._build("6eey", tmp_path, pdb_base_dir)) + + assert torch.all(res_idx[1:] >= res_idx[:-1]) # non-decreasing + assert (torch.bincount(res_idx) > 0).all() + + def test_insertion_codes_do_not_split_residues(self, tmp_path, pdb_base_dir): + """1deu has real insertion codes: they mark distinct residues, but + placeholder codes must not inflate the count.""" + # 1deu's chains sit 45A apart; relax the unrelated interface guard. + dataset = self._build( + "1deu", tmp_path, pdb_base_dir, interface_dist_threshold=100.0 + ) + data = dataset[0] + res_idx = data["protein"].residue_index[: data.num_asu_protein_atoms] + + protein_atoms, _, _ = parse_asu_with_biotite( + str(Path(pdb_base_dir) / "1deu" / "1deu_final.pdb") + ) + sanitized = sanitize_res_names_for_esm(protein_atoms) + for i in range(len(sanitized)): + sanitized.ins_code[i] = normalize_ins_code(sanitized.ins_code[i]) + expected = len(bts.get_residue_starts(sanitized)) + + assert len(torch.unique(res_idx)) == expected + assert data["protein"].num_protein_residues == expected + + # ============== Tests for RBF feature computation ============== diff --git a/tests/test_flow.py b/tests/test_flow.py index 37642d3..9d2f62a 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -10,6 +10,7 @@ import torch import torch.nn.functional as F from torch_geometric.data import Batch, Data, HeteroData +from torch_geometric.nn import knn from src.flow import ( _batch_from_counts, @@ -159,6 +160,74 @@ def test_with_batch(self, device): assert edges.shape[1] > 0 +@pytest.mark.unit +class TestBuildKnnEdgesDirection: + """Exact-set KNN direction tests on asymmetric geometry. + + srcs are spread out, both dsts sit near src[0], so "k nearest srcs per dst" + (correct) and "k nearest dsts per src" (the x/y swap) give different edge + sets -- no distance ties to mask a mix-up. + """ + + SRC = [[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [20.0, 0.0, 0.0]] + DST = [[1.0, 0.0, 0.0], [2.0, 0.0, 0.0]] + + def test_exact_edge_set_is_per_destination(self, device): + """Both dsts are nearest src[0], so src[1]/src[2] must not appear. An x/y + swap gives {(0,0), (1,1), (2,1)} instead.""" + src = torch.tensor(self.SRC, device=device) + dst = torch.tensor(self.DST, device=device) + + edges = build_knn_edges(src, dst, k=1) + + edge_set = set(zip(edges[0].tolist(), edges[1].tolist())) + assert edge_set == {(0, 0), (0, 1)}, f"got {sorted(edge_set)}" + + def test_every_destination_is_covered(self, device): + """Coverage is per-destination: every dst gets k in-edges; an unneeded src + may be absent.""" + src = torch.tensor(self.SRC, device=device) + dst = torch.tensor(self.DST, device=device) + k = 2 + + edges = build_knn_edges(src, dst, k=k) + + dst_row = edges[1] + for d in range(len(self.DST)): + assert (dst_row == d).sum().item() == k, f"dst {d} lacks {k} in-edges" + # src[2] (x=20) is not among the 2 nearest srcs of either dst + assert 2 not in set(edges[0].tolist()) + + def test_output_rows_are_src_then_dst(self, device): + """Row 0 = src, row 1 = dst, pinned by index range. + + Own geometry: both dsts are nearest src[2], so row 0 must reach index 2 + while row 1 only reaches 1. Swapping the rows puts 2 in row 1, which is + out of range for two dsts. (The class fixture can't pin this -- its row 0 + is all zeros, so both ranges hold either way round.) + """ + src = torch.tensor( + [[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [20.0, 0.0, 0.0]], device=device + ) + dst = torch.tensor([[19.0, 0.0, 0.0], [21.0, 0.0, 0.0]], device=device) + + edges = build_knn_edges(src, dst, k=1) + + assert edges[0].max().item() == 2 # src[2]; >= len(dst), so a swap breaks + assert edges[1].max().item() < len(dst) + + def test_torch_geometric_knn_row_convention_unchanged(self, device): + """Pin knn's undocumented rows: row 0 = y (query), row 1 = x (neighbor). + build_knn_edges swaps them, so a flip here would reverse every edge.""" + x = torch.tensor([[0.0, 0.0], [10.0, 0.0], [20.0, 0.0]], device=device) # N=3 + y = torch.tensor([[0.1, 0.0], [19.9, 0.0]], device=device) # M=2 + + out = knn(x, y, k=1) + + assert out[0].tolist() == [0, 1] # queries (y), in order + assert out[1].tolist() == [0, 2] # nearest x: y[0]->x[0], y[1]->x[2] + + @pytest.mark.unit class TestMakeEncoderData: def test_basic_output(self, simple_hetero_data): @@ -895,8 +964,11 @@ def test_all_waters_have_water_edges(self, simple_hetero_data): n_water = simple_hetero_data["water"].num_nodes if n_water > 1: - # Check that all water nodes appear in the water-water edges - water_nodes_with_edges = torch.unique(ww_edges[0]) + # WW edges are built per destination (knn query per water), so every + # water is guaranteed to appear as a destination (row 1); a water that + # is no other water's nearest neighbor would be missing from the source + # row (row 0). Assert coverage on the destination/query row. + water_nodes_with_edges = torch.unique(ww_edges[1]) assert len(water_nodes_with_edges) == n_water, ( f"Only {len(water_nodes_with_edges)}/{n_water} waters have water-water edges" ) @@ -917,9 +989,11 @@ def test_batched_waters_have_edges(self, batched_hetero_data): f"Only {len(water_nodes_with_pw_edges)}/{n_water} waters have protein edges in batched data" ) - # Check water-water edges + # Check water-water edges. WW edges are built per destination, so every + # water appears as a destination (row 1); assert coverage on the + # destination/query row rather than the source row. if n_water > 1: - water_nodes_with_ww_edges = torch.unique(ww_edges[0]) + water_nodes_with_ww_edges = torch.unique(ww_edges[1]) assert len(water_nodes_with_ww_edges) == n_water, ( f"Only {len(water_nodes_with_ww_edges)}/{n_water} waters have water-water edges in batched data" ) diff --git a/tests/test_utils.py b/tests/test_utils.py index 83f878c..45c1145 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -4,7 +4,7 @@ Tests for src/utils.py utility functions. Organized by category to match utils.py structure: -1. Feature encoding (rbf, atom37_to_atoms) +1. Feature encoding (rbf) 2. Optimal transport (ot_coupling) 3. Metrics (recall_precision, compute_rmsd, compute_placement_metrics) 4. Visualization (plot_3d_frame, save_protein_plot, create_trajectory_gif) @@ -14,18 +14,18 @@ from pathlib import Path +import biotite.structure as bts import matplotlib import numpy as np import pytest import torch +from biotite.structure import array, Atom matplotlib.use("Agg") from src.utils import ( - ATOM37_FILL, - atom37_to_atoms, compute_edge_features, compute_edge_geometry, compute_placement_metrics, @@ -39,6 +39,7 @@ rbf, # Metrics recall_precision, + sanitize_res_names_for_esm, save_protein_plot, ) @@ -134,6 +135,88 @@ def test_normalize_valid_code(self): assert normalize_ins_code(" B ") == "B" +@pytest.mark.unit +class TestSanitizeResNamesForEsm: + """Tests for ESM residue-name sanitization and residue-count alignment. + + These guard the contract that src/dataset.py's residue counting (via + biotite.get_residue_starts on a sanitized array) stays in lockstep with + scripts/generate_esm_embeddings.py's residue keys, which are built with + THREE_TO_ONE/normalize_ins_code. A drift here desyncs protein_res_idx from + the cached ESM embeddings. + """ + + @staticmethod + def _make_atoms(residues): + """Build a single-atom-per-row AtomArray from (res_id, res_name, ins) tuples.""" + return array( + [ + Atom( + [0.0, 0.0, 0.0], + chain_id="A", + res_id=res_id, + res_name=res_name, + ins_code=ins, + atom_name="CA", + element="C", + ) + for (res_id, res_name, ins) in residues + ] + ) + + @staticmethod + def _esm_key_count(atoms): + """Replicate the ESM script's residue-key counting.""" + keys = [] + for i in range(len(atoms)): + key = ( + atoms.chain_id[i], + atoms.res_id[i], + normalize_ins_code(atoms.ins_code[i]), + ) + if key not in keys: + keys.append(key) + return len(keys) + + def test_sanitize_canonicalizes_modified_and_unknown(self): + atoms = self._make_atoms([(1, "MSE", ""), (2, "ALA", ""), (3, "Q2K", "")]) + sanitized = sanitize_res_names_for_esm(atoms) + # MSE -> MET (canonical parent), ALA unchanged, Q2K -> UNK (unknown) + assert list(sanitized.res_name) == ["MET", "ALA", "UNK"] + # original array must be untouched (helper returns a copy) + assert list(atoms.res_name) == ["MSE", "ALA", "Q2K"] + + def test_placeholder_ins_code_desync_is_fixed(self): + # Two atoms that share (chain, res_id, res_name) and differ only by a + # placeholder insertion code ('' vs '.'). ESM keys treat them as one + # residue; raw get_residue_starts would split them into two. + atoms = self._make_atoms([(5, "GLY", ""), (5, "GLY", ".")]) + assert self._esm_key_count(atoms) == 1 + + sanitized = sanitize_res_names_for_esm(atoms) + for i in range(len(sanitized)): + sanitized.ins_code[i] = normalize_ins_code(sanitized.ins_code[i]) + assert len(bts.get_residue_starts(sanitized)) == 1 + + def test_residue_count_matches_esm_keys(self): + # Mix of canonical, modified, unknown residues with placeholder and real + # insertion codes; the dataset count must equal the ESM key count. + atoms = self._make_atoms( + [ + (1, "ALA", ""), + (1, "ALA", ""), + (2, "MSE", "."), + (3, "GLY", "?"), + (3, "GLY", "A"), # real insertion code -> distinct residue + (4, "Q2K", ""), + ] + ) + sanitized = sanitize_res_names_for_esm(atoms) + for i in range(len(sanitized)): + sanitized.ins_code[i] = normalize_ins_code(sanitized.ins_code[i]) + assert len(bts.get_residue_starts(sanitized)) == self._esm_key_count(atoms) + + @pytest.mark.unit class TestSplitFileParsing: """Tests for parse_split_file CIF/PDB structure-path resolution.""" @@ -174,83 +257,6 @@ def test_raises_when_structure_missing(self, tmp_path): parse_split_file(split_file, base_dir) -@pytest.mark.unit -class TestAtom37ToAtoms: - """Tests for atom37 representation conversion.""" - - def test_basic_conversion(self): - """Basic conversion from atom37 to flat atoms.""" - # Create atom37 tensor with some present atoms - atom_tensor = torch.full((3, 37, 3), ATOM37_FILL) - # Place atoms at specific positions - atom_tensor[0, 0, :] = torch.tensor([1.0, 2.0, 3.0]) # CA of residue 0 - atom_tensor[0, 1, :] = torch.tensor([1.5, 2.5, 3.5]) # C of residue 0 - atom_tensor[1, 0, :] = torch.tensor([4.0, 5.0, 6.0]) # CA of residue 1 - - coords, residue_idx, atom_type = atom37_to_atoms(atom_tensor) - - assert coords.shape == (3, 3) - assert residue_idx.shape == (3,) - assert atom_type.shape == (3,) - - def test_residue_indices_correct(self): - """Residue indices should match the original residue.""" - atom_tensor = torch.full((2, 37, 3), ATOM37_FILL) - atom_tensor[0, 0, :] = torch.tensor([1.0, 0.0, 0.0]) - atom_tensor[0, 1, :] = torch.tensor([2.0, 0.0, 0.0]) - atom_tensor[1, 5, :] = torch.tensor([3.0, 0.0, 0.0]) - - _, residue_idx, atom_type = atom37_to_atoms(atom_tensor) - - assert residue_idx[0] == 0 - assert residue_idx[1] == 0 - assert residue_idx[2] == 1 - - def test_atom_types_correct(self): - """Atom types should match the slot index.""" - atom_tensor = torch.full((1, 37, 3), ATOM37_FILL) - atom_tensor[0, 0, :] = torch.tensor([1.0, 0.0, 0.0]) # slot 0 - atom_tensor[0, 5, :] = torch.tensor([2.0, 0.0, 0.0]) # slot 5 - atom_tensor[0, 10, :] = torch.tensor([3.0, 0.0, 0.0]) # slot 10 - - _, _, atom_type = atom37_to_atoms(atom_tensor) - - assert atom_type[0] == 0 - assert atom_type[1] == 5 - assert atom_type[2] == 10 - - def test_empty_residues(self): - """Empty residues should not contribute atoms.""" - atom_tensor = torch.full((3, 37, 3), ATOM37_FILL) - # Only residue 0 has atoms - atom_tensor[0, 0, :] = torch.tensor([1.0, 0.0, 0.0]) - - coords, residue_idx, _ = atom37_to_atoms(atom_tensor) - - assert coords.shape == (1, 3) - assert residue_idx[0] == 0 - - def test_all_empty(self): - """All-empty tensor should return empty outputs.""" - atom_tensor = torch.full((5, 37, 3), ATOM37_FILL) - - coords, residue_idx, atom_type = atom37_to_atoms(atom_tensor) - - assert coords.shape == (0, 3) - assert residue_idx.shape == (0,) - assert atom_type.shape == (0,) - - def test_coordinates_preserved(self): - """Coordinates should be preserved exactly.""" - atom_tensor = torch.full((1, 37, 3), ATOM37_FILL) - expected_coord = torch.tensor([1.234, 5.678, 9.012]) - atom_tensor[0, 0, :] = expected_coord - - coords, _, _ = atom37_to_atoms(atom_tensor) - - assert torch.allclose(coords[0], expected_coord) - - @pytest.mark.unit class TestOTCoupling: """Tests for OT coupling in flow matching.""" diff --git a/uv.lock b/uv.lock index ace0e10..0017e0e 100644 --- a/uv.lock +++ b/uv.lock @@ -1307,11 +1307,11 @@ wheels = [ [[package]] name = "pyg-lib" -version = "0.5.0+pt28cu126" +version = "0.6.0+pt28cu126" source = { registry = "https://data.pyg.org/whl/torch-2.8.0+cu126.html" } wheels = [ - { url = "https://data.pyg.org/whl/torch-2.8.0%2Bcu126/pyg_lib-0.5.0%2Bpt28cu126-cp312-cp312-linux_x86_64.whl" }, - { url = "https://data.pyg.org/whl/torch-2.8.0%2Bcu126/pyg_lib-0.5.0%2Bpt28cu126-cp312-cp312-win_amd64.whl" }, + { url = "https://data.pyg.org/whl/torch-2.8.0%2Bcu126/pyg_lib-0.6.0%2Bpt28cu126-cp312-cp312-manylinux_2_28_x86_64.whl" }, + { url = "https://data.pyg.org/whl/torch-2.8.0%2Bcu126/pyg_lib-0.6.0%2Bpt28cu126-cp312-cp312-win_amd64.whl" }, ] [[package]] @@ -1718,7 +1718,7 @@ wheels = [ [[package]] name = "torch-geometric" -version = "2.7.0" +version = "2.8.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohttp" }, @@ -1731,9 +1731,8 @@ dependencies = [ { name = "tqdm" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/75/63/b210152635902da7fe79fcdd16517fae108f457a0ed22c737e702a9afbae/torch_geometric-2.7.0.tar.gz", hash = "sha256:f9099e4aece1a9f618c84dbaac33a77f43139736698c7e8bddf3301ef1f2e8d4", size = 876725, upload-time = "2025-10-15T20:48:03.443Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/d3/4dffd7300500465e0b4a2ae917dcb2ce771de0b9a772670365799a27c024/torch_geometric-2.7.0-py3-none-any.whl", hash = "sha256:6e0cd3ad824d484651ef5d308fc66c687bfcf5ba040d56d1e0fe0f81f365e292", size = 1275346, upload-time = "2025-10-15T20:48:01.949Z" }, + { url = "https://files.pythonhosted.org/packages/ab/68/71b142d713c5449b4b3446233d85dc711f63c42d7768b28c690ff55bc181/torch_geometric-2.8.0-py3-none-any.whl", hash = "sha256:1f62e415a2e9ee69d34617d1b0b230e9d9040f51809b96e801e742770fd4dada", size = 1325330, upload-time = "2026-06-05T21:13:18.257Z" }, ] [[package]] @@ -1976,12 +1975,12 @@ requires-dist = [ { name = "numpy" }, { name = "pandas" }, { name = "pillow" }, - { name = "pyg-lib", index = "https://data.pyg.org/whl/torch-2.8.0+cu126.html" }, + { name = "pyg-lib", specifier = ">=0.6.0", index = "https://data.pyg.org/whl/torch-2.8.0+cu126.html" }, { name = "pymol-open-source-whl", specifier = ">=3.1.0.4" }, { name = "scipy" }, { name = "torch", specifier = "==2.8.0", index = "https://download.pytorch.org/whl/cu126" }, { name = "torch-cluster", index = "https://data.pyg.org/whl/torch-2.8.0+cu126.html" }, - { name = "torch-geometric" }, + { name = "torch-geometric", specifier = ">=2.8.0" }, { name = "torch-scatter", index = "https://data.pyg.org/whl/torch-2.8.0+cu126.html" }, { name = "tqdm" }, { name = "wandb" },