diff --git a/README.md b/README.md index 4724d45..64ce6f9 100644 --- a/README.md +++ b/README.md @@ -72,14 +72,16 @@ WaterFlow processes structure files through several stages to create training-re - For atoms with alternate conformations, the highest-occupancy conformer is selected **Crystal Contact Detection** -- Uses PyMOL's `symexp` to generate symmetry mates within 5.0Å cutoff -- Symmetry mate atoms are included as additional protein context when `include_mates=True` -- Mate atoms are stored separately for proper handling during training +- Uses PyMOL's `symexp` to generate symmetry mates, keeping whole residues and whole ligand entities with any atom within the cutoff of the ASU. Runs only when `include_mates=True`; a no-mates cache never invokes PyMOL +- Protein mates and ligand mates are selected separately by PyMOL's own classifiers, so `is_ligand` stays exact for mate nodes too +- **Mate waters are never selected.** A mate water is a symmetry image of an ASU water, which is what the model predicts, so keeping it as context leaks the label +- Symmetry also maps atoms onto themselves (special positions) and reaches one residue through two operators. Mate atoms within 0.3Å of an ASU atom, a target water, or an already-kept mate atom are dropped (`dedup_mate_atoms`); mate ligands are judged whole, so a ligand is never fragmented (`dedup_mate_ligands_by_residue`) +- A mate keeps its source residue's `(chain, res_id, ins_code)`, so it inherits that residue's ESM row through `emb_res_idx` instead of a zero vector, and it joins the distance-filter reference so a water in a crystal contact — near a neighbour surface but far from the ASU — is not dropped as solvent-far **Graph Representation** - Node types: `protein` (ASU + symmetry mates + ligands), `water` (ground truth) -- ASU ligand atoms are appended after ASU and mate atoms and carry the boolean `is_ligand` mask plus `residue_index = -1` (they have no residue embedding, so residue pooling masks them out) -- `is_ligand` marks **ASU ligands only**. Symmetry-mate generation is currently unfiltered, so mate nodes can include HETATM and water atoms that `is_ligand` does not mark — see `TODO(mates)` in `ProteinWaterDataset._preprocess_one`. Don't treat `is_ligand` as an exhaustive ligand selector +- Ligand atoms are appended after ASU and mate atoms and carry the boolean `is_ligand` mask plus `residue_index = -1` (they have no residue embedding, so residue pooling masks them out) +- `is_mate` marks every non-ASU node, protein or ligand. The flow prior anchors on `~is_mate` so sampled waters start where the targets live - Edge types (defined in `src/constants.py`): - `('protein', 'pp', 'protein')`: protein-protein edges - `('protein', 'pw', 'water')`: protein to water @@ -111,23 +113,29 @@ Preprocessed data is cached under `--processed_dir` in a three-layer architectur / ├── geometry/ # Graph structures; see cache directory naming below │ └── _final.pt -│ - protein_pos: centered protein coordinates (N, 3) +│ - protein_pos: centered node coordinates (N, 3) │ - protein_x: element one-hot encoding (N, 16) │ - protein_res_idx: residue indices for grouping -│ - is_ligand: bool mask marking the appended ASU ligand atoms (N,) +│ - is_ligand: bool mask marking the ligand atoms (N,) +│ - is_mate: bool mask marking the symmetry-mate atoms (N,) +│ - emb_res_idx: embedding row per atom; -1 means no row (N,) │ - water_pos, water_x: water coordinates and features -│ - num_asu_protein: ASU atom count (mate boundary metadata) -│ # Note: When include_mates=True, mate atoms are concatenated into -│ # protein_pos/protein_x, and ASU ligand atoms are appended after those. -│ # Node order is [ASU protein | mates | ASU ligands]. Recover blocks via: -│ # ASU protein atoms = protein_pos[:num_asu_protein] -│ # ASU ligand atoms = protein_pos[is_ligand] # always last -│ # Mate atoms = protein_pos[num_asu_protein:][~is_ligand[num_asu_protein:]] +│ - num_asu_protein: ASU protein atom count (mate boundary metadata) +│ # The protein_* names predate mates and ligands: N is the total node +│ # count and these arrays hold every node, not just protein atoms (same +│ # for the data["protein"] node type). Select blocks with the masks. │ # -│ # is_ligand marks ASU ligands ONLY -- it is not an exhaustive ligand -│ # selector. The mate block is unfiltered (see TODO(mates) in -│ # _preprocess_one), so mate atoms may include HETATM/ligand/water atoms -│ # that are NOT marked by is_ligand. +│ # Node order is [ASU protein | mate protein | ASU ligand | mate ligand], +│ # so the two masks recover every block: +│ # ASU protein = ~is_mate & ~is_ligand (== the first num_asu_protein) +│ # mate protein = is_mate & ~is_ligand +│ # ASU ligand = ~is_mate & is_ligand +│ # mate ligand = is_mate & is_ligand +│ # +│ # emb_res_idx indexes the ESM table: mate atoms carry the row of the ASU +│ # residue they are a symmetry image of, and every ligand carries -1, +│ # which reads as a zero row. +├── /_filter_meta.json # settings this directory was built with ├── esm/ # ESM embeddings (per-residue) │ └── _final.pt │ - residue_embeddings: ESM3 embeddings (N_res, embed_dim) @@ -153,13 +161,29 @@ configs that produce different graphs never share a directory: The base name comes from `--geometry_cache_name` (default `geometry`). +**Filter Provenance:** + +Filtering happens *before* the cache is written, so the thresholds are a property of the +directory, not of the run reading it — and the `.pt` files record none of them. Each geometry +directory therefore carries a `_filter_meta.json` sidecar holding the per-water filters and +their toggles, the structure-level checks that decide which entries exist at all +(`min_water_residue_ratio`, `max_com_dist`, `max_clash_fraction`, `clash_dist`, +`interface_dist_threshold`), and the graph parameters behind the cached PP edges (`cutoff`, +`max_neighbors`). + +The first run with `preprocess=True` writes it; every later run compares against it and +**refuses to start** on a mismatch rather than appending differently filtered entries to the +same directory. A disabled filter records `null` for its threshold, which cannot have changed +the cached waters. Directories built before this existed have no sidecar: they load, and warn +that their provenance is unverifiable, until a preprocessing run stamps them — so check your +thresholds match the cache before that first run. + **Cache Generation Notes:** - Geometry cache is generated automatically when `preprocess=True` (default) - ESM/SLAE caches require running the respective `generate_*_embeddings.py` scripts first - Preprocessing failures are logged to `/preprocessing_failures.log` -- Geometry caches built before ligand support lack the `is_ligand` field and will fail to - load with a `KeyError`. Delete the geometry cache directory and let it regenerate — the - cached graphs are stale, not merely missing a field +- A cache file missing any field the loader reads (`is_ligand`, `is_mate`, `emb_res_idx`, …) + raises `KeyError`. Delete the geometry cache directory and let it regenerate ## Environment Setup @@ -303,7 +327,7 @@ These checks determine whether a structure is included in training: | `--max_com_dist` | `25.0` | Max protein-water center-of-mass distance (A) | | `--max_clash_fraction` | `0.05` | Max fraction of waters clashing with protein | | `--clash_dist` | `2.0` | Distance threshold for clash detection (A) | -| `--min_water_residue_ratio` | `0.6` | Minimum waters per residue ratio | +| `--min_water_residue_ratio` | `0.1` | Minimum waters per residue ratio | ### Per-Water Quality Filters @@ -313,7 +337,7 @@ These filters remove individual low-quality waters (can be toggled): |-----------|---------|-------------|-------------| | `--max_protein_dist` | `5.0` | `--no_filter_by_distance` | Remove waters far from protein | | `--min_edia` | `0.4` | `--no_filter_by_edia` | Remove waters with low EDIA scores | -| `--max_bfactor_zscore` | `1.5` | `--no_filter_by_bfactor` | Remove waters with high B-factor | +| `--max_bfactor_zscore` | `2.0` | `--no_filter_by_bfactor` | Remove waters with high B-factor |
About EDIA Scores @@ -356,6 +380,13 @@ uv run python -m scripts.inference \ | `--water_ratio` | `None` | Sample `num_residues * ratio` waters (if not set, uses ground truth count) | | `--use_sc` | `false` | Use self-conditioning during integration | +> **`--water_ratio` counts mate residues too.** `num_residues` covers ASU *and* mate +> residues, so `--include_mates` emits ~1.7x more waters at the same ratio (~440 vs +> ~263 particles at ratio 1, against ~238 true waters). Two runs share a sampling +> budget only if their mate settings match; compare density-sensitive metrics at +> parity, not at equal ratio. `--include_mates` is inherited from the training config +> when the flag is absent. + ### Output Structure ``` diff --git a/scripts/generate_slae_embeddings.py b/scripts/generate_slae_embeddings.py index 2da2926..8e42f61 100644 --- a/scripts/generate_slae_embeddings.py +++ b/scripts/generate_slae_embeddings.py @@ -1,9 +1,8 @@ """ Precompute SLAE embeddings for protein structures and save to separate cache files. -NOTE: This SLAE encoder is legacy and is NOT currently used. We primarily use the -ESM encoder (see scripts/generate_esm_embeddings.py). This script is retained for -reference/reproducibility only. +NOTE: The SLAE encoder is NOT currently used. We primarily use the ESM encoder +(see scripts/generate_esm_embeddings.py); this script is kept for reproducibility. This script: 1. Reads a split file containing PDB entries diff --git a/scripts/inference.py b/scripts/inference.py index e73b919..1de8897 100644 --- a/scripts/inference.py +++ b/scripts/inference.py @@ -168,8 +168,10 @@ def parse_args(): "--water_ratio", type=float, default=None, - help="Sample num_residues * water_ratio waters instead of using ground truth count. " - "E.g., --water_ratio 0.5 samples 50 waters for a 100-residue protein.", + help="Sample num_residues * water_ratio waters instead of using ground truth " + "count. num_residues counts ASU and symmetry-mate residues, so with " + "--include_mates the same ratio yields ~1.7x more waters than without: two " + "runs share a sampling budget only if their mate settings match.", ) p.add_argument( @@ -215,10 +217,10 @@ def _extract_dataset_filter_config(config: dict) -> dict: "max_clash_fraction": config.get("max_clash_fraction", 0.05), "clash_dist": config.get("clash_dist", 2.0), "interface_dist_threshold": config.get("interface_dist_threshold", 4.0), - "min_water_residue_ratio": config.get("min_water_residue_ratio", 0.6), + "min_water_residue_ratio": config.get("min_water_residue_ratio", 0.1), "max_protein_dist": config.get("max_protein_dist", 5.0), "min_edia": config.get("min_edia", 0.4), - "max_bfactor_zscore": config.get("max_bfactor_zscore", 1.5), + "max_bfactor_zscore": config.get("max_bfactor_zscore", 2.0), "filter_by_distance": config.get("filter_by_distance", True), "filter_by_edia": config.get("filter_by_edia", True), "filter_by_bfactor": config.get("filter_by_bfactor", True), diff --git a/scripts/train.py b/scripts/train.py index f5bc0be..7cbad8f 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -149,8 +149,11 @@ def parse_args(): p.add_argument( "--min_water_residue_ratio", type=float, - default=0.6, - help="Quality: minimum waters/residue ratio required per structure.", + default=0.1, + help=( + "Quality: minimum waters/residue ratio required per structure. Applied " + "at cache-write time, so it decides which structures the cache holds." + ), ) # per-water filtering (toggleable) @@ -169,8 +172,12 @@ def parse_args(): p.add_argument( "--max_bfactor_zscore", type=float, - default=1.5, - help="Water filter: remove waters with normalized B-factor above this threshold.", + default=2.0, + help=( + "Water filter: remove waters with normalized B-factor above this " + "threshold. Baked in at cache-write time, so a warm cache built at a " + "different value is refused rather than extended." + ), ) p.add_argument( "--no_filter_by_distance", diff --git a/src/constants.py b/src/constants.py index 5e2604f..9a06350 100644 --- a/src/constants.py +++ b/src/constants.py @@ -7,9 +7,6 @@ NODE_FEATURE_DIM = 16 # Default node scalar feature dimension # Native widths of the cached embeddings produced by scripts/generate_*_embeddings.py. -# These are fixed by the upstream models, not tunable: ESM3-open emits 1536-wide -# per-residue vectors, SLAE emits 128-wide per-atom vectors. Cached encoders take the -# width as a required config key (embedding_dim); these are the values to pass. ESM_EMBEDDING_DIM = 1536 SLAE_EMBEDDING_DIM = 128 diff --git a/src/dataset.py b/src/dataset.py index 4a6ba0b..2aef7cf 100644 --- a/src/dataset.py +++ b/src/dataset.py @@ -13,6 +13,8 @@ import itertools import json +import os +import re from collections import OrderedDict from pathlib import Path @@ -24,6 +26,7 @@ from biotite.structure.io.pdb import get_structure, PDBFile from biotite.structure.io.pdbx import CIFFile, get_structure as get_structure_cif from loguru import logger +from scipy.spatial import cKDTree from scipy.spatial.distance import cdist from torch import Tensor from torch.utils.data import DataLoader, Dataset @@ -44,6 +47,10 @@ ) +# Per-directory record of the settings a geometry cache was built with. +FILTER_META_FILENAME = "_filter_meta.json" + + def element_onehot(symbols: list[str]) -> Tensor: """One-hot encoding with 'other' bucket at end.""" other_idx = len(ELEMENT_VOCAB) @@ -112,24 +119,38 @@ def parse_asu_with_biotite( def get_crystal_contacts_pymol( - struc_path: str, cutoff: float = 5.0 + struc_path: str, + cutoff: float = 5.0, + include_ligands: bool = False, ) -> dict[str, np.ndarray | list]: """ - Extract ASU and symmetry mate atoms within crystal contact distance. + Extract ASU and symmetry-mate atoms within crystal contact distance. - Uses PyMOL's symexp command to generate symmetry mates and selects - interface atoms within the specified cutoff distance. + PyMOL's symexp generates the mates; `byres` then keeps whole residues and + whole ligand entities with any atom within `cutoff` of the ASU. Protein and + ligand mates come back under separate keys, classified by PyMOL itself + (`polymer.protein` vs the non-protein, non-solvent remainder), so nothing + downstream needs residue-name heuristics. + + Mate waters are never selected: a mate water is a symmetry image of an ASU + water, which is a prediction target, so keeping it as context is a label + leak. Protein and ligand contact surfaces are genuine context. Args: - struc_path: Path to structure file (PDB/CIF) with crystal symmetry information - cutoff: Distance cutoff in Angstroms for interface detection + struc_path: Structure file (PDB/CIF) carrying crystal symmetry. + cutoff: Interface distance cutoff in Angstroms. + include_ligands: Also collect whole ligand/ion/cofactor/nucleic-acid mate + entities (still never waters). Off by default: protein mates alone. Returns: Dict with keys: - 'asu_coords': (N_asu, 3) ASU atom coordinates - - 'mate_coords': (N_mate, 3) symmetry mate atom coordinates - - 'asu_atoms': List of PyMOL atom objects for ASU - - 'mate_atoms': List of PyMOL atom objects for mates + - 'asu_atoms': List of PyMOL atom objects for the ASU + - 'mate_coords': (N_mate, 3) whole protein-mate residues + - 'mate_atoms': List of PyMOL atom objects for protein mates + - 'mate_ligand_coords': (M, 3) whole ligand-mate entities, empty + unless include_ligands + - 'mate_ligand_atoms': List of PyMOL atom objects for ligand mates """ with pymol2.PyMOL() as pm: cmd = pm.cmd @@ -138,24 +159,43 @@ def get_crystal_contacts_pymol( obj = "struct" cmd.load(struc_path, obj) cmd.symexp("sym", obj, obj, cutoff) - cmd.select("interface", f"byres (sym* within {cutoff} of {obj})") - - asu_coords = cmd.get_coords(obj, state=1) - mate_coords = cmd.get_coords("sym* and interface", state=1) - asu_atoms = cmd.get_model(obj, state=1).atom - mate_atoms = cmd.get_model("sym* and interface", state=1).atom - asu_coords = ( - asu_coords if asu_coords is not None else np.zeros((0, 3), dtype=float) - ) - mate_coords = ( - mate_coords if mate_coords is not None else np.zeros((0, 3), dtype=float) + def _coords(selection: str) -> np.ndarray: + coords = cmd.get_coords(selection, state=1) + return coords if coords is not None else np.zeros((0, 3), dtype=float) + + # Whole protein-mate residues with any atom within cutoff of the ASU. + # `not hydro` last: byres would otherwise re-add the residue's hydrogens, + # and mates must stay heavy-atom-only like the ASU. + cmd.select( + "iface_prot", + f"(byres ((sym* and polymer.protein) within {cutoff} of {obj})) " + f"and not hydro", ) + mate_coords = _coords("iface_prot") + mate_atoms = cmd.get_model("iface_prot", state=1).atom + + # Whole ligand-mate entities (non-protein, non-water het atoms: ligands, + # ions, cofactors, nucleic acids). + if include_ligands: + cmd.select( + "iface_lig", + f"(byres ((sym* and (not polymer.protein) and (not solvent)) " + f"within {cutoff} of {obj})) and not hydro", + ) + mate_ligand_coords = _coords("iface_lig") + mate_ligand_atoms = cmd.get_model("iface_lig", state=1).atom + else: + mate_ligand_coords = np.zeros((0, 3), dtype=float) + mate_ligand_atoms = [] + return { - "asu_coords": asu_coords, + "asu_coords": _coords(obj), + "asu_atoms": cmd.get_model(obj, state=1).atom, "mate_coords": mate_coords, - "asu_atoms": asu_atoms, "mate_atoms": mate_atoms, + "mate_ligand_coords": mate_ligand_coords, + "mate_ligand_atoms": mate_ligand_atoms, } @@ -182,8 +222,6 @@ def match_atoms_to_coords( if target_coords.shape[0] == 0 or len(atoms) == 0: return [] - 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)) @@ -199,6 +237,118 @@ def match_atoms_to_coords( return matched +def dedup_mate_atoms( + mate_coords: np.ndarray, + mate_atoms: list, + reference_coords: np.ndarray, + tol: float = 0.3, +) -> tuple[np.ndarray, list]: + """ + Drop mate atoms coincident with a reference atom or an already-kept mate atom. + + Crystal symmetry creates coincidences: an atom on a rotation or screw axis + maps onto itself, and one residue can be reached through two operators. Left + alone each becomes an independent node, giving duplicates joined by ~0 A + edges -- and, for a target water on a special position, a label leak. + + Per atom, unlike `dedup_mate_ligands_by_residue`: a special position is an + atom-level accident, so only the coincident atom is dropped. The self sweep + also catches mates coincident with each other, which the reference tree cannot. + + Args: + mate_coords: (N, 3) mate atom coordinates, uncentered. + mate_atoms: Parallel list of mate atom objects, kept in lockstep. + reference_coords: (M, 3) uncentered ASU coordinates. + tol: Coincidence radius in Angstroms. + + Returns: + (kept_coords, kept_atoms). The first atom of a coincident group is the + one kept, so the result depends on input order. + """ + n = mate_coords.shape[0] + if n == 0: + return mate_coords, mate_atoms + + # An empty reference tree answers inf, so no guard is needed here. + drop = cKDTree(reference_coords).query(mate_coords, k=1)[0] < tol + + # Self-dedup is a first-win sweep. One tree answers every lookup, so the + # sweep only walks each atom's coincident neighbors. + neighbors = cKDTree(mate_coords).query_ball_point(mate_coords, r=tol) + kept = np.zeros(n, dtype=bool) + for i in range(n): + if drop[i]: + continue + earlier = [j for j in neighbors[i] if j < i and kept[j]] + # query_ball_point includes r, this sweep is strict. + dists = np.linalg.norm(mate_coords[earlier] - mate_coords[i], axis=1) + kept[i] = not (dists < tol).any() + + keep_idx = np.flatnonzero(kept) + return mate_coords[keep_idx], [mate_atoms[i] for i in keep_idx] + + +def dedup_mate_ligands_by_residue( + lig_coords: np.ndarray, + lig_atoms: list, + reference_coords: np.ndarray, + tol: float = 0.3, + image_frac: float = 0.5, +) -> tuple[np.ndarray, list]: + """ + Drop whole mate-ligand entities that are symmetry images of ASU atoms. + + Unlike `dedup_mate_atoms`, which works per atom, this works per entity so a + ligand is never fragmented: it goes only when the whole ligand is a redundant + copy. Genuine neighbor-cell ligands are kept whole. + + Args: + lig_coords: (M, 3) mate-ligand atom coordinates, uncentered. + lig_atoms: Parallel list of mate-ligand atom objects. + reference_coords: Uncentered ASU coordinates. + tol: Coincidence radius in Angstroms. + image_frac: Drop a ligand when more than this fraction of its atoms are + coincident with the reference. + + Returns: + (kept_coords, kept_atoms) with whole symmetry-image ligands removed. + """ + if len(lig_atoms) == 0: + return lig_coords, lig_atoms + + ref_tree = cKDTree(reference_coords) + # Group atom indices by ligand entity (chain, residue id, segment). + groups = {} + for i, atom in enumerate(lig_atoms): + key = (atom.chain, atom.resi, getattr(atom, "segi", "")) + groups.setdefault(key, []).append(i) + + keep_idx: list[int] = [] + for idxs in groups.values(): + if np.mean(ref_tree.query(lig_coords[idxs], k=1)[0] < tol) <= image_frac: + keep_idx.extend(idxs) # genuine neighbor ligand: keep whole + keep_idx.sort() + + return lig_coords[keep_idx], [lig_atoms[i] for i in keep_idx] + + +def _parse_pdb_resi(resi) -> tuple[int, str] | None: + """ + Parse a PyMOL residue identifier, which may carry an insertion code. + + Args: + resi: Residue id as PyMOL exposes it, e.g. "52", "-3", "52A". + + Returns: + (res_id, ins_code), or None when there is no integer part, which the + caller scores with a zero embedding rather than crashing on. + """ + match = re.match(r"^\s*(-?\d+)\s*([A-Za-z]?)\s*$", str(resi)) + if match is None: + return None + return int(match.group(1)), match.group(2).strip() + + def _make_undirected(edge_index: torch.Tensor) -> torch.Tensor: """ Convert directed edges to undirected by adding reverse edges. @@ -741,10 +891,10 @@ def __init__( max_clash_fraction: float = 0.05, clash_dist: float = 2.0, interface_dist_threshold: float = 4.0, - min_water_residue_ratio: float = 0.6, + min_water_residue_ratio: float = 0.1, max_protein_dist: float = 5.0, min_edia: float = 0.4, - max_bfactor_zscore: float = 1.5, + max_bfactor_zscore: float = 2.0, filter_by_distance: bool = True, filter_by_edia: bool = True, filter_by_bfactor: bool = True, @@ -850,6 +1000,8 @@ def __init__( self.entries = self._parse_pdb_list(pdb_list_file) + self._sync_filter_meta(write=preprocess) + if preprocess: self._preprocess_all() @@ -914,6 +1066,79 @@ def _parse_pdb_list(self, pdb_list_file: str) -> list[dict]: logger.info(f"Loaded {len(entries)} entries from {pdb_list_file}") return entries + def _sync_filter_meta(self, write: bool) -> None: + """ + Refuse to read or extend a cache built under different settings. + + Filtering happens before the cache is written, so these are properties of + the directory rather than of the run reading it -- and the .pt files + record none of them. Writing entries under different settings would leave + one directory holding two populations no later reader can tell apart. + + Args: + write: Create when it is absent. Only runs that may add + entries (preprocess=True) claim a directory this way. + + Raises: + ValueError: If the recorded settings differ from this run's. + """ + meta_path = self.geometry_dir / FILTER_META_FILENAME + # A disabled filter's threshold is None: it never touched the cached + # waters, so it must not make two identical caches look incompatible. + current = { + "filter_by_distance": self.filter_by_distance, + "filter_by_edia": self.filter_by_edia, + "filter_by_bfactor": self.filter_by_bfactor, + "max_protein_dist": self.max_protein_dist + if self.filter_by_distance + else None, + "min_edia": self.min_edia if self.filter_by_edia else None, + "max_bfactor_zscore": self.max_bfactor_zscore + if self.filter_by_bfactor + else None, + "min_water_residue_ratio": self.min_water_residue_ratio, + "max_com_dist": self.max_com_dist, + "max_clash_fraction": self.max_clash_fraction, + "clash_dist": self.clash_dist, + "interface_dist_threshold": self.interface_dist_threshold, + "cutoff": self.cutoff, + "max_neighbors": self.max_neighbors, + } + + if meta_path.is_file(): + with open(meta_path) as f: + recorded = json.load(f) + differing = [ + f"{name}: cache={recorded.get(name)!r} run={value!r}" + for name, value in current.items() + if recorded.get(name) != value + ] + if differing: + raise ValueError( + f"Settings disagree with {meta_path}: {', '.join(differing)}. " + "The cache was filtered at write time, so one directory cannot " + "hold both. Match the recorded values or point " + "geometry_cache_name at a different directory." + ) + return + + # Warn before writing too: stamping pre-existing entries labels them with + # settings they were never checked against, and later runs trust the label. + if any(self.geometry_dir.glob("*.pt")): + logger.warning( + f"{self.geometry_dir} has no {FILTER_META_FILENAME}; the settings " + "its entries were built with cannot be verified." + ) + + if write: + self.geometry_dir.mkdir(parents=True, exist_ok=True) + # Written through a temp file: cache builds fan out over processes, + # and a reader must never catch a half-written sidecar. + tmp_path = meta_path.with_suffix(f".{os.getpid()}.tmp") + with open(tmp_path, "w") as f: + json.dump(current, f, indent=2) + tmp_path.replace(meta_path) + def _preprocess_all(self): """ Preprocess all PDB files that don't have cached geometry results. @@ -986,20 +1211,31 @@ def _preprocess_one(self, entry: dict, cache_path: Path): if not chain_valid: raise ValueError(f"Quality filter failed: {chain_reason}") - crystal_data = get_crystal_contacts_pymol(struc_path, self.cutoff) + # PyMOL is only needed for symmetry expansion, so a no-mates cache skips + # it (and with it the water cross-check below) entirely. + if self.include_mates: + crystal_data = get_crystal_contacts_pymol( + struc_path, self.cutoff, include_ligands=self.include_ligands + ) - # 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"] - ) - if asu_water_indices: - asu_water_mask = np.zeros(len(water_atoms), dtype=bool) - asu_water_mask[asu_water_indices] = True - water_atoms = water_atoms[asu_water_mask] - else: - water_atoms = water_atoms[:0] + # 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"] + ) + if asu_water_indices: + asu_water_mask = np.zeros(len(water_atoms), dtype=bool) + asu_water_mask[asu_water_indices] = True + water_atoms = water_atoms[asu_water_mask] + else: + if len(water_atoms) > 0: + logger.warning( + f"{entry['pdb_id']}: no waters survived the biotite/PyMOL " + f"cross-check (had {len(water_atoms)})" + ) + water_atoms = water_atoms[:0] # Per-water filtering is optional; structure-level quality checks below always run. use_distance_filter = self.filter_by_distance @@ -1039,11 +1275,19 @@ def _preprocess_one(self, entry: dict, cache_path: Path): ) ) - # apply quality filters + # Apply quality filters. Mate protein atoms join the distance + # reference so a genuine crystal-contact water is not dropped as solvent-far. + if self.include_mates and crystal_data["mate_coords"].shape[0] > 0: + filter_protein_coords = np.concatenate( + [protein_atoms.coord, crystal_data["mate_coords"]], axis=0 + ) + else: + filter_protein_coords = protein_atoms.coord + keep_mask = filter_waters_by_quality( water_atoms.coord, water_keys, - protein_atoms.coord if use_distance_filter else None, + filter_protein_coords if use_distance_filter else None, edia_lookup, bfactor_lookup, max_protein_dist=self.max_protein_dist, @@ -1100,6 +1344,21 @@ def _preprocess_one(self, entry: dict, cache_path: Path): protein_res_idx = torch.from_numpy( bts.spread_residue_wise(sanitized_for_idx, np.arange(num_residues)) ).long() + + # (chain, res_id, ins_code) -> residue index, so a symmetry mate can + # inherit the embedding row of the ASU residue it is an image of. Keyed + # off the same sanitized parse that defines protein_res_idx, so the index + # lines up with the stored ESM rows. + asu_reskey_to_residx: dict[tuple[str, int, str], int] = {} + for res_i, start in enumerate(bts.get_residue_starts(sanitized_for_idx)): + # ins_code already normalized in place above + key = ( + str(sanitized_for_idx.chain_id[start]).strip(), + int(sanitized_for_idx.res_id[start]), + str(sanitized_for_idx.ins_code[start]), + ) + asu_reskey_to_residx.setdefault(key, res_i) + num_waters = len(water_atoms) ratio_valid, ratio_reason = check_water_residue_ratio( num_waters, @@ -1118,33 +1377,75 @@ def _preprocess_one(self, entry: dict, cache_path: Path): water_pos = torch.zeros((0, 3), dtype=torch.float32) water_x = torch.zeros((0, len(ELEMENT_VOCAB) + 1), dtype=torch.float32) - # process symmetry mate atoms - # - # TODO(mates): the mate atom set is unfiltered and inconsistent with the ASU - # path. get_crystal_contacts_pymol runs symexp over the whole object and - # selects "sym* and interface" with no polymer filter, so mate_atoms carries - # het atoms and waters as well as protein, and every one of them becomes a - # protein-type node below. Consequences: mate ligands are already included - # but never marked in is_ligand (unlike ASU ligands) and are not gated by - # include_ligands; mate waters -- symmetry images of the prediction target -- - # enter as protein context. Fix in dev_crystal_mates. - mate_coords = crystal_data["mate_coords"] - if mate_coords.shape[0] > 0: - mate_pos = torch.tensor(mate_coords, dtype=torch.float32) - center - mate_elements = [a.symbol.upper() for a in crystal_data["mate_atoms"]] - mate_x = element_onehot(mate_elements) - - # compute mate residue indices (group atoms by actual residue) - mate_residue_keys = [(a.chain, a.resi) for a in crystal_data["mate_atoms"]] - unique_mate_res = list(dict.fromkeys(mate_residue_keys)) # preserves order - mate_res_map = {k: i for i, k in enumerate(unique_mate_res)} - mate_res_idx = torch.tensor( - [mate_res_map[k] for k in mate_residue_keys], dtype=torch.long + # Mate blocks stay empty unless include_mates (and, for ligands, + # include_ligands) filled them in below. + mate_pos = torch.zeros((0, 3), dtype=torch.float32) + mate_x = torch.zeros((0, len(ELEMENT_VOCAB) + 1), dtype=torch.float32) + mate_res_idx = torch.empty(0, dtype=torch.long) + mate_emb_res_idx = torch.empty(0, dtype=torch.long) + mate_lig_coords = np.zeros((0, 3), dtype=float) + mate_lig_atoms: list = [] + + if self.include_mates: + # Drop mate atoms coincident with an ASU atom, a target water, or an + # already-kept mate atom: special positions and redundant symmetry + # images. Uncentered coords; mates are centered below. + ref_parts = [protein_atoms.coord] + if len(water_atoms): + ref_parts.append(water_atoms.coord) + # ASU ligands join the reference so a mate ligand that is only their + # symmetry image goes too; neighbor-cell ligands stay. + if self.include_ligands and len(ligand_atoms) > 0: + ref_parts.append(ligand_atoms.coord) + reference = np.concatenate(ref_parts, axis=0) + mate_coords, mate_atoms = dedup_mate_atoms( + crystal_data["mate_coords"], crystal_data["mate_atoms"], reference ) - else: - mate_pos = torch.zeros((0, 3), dtype=torch.float32) - mate_x = torch.zeros((0, len(ELEMENT_VOCAB) + 1), dtype=torch.float32) - mate_res_idx = torch.empty(0, dtype=torch.long) + # Ligand mates dedup at entity granularity, so a ligand is never + # fragmented: whole symmetry images go, genuine neighbors stay. + if self.include_ligands: + mate_lig_coords, mate_lig_atoms = dedup_mate_ligands_by_residue( + crystal_data["mate_ligand_coords"], + crystal_data["mate_ligand_atoms"], + reference, + ) + + if mate_coords.shape[0] > 0: + mate_pos = torch.tensor(mate_coords, dtype=torch.float32) - center + mate_x = element_onehot([a.symbol.upper() for a in mate_atoms]) + + # Group mate atoms by residue. The key omits the symmetry-object id + # (atom.model), so two images of one residue share a group. Harmless + # today; add atom.model before enabling GVPEncoder's pool_residue, + # which would otherwise merge the images into one residue. + mate_residue_keys = [(a.chain, a.resi) for a in mate_atoms] + unique_mate_res = list(dict.fromkeys(mate_residue_keys)) # keeps order + mate_res_map = {k: i for i, k in enumerate(unique_mate_res)} + mate_res_idx = torch.tensor( + [mate_res_map[k] for k in mate_residue_keys], dtype=torch.long + ) + + # A mate inherits its ASU residue's ESM row via (chain, resi); + # embeddings are coordinate-free, so image and source share it. -1 + # (no match) reads as a zero embedding: the atom keeps geometry and + # element, losing only its sequence signal, so a miss warns rather + # than raises. Misses come from PyMOL's polymer.protein admitting a + # residue biotite dropped, or an unparseable resi. + mate_emb_idx = [] + for atom in mate_atoms: + parsed = _parse_pdb_resi(atom.resi) + mate_emb_idx.append( + asu_reskey_to_residx.get((str(atom.chain).strip(), *parsed), -1) + if parsed is not None + else -1 + ) + mate_emb_res_idx = torch.tensor(mate_emb_idx, dtype=torch.long) + unmatched = int((mate_emb_res_idx < 0).sum()) + if unmatched: + logger.warning( + f"{entry['cache_key']}: {unmatched}/{len(mate_emb_idx)} mate " + "atoms unmatched to an ASU residue (zero embedding for those)" + ) # Compute final protein data based on include_mates flag num_asu_protein = protein_pos.size(0) @@ -1164,29 +1465,48 @@ def _preprocess_one(self, entry: dict, cache_path: Path): final_protein_x = protein_x final_protein_res_idx = protein_res_idx - # Append ASU ligand atoms after protein (and mate) atoms when enabled. - # is_ligand mask marks which protein-type nodes are ligand atoms. - # Ligands always go last so num_asu_protein and mate counts are unaffected, - # preserving ESM/SLAE embedding alignment via _pad_atom_embeddings_for_mates. - # Only ASU ligands are handled here -- mate het atoms come in unfiltered via - # the mate block above, see TODO(mates) there. + # Append ligand atoms last, giving the node order ASU protein -> mate + # protein -> ASU ligand -> mate ligand: num_asu_protein and the mate count + # stay meaningful, which is what keeps ESM/SLAE aligned. Ligands get + # residue_index = emb_res_idx = -1 (no residue embedding); residue pooling + # masks out those negatives before any scatter (GVPEncoder._pool_by_residue). + ligand_blocks = [] if self.include_ligands and len(ligand_atoms) > 0: - ligand_pos = torch.tensor(ligand_atoms.coord, dtype=torch.float32) - center - ligand_elements = [str(e).upper() for e in ligand_atoms.element] - ligand_x = element_onehot(ligand_elements) - final_protein_pos = torch.cat([final_protein_pos, ligand_pos], dim=0) - final_protein_x = torch.cat([final_protein_x, ligand_x], dim=0) - # Ligand atoms get residue_index = -1 (sentinel; no residue embedding). - # The is_ligand mask identifies them; residue-pooling masks out these - # negative indices before any scatter (see GVPEncoder._pool_by_residue). - ligand_res_idx = torch.full((len(ligand_atoms),), -1, dtype=torch.long) - final_protein_res_idx = torch.cat( - [final_protein_res_idx, ligand_res_idx], dim=0 + ligand_blocks.append( + ( + ligand_atoms.coord, + [str(e).upper() for e in ligand_atoms.element], + False, + ) + ) + if len(mate_lig_atoms) > 0: + ligand_blocks.append( + (mate_lig_coords, [a.symbol.upper() for a in mate_lig_atoms], True) ) - is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool) - is_ligand[-len(ligand_atoms) :] = True - else: - is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool) + + # Mate proteins inherit their source ASU residue's embedding row; ligands + # get -1 whichever cell they came from. + n_protein = final_protein_pos.size(0) + emb_res_idx = torch.cat([protein_res_idx, mate_emb_res_idx], dim=0) + is_mate = torch.zeros(n_protein, dtype=torch.bool) + is_mate[num_asu_protein:] = True + + for coords, elements, from_mate in ligand_blocks: + n_lig = len(elements) + pos = torch.tensor(coords, dtype=torch.float32) - center + final_protein_pos = torch.cat([final_protein_pos, pos], dim=0) + final_protein_x = torch.cat( + [final_protein_x, element_onehot(elements)], dim=0 + ) + sentinel = torch.full((n_lig,), -1, dtype=torch.long) + final_protein_res_idx = torch.cat([final_protein_res_idx, sentinel], dim=0) + emb_res_idx = torch.cat([emb_res_idx, sentinel], dim=0) + is_mate = torch.cat( + [is_mate, torch.full((n_lig,), from_mate, dtype=torch.bool)], dim=0 + ) + + is_ligand = torch.zeros(final_protein_pos.size(0), dtype=torch.bool) + is_ligand[n_protein:] = True # Compute PP edges and features if final_protein_pos.size(0) > 0: @@ -1219,6 +1539,8 @@ def _preprocess_one(self, entry: dict, cache_path: Path): "protein_x": final_protein_x, "protein_res_idx": final_protein_res_idx, "is_ligand": is_ligand, + "is_mate": is_mate, + "emb_res_idx": emb_res_idx, "water_pos": water_pos, "water_x": water_x, # PP topology and features (precomputed) @@ -1240,9 +1562,9 @@ def _annotate_data_with_embeddings( self, data: HeteroData, cache_key: str, - asu_protein_res_idx: torch.Tensor, num_asu_protein: int, num_protein_residues: int, + emb_res_idx: torch.Tensor, ) -> None: """ Load encoder-specific embeddings and attach to data object. @@ -1255,9 +1577,11 @@ def _annotate_data_with_embeddings( Args: data: HeteroData object to attach embeddings to (modified in-place) cache_key: Identifier for cached embedding files - asu_protein_res_idx: (N_asu,) residue index per ASU atom num_asu_protein: Number of ASU protein atoms num_protein_residues: Number of unique protein residues + emb_res_idx: (N_total,) embedding row per atom -- mates inherit their + source ASU residue's row; ligands and unmatched atoms are -1 and + get a zero row. """ if self.encoder_type == "slae": data["protein"].embedding = load_slae_embedding( @@ -1276,10 +1600,15 @@ def _annotate_data_with_embeddings( num_protein_residues=num_protein_residues, cache_load_mmap=self.cache_load_mmap, ) - esm_atom_emb = residue_embeddings[asu_protein_res_idx] - data["protein"].embedding = _pad_atom_embeddings_for_mates( - esm_atom_emb, data["protein"].num_nodes + # Per-atom inheritance: a mate atom takes the row of the ASU residue + # it images; ligands and unmatched atoms (-1) stay zero. + atom_emb = residue_embeddings.new_zeros( + data["protein"].num_nodes, residue_embeddings.size(1) ) + valid = emb_res_idx >= 0 + if valid.any(): + atom_emb[valid] = residue_embeddings[emb_res_idx[valid]] + data["protein"].embedding = atom_emb data["protein"].embedding_type = "esm" def __getitem__(self, idx: int) -> HeteroData: @@ -1323,6 +1652,8 @@ def __getitem__(self, idx: int) -> HeteroData: protein_x = cached["protein_x"] protein_res_idx = cached["protein_res_idx"] is_ligand = cached["is_ligand"] + is_mate = cached["is_mate"] + emb_res_idx = cached["emb_res_idx"] pp_edge_index = cached["pp_edge_index"] pp_edge_unit_vectors = cached["pp_edge_unit_vectors"] pp_edge_rbf = cached["pp_edge_rbf"] @@ -1331,9 +1662,6 @@ def __getitem__(self, idx: int) -> HeteroData: water_pos = cached["water_pos"] water_x = cached["water_x"] - # extract ASU protein residue indices for embedding loading - asu_protein_res_idx = protein_res_idx[:num_asu_protein] - data = HeteroData() # compute total num_residues (protein + mates) @@ -1345,6 +1673,7 @@ def __getitem__(self, idx: int) -> HeteroData: data["protein"].pos = protein_pos data["protein"].residue_index = protein_res_idx data["protein"].is_ligand = is_ligand + data["protein"].is_mate = is_mate data["protein"].num_nodes = protein_pos.size(0) data["protein"].num_residues = num_residues data["protein"].num_protein_residues = num_protein_residues @@ -1352,9 +1681,9 @@ def __getitem__(self, idx: int) -> HeteroData: self._annotate_data_with_embeddings( data=data, cache_key=entry["embedding_key"], # use base key for embeddings - asu_protein_res_idx=asu_protein_res_idx, num_asu_protein=num_asu_protein, num_protein_residues=num_protein_residues, + emb_res_idx=emb_res_idx, ) data["water"].x = water_x diff --git a/src/flow.py b/src/flow.py index 6a4f151..9a5f89f 100644 --- a/src/flow.py +++ b/src/flow.py @@ -14,6 +14,7 @@ import numpy as np import torch import torch.nn.functional as F +from loguru import logger from torch import nn, Tensor from torch_geometric.data import Batch, HeteroData from torch_geometric.nn import knn @@ -57,6 +58,7 @@ def sample_waters_uniform_ball( batch_w: Tensor, cutoff: float = 8.0, device: torch.device | None = None, + anchor_mask: Tensor | None = None, ) -> Tensor: """ Sample water positions uniformly inside balls of radius *cutoff* centred @@ -73,6 +75,11 @@ def sample_waters_uniform_ball( batch vector and get samples aligned to it. cutoff: Ball radius in Angstroms device: Optional output device (defaults to protein_pos.device) + anchor_mask: Optional (N_protein,) bool selecting eligible anchors. Used to + anchor on ASU atoms only, so the prior spawns where the targets live + instead of dispersing onto symmetry mates that OT must then transport + back. Every structure keeps >=1 ASU atom, so the mask never starves a + graph; if it somehow does, the batch anchors on all atoms and warns. Returns: water_pos: (N_water, 3) sampled positions, one per entry of batch_w @@ -97,6 +104,23 @@ def sample_waters_uniform_ball( if batch_p.numel() > 0: num_graphs = max(num_graphs, int(batch_p.max().item()) + 1) + protein_pos = protein_pos.to(device) + + # Drop to the eligible anchors (ASU-only). Skip the mask, rather than starve a + # graph, if it would leave a water-requesting graph with no anchor -- an + # invariant violation the dataset should never produce, so warn if it happens. + if anchor_mask is not None: + eligible = anchor_mask.to(device).bool() + counts = torch.bincount(batch_p[eligible], minlength=num_graphs) + if (counts[batch_w] == 0).any(): + logger.warning( + "sample_waters_uniform_ball: anchor mask leaves a water-requesting " + "graph with no anchor; anchoring the batch on all protein atoms." + ) + else: + protein_pos = protein_pos[eligible] + batch_p = batch_p[eligible] + # per-graph protein atom counts and cumulative offsets num_p_per_graph = scatter( torch.ones(batch_p.size(0), device=device, dtype=torch.long), @@ -123,7 +147,7 @@ def sample_waters_uniform_ball( # pick a random protein atom per water (uniform with replacement) graph_offsets = offsets[batch_w] local_idx = (torch.rand(total_waters, device=device) * graph_sizes.float()).long() - anchors = protein_pos.to(device)[graph_offsets + local_idx] + anchors = protein_pos[graph_offsets + local_idx] # uniform direction on the unit sphere direction = torch.randn(total_waters, 3, device=device, dtype=protein_pos.dtype) @@ -405,15 +429,14 @@ def forward( if EDGE_PP in data.edge_types: pp_edge = data[EDGE_PP] - # V_edge fallback is for backward compatibility with datasets - # that don't have cached edge features. A given model only sees one or the other. + # A given model sees one source or the other, never both. if pp_edge_attr is not None: # Use encoder-learned scalar features (s_edge) with unit vectors s_edge, V_edge = pp_edge_attr if hasattr(pp_edge, "edge_unit_vectors"): cached_edge_attr_dict[EDGE_PP] = (s_edge, pp_edge.edge_unit_vectors) else: - # Fallback for datasets without cached unit vectors + # Graphs built outside the dataset carry vectors on the encoder side cached_edge_attr_dict[EDGE_PP] = (s_edge, V_edge.squeeze(1)) elif hasattr(pp_edge, "edge_rbf") and hasattr(pp_edge, "edge_unit_vectors"): # No encoder edge features (e.g., SLAE/ESM) - use cached geometric features @@ -707,6 +730,20 @@ def _num_graphs(data: HeteroData | Batch) -> int: return 0 return int(batch_p.max().item()) + 1 + @staticmethod + def _asu_mask(data: HeteroData | Batch) -> Tensor | None: + """ + Mask over protein nodes selecting ASU (non-mate) atoms, or None when the + batch carries no mate annotation. + + No ``.any()`` short-circuit: it would force a host sync every step, and both + consumers treat an all-True mask as a no-op. + """ + is_mate = getattr(data["protein"], "is_mate", None) + if is_mate is None: + return None + return ~is_mate.bool() + def _sample_waters( self, batch_data: HeteroData | Batch, @@ -715,6 +752,10 @@ def _sample_waters( ) -> Tensor: """Dispatch to the configured sampling strategy, sampling one water per entry of batch_w and returning them in that order.""" + # Targets are ASU-only, so mate atoms disperse the prior (uniform_ball) and + # inflate its scale (scaled_gaussian). No mates -> no mask -> unchanged. + asu_mask = self._asu_mask(batch_data) + if self.sampling_strategy == "uniform_ball": return sample_waters_uniform_ball( protein_pos=batch_data["protein"].pos, @@ -722,9 +763,12 @@ def _sample_waters( batch_w=batch_w, cutoff=self.graph_cutoff, device=device, + anchor_mask=asu_mask, ) # scaled_gaussian - sigma_per_graph = self.compute_sigma_per_graph(batch_data, device) + sigma_per_graph = self.compute_sigma_per_graph( + batch_data, device, node_mask=asu_mask + ) return sample_waters_scaled_gaussian( batch_w=batch_w, sigma_per_graph=sigma_per_graph, @@ -750,17 +794,31 @@ def compute_sigma(data: HeteroData) -> float: Returns: Scalar sigma value (standard deviation across all protein coordinates) + + Note: + Diagnostic only, no production caller, and does not exclude mates. + Training and inference use compute_sigma_per_graph, which does. """ pos = data["protein"].pos return float(pos.std().item()) @staticmethod def compute_sigma_per_graph( - data: HeteroData | Batch, device: torch.device + data: HeteroData | Batch, + device: torch.device, + node_mask: Tensor | None = None, ) -> torch.Tensor: """ Compute sigma (std of protein coordinates) per graph in a batch. + Args: + data: Batch carrying protein positions and a protein batch vector + device: Unused for the computation; kept for call-site symmetry + node_mask: Optional (N_protein,) bool selecting the atoms that define + the scale. Mates sit far from the ASU, so counting them inflates + sigma and pushes the prior out past the targets. Skipped for the + whole batch (with a warning) if it would empty any graph. + Returns: sigma: (num_graphs,) tensor of sigma values per graph """ @@ -769,7 +827,8 @@ def compute_sigma_per_graph( num_graphs = FlowMatcher._num_graphs(data) # an empty graph would otherwise shorten the output or yield a degenerate - # sigma that silently places its waters at the origin + # sigma that silently places its waters at the origin. Checked against the + # unmasked atoms so the error keeps its original meaning. empty = torch.bincount(batch_p, minlength=num_graphs) == 0 if empty.any(): raise ValueError( @@ -777,6 +836,21 @@ def compute_sigma_per_graph( "they have zero protein atoms." ) + # Restrict to the masked atoms (ASU-only). Skip the mask, rather than yield + # a degenerate sigma, if it would leave a graph with no atoms -- which the + # dataset should never produce, so warn if it happens. + if node_mask is not None: + eligible = node_mask.to(pos.device).bool() + counts = torch.bincount(batch_p[eligible], minlength=num_graphs) + if (counts == 0).any(): + logger.warning( + "compute_sigma_per_graph: node mask leaves a graph with no " + "atoms; computing sigma over all protein atoms." + ) + else: + pos = pos[eligible] + batch_p = batch_p[eligible] + # Var(X) = E[X^2] - E[X]^2 mean_pos = scatter_mean(pos, batch_p, dim=0, dim_size=num_graphs) mean_sq = scatter_mean(pos**2, batch_p, dim=0, dim_size=num_graphs) @@ -828,7 +902,10 @@ def training_step( batch_w = batch["water"].batch num_graphs = self._num_graphs(batch) - sigma_per_graph = self.compute_sigma_per_graph(batch, device) + # same restriction the sampler uses, so the logged sigma matches it + sigma_per_graph = self.compute_sigma_per_graph( + batch, device, node_mask=self._asu_mask(batch) + ) # sampling against the batch's own water order keeps x0 aligned with x1, so # ot_coupling's per-graph mask selects the same nodes from both x0 = self._sample_waters(batch, batch_w, device) diff --git a/tests/conftest.py b/tests/conftest.py index 7fa5117..9be3d92 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -83,7 +83,8 @@ def pdb_1deu(): @pytest.fixture def pdb_4h0b(): - """4h0b - has non-water ligand HETATMs for ligand support tests.""" + """4h0b - has non-water ligand HETATMs for ligand support tests. P6 space group, + so a water on the 6-fold axis has a symmetry copy ~0A away (special position).""" return _resolve_test_path("4h0b", ".pdb") diff --git a/tests/test_dataset.py b/tests/test_dataset.py index b58f2f2..5e56eef 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -36,13 +36,17 @@ from src.dataset import ( _make_undirected, _pad_atom_embeddings_for_mates, + _parse_pdb_resi, apply_threshold_filter, check_chain_interactions, check_com_distance, check_water_clashes, check_water_residue_ratio, compute_normalized_bfactors, + dedup_mate_atoms, + dedup_mate_ligands_by_residue, element_onehot, + FILTER_META_FILENAME, filter_waters_by_quality, get_crystal_contacts_pymol, get_dataloader, @@ -305,6 +309,258 @@ def test_warns_on_odd_counts_below_half(self, warning_log, n_atoms, n_matched): assert f"{n_matched}/{n_atoms} atoms matched" in warning_log[0] +@pytest.mark.unit +class TestDedupMateAtoms: + """Tests for symmetry-mate coordinate deduplication.""" + + @staticmethod + def _atoms(n): + return [object() for _ in range(n)] + + def test_empty_passthrough(self): + coords = np.zeros((0, 3)) + out_coords, out_atoms = dedup_mate_atoms(coords, [], np.zeros((0, 3))) + + assert out_coords.shape == (0, 3) + assert out_atoms == [] + + def test_drops_atoms_coincident_with_reference(self): + """A mate atom sitting on an ASU/target atom is a leak and is removed.""" + mate_coords = np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]]) + reference = np.array([[0.0, 0.0, 0.0]]) + + out_coords, out_atoms = dedup_mate_atoms( + mate_coords, self._atoms(2), reference, tol=0.3 + ) + + assert out_coords.shape[0] == 1 + assert len(out_atoms) == 1 + np.testing.assert_allclose(out_coords[0], [10.0, 0.0, 0.0]) + + def test_keeps_atoms_aligned(self): + """Returned coords and atom objects stay in lockstep.""" + mate_coords = np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]]) + first, second = object(), object() + reference = np.array([[0.0, 0.0, 0.0]]) + + out_coords, out_atoms = dedup_mate_atoms( + mate_coords, [first, second], reference, tol=0.3 + ) + + assert out_atoms == [second] + np.testing.assert_allclose(out_coords[0], [10.0, 0.0, 0.0]) + + def test_keeps_atoms_beyond_tolerance(self): + """Separations at or past tol are distinct atoms, not duplicates.""" + mate_coords = np.array([[0.0, 0.0, 0.0], [0.3, 0.0, 0.0], [0.6, 0.0, 0.0]]) + + out_coords, _ = dedup_mate_atoms( + mate_coords, self._atoms(3), np.zeros((0, 3)), tol=0.3 + ) + + assert out_coords.shape[0] == 3 + + def test_self_dedup_is_first_wins(self): + """A chain of near-coincident mate atoms collapses onto the earliest.""" + mate_coords = np.array([[0.0, 0.0, 0.0], [0.1, 0.0, 0.0], [0.2, 0.0, 0.0]]) + first, second, third = object(), object(), object() + + out_coords, out_atoms = dedup_mate_atoms( + mate_coords, [first, second, third], np.zeros((0, 3)), tol=0.3 + ) + + assert out_atoms == [first] + np.testing.assert_allclose(out_coords[0], [0.0, 0.0, 0.0]) + + @pytest.mark.integration + def test_real_structure_drops_special_position_atoms(self, pdb_8dzt): + """8dzt is P 61: its screw axis maps real atoms onto the ASU, so the dedup + genuinely fires here rather than on hand-placed coordinates.""" + from scipy.spatial import cKDTree + + protein_atoms, water_atoms, ligand_atoms = parse_asu_with_biotite(pdb_8dzt) + crystal = get_crystal_contacts_pymol( + str(pdb_8dzt), cutoff=8.0, include_ligands=True + ) + reference = np.concatenate( + [protein_atoms.coord, water_atoms.coord, ligand_atoms.coord], axis=0 + ) + + kept_coords, kept_atoms = dedup_mate_atoms( + crystal["mate_coords"], crystal["mate_atoms"], reference, tol=0.3 + ) + + n_in = crystal["mate_coords"].shape[0] + assert 0 < kept_coords.shape[0] < n_in, ( + "expected a real structure to have coincident mate atoms to drop" + ) + assert len(kept_atoms) == kept_coords.shape[0] + # nothing coincident with the ASU survives: that is the label-leak guard + assert (cKDTree(reference).query(kept_coords, k=1)[0] >= 0.3).all() + # ...and no two survivors are coincident with each other either + assert (cKDTree(kept_coords).query(kept_coords, k=2)[0][:, 1] >= 0.3).all() + + +class _FakeLigandAtom: + """Stand-in for a PyMOL atom object with the fields the dedup reads.""" + + def __init__(self, chain, resi, segi=""): + self.chain = chain + self.resi = resi + self.segi = segi + + +@pytest.mark.unit +class TestDedupMateLigandsByResidue: + """Tests for whole-entity symmetry-image ligand removal.""" + + def test_empty_passthrough(self): + coords = np.zeros((0, 3)) + out_coords, out_atoms = dedup_mate_ligands_by_residue( + coords, [], np.zeros((0, 3)) + ) + + assert out_coords.shape == (0, 3) + assert out_atoms == [] + + def test_drops_whole_symmetry_image_ligand(self): + """A ligand whose atoms mostly land on ASU atoms is dropped entirely.""" + lig_coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]) + lig_atoms = [_FakeLigandAtom("A", "1") for _ in range(3)] + reference = lig_coords.copy() + + out_coords, out_atoms = dedup_mate_ligands_by_residue( + lig_coords, lig_atoms, reference, tol=0.3 + ) + + assert out_coords.shape[0] == 0 + assert out_atoms == [] + + def test_keeps_neighbour_ligand_whole(self): + """A genuine neighbour-cell ligand keeps every atom, including any that + happen to coincide with the ASU.""" + lig_coords = np.array([[0.0, 0.0, 0.0], [9.0, 0.0, 0.0], [10.0, 0.0, 0.0]]) + lig_atoms = [_FakeLigandAtom("B", "7") for _ in range(3)] + reference = np.array([[0.0, 0.0, 0.0]]) # only one atom coincides + + out_coords, out_atoms = dedup_mate_ligands_by_residue( + lig_coords, lig_atoms, reference, tol=0.3 + ) + + assert out_coords.shape[0] == 3 + assert len(out_atoms) == 3 + + def test_entities_are_judged_independently(self): + """One ligand being an image does not remove its neighbours.""" + lig_coords = np.array([[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]]) + lig_atoms = [_FakeLigandAtom("A", "1"), _FakeLigandAtom("A", "2")] + reference = np.array([[0.0, 0.0, 0.0]]) + + out_coords, out_atoms = dedup_mate_ligands_by_residue( + lig_coords, lig_atoms, reference, tol=0.3 + ) + + assert out_coords.shape[0] == 1 + np.testing.assert_allclose(out_coords[0], [20.0, 0.0, 0.0]) + assert out_atoms[0].resi == "2" + + def test_segment_separates_entities(self): + """Two ligands sharing (chain, resi) but not segi stay independent.""" + lig_coords = np.array([[0.0, 0.0, 0.0], [20.0, 0.0, 0.0]]) + lig_atoms = [ + _FakeLigandAtom("A", "1", segi="X"), + _FakeLigandAtom("A", "1", segi="Y"), + ] + reference = np.array([[0.0, 0.0, 0.0]]) + + out_coords, _ = dedup_mate_ligands_by_residue( + lig_coords, lig_atoms, reference, tol=0.3 + ) + + assert out_coords.shape[0] == 1 + np.testing.assert_allclose(out_coords[0], [20.0, 0.0, 0.0]) + + @pytest.mark.integration + def test_real_mate_ligands_are_kept_or_dropped_whole(self, pdb_4h0b): + """Real mate ligands from 4h0b, judged as entities. + + No fixture puts a ligand on a special position, so against the true ASU + reference nothing drops: every mate ligand here is a genuine neighbour-cell + copy. Referencing one entity against itself builds the drop case from real + atoms -- it goes whole, and the others are not fragmented. + """ + protein_atoms, water_atoms, ligand_atoms = parse_asu_with_biotite(pdb_4h0b) + crystal = get_crystal_contacts_pymol( + str(pdb_4h0b), cutoff=8.0, include_ligands=True + ) + lig_coords = crystal["mate_ligand_coords"] + lig_atoms = crystal["mate_ligand_atoms"] + assert len(lig_atoms) > 0, "4h0b must produce mate ligands to test" + + asu_reference = np.concatenate( + [protein_atoms.coord, water_atoms.coord, ligand_atoms.coord], axis=0 + ) + kept, kept_atoms = dedup_mate_ligands_by_residue( + lig_coords, lig_atoms, asu_reference, tol=0.3 + ) + assert kept.shape[0] == lig_coords.shape[0], ( + "genuine neighbour-cell ligands must all survive the ASU reference" + ) + + # now make one entity a symmetry image by referencing it against itself + groups = {} + for i, atom in enumerate(lig_atoms): + groups.setdefault( + (atom.chain, atom.resi, getattr(atom, "segi", "")), [] + ).append(i) + assert len(groups) > 1, "need more than one entity to show the others survive" + target_key, target_idx = next(iter(groups.items())) + + kept, kept_atoms = dedup_mate_ligands_by_residue( + lig_coords, + lig_atoms, + np.concatenate([asu_reference, lig_coords[target_idx]], axis=0), + tol=0.3, + ) + + # the imaged entity went whole, and nothing else went with it + assert kept.shape[0] == lig_coords.shape[0] - len(target_idx) + surviving = {(a.chain, a.resi, getattr(a, "segi", "")) for a in kept_atoms} + assert target_key not in surviving + assert surviving == set(groups) - {target_key} + # every survivor kept all of its atoms: entities are never fragmented + for key, idxs in groups.items(): + if key == target_key: + continue + assert sum( + 1 + for a in kept_atoms + if (a.chain, a.resi, getattr(a, "segi", "")) == key + ) == len(idxs) + + +@pytest.mark.unit +class TestParsePdbResi: + """Tests for PyMOL residue-identifier parsing.""" + + @pytest.mark.parametrize( + "resi,expected", + [ + ("52", (52, "")), + ("-3", (-3, "")), + ("52A", (52, "A")), + (" 52 ", (52, "")), + (7, (7, "")), + ], + ) + def test_parses(self, resi, expected): + assert _parse_pdb_resi(resi) == expected + + @pytest.mark.parametrize("resi", ["", "A", "52AB", "5.2"]) + def test_returns_none_when_unparseable(self, resi): + assert _parse_pdb_resi(resi) is None + + @pytest.mark.unit class TestCheckComDistance: """Tests for center of mass distance quality filter.""" @@ -768,6 +1024,13 @@ def test_is_ligand_mask_shape(self, pdb_4h0b, tmp_path): assert data["protein"].is_ligand.shape == (data["protein"].num_nodes,) assert data["protein"].is_ligand.dtype == torch.bool + # No residue owns a ligand, so it carries the -1 embedding sentinel. + cached = torch.load( + tmp_path / "processed" / "geometry" / "4h0b_final.pt", weights_only=False + ) + assert cached["is_ligand"].any() + assert (cached["emb_res_idx"][cached["is_ligand"]] == -1).all() + def test_protein_x_dim_unchanged(self, pdb_4h0b, tmp_path): """protein.x should still be 16-dim one-hot for both protein and ligand atoms.""" list_file = tmp_path / "list.txt" @@ -843,6 +1106,45 @@ def test_different_cutoffs(self, pdb_6eey): result_large["mate_coords"].shape[0] >= result_small["mate_coords"].shape[0] ) + def test_mates_never_include_solvent(self, pdb_8dzt): + """No mate atom, protein or ligand, is a water: a mate water is a symmetry + image of an ASU water, which is a prediction target.""" + result = get_crystal_contacts_pymol(pdb_8dzt, cutoff=5.0, include_ligands=True) + + protein_resns = {str(a.resn).upper() for a in result["mate_atoms"]} + ligand_resns = {str(a.resn).upper() for a in result["mate_ligand_atoms"]} + assert not {"HOH", "WAT"} & protein_resns + assert not {"HOH", "WAT"} & ligand_resns + + def test_ligand_mates_are_gated_and_separate(self, pdb_8dzt): + """include_ligands=False suppresses ligand mates and leaves protein mates + alone -- the two sets come back under separate keys.""" + full = get_crystal_contacts_pymol(pdb_8dzt, cutoff=5.0, include_ligands=True) + protein_only = get_crystal_contacts_pymol(pdb_8dzt, cutoff=5.0) + + assert len(protein_only["mate_ligand_atoms"]) == 0 + assert protein_only["mate_ligand_coords"].shape[0] == 0 + assert len(full["mate_ligand_atoms"]) > 0 + assert protein_only["mate_coords"].shape[0] == full["mate_coords"].shape[0] + + def test_special_position_water_never_selected(self, pdb_4h0b): + """4h0b has a target water on the 6-fold axis whose symmetry copy lands + ~0 A away. Since mate waters are never selected, no mate of any kind may + coincide with a target water -- the special-position leak cannot happen.""" + from scipy.spatial import cKDTree + + _, water_atoms, _ = parse_asu_with_biotite(pdb_4h0b) + result = get_crystal_contacts_pymol(pdb_4h0b, cutoff=5.0, include_ligands=True) + + mate_coords = result["mate_coords"] + if result["mate_ligand_coords"].shape[0]: + mate_coords = np.concatenate( + [mate_coords, result["mate_ligand_coords"]], axis=0 + ) + assert mate_coords.shape[0] > 0 + nearest = cKDTree(mate_coords).query(water_atoms.coord, k=1)[0] + assert (nearest < 0.3).sum() == 0 + @pytest.mark.integration class TestProteinWaterDataset: @@ -952,6 +1254,8 @@ def test_getitem_passes_mmap_flag_to_geometry_loader(self, tmp_path, monkeypatch "protein_x": torch.zeros((1, len(ELEMENT_VOCAB) + 1), dtype=torch.float32), "protein_res_idx": torch.zeros(1, dtype=torch.long), "is_ligand": torch.zeros(1, dtype=torch.bool), + "is_mate": torch.zeros(1, dtype=torch.bool), + "emb_res_idx": torch.zeros(1, dtype=torch.long), "pp_edge_index": torch.empty((2, 0), dtype=torch.long), "pp_edge_unit_vectors": torch.empty((0, 3), dtype=torch.float32), "pp_edge_rbf": torch.empty((0, 16), dtype=torch.float32), @@ -2138,9 +2442,9 @@ def test_gvp_encoder_no_embeddings(self, tmp_path, pdb_base_dir): dataset._annotate_data_with_embeddings( data=data, cache_key="test", - asu_protein_res_idx=torch.tensor([0]), num_asu_protein=100, num_protein_residues=50, + emb_res_idx=torch.zeros(100, dtype=torch.long), ) # Should not have added any embedding attributes @@ -2176,9 +2480,9 @@ def test_slae_encoder_loads_slae(self, tmp_path, pdb_base_dir): dataset._annotate_data_with_embeddings( data=data, cache_key="test_final", - asu_protein_res_idx=torch.tensor([0]), num_asu_protein=100, num_protein_residues=50, + emb_res_idx=torch.zeros(100, dtype=torch.long), ) assert hasattr(data["protein"], "embedding") @@ -2217,9 +2521,9 @@ def test_esm_encoder_loads_esm(self, tmp_path, pdb_base_dir): dataset._annotate_data_with_embeddings( data=data, cache_key="test_final", - asu_protein_res_idx=asu_res_idx, num_asu_protein=50, num_protein_residues=10, + emb_res_idx=asu_res_idx, ) assert hasattr(data["protein"], "embedding") @@ -2254,9 +2558,9 @@ def test_slae_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir): dataset._annotate_data_with_embeddings( data=data, cache_key="test_final", - asu_protein_res_idx=torch.zeros(num_asu, dtype=torch.long), num_asu_protein=num_asu, num_protein_residues=1, + emb_res_idx=torch.zeros(num_asu + num_mate + num_ligand, dtype=torch.long), ) emb = data["protein"].embedding @@ -2264,10 +2568,10 @@ def test_slae_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir): assert torch.equal(emb[:num_asu], asu_emb), "ASU rows must be left untouched" assert (emb[num_asu:] == 0).all(), "mate and ligand rows must be zero-padded" - def test_esm_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir): - """ESM residue embeddings broadcast to ASU atoms only. Mate and ligand atoms - are zero-padded -- ligands carry residue_index=-1 and must never be used to - index the residue embedding table.""" + def test_esm_mates_inherit_and_ligands_zero(self, tmp_path, pdb_base_dir): + """ESM rows broadcast to ASU atoms, and mate atoms inherit the row of the + ASU residue they image. Ligands carry -1 and must never index the residue + table, so they stay zero.""" from torch_geometric.data import HeteroData num_residues = 4 @@ -2293,16 +2597,19 @@ def test_esm_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir): data = HeteroData() data["protein"].num_nodes = num_asu + num_mate + num_ligand - # ASU res idx only -- ligand sentinels (-1) live past num_asu_protein and are - # sliced off by __getitem__ before this call. asu_res_idx = torch.arange(num_residues).repeat_interleave(atoms_per_residue) + # Mates image residue 0; ligands get the -1 sentinel. + mate_res_idx = torch.zeros(num_mate, dtype=torch.long) + emb_res_idx = torch.cat( + [asu_res_idx, mate_res_idx, torch.full((num_ligand,), -1)] + ) dataset._annotate_data_with_embeddings( data=data, cache_key="test_final", - asu_protein_res_idx=asu_res_idx, num_asu_protein=num_asu, num_protein_residues=num_residues, + emb_res_idx=emb_res_idx, ) emb = data["protein"].embedding @@ -2310,7 +2617,10 @@ def test_esm_zero_pads_mate_and_ligand_atoms(self, tmp_path, pdb_base_dir): assert torch.equal(emb[:num_asu], residue_emb[asu_res_idx]), ( "each ASU atom must carry its own residue's embedding" ) - assert (emb[num_asu:] == 0).all(), "mate and ligand rows must be zero-padded" + assert torch.equal( + emb[num_asu : num_asu + num_mate], residue_emb[mate_res_idx] + ), "mate atoms must inherit their source residue's embedding" + assert (emb[num_asu + num_mate :] == 0).all(), "ligand rows must stay zero" # ============== Tests for caching behavior ============== @@ -2665,6 +2975,369 @@ def test_num_asu_protein_metadata_correct( assert data.num_asu_protein_atoms <= data["protein"].num_nodes assert data.num_asu_protein_atoms > 0 + def _mate_dataset(self, single_pdb_list_file, tmp_path, pdb_base_dir, **kwargs): + return ProteinWaterDataset( + pdb_list_file=single_pdb_list_file, + processed_dir=str(tmp_path), + base_pdb_dir=str(pdb_base_dir), + include_mates=True, + preprocess=True, + **kwargs, + ) + + def test_mate_provenance_fields(self, single_pdb_list_file, tmp_path, pdb_base_dir): + """is_mate splits ASU from mate at num_asu_protein, and the mate atoms carry + the row of the ASU residue they image, not the -1 that reads as zero.""" + dataset = self._mate_dataset(single_pdb_list_file, tmp_path, pdb_base_dir) + data = dataset[0] + cached = torch.load( + tmp_path / "geometry_mates" / "6eey_final.pt", weights_only=False + ) + is_mate, emb_res_idx = data["protein"].is_mate, cached["emb_res_idx"] + num_asu = data.num_asu_protein_atoms + + assert is_mate.shape == emb_res_idx.shape == (data["protein"].num_nodes,) + assert not is_mate[:num_asu].any() + assert is_mate.sum().item() > 0 + + # is_mate is not contiguous: the ASU ligand block sits between its two True + # runs. The mate *protein* block is, and it starts at num_asu. + mate_protein = (is_mate & ~cached["is_ligand"]).nonzero().flatten() + assert mate_protein.numel() > 0 + assert mate_protein[0].item() == num_asu + assert (mate_protein.diff() == 1).all() + + mate_protein_mask = is_mate & ~cached["is_ligand"] + assert (emb_res_idx[mate_protein_mask] >= 0).all() + assert ( + emb_res_idx[mate_protein_mask] < data["protein"].num_protein_residues + ).all() + + def test_mate_emb_res_idx_equals_its_asu_residue_row( + self, single_pdb_list_file, tmp_path, pdb_base_dir + ): + """Every mate atom points at the exact ESM row of the ASU residue it images, + not merely at some row in range: an off-by-one passes a range check but + gives every mate the wrong embedding. + + Water filtering is off so the dedup reference, and so the surviving mate + atom list, is reproducible atom for atom. + """ + dataset = self._mate_dataset( + single_pdb_list_file, + tmp_path, + pdb_base_dir, + filter_by_distance=False, + filter_by_edia=False, + filter_by_bfactor=False, + ) + data = dataset[0] + cached = torch.load( + tmp_path / "geometry_mates" / "6eey_final.pt", weights_only=False + ) + + pdb_path = Path(pdb_base_dir) / "6eey" / "6eey_final.pdb" + protein_atoms, water_atoms, ligand_atoms = parse_asu_with_biotite(pdb_path) + + # rebuild the ASU (chain, res_id, ins_code) -> ESM row map exactly as + # _preprocess_one does, off the same sanitized parse + 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]) + asu_key_to_row = {} + for res_i, start in enumerate(bts.get_residue_starts(sanitized)): + key = ( + str(sanitized.chain_id[start]).strip(), + int(sanitized.res_id[start]), + str(sanitized.ins_code[start]), + ) + asu_key_to_row.setdefault(key, res_i) + + # reproduce the surviving mate atom list, in order + crystal = get_crystal_contacts_pymol( + str(pdb_path), dataset.cutoff, include_ligands=dataset.include_ligands + ) + water_atoms = water_atoms[ + match_atoms_to_coords(water_atoms, crystal["asu_coords"]) + ] + reference = [protein_atoms.coord] + if len(water_atoms): + reference.append(water_atoms.coord) + if dataset.include_ligands and len(ligand_atoms) > 0: + reference.append(ligand_atoms.coord) + _, mate_atoms = dedup_mate_atoms( + crystal["mate_coords"], + crystal["mate_atoms"], + np.concatenate(reference, axis=0), + ) + + num_asu = data.num_asu_protein_atoms + emb_res_idx = cached["emb_res_idx"] + mate_protein_count = int((cached["is_mate"] & ~cached["is_ligand"]).sum()) + assert len(mate_atoms) == mate_protein_count > 0 + + expected = [] + for atom in mate_atoms: + parsed = _parse_pdb_resi(atom.resi) + assert parsed is not None, f"unparseable mate resi {atom.resi!r}" + expected.append(asu_key_to_row[(str(atom.chain).strip(), *parsed)]) + + actual = emb_res_idx[num_asu : num_asu + mate_protein_count] + assert torch.equal(actual, torch.tensor(expected, dtype=actual.dtype)) + + # mates reuse ASU rows rather than introducing rows of their own + asu_rows = set(emb_res_idx[:num_asu].tolist()) + assert set(actual.tolist()) <= asu_rows + + def test_no_mates_run_marks_nothing_as_mate( + self, single_pdb_list_file, tmp_path, pdb_base_dir + ): + """Without mates every node is ASU, so the prior's anchor mask is inert.""" + dataset = ProteinWaterDataset( + pdb_list_file=single_pdb_list_file, + processed_dir=str(tmp_path), + base_pdb_dir=str(pdb_base_dir), + include_mates=False, + preprocess=True, + ) + + data = dataset[0] + assert not data["protein"].is_mate.any() + + def test_mate_waters_never_enter_the_graph( + self, single_pdb_list_file, tmp_path, pdb_base_dir + ): + """No protein node may sit on a target water: that is the label leak the + mate selection and the dedup pass exist to prevent.""" + dataset = self._mate_dataset(single_pdb_list_file, tmp_path, pdb_base_dir) + data = dataset[0] + + waters = data["water"].pos + if waters.size(0) == 0: + pytest.skip("structure has no waters after filtering") + nearest = torch.cdist(waters, data["protein"].pos).min(dim=1).values + assert nearest.min().item() > 0.3 + + +@pytest.mark.integration +class TestMatesWithLigands: + """Mates and ligands turned on together, end to end. + + Other mates tests use 6eey (no ligands) and other ligand tests set + include_mates=False, so the four-block layout was never exercised on a + structure with all four. 4h0b has ligands and crystal contacts. + """ + + def _dataset(self, tmp_path, pdb_base_dir, **kwargs): + tmp_path.mkdir(parents=True, exist_ok=True) + list_file = tmp_path / "list.txt" + list_file.write_text("4h0b_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=True, + include_ligands=True, + preprocess=True, + # 4h0b ships no EDIA sidecar; this class is about the node layout + filter_by_distance=False, + filter_by_edia=False, + filter_by_bfactor=False, + **kwargs, + ) + + @staticmethod + def _blocks(is_mate, is_ligand): + return { + "asu_protein": ~is_mate & ~is_ligand, + "mate_protein": is_mate & ~is_ligand, + "asu_ligand": ~is_mate & is_ligand, + "mate_ligand": is_mate & is_ligand, + } + + def test_four_blocks_appear_in_documented_order(self, tmp_path, pdb_base_dir): + """Node order is [ASU protein | mate protein | ASU ligand | mate ligand]. + Each block is contiguous and they follow one another in that order.""" + data = self._dataset(tmp_path, pdb_base_dir)[0] + cached = torch.load( + tmp_path / "processed" / "geometry_mates" / "4h0b_final.pt", + weights_only=False, + ) + blocks = self._blocks(data["protein"].is_mate, cached["is_ligand"]) + + order = ["asu_protein", "mate_protein", "asu_ligand", "mate_ligand"] + spans = [] + for name in order: + idx = blocks[name].nonzero().flatten() + assert idx.numel() > 0, f"4h0b should populate the {name} block" + assert (idx.diff() == 1).all(), f"{name} block is not contiguous" + spans.append((name, idx[0].item(), idx[-1].item())) + + # blocks tile the node range back to back, in the documented order + assert spans[0][1] == 0 + for (_, _, prev_end), (name, start, _) in zip(spans, spans[1:]): + assert start == prev_end + 1, f"{name} does not follow the previous block" + assert spans[-1][2] == data["protein"].num_nodes - 1 + + def test_num_asu_protein_excludes_asu_ligands(self, tmp_path, pdb_base_dir): + """num_asu_protein is the ASU *protein* count. ASU ligands sit behind the + mates, so counting them here would misalign every SLAE/ESM lookup.""" + data = self._dataset(tmp_path, pdb_base_dir)[0] + cached = torch.load( + tmp_path / "processed" / "geometry_mates" / "4h0b_final.pt", + weights_only=False, + ) + blocks = self._blocks(data["protein"].is_mate, cached["is_ligand"]) + + assert data.num_asu_protein_atoms == int(blocks["asu_protein"].sum()) + assert int(blocks["asu_ligand"].sum()) > 0 + assert data.num_asu_protein_atoms < data["protein"].num_nodes + + def test_emb_res_idx_splits_ligands_from_mate_protein(self, tmp_path, pdb_base_dir): + """Ligands carry the -1 sentinel whichever cell they came from; mate protein + atoms carry a real ASU row.""" + data = self._dataset(tmp_path, pdb_base_dir)[0] + cached = torch.load( + tmp_path / "processed" / "geometry_mates" / "4h0b_final.pt", + weights_only=False, + ) + emb_res_idx = cached["emb_res_idx"] + blocks = self._blocks(data["protein"].is_mate, cached["is_ligand"]) + + assert (emb_res_idx[blocks["asu_ligand"]] == -1).all() + assert (emb_res_idx[blocks["mate_ligand"]] == -1).all() + + mate_rows = emb_res_idx[blocks["mate_protein"]] + assert (mate_rows >= 0).all() + assert (mate_rows < data["protein"].num_protein_residues).all() + # mates reuse ASU rows rather than introducing rows of their own + assert set(mate_rows.tolist()) <= set( + emb_res_idx[blocks["asu_protein"]].tolist() + ) + + def test_esm_encoder_over_the_real_preprocessing_path(self, tmp_path, pdb_base_dir): + """Other integration tests run GVP, and the ESM coverage is over hand-built + graphs. This drives the real preprocessing path with encoder_type='esm'. + The ESM cache is synthesised so the broadcast is checkable atom by atom. + """ + # a first pass tells us how many residues the ESM cache must cover + probe = self._dataset(tmp_path / "probe", pdb_base_dir)[0] + num_residues = int(probe["protein"].num_protein_residues) + + esm_dir = tmp_path / "processed" / "esm" + esm_dir.mkdir(parents=True, exist_ok=True) + residue_emb = torch.randn(num_residues, ESM_EMBEDDING_DIM) + torch.save({"residue_embeddings": residue_emb}, esm_dir / "4h0b_final.pt") + + data = self._dataset(tmp_path, pdb_base_dir, encoder_type="esm")[0] + cached = torch.load( + tmp_path / "processed" / "geometry_mates" / "4h0b_final.pt", + weights_only=False, + ) + emb = data["protein"].embedding + emb_res_idx = cached["emb_res_idx"] + blocks = self._blocks(data["protein"].is_mate, cached["is_ligand"]) + + assert emb.shape == (data["protein"].num_nodes, ESM_EMBEDDING_DIM) + + protein_nodes = blocks["asu_protein"] | blocks["mate_protein"] + assert torch.equal(emb[protein_nodes], residue_emb[emb_res_idx[protein_nodes]]) + # every ligand, ASU or mate, reads as a zero row + ligand_nodes = blocks["asu_ligand"] | blocks["mate_ligand"] + assert (emb[ligand_nodes] == 0).all() + + +@pytest.mark.unit +class TestFilterMetaSidecar: + """A geometry directory records the settings its entries were built with.""" + + def _dataset(self, tmp_path, *, preprocess=True, **kwargs): + """Dataset over an empty list: claims the directory, preprocesses nothing.""" + list_file = tmp_path / "empty.txt" + list_file.write_text("") + return ProteinWaterDataset( + pdb_list_file=str(list_file), + processed_dir=str(tmp_path / "processed"), + base_pdb_dir=str(tmp_path), + preprocess=preprocess, + **kwargs, + ) + + def _meta_path(self, tmp_path): + return tmp_path / "processed" / "geometry_mates" / FILTER_META_FILENAME + + def test_written_on_preprocess(self, tmp_path): + self._dataset(tmp_path, max_bfactor_zscore=2.0) + + recorded = json.loads(self._meta_path(tmp_path).read_text()) + assert recorded["max_bfactor_zscore"] == 2.0 + assert recorded["min_edia"] == 0.4 + assert recorded["filter_by_bfactor"] is True + # Structure-level checks decide which entries exist at all, and the graph + # parameters decide the cached edges: both belong to the directory too. + assert recorded["min_water_residue_ratio"] == 0.1 + assert recorded["cutoff"] == 8.0 + assert recorded["max_neighbors"] == 256 + + def test_matching_settings_accepted(self, tmp_path): + self._dataset(tmp_path, max_bfactor_zscore=2.0) + self._dataset(tmp_path, max_bfactor_zscore=2.0) # must not raise + + @pytest.mark.parametrize( + "changed,preprocess", + [ + ({"max_bfactor_zscore": 1.5}, True), # water threshold + ({"filter_by_edia": False}, True), # water filter toggle + ({"min_water_residue_ratio": 0.6}, True), # which entries exist + ({"cutoff": 6.0}, True), # which PP edges were cached + # A read-only run is refused too: it would report metrics over waters + # filtered differently than it asked for. + ({"min_edia": 0.6}, False), + ], + ) + def test_mismatch_refused(self, tmp_path, changed, preprocess): + self._dataset(tmp_path) + + with pytest.raises(ValueError, match=next(iter(changed))): + self._dataset(tmp_path, preprocess=preprocess, **changed) + + def test_disabled_filter_ignores_its_threshold(self, tmp_path): + """A disabled filter never touched the cached waters, so its threshold + must not make two identical caches look incompatible.""" + self._dataset(tmp_path, filter_by_bfactor=False, max_bfactor_zscore=2.0) + + assert ( + json.loads(self._meta_path(tmp_path).read_text())["max_bfactor_zscore"] + is None + ) + self._dataset(tmp_path, filter_by_bfactor=False, max_bfactor_zscore=1.5) + + def test_directories_are_claimed_independently(self, tmp_path): + """Mates and no-mates are separate directories, so they may disagree.""" + self._dataset(tmp_path, include_mates=True, max_bfactor_zscore=2.0) + self._dataset(tmp_path, include_mates=False, max_bfactor_zscore=1.5) + + @pytest.mark.parametrize("preprocess", [False, True]) + def test_unlabelled_cache_warns(self, tmp_path, warning_log, preprocess): + """An existing directory with no sidecar is unverifiable whether or not this + run also claims it.""" + geometry_dir = tmp_path / "processed" / "geometry_mates" + geometry_dir.mkdir(parents=True) + (geometry_dir / "6eey_final.pt").write_bytes(b"cache") + + self._dataset(tmp_path, preprocess=preprocess) + + assert any(FILTER_META_FILENAME in message for message in warning_log) + # the writing run still claims the directory; the read-only one must not + assert self._meta_path(tmp_path).is_file() is preprocess + + def test_empty_directory_does_not_warn(self, tmp_path, warning_log): + (tmp_path / "processed" / "geometry_mates").mkdir(parents=True) + + self._dataset(tmp_path, preprocess=False) + + assert not any(FILTER_META_FILENAME in message for message in warning_log) + # ============== Tests for residue index assignment ============== diff --git a/tests/test_flow.py b/tests/test_flow.py index 9d2f62a..be4d2dc 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -24,6 +24,17 @@ from src.gvp_encoder import GVPEncoder, make_gvp_encoder_data, ProteinGVPEncoder +@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.fixture def simple_hetero_data(device): """Minimal HeteroData with protein and water nodes.""" @@ -633,6 +644,101 @@ def test_zero_protein_graph_raises(self, device): device=device, ) + def test_anchor_mask_is_per_graph(self, device): + """Masked-out atoms are never ball centres, and each graph stays within + its own eligible atoms rather than its neighbour's.""" + torch.manual_seed(0) + protein_pos = torch.tensor( + [[0.0, 0.0, 0.0], [500.0, 0.0, 0.0], [100.0, 0.0, 0.0], [900.0, 0.0, 0.0]], + device=device, + ) + batch_p = torch.tensor([0, 0, 1, 1], dtype=torch.long, device=device) + batch_w = _batch_from_counts( + torch.tensor([50, 50], dtype=torch.long, device=device), device + ) + anchor_mask = torch.tensor([True, False, True, False], device=device) + + pos = sample_waters_uniform_ball( + protein_pos=protein_pos, + batch_p=batch_p, + batch_w=batch_w, + cutoff=2.0, + device=device, + anchor_mask=anchor_mask, + ) + + assert pos[batch_w == 0][:, 0].abs().max().item() < 5.0 + assert (pos[batch_w == 1][:, 0] - 100.0).abs().max().item() < 5.0 + + def test_anchor_mask_skipped_for_batch_when_a_graph_starves( + self, device, warning_log + ): + """If the mask would leave a water-requesting graph with no anchor, the + whole batch anchors on all atoms (as local_flow does) and warns.""" + torch.manual_seed(0) + protein_pos = torch.tensor( + [ + [0.0, 0.0, 0.0], # graph 0, masked out + [10.0, 0.0, 0.0], # graph 0, masked out -> graph 0 is starved + [1000.0, 0.0, 0.0], # graph 1, eligible + [2000.0, 0.0, 0.0], # graph 1, masked out + ], + device=device, + ) + batch_p = torch.tensor([0, 0, 1, 1], dtype=torch.long, device=device) + batch_w = _batch_from_counts( + torch.tensor([100, 100], dtype=torch.long, device=device), device + ) + anchor_mask = torch.tensor([False, False, True, False], device=device) + cutoff = 2.0 + + pos = sample_waters_uniform_ball( + protein_pos=protein_pos, + batch_p=batch_p, + batch_w=batch_w, + cutoff=cutoff, + device=device, + anchor_mask=anchor_mask, + ) + + # each water still sits within cutoff of one of its own graph's atoms + d0 = torch.cdist(pos[batch_w == 0], protein_pos[:2]).min(dim=1).values + assert d0.max().item() <= cutoff + 1e-5 + # the mask was dropped for the whole batch, so graph 1 uses its masked-out + # atom at x=2000 too, not only the eligible one at x=1000 + g1_x = pos[batch_w == 1][:, 0] + assert (g1_x > 1500.0).any() + + assert any("all protein atoms" in message for message in warning_log) + + def test_anchor_mask_none_matches_all_true_mask(self, device): + """An all-True mask changes neither the draws nor their order.""" + protein_pos = torch.randn(12, 3, device=device) * 10 + batch_p = torch.cat([torch.zeros(6), torch.ones(6)]).long().to(device) + batch_w = _batch_from_counts( + torch.tensor([20, 15], dtype=torch.long, device=device), device + ) + + torch.manual_seed(7) + without = sample_waters_uniform_ball( + protein_pos=protein_pos, + batch_p=batch_p, + batch_w=batch_w, + cutoff=8.0, + device=device, + ) + torch.manual_seed(7) + with_mask = sample_waters_uniform_ball( + protein_pos=protein_pos, + batch_p=batch_p, + batch_w=batch_w, + cutoff=8.0, + device=device, + anchor_mask=torch.ones(12, dtype=torch.bool, device=device), + ) + + assert torch.equal(without, with_mask) + def test_large_spread_protein_succeeds(self, device): """The scenario that crashes truncated Gaussian (sigma~50) works here.""" torch.manual_seed(0) @@ -762,6 +868,144 @@ def test_empty_waters(self, device): assert pos.shape == (0, 3) +@pytest.mark.unit +class TestCrystalMateAwareSampling: + """Both samplers work from ASU atoms only when the batch carries is_mate.""" + + @staticmethod + def _matcher(strategy): + return FlowMatcher(model=Mock(cutoff=8.0), sampling_strategy=strategy) + + @staticmethod + def _graph(pos, batch, is_mate=None): + data = HeteroData() + data["protein"].pos = pos + data["protein"].batch = batch + if is_mate is not None: + data["protein"].is_mate = is_mate + return data + + def test_uniform_ball_anchors_on_asu_only(self, device): + """Waters spawn around ASU atoms, not around the distant symmetry mates.""" + torch.manual_seed(0) + data = self._graph( + torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [500.0, 0.0, 0.0], + [900.0, 0.0, 0.0], + ], + device=device, + ), + torch.zeros(4, dtype=torch.long, device=device), + torch.tensor([False, False, True, True], device=device), + ) + batch_w = torch.zeros(200, dtype=torch.long, device=device) + + matcher = self._matcher("uniform_ball") + pos = matcher._sample_waters(data, batch_w, device) + + # ASU spans x in [0, 1]; a mate anchor would land a water near 500 or 900 + assert pos[:, 0].max().item() < 1.0 + matcher.graph_cutoff + 1e-5 + assert pos[:, 0].min().item() > -matcher.graph_cutoff - 1e-5 + + def test_uniform_ball_skips_mask_for_all_mate_graph(self, device, warning_log): + """An all-mate graph has no ASU anchor, so the batch anchors on all atoms + and warns; every graph still draws from its own atoms.""" + torch.manual_seed(0) + data = self._graph( + torch.tensor([[0.0, 0.0, 0.0], [500.0, 0.0, 0.0]], device=device), + torch.tensor([0, 1], dtype=torch.long, device=device), + torch.tensor([False, True], device=device), + ) + batch_w = _batch_from_counts( + torch.tensor([20, 20], dtype=torch.long, device=device), device + ) + + pos = self._matcher("uniform_ball")._sample_waters(data, batch_w, device) + + assert pos[batch_w == 0][:, 0].abs().max().item() < 10.0 + assert (pos[batch_w == 1][:, 0] - 500.0).abs().max().item() < 10.0 + assert any("all protein atoms" in message for message in warning_log) + + def test_all_false_is_mate_matches_no_attribute(self, device): + """Dropping the .any() guard is safe: an all-True mask draws exactly what + no mask draws, bit for bit.""" + protein_pos = torch.randn(12, 3, device=device) * 10 + batch_p = torch.cat([torch.zeros(6), torch.ones(6)]).long().to(device) + batch_w = _batch_from_counts( + torch.tensor([20, 15], dtype=torch.long, device=device), device + ) + matcher = self._matcher("uniform_ball") + + plain = self._graph(protein_pos, batch_p) + flagged = self._graph( + protein_pos, batch_p, torch.zeros(12, dtype=torch.bool, device=device) + ) + + torch.manual_seed(7) + without = matcher._sample_waters(plain, batch_w, device) + torch.manual_seed(7) + with_mask = matcher._sample_waters(flagged, batch_w, device) + + assert torch.equal(without, with_mask) + + def test_sigma_ignores_distant_mates(self, device): + """Counting distant mates would inflate sigma and push the prior out past + the targets.""" + torch.manual_seed(0) + asu = torch.randn(30, 3, device=device) + mates = torch.randn(30, 3, device=device) + 500.0 + + asu_only = self._graph(asu, torch.zeros(30, dtype=torch.long, device=device)) + with_mates = self._graph( + torch.cat([asu, mates], dim=0), + torch.zeros(60, dtype=torch.long, device=device), + torch.cat( + [torch.zeros(30, dtype=torch.bool), torch.ones(30, dtype=torch.bool)] + ).to(device), + ) + + reference = FlowMatcher.compute_sigma_per_graph(asu_only, device) + masked = FlowMatcher.compute_sigma_per_graph( + with_mates, device, node_mask=FlowMatcher._asu_mask(with_mates) + ) + unmasked = FlowMatcher.compute_sigma_per_graph(with_mates, device) + + # adding the mates leaves sigma untouched once they are masked out + assert torch.allclose(masked, reference, atol=1e-4) + # ...and would otherwise blow it up by two orders of magnitude + assert unmasked.item() > 50 * reference.item() + + def test_sigma_skips_mask_for_all_mate_graph(self, device, warning_log): + """If any graph has no ASU atom, sigma is computed over all atoms for the + whole batch (as local_flow does) with a warning, not a degenerate value.""" + torch.manual_seed(0) + g0 = HeteroData() + g0["protein"].pos = torch.cat( + [torch.randn(20, 3), torch.randn(20, 3) + 500.0] + ).to(device) + g0["protein"].is_mate = torch.cat( + [torch.zeros(20, dtype=torch.bool), torch.ones(20, dtype=torch.bool)] + ).to(device) + g1 = HeteroData() + g1["protein"].pos = torch.randn(20, 3).to(device) + g1["protein"].is_mate = torch.ones(20, dtype=torch.bool, device=device) + + batch = Batch.from_data_list([g0, g1]) + masked = FlowMatcher.compute_sigma_per_graph( + batch, device, node_mask=FlowMatcher._asu_mask(batch) + ) + full = FlowMatcher.compute_sigma_per_graph(batch, device) + + # graph 1 has no ASU atom, so the mask is skipped for the whole batch: + # sigma matches the unmasked computation everywhere, non-degenerate + assert torch.allclose(masked, full) + assert masked[1].item() > 0.5 + assert any("all protein atoms" in message for message in warning_log) + + @pytest.mark.unit class TestSamplingHonoursNodeOrder: """Samplers return one water per batch_w entry, in batch_w's own order."""