Edge type flags

README.md

@@ -83,10 +83,13 @@ WaterFlow processes structure files through several stages to create training-re
 - 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
-  - `('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**
@@ -223,6 +226,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 |
+|-------------------------|-----------|
+| `auto` (default) | Resolves off the prior: `radius` under `uniform_ball`, `knn_if_isolated` under `scaled_gaussian` |
+| `radius` | 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`) |
+| `knn_if_isolated` | A `radius` graph plus a KNN rescue for any node the cutoff stranded |
+
+`radius` and `knn` 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_if_isolated` repairs that: any water the radius query stranded is
+reconnected to its `--knn_fallback_k` nearest protein atoms regardless of
+distance (`0` disables the rescue). Plain `radius` does *not* rescue, and the
+flag has no effect under `knn`, which cannot strand a node. `auto` picks
+`knn_if_isolated` for `scaled_gaussian` precisely because Gaussian samples can
+land outside every cutoff, whereas uniform-ball samples cannot.
+
+Set `--disable_ww` / `--disable_wp` to ablate those edge types; PW and PP are
+always active.
+
 ## Embedding Generation
 
 For `esm` and `slae` encoder types, you must precompute embeddings before training or inference.
@@ -296,6 +327,12 @@ 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 |
+| `--sampling_strategy` | `uniform_ball` | Flow prior: `uniform_ball` or `scaled_gaussian`; also resolves `--dynamic_edge_policy auto` |
+| `--dynamic_edge_policy` | `auto` | How water-touching edges are built: `auto`, `radius`, `knn`, or `knn_if_isolated` (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 under `knn_if_isolated`; `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

@@ -275,8 +275,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
@@ -433,6 +442,7 @@ def main():
     flow_matcher = FlowMatcher(
         model=model,
         p_self_cond=config.get("p_self_cond", 0.5),
+        sampling_strategy=config.get("sampling_strategy", "uniform_ball"),
     )
 
     # Load dataset

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,
@@ -226,8 +226,82 @@ 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)
+    # flow-matching prior
+    p.add_argument(
+        "--sampling_strategy",
+        type=str,
+        default="uniform_ball",
+        choices=["uniform_ball", "scaled_gaussian"],
+        help=(
+            "Source distribution for the flow prior. Also resolves "
+            "--dynamic_edge_policy auto (default: uniform_ball)"
+        ),
+    )
+
+    # 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(
@@ -570,8 +644,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
@@ -1037,6 +1120,7 @@ def main():
     flow_matcher = FlowMatcher(
         model=model,
         p_self_cond=args.p_self_cond,
+        sampling_strategy=args.sampling_strategy,
         use_distortion=args.use_distortion,
         p_distort=args.p_distort,
         t_distort=args.t_distort,

src/constants.py

@@ -14,6 +14,11 @@
 NUM_RBF = 16  # Number of RBF basis functions
 RBF_CUTOFF = 8.0  # Distance cutoff in Angstroms for RBF encoding
 
+# Default radius for dynamic edge construction. Kept equal to RBF_CUTOFF on
+# purpose: an edge longer than RBF_CUTOFF encodes to ~0 under the RBF, so
+# building edges beyond it would feed the message passing near-zero features.
+DEFAULT_EDGE_CUTOFF = RBF_CUTOFF
+
 # Edge type tuples: (src_node_type, edge_name, dst_node_type)
 EDGE_PP = ("protein", "pp", "protein")  # protein -> protein
 EDGE_WW = ("water", "ww", "water")  # water -> water
@@ -23,6 +28,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'