diff --git a/src/flow.py b/src/flow.py
index ad5f7e4..42f2953 100644
--- a/src/flow.py
+++ b/src/flow.py
@@ -17,15 +17,157 @@
 from torch import nn, Tensor
 from torch_geometric.data import Batch, HeteroData
 from torch_geometric.nn import knn
-from torch_scatter import scatter_mean
+from torch_scatter import scatter, scatter_mean
 from tqdm.auto import tqdm
 
-from src.constants import ALL_EDGE_TYPES, EDGE_PP, EDGE_PW, EDGE_WP, EDGE_WW, NUM_RBF
+from src.constants import (
+    ALL_EDGE_TYPES,
+    EDGE_PP,
+    EDGE_PW,
+    EDGE_WP,
+    EDGE_WW,
+    ELEM_IDX,
+    ELEMENT_VOCAB,
+    NUM_RBF,
+)
 from src.encoder_base import BaseProteinEncoder
 from src.gvp import GVP, GVPMultiEdgeConv
 from src.utils import ot_coupling
 
 
+def _batch_from_counts(num_waters: Tensor, device: torch.device) -> Tensor:
+    """
+    Build a graph-grouped batch vector from per-graph counts.
+
+    Args:
+        num_waters: (num_graphs,) water count per graph
+        device: Output device
+
+    Returns:
+        (sum(num_waters),) graph index per water, non-decreasing
+    """
+    return torch.repeat_interleave(
+        torch.arange(num_waters.numel(), device=device), num_waters.to(device)
+    )
+
+
+def sample_waters_uniform_ball(
+    protein_pos: Tensor,
+    batch_p: Tensor,
+    batch_w: Tensor,
+    cutoff: float = 8.0,
+    device: torch.device | None = None,
+) -> Tensor:
+    """
+    Sample water positions uniformly inside balls of radius *cutoff* centred
+    on randomly chosen protein atoms.
+
+    Every sample is guaranteed within *cutoff* of at least one protein atom.
+    No rejection sampling — runs in O(1) rounds, fully vectorised.
+
+    Args:
+        protein_pos: (N_protein, 3) protein coordinates for all graphs
+        batch_p: (N_protein,) graph indices for protein atoms
+        batch_w: (N_water,) graph index per water to sample. Positions are returned
+            in this order, so callers with existing water nodes can pass their own
+            batch vector and get samples aligned to it.
+        cutoff: Ball radius in Angstroms
+        device: Optional output device (defaults to protein_pos.device)
+
+    Returns:
+        water_pos: (N_water, 3) sampled positions, one per entry of batch_w
+    """
+    if device is None:
+        device = protein_pos.device
+
+    batch_w = batch_w.to(device)
+    total_waters = batch_w.numel()
+
+    if total_waters == 0:
+        return torch.empty(0, 3, dtype=protein_pos.dtype, device=device)
+
+    # offsets below assume protein atoms are grouped contiguously by graph;
+    # interleaved batch_p would pick anchors from the wrong graph.
+    batch_p = batch_p.to(device)
+    if batch_p.numel() > 1 and (batch_p[1:] < batch_p[:-1]).any():
+        raise ValueError("batch_p must be sorted (non-decreasing) by graph index.")
+
+    # cover every graph named by either side so the guard below can see empty ones
+    num_graphs = int(batch_w.max().item()) + 1
+    if batch_p.numel() > 0:
+        num_graphs = max(num_graphs, int(batch_p.max().item()) + 1)
+
+    # per-graph protein atom counts and cumulative offsets
+    num_p_per_graph = scatter(
+        torch.ones(batch_p.size(0), device=device, dtype=torch.long),
+        batch_p,
+        dim=0,
+        dim_size=num_graphs,
+        reduce="sum",
+    )
+
+    # fail fast: a graph that requests waters must have at least one protein atom.
+    # Otherwise graph_sizes is 0 and graph_offsets + local_idx would index into a
+    # neighbouring graph's atoms (or out of bounds) when picking anchors below.
+    graph_sizes = num_p_per_graph[batch_w]
+    if (graph_sizes == 0).any():
+        bad = batch_w[graph_sizes == 0].unique().tolist()
+        raise ValueError(
+            f"Cannot sample waters for graph(s) {bad}: they request waters "
+            "but have zero protein atoms."
+        )
+
+    offsets = torch.zeros(num_graphs + 1, dtype=torch.long, device=device)
+    offsets[1:] = num_p_per_graph.cumsum(dim=0)
+
+    # 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]
+
+    # uniform direction on the unit sphere
+    direction = torch.randn(total_waters, 3, device=device, dtype=protein_pos.dtype)
+    direction = direction / direction.norm(dim=-1, keepdim=True).clamp(min=1e-12)
+
+    # uniform radius inside the ball: r = R * U^(1/3)
+    r = cutoff * torch.rand(
+        total_waters, 1, device=device, dtype=protein_pos.dtype
+    ).pow(1.0 / 3.0)
+
+    return anchors + r * direction
+
+
+def sample_waters_scaled_gaussian(
+    batch_w: Tensor,
+    sigma_per_graph: Tensor,
+    device: torch.device,
+    dtype: torch.dtype = torch.float32,
+) -> Tensor:
+    """
+    Sample water positions from N(0, sigma^2 * I) with no rejection.
+
+    Args:
+        batch_w: (N_water,) graph index per water to sample. Positions are returned
+            in this order, so callers with existing water nodes can pass their own
+            batch vector and get samples aligned to it.
+        sigma_per_graph: (num_graphs,) Gaussian scale per graph
+        device: Output device
+        dtype: Output dtype
+
+    Returns:
+        water_pos: (N_water, 3) sampled positions, one per entry of batch_w
+    """
+    batch_w = batch_w.to(device)
+    total_waters = batch_w.numel()
+
+    if total_waters == 0:
+        return torch.empty(0, 3, dtype=dtype, device=device)
+
+    sigma = sigma_per_graph.to(device=device, dtype=dtype)[batch_w].unsqueeze(-1)
+
+    return torch.randn(total_waters, 3, device=device, dtype=dtype) * sigma
+
+
 def build_knn_edges(
     src_pos: torch.Tensor,
     dst_pos: torch.Tensor,
@@ -477,6 +619,9 @@ class FlowMatcher:
     High level class for flow matching training, validation, and numerical integration
     """
 
+    SAMPLING_STRATEGIES = ("uniform_ball", "scaled_gaussian")
+    DYNAMIC_EDGE_POLICIES = ("auto", "radius", "knn_if_isolated")
+
     def __init__(
         self,
         model,
@@ -486,6 +631,8 @@ def __init__(
         t_distort: float = 0.5,
         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.
@@ -498,7 +645,21 @@ def __init__(
             t_distort: Time threshold after which distortion may be applied
             sigma_distort: Standard deviation of distortion noise
             loss_eps: Small constant for numerical stability in loss weighting
+            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.
         """
+        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
@@ -506,6 +667,53 @@ def __init__(
         self.t_distort = t_distort
         self.sigma_distort = sigma_distort
         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:
+        """Infer graph count without forcing a device sync when batch metadata exists."""
+        num_graphs = getattr(data, "num_graphs", None)
+        if num_graphs is not None:
+            return int(num_graphs)
+        batch_p = data["protein"].batch
+        if batch_p.numel() == 0:
+            return 0
+        return int(batch_p.max().item()) + 1
+
+    def _sample_waters(
+        self,
+        batch_data: HeteroData | Batch,
+        batch_w: Tensor,
+        device: torch.device,
+    ) -> Tensor:
+        """Dispatch to the configured sampling strategy, sampling one water per
+        entry of batch_w and returning them in that order."""
+        if self.sampling_strategy == "uniform_ball":
+            return sample_waters_uniform_ball(
+                protein_pos=batch_data["protein"].pos,
+                batch_p=batch_data["protein"].batch,
+                batch_w=batch_w,
+                cutoff=self.graph_cutoff,
+                device=device,
+            )
+        # scaled_gaussian
+        sigma_per_graph = self.compute_sigma_per_graph(batch_data, device)
+        return sample_waters_scaled_gaussian(
+            batch_w=batch_w,
+            sigma_per_graph=sigma_per_graph,
+            device=device,
+            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:
@@ -533,10 +741,20 @@ def compute_sigma_per_graph(
         """
         pos = data["protein"].pos  # (N_total, 3)
         batch_p = data["protein"].batch  # (N_total,)
+        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
+        empty = torch.bincount(batch_p, minlength=num_graphs) == 0
+        if empty.any():
+            raise ValueError(
+                f"Cannot compute sigma for graph(s) {empty.nonzero().flatten().tolist()}: "
+                "they have zero protein atoms."
+            )
 
         # Var(X) = E[X^2] - E[X]^2
-        mean_pos = scatter_mean(pos, batch_p, dim=0)  # (num_graphs, 3)
-        mean_sq = scatter_mean(pos**2, batch_p, dim=0)  # (num_graphs, 3)
+        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)
         var_per_dim = mean_sq - mean_pos**2  # (num_graphs, 3)
         sigma = torch.sqrt(var_per_dim.mean(dim=-1).clamp(min=1e-8))  # (num_graphs,)
 
@@ -547,7 +765,7 @@ def training_step(
         batch: HeteroData,
         use_self_conditioning: bool = True,
         accumulation_steps: int = 1,
-    ) -> dict[str, float | int | None | dict]:
+    ) -> dict[str, object]:
         """
         Single flow matching training step (forward + backward only).
 
@@ -579,15 +797,16 @@ 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
-        batch_p = batch["protein"].batch
-        num_graphs = int(batch_p.max().item()) + 1
+        num_graphs = self._num_graphs(batch)
 
-        sigma = self.compute_sigma(batch)
-
-        x0 = torch.randn_like(x1) * sigma
+        sigma_per_graph = self.compute_sigma_per_graph(batch, device)
+        # 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)
         x0_star, x1_star = ot_coupling(x1=x1, batch=batch_w, x0=x0)
 
         t = torch.rand(num_graphs, device=device)
@@ -624,6 +843,12 @@ def training_step(
         per_atom_mse = (v_pred - v_target).pow(2).mean(dim=-1, keepdim=True)
         loss = (w * per_atom_mse).sum() / w.sum()
 
+        # training RMSD
+        with torch.no_grad():
+            x1_hat = x_t + (1.0 - t_per_atom) * v_pred
+            diff2 = ((x1_hat - x1_star) ** 2).sum(-1)  # (Nw,)
+            rmsd = torch.sqrt(scatter_mean(diff2, batch_w, dim=0)).mean()
+
         # check for high loss and compute per-sample losses for debugging
         per_sample_info = None
         if loss.item() > 100.0:
@@ -631,28 +856,21 @@ def training_step(
                 from torch_scatter import scatter_add
 
                 weighted_mse = (w * per_atom_mse).squeeze(-1)
-                # compute per-graph loss: sum(weighted_mse) / sum(w) for each graph
-                numerator = scatter_add(weighted_mse, batch_w, dim=0)
-                denominator = scatter_add(w.squeeze(-1), batch_w, dim=0)
+                numerator = scatter_add(
+                    weighted_mse, batch_w, dim=0, dim_size=num_graphs
+                )
+                denominator = scatter_add(
+                    w.squeeze(-1), batch_w, dim=0, dim_size=num_graphs
+                )
                 per_sample_loss = numerator / (denominator + 1e-8)
                 per_sample_info = {"losses": per_sample_loss, "num_graphs": num_graphs}
 
-        # backward (scale loss for gradient accumulation)
         (loss / accumulation_steps).backward()
 
-        # training RMSD
-        with torch.no_grad():
-            x1_hat = x_t + (1.0 - t_per_atom) * v_pred
-            # rmsd = compute_rmsd(x1_hat, x1_star)
-
-            # on-gpu version of rmsd
-            diff2 = ((x1_hat - x1_star) ** 2).sum(-1)  # (Nw,)
-            rmsd = torch.sqrt(scatter_mean(diff2, batch_w, dim=0)).mean().item()
-
         return {
             "loss": loss.item(),
-            "rmsd": rmsd,
-            "sigma": sigma,
+            "rmsd": rmsd.item(),
+            "sigma": sigma_per_graph,
             "per_sample_info": per_sample_info,
         }
 
@@ -674,14 +892,15 @@ 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
-        batch_p = batch["protein"].batch
-        num_graphs = int(batch_p.max().item()) + 1
+        num_graphs = self._num_graphs(batch)
 
-        sigma = self.compute_sigma(batch)
-        x0 = torch.randn_like(x1) * sigma
+        # 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)
         x0_star, x1_star = ot_coupling(x1=x1, batch=batch_w, x0=x0)
 
         t = torch.rand(num_graphs, device=device)
@@ -690,6 +909,7 @@ def validation_step(self, batch: HeteroData) -> dict[str, float]:
 
         batch["water"].pos = x_t
         v_pred = self.model(batch, t, self_cond=None)
+
         v_target = x1_star - x0_star
 
         w = 1.0 / (self.loss_eps + (1.0 - t_per_atom))
@@ -699,9 +919,12 @@ def validation_step(self, batch: HeteroData) -> dict[str, float]:
         # GPU RMSD
         x1_hat = x_t + (1.0 - t_per_atom) * v_pred
         diff2 = ((x1_hat - x1_star) ** 2).sum(-1)  # (Nw,)
-        rmsd = torch.sqrt(scatter_mean(diff2, batch_w, dim=0)).mean().item()
+        rmsd = torch.sqrt(scatter_mean(diff2, batch_w, dim=0)).mean()
 
-        return {"loss": loss.item(), "rmsd": rmsd}
+        return {
+            "loss": loss.item(),
+            "rmsd": rmsd.item(),
+        }
 
     def _setup_water_nodes_from_ratio(
         self,
@@ -722,27 +945,67 @@ def _setup_water_nodes_from_ratio(
             batch_w: (N_water_total,) batch indices
         """
         num_residues = g["protein"].num_residues  # (num_graphs,)
-        num_graphs = num_residues.size(0)
 
         # compute waters per graph: num_residues * ratio, minimum 1
         num_waters = (num_residues.float() * water_ratio).long().clamp(min=1)
 
-        # create batch indices (vectorized)
-        batch_w = torch.repeat_interleave(
-            torch.arange(num_graphs, device=device), num_waters
-        )
+        batch_w = _batch_from_counts(num_waters, device)
+        x = self._sample_waters(g, batch_w, device)
         total_waters = batch_w.size(0)
 
-        # compute sigma per graph and expand to per-water
-        sigma_per_graph = self.compute_sigma_per_graph(g, device)
-        sigma_per_water = sigma_per_graph[batch_w]
+        # create water features (oxygen one-hot; +1 for the trailing 'other' bucket)
+        water_x = torch.zeros(
+            total_waters, len(ELEMENT_VOCAB) + 1, dtype=torch.float32, device=device
+        )
+        water_x[:, ELEM_IDX["O"]] = 1.0
 
-        # sample noise
-        x = torch.randn(total_waters, 3, device=device) * sigma_per_water.unsqueeze(-1)
+        # update graph with new water nodes
+        g["water"].pos = x
+        g["water"].x = water_x
+        g["water"].batch = batch_w
+        g["water"].num_nodes = total_waters
 
-        # create water features (oxygen one-hot, index 2 for 'O' in ELEMENT_VOCAB)
-        water_x = torch.zeros(total_waters, 16, device=device)
-        water_x[:, 2] = 1.0  # oxygen is index 2 in ELEMENT_VOCAB
+        return x, batch_w
+
+    def _setup_water_nodes_from_count(
+        self,
+        g: Batch,
+        water_count: int,
+        device: torch.device,
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Create water node positions and batch indices using a fixed count per protein.
+
+        Args:
+            g: Batched HeteroData graph (modified in-place)
+            water_count: Exact number of waters to sample per protein
+            device: Device to create tensors on
+
+        Returns:
+            x: (N_water_total, 3) initial noise positions
+            batch_w: (N_water_total,) batch indices
+        """
+        if water_count < 0:
+            raise ValueError(f"water_count must be >= 0, got {water_count}")
+
+        num_graphs = self._num_graphs(g)
+
+        num_waters = torch.full(
+            (num_graphs,),
+            water_count,
+            dtype=torch.long,
+            device=device,
+        )
+
+        batch_w = _batch_from_counts(num_waters, device)
+        x = self._sample_waters(g, batch_w, device)
+        total_waters = batch_w.size(0)
+
+        # create water features (oxygen one-hot; +1 for the trailing 'other' bucket)
+        water_x = torch.zeros(
+            total_waters, len(ELEMENT_VOCAB) + 1, dtype=torch.float32, device=device
+        )
+        water_x[:, ELEM_IDX["O"]] = 1.0
 
         # update graph with new water nodes
         g["water"].pos = x
@@ -752,6 +1015,43 @@ def _setup_water_nodes_from_ratio(
 
         return x, batch_w
 
+    def _setup_water_nodes(
+        self,
+        g: Batch,
+        water_ratio: float | None,
+        water_count: int | None,
+        device: torch.device,
+    ) -> tuple[Tensor, Tensor]:
+        """
+        Create the initial water nodes to integrate from.
+
+        Args:
+            g: Batched HeteroData graph (modified in-place)
+            water_ratio: If provided, sample num_residues * water_ratio waters.
+                        Ignored when water_count is also given.
+            water_count: If provided, sample exactly this many waters per protein.
+                        Takes precedence over water_ratio. When neither is given,
+                        the ground-truth water count is resampled from the prior.
+            device: Device to create tensors on
+
+        Returns:
+            x: (N_water_total, 3) initial noise positions
+            batch_w: (N_water_total,) batch indices
+        """
+        if water_count is not None:
+            # sample fixed number of waters per protein
+            return self._setup_water_nodes_from_count(g, water_count, device)
+
+        if water_ratio is not None:
+            # sample waters based on residue count
+            return self._setup_water_nodes_from_ratio(g, water_ratio, device)
+
+        # resample the existing water nodes in place; their batch is unchanged
+        batch_w = g["water"].batch
+        x = self._sample_waters(g, batch_w, device)
+
+        return x, batch_w
+
     @torch.inference_mode()
     def euler_integrate(
         self,
@@ -761,6 +1061,7 @@ def euler_integrate(
         sc_ema_alpha: float = 0.2,
         device: str | torch.device = "cuda",
         water_ratio: float | None = None,
+        water_count: int | None = None,
     ) -> list[dict[str, np.ndarray]]:
         """
         Euler integration from noise to final positions.
@@ -772,12 +1073,17 @@ def euler_integrate(
             sc_ema_alpha: EMA decay for self-conditioning
             device: Device to run on
             water_ratio: If provided, sample num_residues * water_ratio waters
-                        instead of using ground truth water count
+                        instead of using ground truth water count. Ignored when
+                        water_count is also given.
+            water_count: If provided, sample exactly this many waters per protein.
+                        Takes precedence over water_ratio. When neither is given,
+                        the ground-truth water count is resampled from the prior.
 
         Returns:
             List of dicts, one per input graph, each with keys:
                 'protein_pos': (Np, 3) - includes both ASU and mate atoms
-                'water_true': (Nw, 3) - None if water_ratio is used
+                'water_true': (Nw, 3) ground-truth waters (always returned; when
+                        water_ratio/water_count is set its count may differ from water_pred)
                 'water_pred': (Nw, 3) final prediction
                 'pdb_id': PDB identifier
         """
@@ -795,24 +1101,13 @@ def euler_integrate(
         g = Batch.from_data_list([copy.deepcopy(graph) for graph in graphs]).to(device)
 
         batch_p = g["protein"].batch
+        num_graphs = self._num_graphs(g)
 
         # store ground truth water positions and batch indices before modifying
         x1_true = g["water"].pos.clone()
         batch_w_true = g["water"].batch.clone()
 
-        if water_ratio is not None:
-            # sample waters based on residue count
-            x, batch_w = self._setup_water_nodes_from_ratio(g, water_ratio, device)
-            num_graphs = g["protein"].num_residues.size(0)
-        else:
-            # use existing water nodes
-            batch_w = g["water"].batch
-            num_graphs = int(batch_w.max().item()) + 1
-            sigma_per_graph = self.compute_sigma_per_graph(g, device)
-            sigma_per_water = sigma_per_graph[batch_w]
-            x = torch.randn(
-                g["water"].num_nodes, 3, device=device
-            ) * sigma_per_water.unsqueeze(-1)
+        x, batch_w = self._setup_water_nodes(g, water_ratio, water_count, device)
 
         x1_pred_ema = x.clone()
 
@@ -872,6 +1167,7 @@ def rk4_integrate(
         device: str | torch.device = "cuda",
         return_trajectory: bool = True,
         water_ratio: float | None = None,
+        water_count: int | None = None,
     ) -> list[dict[str, np.ndarray]]:
         """
         RK4 integration from noise to final positions.
@@ -884,12 +1180,17 @@ def rk4_integrate(
             device: Device to run on
             return_trajectory: Whether to return full trajectory and metrics
             water_ratio: If provided, sample num_residues * water_ratio waters
-                        instead of using ground truth water count
+                        instead of using ground truth water count. Ignored when
+                        water_count is also given.
+            water_count: If provided, sample exactly this many waters per protein.
+                        Takes precedence over water_ratio. When neither is given,
+                        the ground-truth water count is resampled from the prior.
 
         Returns:
             List of dicts, one per input graph, each with keys:
                 'protein_pos': (Np, 3) - includes both ASU and mate atoms
-                'water_true': (Nw, 3) - None if water_ratio is used
+                'water_true': (Nw, 3) ground-truth waters (always returned; when
+                        water_ratio/water_count is set its count may differ from water_pred)
                 'water_pred': (Nw, 3) final prediction
                 'trajectory': list of (Nw, 3) at each step (if return_trajectory=True)
         """
@@ -907,22 +1208,13 @@ def rk4_integrate(
         g = Batch.from_data_list([copy.deepcopy(graph) for graph in graphs]).to(device)
 
         batch_p = g["protein"].batch
+        num_graphs = self._num_graphs(g)
 
         # store ground truth water positions and batch indices before modifying
         x1_true = g["water"].pos.clone()
         batch_w_true = g["water"].batch.clone()
 
-        if water_ratio is not None:
-            # sample waters based on residue count
-            x, batch_w = self._setup_water_nodes_from_ratio(g, water_ratio, device)
-            num_graphs = g["protein"].num_residues.size(0)
-        else:
-            # use existing water nodes
-            batch_w = g["water"].batch
-            num_graphs = int(batch_w.max().item()) + 1
-            sigma_per_graph = self.compute_sigma_per_graph(g, device)
-            sigma_per_water = sigma_per_graph[batch_w]
-            x = torch.randn_like(x1_true) * sigma_per_water.unsqueeze(-1)
+        x, batch_w = self._setup_water_nodes(g, water_ratio, water_count, device)
 
         x1_pred_ema = x.clone()
 
diff --git a/tests/test_flow.py b/tests/test_flow.py
index a3a519c..37642d3 100644
--- a/tests/test_flow.py
+++ b/tests/test_flow.py
@@ -9,13 +9,16 @@
 import pytest
 import torch
 import torch.nn.functional as F
-from torch_geometric.data import Data, HeteroData
+from torch_geometric.data import Batch, Data, HeteroData
 
 from src.flow import (
+    _batch_from_counts,
     build_knn_edges,
     FlowMatcher,
     FlowWaterGVP,
     ProteinWaterUpdate,
+    sample_waters_scaled_gaussian,
+    sample_waters_uniform_ball,
 )
 from src.gvp_encoder import GVPEncoder, make_gvp_encoder_data, ProteinGVPEncoder
 
@@ -364,6 +367,16 @@ def test_compute_sigma(self, simple_hetero_data):
         assert isinstance(sigma, float)
         assert sigma > 0
 
+    def test_compute_sigma_per_graph_zero_protein_raises(self, device):
+        """A graph with no protein atoms has no meaningful sigma."""
+        g0 = HeteroData()
+        g0["protein"].pos = torch.randn(4, 3, device=device)
+        g1 = HeteroData()
+        g1["protein"].pos = torch.empty(0, 3, device=device)
+
+        with pytest.raises(ValueError, match="zero protein atoms"):
+            FlowMatcher.compute_sigma_per_graph(Batch.from_data_list([g0, g1]), device)
+
     def test_training_step(self, flow_matcher, simple_hetero_data, device):
         optimizer = torch.optim.Adam(flow_matcher.model.parameters(), lr=1e-4)
 
@@ -400,6 +413,23 @@ 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
+    ):
+        model = FlowWaterGVP(
+            encoder=gvp_encoder,
+            hidden_dims=(64, 8),
+            layers=1,
+        ).to(device)
+
+        flow_matcher = FlowMatcher(
+            model,
+            sampling_strategy="scaled_gaussian",
+            dynamic_edge_policy="auto",
+        )
+
+        assert flow_matcher._effective_dynamic_edge_policy() == "knn_if_isolated"
+
     @pytest.mark.slow
     def test_euler_integrate(self, flow_matcher, simple_hetero_data, device):
         results = flow_matcher.euler_integrate(
@@ -452,6 +482,298 @@ def test_sample_rk4(self, flow_matcher, simple_hetero_data, device):
         assert water_pred.shape == (n_water, 3)
 
 
+# ============== Tests for water sampling strategies ==============
+
+
+@pytest.mark.unit
+class TestUniformBallSampling:
+    def test_shapes_and_counts(self, device):
+        torch.manual_seed(0)
+        protein_pos = torch.tensor(
+            [[0.0, 0.0, 0.0], [0.5, 0.0, 0.0], [0.0, 0.0, 0.0], [0.5, 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([4, 3], dtype=torch.long, device=device), device
+        )
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=2.0,
+            device=device,
+        )
+
+        assert pos.shape == (7, 3)
+
+    def test_all_within_cutoff(self, device):
+        torch.manual_seed(42)
+        protein_pos = torch.randn(20, 3, device=device) * 50
+        batch_p = torch.cat([torch.zeros(10), torch.ones(10)]).long().to(device)
+        batch_w = _batch_from_counts(
+            torch.tensor([50, 50], dtype=torch.long, device=device), device
+        )
+        cutoff = 8.0
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=cutoff,
+            device=device,
+        )
+
+        for g in range(2):
+            g_waters = pos[batch_w == g]
+            g_protein = protein_pos[batch_p == g]
+            dists = torch.cdist(g_waters, g_protein)
+            assert dists.min(dim=1).values.max().item() <= cutoff + 1e-5
+
+    def test_empty_waters(self, device):
+        protein_pos = torch.randn(5, 3, device=device)
+        batch_p = torch.zeros(5, dtype=torch.long, device=device)
+        batch_w = torch.empty(0, dtype=torch.long, device=device)
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=8.0,
+            device=device,
+        )
+
+        assert pos.shape == (0, 3)
+
+    def test_zero_protein_graph_raises(self, device):
+        """Requesting waters for a graph with no protein atoms fails fast."""
+        # graph 0 has protein atoms, graph 1 has none but requests waters
+        protein_pos = torch.randn(5, 3, device=device)
+        batch_p = torch.zeros(5, dtype=torch.long, device=device)
+        batch_w = _batch_from_counts(
+            torch.tensor([3, 4], dtype=torch.long, device=device), device
+        )
+
+        with pytest.raises(ValueError, match="zero protein atoms"):
+            sample_waters_uniform_ball(
+                protein_pos=protein_pos,
+                batch_p=batch_p,
+                batch_w=batch_w,
+                cutoff=8.0,
+                device=device,
+            )
+
+    def test_large_spread_protein_succeeds(self, device):
+        """The scenario that crashes truncated Gaussian (sigma~50) works here."""
+        torch.manual_seed(0)
+        protein_pos = torch.randn(500, 3, device=device) * 50
+        batch_p = torch.zeros(500, dtype=torch.long, device=device)
+        batch_w = _batch_from_counts(
+            torch.tensor([301], dtype=torch.long, device=device), device
+        )
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=8.0,
+            device=device,
+        )
+
+        assert pos.shape == (301, 3)
+
+    @pytest.mark.slow
+    def test_real_structure_cutoff_and_batch(self, device, pdb_6eey):
+        """Cutoff guarantee holds on real protein geometry; batch indexing is correct
+        when two structures with different water counts are packed into one call."""
+        import biotite.structure as bts
+        from biotite.structure.io.pdb import get_structure, PDBFile
+
+        torch.manual_seed(0)
+
+        pdb_file = PDBFile.read(pdb_6eey)
+        atoms = get_structure(pdb_file, model=1, altloc="occupancy")
+        atoms = atoms[atoms.element != "H"]
+        protein_atoms = atoms[bts.filter_amino_acids(atoms)]
+
+        protein_pos_np = protein_atoms.coord  # (N, 3) float64
+        n_atoms = len(protein_pos_np)
+
+        # batch two copies: graph 0 gets 50 waters, graph 1 gets 30
+        protein_pos = torch.tensor(protein_pos_np, dtype=torch.float32, device=device)
+        protein_pos_both = torch.cat([protein_pos, protein_pos], dim=0)
+        batch_p = torch.cat(
+            [
+                torch.zeros(n_atoms, dtype=torch.long, device=device),
+                torch.ones(n_atoms, dtype=torch.long, device=device),
+            ]
+        )
+        num_waters = torch.tensor([50, 30], dtype=torch.long, device=device)
+        batch_w = _batch_from_counts(num_waters, device)
+        cutoff = 8.0
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos_both,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=cutoff,
+            device=device,
+        )
+
+        # correct total count and per-graph split
+        assert pos.shape == (80, 3)
+        assert (batch_w == 0).sum().item() == 50
+        assert (batch_w == 1).sum().item() == 30
+
+        # every water must be within cutoff of at least one protein atom in its graph
+        for g, n_w in enumerate(num_waters.tolist()):
+            g_waters = pos[batch_w == g]  # (n_w, 3)
+            g_protein = protein_pos_both[batch_p == g]  # (n_atoms, 3)
+            dists = torch.cdist(g_waters, g_protein)  # (n_w, n_atoms)
+            min_dists = dists.min(dim=1).values  # (n_w,)
+            assert min_dists.max().item() <= cutoff + 1e-4, (
+                f"Graph {g}: water too far from protein "
+                f"(max dist {min_dists.max().item():.4f} > {cutoff})"
+            )
+
+
+@pytest.mark.unit
+class TestScaledGaussianSampling:
+    def test_shapes_and_counts(self, device):
+        torch.manual_seed(0)
+        batch_w = _batch_from_counts(
+            torch.tensor([4, 3], dtype=torch.long, device=device), device
+        )
+        sigma = torch.tensor([1.0, 2.0], device=device)
+
+        pos = sample_waters_scaled_gaussian(
+            batch_w=batch_w,
+            sigma_per_graph=sigma,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        assert pos.shape == (7, 3)
+
+    def test_sigma_broadcasts_per_graph(self, device):
+        """Each graph's waters must be scaled by that graph's own sigma."""
+        batch_w = _batch_from_counts(
+            torch.tensor([4, 3], dtype=torch.long, device=device), device
+        )
+        sigma = torch.tensor([1.0, 2.0], device=device)
+
+        torch.manual_seed(0)
+        pos = sample_waters_scaled_gaussian(
+            batch_w=batch_w,
+            sigma_per_graph=sigma,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        # randn is the sampler's only RNG draw, so the same seed reproduces it
+        torch.manual_seed(0)
+        expected = torch.randn(7, 3, device=device, dtype=torch.float32) * sigma[
+            batch_w
+        ].unsqueeze(-1)
+
+        assert torch.allclose(pos, expected)
+
+    def test_empty_waters(self, device):
+        batch_w = torch.empty(0, dtype=torch.long, device=device)
+        sigma = torch.tensor([1.0], device=device)
+
+        pos = sample_waters_scaled_gaussian(
+            batch_w=batch_w,
+            sigma_per_graph=sigma,
+            device=device,
+            dtype=torch.float32,
+        )
+
+        assert pos.shape == (0, 3)
+
+
+@pytest.mark.unit
+class TestSamplingHonoursNodeOrder:
+    """Samplers return one water per batch_w entry, in batch_w's own order."""
+
+    @staticmethod
+    def _two_graphs_far_apart(device):
+        """Graph 0's protein at the origin, graph 1's 100A away."""
+        protein_pos = torch.cat(
+            [torch.zeros(4, 3, device=device), torch.full((4, 3), 100.0, device=device)]
+        )
+        batch_p = torch.tensor([0] * 4 + [1] * 4, dtype=torch.long, device=device)
+        # water nodes interleaved across graphs rather than grouped
+        batch_w = torch.tensor([0, 1, 0, 1], dtype=torch.long, device=device)
+
+        return protein_pos, batch_p, batch_w
+
+    def test_uniform_ball_follows_interleaved_batch(self, device):
+        """Each water anchors on its own graph even when nodes are interleaved."""
+        torch.manual_seed(0)
+        protein_pos, batch_p, batch_w = self._two_graphs_far_apart(device)
+
+        pos = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=2.0,
+            device=device,
+        )
+
+        assert pos[batch_w == 0].abs().max().item() < 10.0
+        assert (pos[batch_w == 1] - 100.0).abs().max().item() < 10.0
+
+    def test_scaled_gaussian_follows_interleaved_batch(self, device):
+        """Sigma follows each water's own graph when nodes are interleaved."""
+        batch_w = torch.tensor([0, 1, 0, 1], dtype=torch.long, device=device)
+        sigma = torch.tensor([1.0, 2.0], device=device)
+
+        torch.manual_seed(0)
+        pos = sample_waters_scaled_gaussian(
+            batch_w=batch_w, sigma_per_graph=sigma, device=device, dtype=torch.float32
+        )
+
+        torch.manual_seed(0)
+        expected = torch.randn(4, 3, device=device) * sigma[batch_w].unsqueeze(-1)
+
+        assert torch.allclose(pos, expected)
+
+    def test_ot_coupling_pairs_within_graph_when_interleaved(self, device):
+        """A graph's waters pair with its own prior, not a neighbour's."""
+        from src.utils import ot_coupling
+
+        torch.manual_seed(0)
+        protein_pos, batch_p, batch_w = self._two_graphs_far_apart(device)
+        x1 = torch.tensor(
+            [[0.0] * 3, [100.0] * 3, [0.0] * 3, [100.0] * 3], device=device
+        )
+
+        x0 = sample_waters_uniform_ball(
+            protein_pos=protein_pos,
+            batch_p=batch_p,
+            batch_w=batch_w,
+            cutoff=2.0,
+            device=device,
+        )
+        x0_star, x1_star = ot_coupling(x1=x1, batch=batch_w, x0=x0)
+
+        # pairings stay inside their graph, so nothing is dragged 100A across
+        assert (x1_star - x0_star).norm(dim=-1).max().item() < 10.0
+
+
+@pytest.mark.unit
+class TestWaterCountValidation:
+    def test_negative_water_count_raises(self, device):
+        """A negative water_count is rejected before any sampling work."""
+        fm = FlowMatcher(model=Mock(cutoff=8.0))
+        g = HeteroData()  # guard fires before touching graph contents
+
+        with pytest.raises(ValueError, match="water_count must be >= 0"):
+            fm._setup_water_nodes_from_count(g, -1, device)
+
+
 # ============== Tests for distortion ==============