Edge type flags

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 |

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

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

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.<protein___pw___water>``),
+    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'