diff --git a/README.md b/README.md index 4724d45..ce09d94 100644 --- a/README.md +++ b/README.md @@ -81,10 +81,13 @@ WaterFlow processes structure files through several stages to create training-re - 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 - Edge types (defined in `src/constants.py`): - - `('protein', 'pp', 'protein')`: protein-protein edges - - `('protein', 'pw', 'water')`: protein to water - - `('water', 'wp', 'protein')`: water to protein - - `('water', 'ww', 'water')`: water-water edges + - `('protein', 'pp', 'protein')`: protein-protein edges — cached at preprocessing + - `('protein', 'pw', 'water')`: protein to water — built at runtime + - `('water', 'wp', 'protein')`: water to protein — built at runtime, ablatable + - `('water', 'ww', 'water')`: water-water edges — built at runtime, ablatable +- Only PP edges are stored in the geometry cache; every water-touching edge is + rebuilt each forward pass, since water positions move during integration. See + [Edge Construction](#edge-construction) - Default edge cutoff: 8.0Å (`RBF_CUTOFF` in constants.py) **Feature Encoding** @@ -199,6 +202,34 @@ WaterFlow uses a two-stage architecture: | `esm` | Uses ESM3 language model embeddings | Yes (`generate_esm_embeddings.py`) | | `slae` | Uses SLAE ([Strictly Local All-Atom Environment](https://www.biorxiv.org/content/10.1101/2025.10.03.680398v1)) embeddings | Yes (`generate_slae_embeddings.py`) | +### Edge Construction + +Water-touching edges (PW, WW, WP) are rebuilt every forward pass because water +positions change during integration. How they are built is fixed at model +construction, so training and inference always agree: + +| `--dynamic_edge_policy` | Behaviour | +|-------------------------|-----------| +| `radius` (default) | Connect every pair within `--cutoff`, capped at `--max_neighbors` per source | +| `knn` | Connect a fixed number of nearest neighbours (`--k_pw`, `--k_ww`, `--k_wp`) | + +The two differ in which side the neighbour budget applies to. KNN queries *per +destination*, so every destination is guaranteed edges but a source may have +none — coverage checks must read the destination row. Radius guarantees nothing: +a water with no protein atom inside `--cutoff` gets no PW edges at all. + +`--knn_fallback_k` repairs that. Under `radius`, any water the query stranded is +reconnected to that many nearest protein atoms regardless of distance. Set it to +`0` to disable. It has no effect under `knn`, which cannot strand a node. + +Set `--disable_ww` / `--disable_wp` to ablate those edge types; PW and PP are +always active. + +> Configs written before the radius/KNN split recorded a three-valued +> `dynamic_edge_policy` (`auto`, `radius`, `knn_if_isolated`). All three built a +> radius graph, so they load and map to `radius`; whether stranded waters are +> rescued is now `--knn_fallback_k`'s job. + ## Embedding Generation For `esm` and `slae` encoder types, you must precompute embeddings before training or inference. @@ -272,6 +303,11 @@ To resume training from a checkpoint, you can load the model weights and optimiz | `--scheduler` | `cosine` | LR scheduler: `cosine`, `step`, or `none` | | `--warmup_steps` | `0` | Linear warmup steps | | `--processed_dir` | `~/flow_cache/` | Cache directory for preprocessed data | +| `--dynamic_edge_policy` | `radius` | How water-touching edges are built: `radius` or `knn` (see [Edge Construction](#edge-construction)) | +| `--cutoff` | `8.0` | Distance cutoff in Å for radius edges | +| `--knn_fallback_k` | `8` | Nearest neighbours attached to waters stranded by the radius query; `0` disables | +| `--disable_ww` | `false` | Ablate water→water edges | +| `--disable_wp` | `false` | Ablate water→protein edges | | `--include_mates` | `false` | Include symmetry mate atoms as protein nodes | | `--include_ligands` | `true` | Include ligand/ion/cofactor/nucleic acid heavy atoms as protein nodes. Negate with `--no-include_ligands` | | `--save_dir` | `../flow_checkpoints` | Directory to save checkpoints | diff --git a/scripts/inference.py b/scripts/inference.py index e73b919..0dec53e 100644 --- a/scripts/inference.py +++ b/scripts/inference.py @@ -273,8 +273,17 @@ def build_model_from_config(config: dict, device: torch.device) -> nn.Module: drop_rate=config.get("drop_rate", 0.1), n_message_gvps=config.get("n_message_gvps", 2), n_update_gvps=config.get("n_update_gvps", 2), - k_pw=config.get("k_pw") or 16, - k_ww=config.get("k_ww") or 16, + cutoff=config.get("cutoff", 8.0), + max_neighbors=config.get("max_neighbors", 256), + dynamic_edge_policy=config.get("dynamic_edge_policy", "radius"), + # "auto" depends on which prior the run uses, so pass that through. + sampling_strategy=config.get("sampling_strategy", "uniform_ball"), + knn_fallback_k=config.get("knn_fallback_k", 8), + disable_ww=config.get("disable_ww", False), + disable_wp=config.get("disable_wp", False), + k_pw=config.get("k_pw", 12), + k_ww=config.get("k_ww", 8), + k_wp=config.get("k_wp", 8), ).to(device) return model diff --git a/scripts/train.py b/scripts/train.py index f5bc0be..1ab421a 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -36,7 +36,7 @@ from src.dataset import get_dataloader from src.encoder_base import build_encoder -from src.flow import FlowMatcher, FlowWaterGVP +from src.flow import DYNAMIC_EDGE_POLICIES, FlowMatcher, FlowWaterGVP from src.utils import ( compute_placement_metrics, compute_rmsd, @@ -219,8 +219,70 @@ def parse_args(): default=0.1, help="Dropout rate for GVP layers (default: 0.1)", ) - p.add_argument("--k_pw", type=int, default=16) - p.add_argument("--k_ww", type=int, default=16) + # edge construction + p.add_argument( + "--dynamic_edge_policy", + type=str, + default="auto", + choices=["auto", *DYNAMIC_EDGE_POLICIES], + help=( + "How water-touching edges are built: 'radius' connects everything " + "within --cutoff, 'knn' takes a fixed neighbour count, " + "'knn_if_isolated' is radius plus a rescue for stranded waters. " + "'auto' picks radius under uniform_ball and knn_if_isolated under " + "scaled_gaussian (default: auto)" + ), + ) + p.add_argument( + "--cutoff", + type=float, + default=8.0, + help="Distance cutoff in Angstroms for radius edges (default: 8.0)", + ) + p.add_argument( + "--max_neighbors", + type=int, + default=256, + help="Per-source cap on radius query results (default: 256)", + ) + p.add_argument( + "--knn_fallback_k", + type=int, + default=8, + help=( + "Nearest neighbours attached to waters the radius query stranded; " + "0 disables the rescue. Ignored under --dynamic_edge_policy knn " + "(default: 8)" + ), + ) + p.add_argument( + "--disable_ww", + action="store_true", + help="Ablate water->water edges", + ) + p.add_argument( + "--disable_wp", + action="store_true", + help="Ablate water->protein edges", + ) + p.add_argument( + "--k_pw", + type=int, + default=12, + help="Nearest neighbours for protein->water edges under 'knn' (default: 12)", + ) + p.add_argument( + "--k_ww", + type=int, + default=8, + help="Nearest neighbours for water->water edges under 'knn' (default: 8)", + ) + p.add_argument( + "--k_wp", + type=int, + default=8, + help="Nearest neighbours for water->protein edges under 'knn' (default: 8)", + ) # optional cached-embedding override p.add_argument( @@ -563,8 +625,17 @@ def build_model( n_message_gvps=args.n_message_gvps, n_update_gvps=args.n_update_gvps, drop_rate=args.drop_rate, + cutoff=args.cutoff, + max_neighbors=args.max_neighbors, + dynamic_edge_policy=args.dynamic_edge_policy, + # "auto" depends on which prior the run uses, so pass that through. + sampling_strategy=args.sampling_strategy, + knn_fallback_k=args.knn_fallback_k, + disable_ww=args.disable_ww, + disable_wp=args.disable_wp, k_pw=args.k_pw, k_ww=args.k_ww, + k_wp=args.k_wp, ).to(device) return model diff --git a/src/constants.py b/src/constants.py index 5e2604f..4faac55 100644 --- a/src/constants.py +++ b/src/constants.py @@ -26,6 +26,35 @@ # all edge types used in the model (future support: het atoms) ALL_EDGE_TYPES = [EDGE_PW, EDGE_WW, EDGE_PP, EDGE_WP] + +def get_active_edge_types( + disable_ww: bool = False, disable_wp: bool = False +) -> list[tuple[str, str, str]]: + """ + Return the active edge types for a model configuration. + + PW and PP are always active: PW carries protein context onto waters and PP + is read from the geometry cache. WW and WP are ablatable. + + The returned order differs from ``ALL_EDGE_TYPES``. That is safe -- edge + types key ``HeteroConv``'s parameters by name (``convs.``), + not by position, so ordering does not affect state-dict compatibility. + + Args: + disable_ww: Drop water -> water edges. + disable_wp: Drop water -> protein edges. + + Returns: + List of (src_type, relation, dst_type) tuples. + """ + etypes = [EDGE_PW, EDGE_PP] + if not disable_ww: + etypes.append(EDGE_WW) + if not disable_wp: + etypes.append(EDGE_WP) + return etypes + + # Standard 3-letter to 1-letter amino acid mapping # Includes 20 canonical amino acids plus common non-standard residues # Non-canonical residues not in this dict should be mapped to 'X' diff --git a/src/flow.py b/src/flow.py index 6a4f151..3335931 100644 --- a/src/flow.py +++ b/src/flow.py @@ -2,7 +2,7 @@ Flow matching model components for water placement prediction. This module provides: -- ProteinWaterUpdate: Heterogeneous GVP message passing across 4 edge types +- ProteinWaterUpdate: Heterogeneous GVP message passing across the active edge types - FlowWaterGVP: End-to-end flow model combining encoder + GVP updates + vector field head - FlowMatcher: High-level training, validation, and numerical integration interface """ @@ -16,7 +16,7 @@ import torch.nn.functional as F from torch import nn, Tensor from torch_geometric.data import Batch, HeteroData -from torch_geometric.nn import knn +from torch_geometric.nn import knn, radius, radius_graph from torch_scatter import scatter, scatter_mean from tqdm.auto import tqdm @@ -28,6 +28,7 @@ EDGE_WW, ELEM_IDX, ELEMENT_VOCAB, + get_active_edge_types, NUM_RBF, ) from src.encoder_base import BaseProteinEncoder @@ -35,6 +36,40 @@ from src.utils import ot_coupling +# "knn_if_isolated" = radius edges, plus nearest-neighbour edges for nodes left with none. +DYNAMIC_EDGE_POLICIES = ("radius", "knn", "knn_if_isolated") + + +def resolve_edge_policy(policy: str, sampling_strategy: str = "uniform_ball") -> str: + """ + Resolve a recorded `dynamic_edge_policy`, including the "auto" setting. + + "auto" is what every recorded run carries. It reads off the prior: uniform + ball samples land within `cutoff` of a protein atom by construction, so + nothing is stranded and the rescue is not wanted; Gaussian samples carry no + such guarantee, so they get it. + + Args: + policy: Value from a config file or CLI flag. + sampling_strategy: Prior the run uses, consulted only for "auto". + + Returns: + One of DYNAMIC_EDGE_POLICIES. + + Raises: + ValueError: If the value is not "auto" or a member of + DYNAMIC_EDGE_POLICIES. + """ + if policy == "auto": + return "knn_if_isolated" if sampling_strategy == "scaled_gaussian" else "radius" + if policy not in DYNAMIC_EDGE_POLICIES: + raise ValueError( + f"dynamic_edge_policy must be 'auto' or one of {DYNAMIC_EDGE_POLICIES}, " + f"got '{policy}'" + ) + return policy + + def _batch_from_counts(num_waters: Tensor, device: torch.device) -> Tensor: """ Build a graph-grouped batch vector from per-graph counts. @@ -168,30 +203,42 @@ def sample_waters_scaled_gaussian( return torch.randn(total_waters, 3, device=device, dtype=dtype) * sigma -def build_knn_edges( +def build_dynamic_edges( src_pos: torch.Tensor, dst_pos: torch.Tensor, + *, + policy: str, k: int, + r: float, + max_neighbors: int = 256, batch_src: torch.Tensor | None = None, batch_dst: torch.Tensor | None = None, ) -> torch.Tensor: """ - Build KNN edges from src -> dst (source indices in row 0, dest in row 1). + Build 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). + The two policies differ in which side the neighbour budget applies to, which + matters when reading coverage guarantees off the result: + + - ``"knn"`` queries *per destination*: each point in ``dst_pos`` takes its + ``k`` nearest neighbours in ``src_pos``. Every destination is therefore + guaranteed incoming edges and appears in row 1, while a source that is no + destination's nearest neighbour may not appear in row 0 at all. Coverage + checks must be made against row 1. + - ``"radius"`` connects every pair within ``r``, capped at ``max_neighbors`` + *per source*. Nothing is guaranteed: a node with an empty neighbourhood + gets no edges, which is what :meth:`ProteinWaterUpdate._add_knn_fallback` + exists to repair. 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. + policy: One of DYNAMIC_EDGE_POLICIES. + k: Nearest neighbours per destination, used when policy is "knn". + r: Distance cutoff in Angstroms, used when policy is "radius". + max_neighbors: Per-source cap on radius results. batch_src: (N_src,) batch assignment for source nodes, or None. batch_dst: (N_dst,) batch assignment for destination nodes, or None. @@ -202,27 +249,50 @@ def build_knn_edges( if src_pos.numel() == 0 or dst_pos.numel() == 0: return torch.empty(2, 0, dtype=torch.long, device=src_pos.device) - # 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) + # Same object + homogeneous = src_pos is dst_pos - # remove self-edges if homogeneous - if src_pos.data_ptr() == dst_pos.data_ptr(): - mask = idx[0] != idx[1] - idx = idx[:, mask] + if policy == "knn": + # Asked for each destination's nearest sources, so sources come back second. + dst_idx, src_idx = knn( + x=src_pos, y=dst_pos, k=k, batch_x=batch_src, batch_y=batch_dst + ) + edge_index = torch.stack((src_idx, dst_idx), dim=0) + if homogeneous: + edge_index = edge_index[:, edge_index[0] != edge_index[1]] + return edge_index.unique(dim=1) + + # Cap against the number of reachable counterparts, minus the self-edge that + # a homogeneous query would otherwise spend a slot on. + num_candidates = dst_pos.size(0) - 1 if homogeneous else dst_pos.size(0) + cap = max(1, min(num_candidates, max_neighbors)) + + if homogeneous: + return radius_graph( + src_pos, r=r, batch=batch_src, loop=False, max_num_neighbors=cap + ) - return idx.unique(dim=1) + # Asked for each source's neighbours within r, so sources already come back first. + return radius( + x=dst_pos, + y=src_pos, + r=r, + batch_x=batch_dst, + batch_y=batch_src, + max_num_neighbors=cap, + ) class ProteinWaterUpdate(nn.Module): """ - Heterogeneous GVP message passing with all four edge types: - - protein -> water (pw) - - water -> water (ww) - - protein -> protein (pp) - - water -> protein (wp) + Heterogeneous GVP message passing over the active edge types: + - protein -> water (pw) always active + - protein -> protein (pp) always active + - water -> water (ww) ablatable + - water -> protein (wp) ablatable + + Which are active is fixed at construction via `etypes`; see + `constants.get_active_edge_types`. """ def __init__( @@ -236,6 +306,15 @@ def __init__( vector_gate=True, aggr_edges="sum", use_dst_feats=True, + etypes: list[tuple[str, str, str]] | None = None, + cutoff: float = 8.0, + max_neighbors: int = 256, + dynamic_edge_policy: str = "radius", + sampling_strategy: str = "uniform_ball", + knn_fallback_k: int = 8, + k_pw: int = 12, + k_ww: int = 8, + k_wp: int = 8, ): """ Initialize heterogeneous protein-water message passing module. @@ -252,12 +331,56 @@ def __init__( vector_gate: Whether to use vector gating in GVP layers aggr_edges: Edge aggregation method ('sum' or 'mean') use_dst_feats: Whether to include destination features in messages + etypes: Active edge types. Defaults to ALL_EDGE_TYPES. Build with + `constants.get_active_edge_types` to ablate WW/WP. + cutoff: Distance cutoff in Angstroms for radius edges. + max_neighbors: Cap on neighbours per source node in radius queries, + bounding edge count and runtime on dense structures. + dynamic_edge_policy: "auto" or one of DYNAMIC_EDGE_POLICIES. "radius" + connects everything within `cutoff`; "knn" connects a fixed + number of nearest neighbours using k_pw/k_ww/k_wp. + sampling_strategy: Prior the run uses, consulted only to resolve + "auto"; see `resolve_edge_policy`. + knn_fallback_k: Under "knn_if_isolated", attach this many nearest + neighbours to any water the radius query left with no edges. 0 + disables the rescue. Ignored under the other policies. + k_pw: Nearest neighbours for protein -> water edges under "knn". + k_ww: Nearest neighbours for water -> water edges under "knn". + k_wp: Nearest neighbours for water -> protein edges under "knn". + + Raises: + ValueError: If `dynamic_edge_policy` is not a known value, or if + `knn_fallback_k` is negative. """ super().__init__() # Unpack hidden dimensions: s_h = scalar hidden dim, v_h = vector hidden dim s_h, v_h = hidden_dims - etypes = ALL_EDGE_TYPES + if knn_fallback_k < 0: + raise ValueError(f"knn_fallback_k must be >= 0, got {knn_fallback_k}") + + # build_edges only knows how to construct the four known relations, and + # HeteroConv would KeyError on any other, so reject it up front. + unknown = [et for et in (etypes or []) if et not in ALL_EDGE_TYPES] + if unknown: + raise ValueError( + f"etypes must be a subset of {ALL_EDGE_TYPES}, got unknown {unknown}" + ) + + self.cutoff = cutoff + self.max_neighbors = max_neighbors + resolved = resolve_edge_policy(dynamic_edge_policy, sampling_strategy) + # Same edges as "radius"; the difference is the extra pass for nodes left with none. + self.rescue_isolated = resolved == "knn_if_isolated" and knn_fallback_k > 0 + self.dynamic_edge_policy = ( + "radius" if resolved == "knn_if_isolated" else resolved + ) + self.knn_fallback_k = knn_fallback_k + self.k_pw = k_pw + self.k_ww = k_ww + self.k_wp = k_wp + + etypes = ALL_EDGE_TYPES if etypes is None else etypes self.blocks = nn.ModuleList( [ @@ -279,30 +402,92 @@ def __init__( ) self.etypes = etypes - def build_edges( - self, data: HeteroData, k_pw: int = 12, k_ww: int = 8, k_wp: int = 8 - ) -> dict[tuple[str, str, str], torch.Tensor]: + def _add_knn_fallback( + self, + edge_index: torch.Tensor, + src_pos: torch.Tensor, + dst_pos: torch.Tensor, + batch_src: torch.Tensor | None, + batch_dst: torch.Tensor | None, + isolate_axis: int, + ) -> torch.Tensor: + """ + Attach KNN edges for nodes the radius query left with no edges. + + A radius query strands any node with nothing inside `cutoff`. Those nodes + would reach the GVP blocks with no incoming messages, so they are + reconnected to their `knn_fallback_k` nearest counterparts regardless of + distance. + + Args: + edge_index: (2, E) radius edges, source row 0, destination row 1. + src_pos: (N_src, 3) source positions. + dst_pos: (N_dst, 3) destination positions. + batch_src: (N_src,) batch assignment for sources, or None. + batch_dst: (N_dst,) batch assignment for destinations, or None. + isolate_axis: Row to check for stranded nodes -- 0 for sources, + 1 for destinations. + + Returns: + (2, E') edge index with fallback edges merged in and deduplicated. """ - Build KNN edges for protein-water interactions. + device = src_pos.device + num_nodes = dst_pos.size(0) if isolate_axis == 1 else src_pos.size(0) + if num_nodes == 0: + return edge_index + + connected = torch.zeros(num_nodes, dtype=torch.bool, device=device) + if edge_index.numel() > 0: + connected[edge_index[isolate_axis].unique()] = True + isolated = (~connected).nonzero(as_tuple=False).flatten() + if isolated.numel() == 0: + return edge_index + + # Query only the stranded nodes, then lift the returned local indices + # back into the full node set via `isolated`. + if isolate_axis == 1: + fallback = build_dynamic_edges( + src_pos, + dst_pos[isolated], + policy="knn", + k=max(1, min(self.knn_fallback_k, src_pos.size(0))), + r=self.cutoff, + batch_src=batch_src, + batch_dst=batch_dst[isolated] if batch_dst is not None else None, + ) + fallback = torch.stack((fallback[0], isolated[fallback[1]]), dim=0) + else: + fallback = build_dynamic_edges( + dst_pos, + src_pos[isolated], + policy="knn", + k=max(1, min(self.knn_fallback_k, dst_pos.size(0))), + r=self.cutoff, + batch_src=batch_dst, + batch_dst=batch_src[isolated] if batch_src is not None else None, + ) + fallback = torch.stack((isolated[fallback[1]], fallback[0]), dim=0) + + if edge_index.numel() == 0: + return fallback + return torch.cat((edge_index, fallback), dim=1).unique(dim=1) - For protein->water edges, we take the union of: - - KNN(protein -> water): k nearest waters per protein - - KNN(water -> protein) reversed: k nearest proteins per water - This ensures every water has at least k_pw protein neighbors. + def build_edges(self, data: HeteroData) -> dict[tuple[str, str, str], torch.Tensor]: + """ + Build the edge set for one batch under the active policy. - PP edges are read from the dataset (cached at preprocessing time). + PP and PW edges are read from the dataset when cached at preprocessing + time and built on the fly otherwise. WW and WP are always dynamic, and + are skipped entirely when ablated out of `etypes`. Args: - data: HeteroData with 'protein' and 'water' node types containing positions - k_pw: Number of nearest neighbors for protein-water edges - k_ww: Number of nearest neighbors for water-water edges - k_wp: Number of nearest neighbors for water-protein edges + data: HeteroData with 'protein' and 'water' node types containing + positions, optionally carrying cached PP/PW edges. Returns: - Dict mapping edge type tuples to (2, E) edge index tensors + Dict mapping each active edge type to a (2, E) edge index tensor. """ edge_index_dict: dict[tuple[str, str, str], torch.Tensor] = {} - device = data["protein"].pos.device batch_p = data["protein"].batch if "batch" in data["protein"] else None batch_w = data["water"].batch if "batch" in data["water"] else None @@ -310,56 +495,77 @@ def build_edges( pos_p = data["protein"].pos pos_w = data["water"].pos - # protein -> water - if pos_p.numel() > 0 and pos_w.numel() > 0: - # p->w - ei_pw = build_knn_edges( - pos_p, pos_w, k=k_pw, batch_src=batch_p, batch_dst=batch_w - ) - # w->p then reverse - ei_wp = build_knn_edges( - pos_w, pos_p, k=k_pw, batch_src=batch_w, batch_dst=batch_p - ) - ei_wp_reversed = ei_wp.flip(0) - # union - ei_pw_union = torch.cat([ei_pw, ei_wp_reversed], dim=1).unique(dim=1) - edge_index_dict[EDGE_PW] = ei_pw_union - else: - edge_index_dict[EDGE_PW] = torch.empty( - 2, 0, dtype=torch.long, device=device + # Only protein-water edges get the extra pass; waters keep protein context anyway. + rescue = self.rescue_isolated + + # protein -> water (water is the destination, so it is row 1) + if EDGE_PW in self.etypes: + if EDGE_PW in data.edge_types: + edge_index_dict[EDGE_PW] = data[EDGE_PW].edge_index + else: + ei = build_dynamic_edges( + pos_p, + pos_w, + policy=self.dynamic_edge_policy, + k=self.k_pw, + r=self.cutoff, + max_neighbors=self.max_neighbors, + batch_src=batch_p, + batch_dst=batch_w, + ) + if rescue: + ei = self._add_knn_fallback( + ei, pos_p, pos_w, batch_p, batch_w, isolate_axis=1 + ) + edge_index_dict[EDGE_PW] = ei + + # protein -> protein (cached from the dataset in every normal run) + if EDGE_PP in self.etypes: + edge_index_dict[EDGE_PP] = ( + data[EDGE_PP].edge_index + if EDGE_PP in data.edge_types + else build_dynamic_edges( + pos_p, + pos_p, + policy=self.dynamic_edge_policy, + k=self.k_pw, + r=self.cutoff, + max_neighbors=self.max_neighbors, + batch_src=batch_p, + batch_dst=batch_p, + ) ) # water -> water - if pos_w.numel() > 0: - edge_index_dict[EDGE_WW] = build_knn_edges( - pos_w, pos_w, k=k_ww, batch_src=batch_w, batch_dst=batch_w - ) - else: - edge_index_dict[EDGE_WW] = torch.empty( - 2, 0, dtype=torch.long, device=device + if EDGE_WW in self.etypes: + edge_index_dict[EDGE_WW] = build_dynamic_edges( + pos_w, + pos_w, + policy=self.dynamic_edge_policy, + k=self.k_ww, + r=self.cutoff, + max_neighbors=self.max_neighbors, + batch_src=batch_w, + batch_dst=batch_w, ) - # protein-protein edges (cached from dataset) - if EDGE_PP in data.edge_types: - edge_index_dict[EDGE_PP] = data[EDGE_PP].edge_index - else: - edge_index_dict[EDGE_PP] = build_knn_edges( - pos_p, pos_p, k=k_pw, batch_src=batch_p, batch_dst=batch_p - ) - - # water -> protein - if pos_w.numel() > 0 and pos_p.numel() > 0: - edge_index_dict[EDGE_WP] = build_knn_edges( - pos_w, pos_p, k=k_wp, batch_src=batch_w, batch_dst=batch_p - ) - else: - edge_index_dict[EDGE_WP] = torch.empty( - 2, 0, dtype=torch.long, device=device + # water -> protein (water is the source, so it is row 0) + if EDGE_WP in self.etypes: + ei = build_dynamic_edges( + pos_w, + pos_p, + policy=self.dynamic_edge_policy, + k=self.k_wp, + r=self.cutoff, + max_neighbors=self.max_neighbors, + batch_src=batch_w, + batch_dst=batch_p, ) - - for et in self.etypes: - if et not in edge_index_dict: - edge_index_dict[et] = torch.empty(2, 0, dtype=torch.long, device=device) + if rescue: + ei = self._add_knn_fallback( + ei, pos_w, pos_p, batch_w, batch_p, isolate_axis=0 + ) + edge_index_dict[EDGE_WP] = ei return edge_index_dict @@ -367,9 +573,6 @@ def forward( self, x_dict: dict[str, tuple[torch.Tensor, torch.Tensor]], data: HeteroData, - k_pw: int = 12, - k_ww: int = 8, - k_wp: int = 8, pp_edge_attr: tuple | None = None, ): """ @@ -380,9 +583,6 @@ def forward( - 'protein': (s_p, v_p) where s_p is (N_p, scalar_dim), v_p is (N_p, vector_dim, 3) - 'water': (s_w, v_w) where s_w is (N_w, scalar_dim), v_w is (N_w, vector_dim, 3) data: HeteroData with 'protein' and 'water' node positions - k_pw: Number of nearest neighbors for protein-water edges - k_ww: Number of nearest neighbors for water-water edges - k_wp: Number of nearest neighbors for water-protein edges pp_edge_attr: Optional encoder-learned edge features (s_edge, V_edge) for PP edges. If provided, uses encoder-learned scalar features (s_edge) combined with cached edge direction unit vectors (edge_unit_vectors, pre-normalized at preprocessing). @@ -393,12 +593,7 @@ def forward( """ pos_dict = {nt: data[nt].pos for nt in data.node_types if "pos" in data[nt]} - edge_index_dict = self.build_edges( - data, - k_pw=k_pw, - k_ww=k_ww, - k_wp=k_wp, - ) + edge_index_dict = self.build_edges(data) # PP edge features: encoder-provided take priority over cached geometric features cached_edge_attr_dict = {} @@ -448,10 +643,17 @@ def __init__( n_message_gvps: int = 2, n_update_gvps: int = 2, vector_gate: bool = True, + water_input_dim: int = 16, # 1 hot with oxygen, same as encoder + cutoff: float = 8.0, + max_neighbors: int = 256, + dynamic_edge_policy: str = "radius", + sampling_strategy: str = "uniform_ball", + knn_fallback_k: int = 8, + disable_ww: bool = False, + disable_wp: bool = False, k_pw: int = 12, k_ww: int = 8, k_wp: int = 8, - water_input_dim: int = 16, # 1 hot with oxygen, same as encoder ): """ Initialize end-to-end flow model for water placement. @@ -467,10 +669,19 @@ def __init__( n_update_gvps: Number of GVP modules in the node update function (applied after aggregating messages from all edge types). Default: 2 vector_gate: Whether to use vector gating in GVP layers. Default: True - k_pw: K nearest neighbors for protein-water edges. Default: 12 - k_ww: K nearest neighbors for water-water edges. Default: 8 - k_wp: K nearest neighbors for water-protein edges. Default: 8 water_input_dim: Input dimension for water node features. Default: 16 + cutoff: Distance cutoff in Angstroms for radius edges. Default: 8.0 + max_neighbors: Per-source cap on radius results. Default: 256 + dynamic_edge_policy: How water-touching edges are built, one of + DYNAMIC_EDGE_POLICIES. Default: "radius" + knn_fallback_k: Nearest neighbours attached to waters the radius + query stranded; 0 disables the rescue. Default: 8 + disable_ww: Ablate water -> water edges. Default: False + disable_wp: Ablate water -> protein edges. Default: False + k_pw: K nearest neighbors for protein-water edges under the "knn" + policy. Default: 12 + k_ww: K nearest neighbors for water-water edges under "knn". Default: 8 + k_wp: K nearest neighbors for water-protein edges under "knn". Default: 8 """ super().__init__() self.encoder = encoder @@ -481,9 +692,10 @@ def __init__( self.n_message_gvps = n_message_gvps self.n_update_gvps = n_update_gvps self.vector_gate = vector_gate - self.k_pw = k_pw - self.k_ww = k_ww - self.k_wp = k_wp + # Read back by FlowMatcher for the water prior's sampling radius. The rest + # of the edge configuration is owned by `self.updater` and deliberately not + # mirrored here -- two copies would be two things to keep in sync. + self.cutoff = cutoff s_h, v_h = hidden_dims @@ -520,6 +732,15 @@ def __init__( vector_gate=vector_gate, aggr_edges="sum", use_dst_feats=True, + etypes=get_active_edge_types(disable_ww=disable_ww, disable_wp=disable_wp), + cutoff=cutoff, + max_neighbors=max_neighbors, + dynamic_edge_policy=dynamic_edge_policy, + sampling_strategy=sampling_strategy, + knn_fallback_k=knn_fallback_k, + k_pw=k_pw, + k_ww=k_ww, + k_wp=k_wp, ) self.sc_vec_encoder = GVP( @@ -628,9 +849,6 @@ def forward( x_dict = self.updater( x_dict, data, - k_pw=self.k_pw, - k_ww=self.k_ww, - k_wp=self.k_wp, pp_edge_attr=pp_edge_attr, ) @@ -645,7 +863,6 @@ class FlowMatcher: """ SAMPLING_STRATEGIES = ("uniform_ball", "scaled_gaussian") - DYNAMIC_EDGE_POLICIES = ("auto", "radius", "knn_if_isolated") def __init__( self, @@ -657,7 +874,6 @@ def __init__( sigma_distort: float = 0.5, loss_eps: float = 1e-3, sampling_strategy: str = "uniform_ball", - dynamic_edge_policy: str = "auto", ): """ Initialize flow matcher for training and inference. @@ -673,18 +889,16 @@ def __init__( sampling_strategy: Source distribution for flow matching noise. "uniform_ball" samples uniformly in balls around protein atoms. "scaled_gaussian" samples from N(0, sigma^2*I). - dynamic_edge_policy: Runtime policy for dynamic water-edge building. + + Note: + Edge construction is configured on the model, not here, so training + and integration always build edges the same way. """ if sampling_strategy not in self.SAMPLING_STRATEGIES: raise ValueError( f"sampling_strategy must be one of {self.SAMPLING_STRATEGIES}, " f"got '{sampling_strategy}'" ) - if dynamic_edge_policy not in self.DYNAMIC_EDGE_POLICIES: - raise ValueError( - f"dynamic_edge_policy must be one of {self.DYNAMIC_EDGE_POLICIES}, " - f"got '{dynamic_edge_policy}'" - ) self.model = model self.p_self_cond = p_self_cond self.use_distortion = use_distortion @@ -694,7 +908,6 @@ def __init__( self.loss_eps = loss_eps self.graph_cutoff = getattr(model, "cutoff", 8.0) self.sampling_strategy = sampling_strategy - self.dynamic_edge_policy = dynamic_edge_policy @staticmethod def _num_graphs(data: HeteroData | Batch) -> int: @@ -732,14 +945,6 @@ def _sample_waters( dtype=batch_data["protein"].pos.dtype, ) - def _effective_dynamic_edge_policy(self) -> str: - """Resolve the dynamic edge policy for the current sampling strategy.""" - if self.dynamic_edge_policy == "auto": - if self.sampling_strategy == "scaled_gaussian": - return "knn_if_isolated" - return "radius" - return self.dynamic_edge_policy - @staticmethod def compute_sigma(data: HeteroData) -> float: """ @@ -822,7 +1027,6 @@ def training_step( self.model.train() device = batch["protein"].pos.device - batch.dynamic_edge_policy = self._effective_dynamic_edge_policy() x1 = batch["water"].pos batch_w = batch["water"].batch @@ -917,7 +1121,6 @@ def validation_step(self, batch: HeteroData) -> dict[str, float]: """ self.model.eval() device = batch["protein"].pos.device - batch.dynamic_edge_policy = self._effective_dynamic_edge_policy() x1 = batch["water"].pos batch_w = batch["water"].batch diff --git a/tests/test_flow.py b/tests/test_flow.py index 9d2f62a..222333a 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -10,14 +10,22 @@ import torch import torch.nn.functional as F from torch_geometric.data import Batch, Data, HeteroData -from torch_geometric.nn import knn - +from torch_geometric.nn import knn, radius + +from src.constants import ( + ALL_EDGE_TYPES, + EDGE_PP, + EDGE_PW, + EDGE_WW, + get_active_edge_types, +) from src.flow import ( _batch_from_counts, - build_knn_edges, + build_dynamic_edges, FlowMatcher, FlowWaterGVP, ProteinWaterUpdate, + resolve_edge_policy, sample_waters_scaled_gaussian, sample_waters_uniform_ball, ) @@ -103,14 +111,14 @@ def mock_forward(data): @pytest.mark.unit -class TestBuildKnnEdges: +class TestBuildDynamicEdgesKnn: def test_basic_knn(self, device): src = torch.tensor( [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]], device=device ) dst = torch.tensor([[0.5, 0.0, 0.0], [1.5, 0.0, 0.0]], device=device) - edges = build_knn_edges(src, dst, k=2) + edges = build_dynamic_edges(src, dst, k=2, policy="knn", r=8.0) assert edges.shape[0] == 2 assert edges.shape[1] >= 4 # At least 2 dst points × 2 neighbors each @@ -128,7 +136,7 @@ def test_empty_src(self, device): src = torch.empty(0, 3, device=device) dst = torch.randn(5, 3, device=device) - edges = build_knn_edges(src, dst, k=3) + edges = build_dynamic_edges(src, dst, k=3, policy="knn", r=8.0) assert edges.shape == (2, 0) @@ -136,14 +144,14 @@ def test_empty_dst(self, device): src = torch.randn(5, 3, device=device) dst = torch.empty(0, 3, device=device) - edges = build_knn_edges(src, dst, k=3) + edges = build_dynamic_edges(src, dst, k=3, policy="knn", r=8.0) assert edges.shape == (2, 0) def test_self_edges_removed(self, device): pos = torch.randn(10, 3, device=device) - edges = build_knn_edges(pos, pos, k=5) + edges = build_dynamic_edges(pos, pos, k=5, policy="knn", r=8.0) # No self-loops assert (edges[0] != edges[1]).all() @@ -154,19 +162,26 @@ def test_with_batch(self, device): batch_src = torch.cat([torch.zeros(5), torch.ones(5)]).long().to(device) batch_dst = torch.cat([torch.zeros(4), torch.ones(4)]).long().to(device) - edges = build_knn_edges(src, dst, k=3, batch_src=batch_src, batch_dst=batch_dst) + edges = build_dynamic_edges( + src, dst, k=3, batch_src=batch_src, batch_dst=batch_dst, policy="knn", r=8.0 + ) assert edges.shape[0] == 2 assert edges.shape[1] > 0 @pytest.mark.unit -class TestBuildKnnEdgesDirection: - """Exact-set KNN direction tests on asymmetric geometry. +class TestBuildDynamicEdgesDirection: + """Row-convention tests for both policies, 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. + The KNN cases spread srcs out and put both dsts 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. The radius cases + keep the two index ranges different sizes for the same reason. + + Both policies feed row 0 = src, row 1 = dst, but reach it differently: KNN + swaps the rows PyG returns, radius inverts the arguments instead. Each needs + its own pin. """ SRC = [[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [20.0, 0.0, 0.0]] @@ -178,7 +193,7 @@ def test_exact_edge_set_is_per_destination(self, device): src = torch.tensor(self.SRC, device=device) dst = torch.tensor(self.DST, device=device) - edges = build_knn_edges(src, dst, k=1) + edges = build_dynamic_edges(src, dst, k=1, policy="knn", r=8.0) edge_set = set(zip(edges[0].tolist(), edges[1].tolist())) assert edge_set == {(0, 0), (0, 1)}, f"got {sorted(edge_set)}" @@ -190,7 +205,7 @@ def test_every_destination_is_covered(self, device): dst = torch.tensor(self.DST, device=device) k = 2 - edges = build_knn_edges(src, dst, k=k) + edges = build_dynamic_edges(src, dst, k=k, policy="knn", r=8.0) dst_row = edges[1] for d in range(len(self.DST)): @@ -211,14 +226,14 @@ def test_output_rows_are_src_then_dst(self, 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) + edges = build_dynamic_edges(src, dst, k=1, policy="knn", r=8.0) 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.""" + build_dynamic_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 @@ -227,6 +242,78 @@ def test_torch_geometric_knn_row_convention_unchanged(self, device): 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] + def test_torch_geometric_radius_row_convention_unchanged(self, device): + """Pin radius's rows: row 0 = y (query), row 1 = x (neighbor) -- the same + convention as knn. build_dynamic_edges relies on this by passing src as + `y`, so it lands in row 0 with no swap. Index ranges are kept distinct + (3 x's vs 2 y's) so a flip cannot pass silently.""" + x = torch.tensor( + [[0.0, 0.0], [0.1, 0.0], [0.2, 0.0], [10.0, 0.0], [10.1, 0.0]], + device=device, + ) # N=5 + y = torch.tensor([[0.05, 0.0], [10.05, 0.0]], device=device) # M=2 + + out = radius(x, y, r=0.5) + + assert out[0].max().item() <= 1 # queries (y), max index 1 + assert out[1].max().item() == 4 # neighbours (x), reaches index 4 + + def test_radius_rows_are_src_then_dst(self, device): + """Row 0 = src, row 1 = dst under the radius policy too, pinned by index + range: 2 srcs and 5 dsts, so a swap puts an out-of-range index in row 0.""" + src = torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]], device=device) + dst = torch.tensor( + [ + [0.1, 0.0, 0.0], + [0.2, 0.0, 0.0], + [0.3, 0.0, 0.0], + [10.1, 0.0, 0.0], + [10.2, 0.0, 0.0], + ], + device=device, + ) + + edges = build_dynamic_edges(src, dst, policy="radius", k=1, r=1.0) + + assert edges[0].max().item() < len(src) + assert edges[1].max().item() == 4 # dst[4]; >= len(src), so a swap breaks + + def test_radius_self_graph_has_no_self_loops(self, device): + pos = torch.tensor( + [[0.0, 0.0, 0.0], [0.5, 0.0, 0.0], [99.0, 0.0, 0.0]], device=device + ) + + edges = build_dynamic_edges(pos, pos, policy="radius", k=1, r=1.0) + + assert (edges[0] != edges[1]).all() + assert set(zip(edges[0].tolist(), edges[1].tolist())) == {(0, 1), (1, 0)} + + +@pytest.mark.unit +class TestResolveEdgePolicy: + """Replaying a recorded config must reproduce the original run. Every run on + record carries "auto", so rejecting it would strand all of them.""" + + @pytest.mark.parametrize( + "policy,strategy,expected", + [ + # "auto" reads off the prior: the uniform ball never strands a water, + # the Gaussian can, so only the latter earns the rescue. + ("auto", "uniform_ball", "radius"), + ("auto", "scaled_gaussian", "knn_if_isolated"), + # An explicit policy ignores the prior. + ("radius", "scaled_gaussian", "radius"), + ("knn", "uniform_ball", "knn"), + ("knn_if_isolated", "uniform_ball", "knn_if_isolated"), + ], + ) + def test_resolves(self, policy, strategy, expected): + assert resolve_edge_policy(policy, strategy) == expected + + def test_unknown_value_raises(self): + with pytest.raises(ValueError, match="dynamic_edge_policy"): + resolve_edge_policy("knn_if_lonely") + @pytest.mark.unit class TestMakeEncoderData: @@ -278,20 +365,116 @@ def test_init(self): assert ("protein", "pw", "water") in updater.etypes assert ("water", "ww", "water") in updater.etypes - def test_init_always_includes_all_edge_types(self): + def test_init_defaults_to_all_edge_types(self): updater = ProteinWaterUpdate( hidden_dims=(128, 16), rbf_dim=16, layers=2, ) - assert ("protein", "pp", "protein") in updater.etypes - assert ("water", "wp", "protein") in updater.etypes + assert set(updater.etypes) == set(ALL_EDGE_TYPES) + + def test_init_rejects_unknown_etype(self): + with pytest.raises(ValueError, match="subset"): + ProteinWaterUpdate( + hidden_dims=(128, 16), + layers=1, + etypes=[EDGE_PW, ("water", "xx", "water")], + ) + + def test_knn_if_isolated_builds_a_radius_graph_and_rescues(self): + """It differs from plain radius only by the rescue; build_edges is handed + "radius" either way.""" + updater = ProteinWaterUpdate( + hidden_dims=(128, 16), layers=1, dynamic_edge_policy="knn_if_isolated" + ) + + assert updater.dynamic_edge_policy == "radius" + assert updater.rescue_isolated + + def test_plain_radius_does_not_rescue(self): + """The rescue belongs to knn_if_isolated, so a positive knn_fallback_k + must not switch it on by itself.""" + updater = ProteinWaterUpdate( + hidden_dims=(128, 16), + layers=1, + dynamic_edge_policy="radius", + knn_fallback_k=8, + ) + + assert not updater.rescue_isolated + + def test_init_rejects_negative_fallback_k(self): + with pytest.raises(ValueError, match="knn_fallback_k"): + ProteinWaterUpdate(hidden_dims=(128, 16), layers=1, knn_fallback_k=-1) + + def test_ablated_etypes_are_absent_from_edge_dict(self, simple_hetero_data): + updater = ProteinWaterUpdate( + hidden_dims=(128, 16), + layers=1, + etypes=get_active_edge_types(disable_ww=True, disable_wp=True), + ) + + edge_dict = updater.build_edges(simple_hetero_data) + + assert set(edge_dict) == {EDGE_PW, EDGE_PP} + + def test_radius_strands_far_water_and_fallback_rescues_it(self, device): + """A water parked far outside `cutoff` gets no radius PW edges. With the + rescue enabled it is reconnected to its nearest protein atoms anyway.""" + data = HeteroData() + data["protein"].pos = torch.randn(10, 3, device=device) + data["water"].pos = torch.cat( + [torch.randn(4, 3, device=device), torch.full((1, 3), 500.0, device=device)] + ) + + def stranded(knn_fallback_k): + updater = ProteinWaterUpdate( + hidden_dims=(32, 4), + layers=1, + cutoff=8.0, + dynamic_edge_policy="knn_if_isolated", + knn_fallback_k=knn_fallback_k, + ) + edge_index = updater.build_edges(data)[EDGE_PW] + connected = torch.zeros(5, dtype=torch.bool, device=device) + connected[edge_index[1].unique()] = True + return int((~connected).sum()) + + assert stranded(knn_fallback_k=0) == 1 + assert stranded(knn_fallback_k=3) == 0 + + def test_knn_policy_never_strands_a_water(self, device): + """KNN budgets per destination, so distance cannot isolate a node and the + rescue is unnecessary.""" + data = HeteroData() + data["protein"].pos = torch.randn(10, 3, device=device) + data["water"].pos = torch.cat( + [torch.randn(4, 3, device=device), torch.full((1, 3), 500.0, device=device)] + ) + updater = ProteinWaterUpdate( + hidden_dims=(32, 4), layers=1, dynamic_edge_policy="knn", k_pw=3 + ) + + edge_index = updater.build_edges(data)[EDGE_PW] + + assert set(edge_index[1].tolist()) == {0, 1, 2, 3, 4} + + def test_cached_pw_edges_are_used_verbatim(self, simple_hetero_data, device): + """A dataset that precomputes PW edges (the confidence pipeline) must not + have them rebuilt underneath it.""" + cached = torch.tensor([[0, 1], [2, 3]], device=device) + simple_hetero_data[EDGE_PW].edge_index = cached + updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1) + + edge_dict = updater.build_edges(simple_hetero_data) + + assert torch.equal(edge_dict[EDGE_PW], cached) def test_build_edges(self, simple_hetero_data): updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1) - edge_dict = updater.build_edges(simple_hetero_data, k_pw=4, k_ww=3) + edge_dict = updater.build_edges(simple_hetero_data) assert ("protein", "pw", "water") in edge_dict assert ("water", "ww", "water") in edge_dict @@ -482,22 +665,24 @@ def test_validation_step(self, flow_matcher, simple_hetero_data): assert "rmsd" in result assert result["loss"] >= 0 - def test_scaled_gaussian_auto_policy_enables_knn_fallback( - self, device, gvp_encoder - ): + def test_edge_config_propagates_to_updater(self, device, gvp_encoder): + """Edge construction is configured once on the model and used for both + training and integration, so it must reach the updater that builds the + edges. An earlier design resolved it per batch and let the two diverge.""" model = FlowWaterGVP( encoder=gvp_encoder, hidden_dims=(64, 8), layers=1, + dynamic_edge_policy="knn", + cutoff=6.0, + knn_fallback_k=0, + disable_ww=True, ).to(device) - flow_matcher = FlowMatcher( - model, - sampling_strategy="scaled_gaussian", - dynamic_edge_policy="auto", - ) - - assert flow_matcher._effective_dynamic_edge_policy() == "knn_if_isolated" + assert model.updater.dynamic_edge_policy == "knn" + assert model.updater.cutoff == 6.0 + assert model.updater.knn_fallback_k == 0 + assert EDGE_WW not in model.updater.etypes @pytest.mark.slow def test_euler_integrate(self, flow_matcher, simple_hetero_data, device): @@ -943,7 +1128,7 @@ def test_all_waters_have_protein_edges(self, simple_hetero_data): """Ensure every water has at least one protein-water edge.""" updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1) - edge_dict = updater.build_edges(simple_hetero_data, k_pw=4, k_ww=3) + edge_dict = updater.build_edges(simple_hetero_data) pw_edges = edge_dict[("protein", "pw", "water")] n_water = simple_hetero_data["water"].num_nodes @@ -958,7 +1143,7 @@ def test_all_waters_have_water_edges(self, simple_hetero_data): """Ensure every water has at least one water-water edge (if multiple waters exist).""" updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1) - edge_dict = updater.build_edges(simple_hetero_data, k_pw=4, k_ww=3) + edge_dict = updater.build_edges(simple_hetero_data) ww_edges = edge_dict[("water", "ww", "water")] n_water = simple_hetero_data["water"].num_nodes @@ -977,7 +1162,7 @@ def test_batched_waters_have_edges(self, batched_hetero_data): """Ensure all waters in a batched graph have edges.""" updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1) - edge_dict = updater.build_edges(batched_hetero_data, k_pw=4, k_ww=3) + edge_dict = updater.build_edges(batched_hetero_data) pw_edges = edge_dict[("protein", "pw", "water")] ww_edges = edge_dict[("water", "ww", "water")] @@ -1012,7 +1197,7 @@ def test_single_water_has_protein_edges_no_water_edges(self, device): ) updater = ProteinWaterUpdate(hidden_dims=(128, 16), layers=1) - edge_dict = updater.build_edges(data, k_pw=4, k_ww=3) + edge_dict = updater.build_edges(data) pw_edges = edge_dict[("protein", "pw", "water")] ww_edges = edge_dict[("water", "ww", "water")] diff --git a/tests/test_forward.py b/tests/test_forward.py index 1653bb0..9a9c893 100644 --- a/tests/test_forward.py +++ b/tests/test_forward.py @@ -183,8 +183,6 @@ def test_forward_pass_no_nan_with_module_hooks(device): encoder=encoder, hidden_dims=(64, 8), layers=2, - k_pw=8, # keep <= n_water_per - k_ww=8, # keep <= n_water_per ).to(device) # Quick pre-check: protein encoder input features created from pp edges @@ -193,8 +191,8 @@ def test_forward_pass_no_nan_with_module_hooks(device): enc_data.edge_index, enc_data.x.size(0), enc_data.x.size(0), "pp edge_index" ) - # Also validate knn edges are sane (catches orientation / k issues) - edge_dict = model.updater.build_edges(data, k_pw=model.k_pw, k_ww=model.k_ww) + # Also validate dynamic edges are sane (catches orientation / cutoff issues) + edge_dict = model.updater.build_edges(data) assert_edge_index_in_range( edge_dict[("protein", "pw", "water")], data["protein"].pos.size(0), diff --git a/tests/test_train_config.py b/tests/test_train_config.py index b59a1a4..b7e3f2e 100644 --- a/tests/test_train_config.py +++ b/tests/test_train_config.py @@ -131,6 +131,33 @@ def test_inference_build_model_from_config_uses_embedding_dim(device): assert model.encoder.output_dims == (128, 0) +def test_inference_build_model_from_config_replays_recorded_edge_policy(device): + """Every recorded config carries "auto". Replaying one must build a model, + not raise, and must land on the radius path those runs actually used.""" + config = { + "encoder_type": "slae", + "hidden_s": 128, + "hidden_v": 32, + "flow_layers": 2, + "node_scalar_in": 16, + "embedding_dim": 128, + "dynamic_edge_policy": "auto", + "knn_fallback_k": 8, + "cutoff": 8.0, + "max_neighbors": 256, + "disable_ww": True, + "disable_wp": True, + } + + model = build_model_from_config(config, device) + + assert model.updater.dynamic_edge_policy == "radius" + assert set(model.updater.etypes) == { + ("protein", "pw", "water"), + ("protein", "pp", "protein"), + } + + def test_parse_args_rejects_embedding_dim_for_gvp(monkeypatch): monkeypatch.setattr( "sys.argv",