diff --git a/README.md b/README.md
index ce09d94..760f93c 100644
--- a/README.md
+++ b/README.md
@@ -9,6 +9,7 @@ WaterFlow/
 ├── src/                    # Core library code
 │   ├── dataset.py          # ProteinWaterDataset and data loading
 │   ├── flow.py             # FlowMatcher and FlowWaterGVP model
+│   ├── confidence.py       # ConfidenceGVP scorer, targets, vdW clustering
 │   ├── gvp.py              # Geometric Vector Perceptron layers
 │   ├── gvp_encoder.py      # GVP-based protein encoder
 │   ├── encoder_base.py     # Encoder registry and factory (includes ESM/SLAE)
@@ -21,6 +22,7 @@ WaterFlow/
 │   └── generate_slae_embeddings.py  # Precompute SLAE embeddings
 ├── tests/                  # Test suite
 │   ├── test_dataset.py     # Dataset and preprocessing tests
+│   ├── test_confidence.py  # Confidence scorer, target and clustering tests
 │   ├── test_flow.py        # Flow matching tests
 │   ├── test_encoder.py     # Encoder tests
 │   ├── test_forward.py     # End-to-end forward pass tests
diff --git a/src/confidence.py b/src/confidence.py
new file mode 100644
index 0000000..4d45923
--- /dev/null
+++ b/src/confidence.py
@@ -0,0 +1,411 @@
+"""
+Confidence model and post-processing for candidate waters.
+
+Mirrors the SuperWater (Nature Comm. Chem. s42004-025-01789-4) confidence +
+clustering stage that sits after the generator in the DiffDock-style two-stage
+pipeline.
+
+This module provides:
+- smootherstep_target / smootherstep_confidence: per-candidate supervision from
+  a soft cutoff of the nearest-GT-water distance
+- cluster_waters_vdw: vdW-radius clustering with confidence-weighted centroids
+  and NMS over those centroids
+- ConfidenceGVP: scores candidate waters through the same GVP backbone as
+  FlowWaterGVP, minus time conditioning
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+from torch_geometric.data import HeteroData
+
+from src.constants import EDGE_PP, EDGE_PW, NUM_RBF
+from src.encoder_base import BaseProteinEncoder
+from src.flow import ProteinWaterUpdate
+from src.gvp import GVP
+
+
+# ---------------------------------------------------------------------------
+# Targets
+# ---------------------------------------------------------------------------
+
+
+def _nearest_gt_distance(candidate_pos: Tensor, gt_pos: Tensor) -> Tensor:
+    """
+    Distance from each candidate to its nearest ground-truth water.
+
+    Single-protein (no batch dim); the caller guards the empty cases.
+
+    Args:
+        candidate_pos: (Nc, 3) candidate positions.
+        gt_pos: (Ng, 3) ground-truth positions.
+
+    Returns:
+        (Nc,) Euclidean distances.
+    """
+    diffs = candidate_pos.unsqueeze(1) - gt_pos.unsqueeze(0)  # (Nc, Ng, 3)
+    return diffs.norm(dim=-1).min(dim=1).values
+
+
+def smootherstep(x: Tensor) -> Tensor:
+    """
+    Perlin smootherstep S(x) = 6x^5 - 15x^4 + 10x^3 on x in [0, 1].
+
+    C2-continuous: S, S' and S'' all vanish at both ends, so it joins flat
+    plateaus with no kink in value, slope, or curvature.
+
+    Args:
+        x: Input, already clamped to [0, 1] by the caller.
+
+    Returns:
+        S(x), same shape as `x`.
+    """
+    # Horner form: degree-5 polynomial, no transcendentals.
+    return x * x * x * (10.0 + x * (x * 6.0 - 15.0))
+
+
+def smootherstep_confidence(
+    d: Tensor,
+    r_in: float = 0.5,
+    r_out: float = 1.5,
+) -> Tensor:
+    """
+    Soft-cutoff confidence from a nearest-GT distance.
+
+        conf = 1                    for d <= r_in
+        conf = 1 - smootherstep(u)  for r_in < d < r_out, u = (d-r_in)/(r_out-r_in)
+        conf = 0                    for d >= r_out
+
+    A candidate within `r_in` of a GT water is on the site within experimental
+    error; past `r_out` it is background. The 0.5-crossing sits at the midpoint
+    and the band width sets the steepness, so location and sharpness are
+    decoupled.
+
+    Args:
+        d: (N,) nearest-GT distances in Angstroms.
+        r_in: Plateau radius -- at or under this, confidence is 1.
+        r_out: Floor radius -- at or over this, confidence is 0.
+
+    Returns:
+        (N,) confidences in [0, 1].
+
+    Raises:
+        ValueError: If `r_out` does not exceed `r_in`.
+    """
+    if r_out <= r_in:
+        raise ValueError(f"r_out ({r_out}) must exceed r_in ({r_in}).")
+    u = ((d - r_in) / (r_out - r_in)).clamp(0.0, 1.0)
+    return 1.0 - smootherstep(u)
+
+
+def smootherstep_target(
+    candidate_pos: Tensor,
+    gt_pos: Tensor,
+    r_in: float = 0.5,
+    r_out: float = 1.5,
+) -> Tensor:
+    """
+    Confidence target for each candidate: higher means closer to a GT water.
+
+    Single-protein (no batch dim).
+
+    Args:
+        candidate_pos: (Nc, 3) candidate water positions.
+        gt_pos: (Ng, 3) ground-truth water positions, Ng >= 1.
+        r_in: Plateau radius (A). See `smootherstep_confidence`.
+        r_out: Floor radius (A). See `smootherstep_confidence`.
+
+    Returns:
+        (Nc,) target confidences in [0, 1]; empty when there are no candidates.
+
+    Raises:
+        ValueError: If `gt_pos` is empty.
+    """
+    if gt_pos.numel() == 0:
+        raise ValueError("smootherstep_target requires at least one GT water.")
+    if candidate_pos.numel() == 0:
+        return candidate_pos.new_empty(0)
+    return smootherstep_confidence(
+        _nearest_gt_distance(candidate_pos, gt_pos), r_in=r_in, r_out=r_out
+    )
+
+
+# ---------------------------------------------------------------------------
+# Post-processor
+# ---------------------------------------------------------------------------
+
+
+def cluster_waters_vdw(
+    positions: Tensor,
+    confidences: Tensor,
+    radius: float = 1.52,
+    threshold: float | None = None,
+) -> tuple[Tensor, Tensor]:
+    """
+    Two-pass vdW clustering of scored candidates (SuperWater Fig. 4 / Methods).
+
+    Round 1 absorbs: seed a cluster with the highest-confidence unassigned
+    water, absorb every unassigned water within `radius`, and emit a
+    confidence-weighted centroid carrying the cluster's max confidence.
+    Round 2 runs NMS over those centroids, dropping the lower-confidence member
+    of any pair still within `radius`.
+
+    Args:
+        positions: (N, 3) candidate water positions.
+        confidences: (N,) scalar confidences, higher is better.
+        radius: Absorption and NMS radius in Angstroms. Default 1.52, the vdW
+            radius of oxygen.
+        threshold: Drop candidates scoring below this before clustering. None
+            keeps all.
+
+    Returns:
+        ((M, 3) positions, (M,) confidences), in descending confidence order.
+
+    Raises:
+        ValueError: If the input shapes disagree or are not (N, 3) / (N,).
+    """
+    if positions.dim() != 2 or positions.size(-1) != 3:
+        raise ValueError(f"positions must be (N, 3), got {tuple(positions.shape)}")
+    if confidences.dim() != 1 or confidences.size(0) != positions.size(0):
+        raise ValueError(
+            f"confidences must be (N,), got {tuple(confidences.shape)} "
+            f"for positions {tuple(positions.shape)}"
+        )
+
+    if threshold is not None:
+        keep = confidences >= threshold
+        positions = positions[keep]
+        confidences = confidences[keep]
+
+    if positions.numel() == 0:
+        return positions, confidences
+
+    # --- Round 1: absorb into confidence-weighted centroids ---
+    order = torch.argsort(confidences, descending=True)
+    pos_sorted = positions[order]
+    conf_sorted = confidences[order]
+
+    n = pos_sorted.size(0)
+    assigned = torch.zeros(n, dtype=torch.bool, device=pos_sorted.device)
+    r2 = float(radius) ** 2
+
+    centroid_positions: list[Tensor] = []
+    centroid_confidences: list[Tensor] = []
+
+    for i in range(n):
+        if assigned[i]:
+            continue
+        seed_pos = pos_sorted[i]
+        d2 = ((pos_sorted - seed_pos) ** 2).sum(dim=-1)
+        # Always includes i: d2[i] == 0 and assigned[i] is False.
+        within = (d2 <= r2) & (~assigned)
+        cluster_pos = pos_sorted[within]
+        cluster_conf = conf_sorted[within]
+
+        # Fall back to an unweighted mean when the weights sum to 0, so an
+        # all-zero-confidence cluster yields a position rather than NaN.
+        w_sum = cluster_conf.sum()
+        if w_sum.abs() > 0:
+            centroid = (cluster_pos * cluster_conf.unsqueeze(-1)).sum(dim=0) / w_sum
+        else:
+            centroid = cluster_pos.mean(dim=0)
+
+        centroid_positions.append(centroid)
+        centroid_confidences.append(cluster_conf.max())
+        assigned = assigned | within
+
+    if not centroid_positions:
+        return pos_sorted.new_empty(0, 3), conf_sorted.new_empty(0)
+
+    cent_pos = torch.stack(centroid_positions, dim=0)
+    cent_conf = torch.stack(centroid_confidences, dim=0)
+
+    # --- Round 2: NMS between centroids ---
+    keep_mask = torch.ones(cent_pos.size(0), dtype=torch.bool, device=cent_pos.device)
+    # Already in descending seed order, but sort explicitly rather than rely on it.
+    cent_order = torch.argsort(cent_conf, descending=True)
+    for idx in cent_order.tolist():
+        if not keep_mask[idx]:
+            continue
+        d2 = ((cent_pos - cent_pos[idx]) ** 2).sum(dim=-1)
+        collide = (d2 <= r2) & keep_mask
+        collide[idx] = False  # never drop self
+        keep_mask[collide] = False
+
+    cent_pos = cent_pos[keep_mask]
+    cent_conf = cent_conf[keep_mask]
+
+    # Descending order gives downstream thresholding a stable prefix.
+    final_order = torch.argsort(cent_conf, descending=True)
+    return cent_pos[final_order], cent_conf[final_order]
+
+
+# ---------------------------------------------------------------------------
+# Model
+# ---------------------------------------------------------------------------
+
+
+class ConfidenceGVP(nn.Module):
+    """
+    DiffDock-style confidence model scoring candidate waters.
+
+    Mirrors `FlowWaterGVP`'s structure but drops time and self conditioning
+    (inputs are clean samples), emits one scalar per candidate water, and runs
+    on PW and PP edges only -- candidates come from inference and are not
+    refined here, so WW and WP carry nothing.
+
+    PW edges come from the dataset when it supplies them, otherwise from a
+    radius query plus a nearest-neighbour pass for candidates it left with no
+    edges -- candidates can land in sparse regions the generator never visits.
+
+    Keeping the backbone identical to the generator is what lets the confidence
+    network warm-start from a flow checkpoint.
+    """
+
+    def __init__(
+        self,
+        encoder: BaseProteinEncoder,
+        hidden_dims: tuple[int, int] = (256, 32),
+        edge_scalar_dim: int = NUM_RBF,
+        layers: int = 4,
+        drop_rate: float = 0.1,
+        n_message_gvps: int = 2,
+        n_update_gvps: int = 2,
+        vector_gate: bool = True,
+        cutoff: float = 8.0,
+        max_neighbors: int = 256,
+        dynamic_edge_policy: str = "knn_if_isolated",
+        knn_fallback_k: int = 8,
+        water_input_dim: int = 16,
+    ):
+        """
+        Initialize the confidence scorer.
+
+        Args mirror `FlowWaterGVP`; see that class for the shared ones. The
+        differences: no time dimension is added to the scalar encoders, and the
+        head projects hidden -> scalar logit (vo=0) instead of a vector field.
+
+        Args:
+            encoder: Protein encoder implementing BaseProteinEncoder.
+            hidden_dims: (scalar_dim, vector_dim) hidden dimensions.
+            edge_scalar_dim: Dimension of edge scalar features.
+            layers: Number of heterogeneous GVP message passing layers.
+            drop_rate: Dropout rate.
+            n_message_gvps: GVP modules in each edge type's message function.
+            n_update_gvps: GVP modules in the node update function.
+            vector_gate: Whether to use vector gating in GVP layers.
+            cutoff: Distance cutoff in Angstroms for radius PW edges.
+            max_neighbors: Per-source cap on radius results.
+            dynamic_edge_policy: How PW edges are built when not cached.
+            knn_fallback_k: Neighbours attached to candidates the radius query
+                left with no edges; 0 disables that pass.
+            water_input_dim: Input dimension for water node features.
+        """
+        super().__init__()
+        self.encoder = encoder
+        self.hidden_dims = hidden_dims
+        self.edge_scalar_dim = edge_scalar_dim
+        self.layers = layers
+        self.drop_rate = drop_rate
+        self.n_message_gvps = n_message_gvps
+        self.n_update_gvps = n_update_gvps
+        self.vector_gate = vector_gate
+
+        s_h, v_h = hidden_dims
+
+        self.encoder_to_flow = GVP(
+            in_dims=encoder.output_dims,
+            out_dims=hidden_dims,
+            activations=(F.relu, torch.sigmoid),
+            vector_gate=True,
+        )
+
+        # No time concat, unlike FlowWaterGVP: these are clean samples.
+        self.protein_scalar_encoder = nn.Sequential(
+            nn.Linear(s_h, s_h),
+            nn.GELU(),
+            nn.LayerNorm(s_h),
+        )
+        self.water_scalar_encoder = nn.Sequential(
+            nn.Linear(water_input_dim, s_h),
+            nn.GELU(),
+            nn.LayerNorm(s_h),
+        )
+
+        self.updater = ProteinWaterUpdate(
+            hidden_dims=hidden_dims,
+            rbf_dim=edge_scalar_dim,
+            layers=layers,
+            drop_rate=drop_rate,
+            n_message_gvps=n_message_gvps,
+            n_update_gvps=n_update_gvps,
+            vector_gate=vector_gate,
+            aggr_edges="sum",
+            use_dst_feats=True,
+            etypes=[EDGE_PW, EDGE_PP],
+            cutoff=cutoff,
+            max_neighbors=max_neighbors,
+            dynamic_edge_policy=dynamic_edge_policy,
+            knn_fallback_k=knn_fallback_k,
+        )
+
+        # vo=0 makes GVP return a tensor rather than a (scalar, vector) tuple.
+        self.score_head = GVP(
+            in_dims=hidden_dims,
+            out_dims=(1, 0),
+            activations=(None, None),
+            vector_gate=False,
+        )
+
+    def forward(
+        self,
+        data: HeteroData,
+        return_logits: bool = False,
+    ) -> Tensor:
+        """
+        Score every water node in `data`.
+
+        Args:
+            data: HeteroData with 'protein' and 'water' node types, where
+                `water.pos` holds the candidate positions to score.
+            return_logits: Return raw pre-sigmoid logits, for callers using
+                BCEWithLogits instead of MSE.
+
+        Returns:
+            (N_w,) scores, in [0, 1] unless `return_logits` is set.
+        """
+        device = data["protein"].pos.device
+
+        s_all, v_all, pp_edge_attr = self.encoder(data)
+        encoder_input = (s_all, v_all) if self.encoder.output_dims[1] > 0 else s_all
+        s_p_latent, v_p_latent = self.encoder_to_flow(encoder_input)
+
+        if "water" not in data.node_types or data["water"].num_nodes == 0:
+            return torch.zeros(0, device=device)
+
+        s_p = self.protein_scalar_encoder(s_p_latent)
+        s_w = self.water_scalar_encoder(data["water"].x)
+
+        v_w = torch.zeros(
+            data["water"].num_nodes,
+            self.hidden_dims[1],
+            3,
+            device=device,
+        )
+
+        x_dict = {
+            "protein": (s_p, v_p_latent),
+            "water": (s_w, v_w),
+        }
+        x_dict = self.updater(
+            x_dict,
+            data,
+            pp_edge_attr=pp_edge_attr,
+        )
+
+        logits = self.score_head(x_dict["water"]).squeeze(-1)  # (N_w,)
+        if return_logits:
+            return logits
+        return torch.sigmoid(logits)
diff --git a/tests/test_confidence.py b/tests/test_confidence.py
new file mode 100644
index 0000000..d514a4b
--- /dev/null
+++ b/tests/test_confidence.py
@@ -0,0 +1,420 @@
+"""Unit tests for src/confidence.py -- targets, clustering, and ConfidenceGVP."""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch_geometric.data import HeteroData
+
+from src.confidence import (
+    cluster_waters_vdw,
+    ConfidenceGVP,
+    smootherstep_confidence,
+    smootherstep_target,
+)
+from src.constants import EDGE_PP, EDGE_PW, NUM_RBF
+from src.utils import compute_edge_features
+
+
+# ============== smootherstep target ==============
+
+
+@pytest.mark.unit
+class TestSmootherstepTarget:
+    def test_plateau_and_floor(self, device):
+        gt = torch.zeros(1, 3, device=device)
+        ds = torch.tensor([0.0, 0.4, 2.0, 5.0], device=device)
+        cand = torch.stack([ds, torch.zeros_like(ds), torch.zeros_like(ds)], dim=1)
+        out = smootherstep_target(cand, gt, r_in=0.4, r_out=2.0)
+        assert out[0].item() == pytest.approx(1.0, abs=1e-6)  # d < r_in -> 1
+        assert out[1].item() == pytest.approx(1.0, abs=1e-6)  # d == r_in -> 1
+        assert out[2].item() == pytest.approx(0.0, abs=1e-6)  # d == r_out -> 0
+        assert out[3].item() == pytest.approx(0.0, abs=1e-6)  # far -> 0
+
+    def test_midpoint_is_half(self, device):
+        gt = torch.zeros(1, 3, device=device)
+        cand = torch.tensor([[1.2, 0.0, 0.0]], device=device)  # midpoint of [0.4, 2.0]
+        out = smootherstep_target(cand, gt, r_in=0.4, r_out=2.0)
+        assert out.item() == pytest.approx(0.5, abs=1e-6)
+
+    def test_monotone_decreasing(self, device):
+        gt = torch.zeros(1, 3, device=device)
+        ds = torch.linspace(0.0, 2.5, 12, device=device)
+        cand = torch.stack([ds, torch.zeros_like(ds), torch.zeros_like(ds)], dim=1)
+        out = smootherstep_target(cand, gt, r_in=0.4, r_out=2.0)
+        diffs = out[1:] - out[:-1]
+        assert (diffs <= 1e-7).all()  # non-increasing in distance
+
+    def test_values_in_unit_interval(self, device):
+        gt = torch.randn(3, 3, device=device)
+        cand = torch.randn(20, 3, device=device) * 3.0
+        out = smootherstep_target(cand, gt, r_in=0.4, r_out=2.0)
+        assert (out >= 0).all() and (out <= 1.0).all()
+
+    def test_nearest_gt_is_used(self, device):
+        # The nearest GT sets the target; the far one must not dilute it.
+        gt = torch.tensor([[0.0, 0.0, 0.0], [100.0, 0.0, 0.0]], device=device)
+        cand = torch.tensor([[1.2, 0.0, 0.0]], device=device)
+        out = smootherstep_target(cand, gt, r_in=0.4, r_out=2.0)
+        assert out.item() == pytest.approx(0.5, abs=1e-6)
+
+    def test_confidence_matches_hand_values(self, device):
+        d = torch.tensor([0.5, 0.8, 1.0, 1.5], device=device)
+        out = smootherstep_confidence(d, r_in=0.4, r_out=2.0)
+        expected = torch.tensor([0.998, 0.896, 0.725, 0.179], device=device)
+        assert torch.allclose(out, expected, atol=2e-3)
+
+    def test_default_radii_are_half_and_one_and_a_half(self, device):
+        """The 0.5-crossing of the shipped defaults sits at 1.0 A."""
+        gt = torch.zeros(1, 3, device=device)
+        cand = torch.tensor([[1.0, 0.0, 0.0]], device=device)
+        assert smootherstep_target(cand, gt).item() == pytest.approx(0.5, abs=1e-6)
+
+    def test_empty_candidates_ok(self, device):
+        gt = torch.randn(3, 3, device=device)
+        cand = torch.empty(0, 3, device=device)
+        out = smootherstep_target(cand, gt)
+        assert out.shape == (0,)
+
+    def test_empty_gt_raises(self, device):
+        cand = torch.randn(3, 3, device=device)
+        gt = torch.empty(0, 3, device=device)
+        with pytest.raises(ValueError, match="at least one GT"):
+            smootherstep_target(cand, gt)
+
+    def test_bad_radii_raises(self, device):
+        gt = torch.zeros(1, 3, device=device)
+        cand = torch.zeros(1, 3, device=device)
+        with pytest.raises(ValueError, match="r_out"):
+            smootherstep_target(cand, gt, r_in=2.0, r_out=1.0)
+
+
+# ============== vdW clustering ==============
+
+
+@pytest.mark.unit
+class TestClusterWatersVdw:
+    def test_two_close_waters_collapse(self, device):
+        # 1 A apart with r=1.52 => absorbed.
+        pos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], device=device)
+        conf = torch.tensor([0.9, 0.8], device=device)
+        out_pos, out_conf = cluster_waters_vdw(pos, conf, radius=1.52)
+        assert out_pos.size(0) == 1
+        expected_x = (0.9 * 0.0 + 0.8 * 1.0) / (0.9 + 0.8)
+        assert out_pos[0, 0].item() == pytest.approx(expected_x, abs=1e-5)
+        assert out_conf[0].item() == pytest.approx(0.9, abs=1e-6)  # cluster max
+
+    def test_two_far_waters_preserved(self, device):
+        # 2 A apart with r=1.52 => kept separate.
+        pos = torch.tensor([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], device=device)
+        conf = torch.tensor([0.9, 0.8], device=device)
+        out_pos, out_conf = cluster_waters_vdw(pos, conf, radius=1.52)
+        assert out_pos.size(0) == 2
+        assert out_conf[0].item() >= out_conf[1].item()
+
+    def test_threshold_filters_pre_cluster(self, device):
+        pos = torch.tensor(
+            [[0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [10.0, 0.0, 0.0]], device=device
+        )
+        conf = torch.tensor([0.1, 0.6, 0.9], device=device)
+        out_pos, out_conf = cluster_waters_vdw(pos, conf, radius=1.52, threshold=0.5)
+        assert out_pos.size(0) == 2  # 0.1 dropped
+        assert (out_conf >= 0.5).all()
+
+    def test_threshold_runs_before_clustering(self, device):
+        """
+        A sub-threshold neighbour must not pull the centroid.
+
+        Filtering after clustering would let it vote in the weighted mean, so
+        the surviving centroid would land off its own position.
+        """
+        pos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], device=device)
+        conf = torch.tensor([0.9, 0.1], device=device)
+        out_pos, _ = cluster_waters_vdw(pos, conf, radius=1.52, threshold=0.5)
+        assert out_pos.size(0) == 1
+        assert out_pos[0, 0].item() == pytest.approx(0.0, abs=1e-6)
+
+    def test_output_sorted_descending_confidence(self, device):
+        torch.manual_seed(0)
+        pos = torch.randn(30, 3, device=device) * 5.0
+        conf = torch.rand(30, device=device)
+        _, out_conf = cluster_waters_vdw(pos, conf, radius=1.52)
+        assert (out_conf[:-1] >= out_conf[1:]).all()
+
+    def test_no_pair_within_radius_after_clustering(self, device):
+        torch.manual_seed(1)
+        pos = torch.randn(50, 3, device=device) * 3.0
+        conf = torch.rand(50, device=device)
+        out_pos, _ = cluster_waters_vdw(pos, conf, radius=1.52)
+        if out_pos.size(0) > 1:
+            diffs = out_pos.unsqueeze(0) - out_pos.unsqueeze(1)
+            dists = diffs.norm(dim=-1)
+            mask = ~torch.eye(out_pos.size(0), dtype=torch.bool, device=device)
+            assert (dists[mask] > 1.52 - 1e-5).all()
+
+    def test_nms_operates_on_centroids_not_seeds(self, device):
+        """
+        Round 2 must compare centroids, not the seeds they came from.
+
+        Seeds at 0.0 and 1.6 are 1.6 A apart, so round 1 keeps them separate.
+        Absorbing the 1.5 A neighbour drags the first centroid to ~0.54, which
+        brings the pair inside the radius -- a collision that only exists after
+        the weighting.
+        """
+        pos = torch.tensor(
+            [[0.0, 0.0, 0.0], [1.6, 0.0, 0.0], [10.0, 0.0, 0.0], [1.5, 0.0, 0.0]],
+            device=device,
+        )
+        conf = torch.tensor([0.9, 0.8, 0.7, 0.5], device=device)
+        out_pos, out_conf = cluster_waters_vdw(pos, conf, radius=1.52)
+        assert out_pos.size(0) == 2
+        assert out_conf.tolist() == pytest.approx([0.9, 0.7], abs=1e-6)
+        assert out_pos[0, 0].item() == pytest.approx(0.75 / 1.4, abs=1e-5)
+        assert out_pos[1, 0].item() == pytest.approx(10.0, abs=1e-6)
+
+    def test_cluster_confidence_is_the_max_not_the_mean(self, device):
+        pos = torch.zeros(3, 3, device=device)
+        conf = torch.tensor([0.9, 0.5, 0.1], device=device)
+        _, out_conf = cluster_waters_vdw(pos, conf, radius=1.52)
+        assert out_conf.shape == (1,)
+        assert out_conf[0].item() == pytest.approx(0.9, abs=1e-6)
+
+    def test_all_zero_confidence_falls_back_to_unweighted_mean(self, device):
+        """Weights summing to zero must yield a position, not NaN."""
+        pos = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], device=device)
+        conf = torch.zeros(2, device=device)
+        out_pos, out_conf = cluster_waters_vdw(pos, conf, radius=1.52)
+        assert out_pos.size(0) == 1
+        assert torch.isfinite(out_pos).all()
+        assert out_pos[0, 0].item() == pytest.approx(0.5, abs=1e-6)
+        assert out_conf[0].item() == pytest.approx(0.0, abs=1e-6)
+
+    def test_empty_input(self, device):
+        pos = torch.empty(0, 3, device=device)
+        conf = torch.empty(0, device=device)
+        out_pos, out_conf = cluster_waters_vdw(pos, conf, radius=1.52)
+        assert out_pos.shape == (0, 3)
+        assert out_conf.shape == (0,)
+
+    def test_threshold_drops_everything(self, device):
+        pos = torch.randn(5, 3, device=device)
+        conf = torch.zeros(5, device=device)
+        out_pos, out_conf = cluster_waters_vdw(pos, conf, radius=1.52, threshold=0.5)
+        assert out_pos.shape == (0, 3)
+        assert out_conf.shape == (0,)
+
+    def test_shape_validation(self, device):
+        pos = torch.randn(5, 3, device=device)
+        conf = torch.randn(4, device=device)  # wrong size
+        with pytest.raises(ValueError):
+            cluster_waters_vdw(pos, conf, radius=1.52)
+
+    def test_positions_shape_validation(self, device):
+        pos = torch.randn(5, 2, device=device)
+        conf = torch.randn(5, device=device)
+        with pytest.raises(ValueError, match=r"positions must be \(N, 3\)"):
+            cluster_waters_vdw(pos, conf, radius=1.52)
+
+
+# ============== ConfidenceGVP ==============
+
+
+def _make_hetero(device, n_prot=10, n_wat=5, cached_pw=False):
+    """
+    A minimal single-graph HeteroData with cached PP edges.
+
+    Args:
+        device: Device to build on.
+        n_prot: Number of protein atoms.
+        n_wat: Number of candidate waters.
+        cached_pw: Also attach a PW edge_index, standing in for a dataset that
+            supplies pre-built protein->water edges.
+    """
+    data = HeteroData()
+    data["protein"].pos = torch.randn(n_prot, 3, device=device)
+    idx = torch.randint(0, 16, (n_prot,), device=device)
+    data["protein"].x = F.one_hot(idx, num_classes=16).float()
+    data["protein"].batch = torch.zeros(n_prot, dtype=torch.long, device=device)
+
+    data["water"].pos = torch.randn(n_wat, 3, device=device)
+    wat_idx = torch.full((n_wat,), 2, dtype=torch.long, device=device)
+    data["water"].x = F.one_hot(wat_idx, num_classes=16).float()
+    data["water"].batch = torch.zeros(n_wat, dtype=torch.long, device=device)
+
+    pp_edge_index = torch.tensor(
+        [[0, 1, 2, 3], [1, 2, 3, 4]], dtype=torch.long, device=device
+    )
+    edge_unit_vectors, edge_rbf = compute_edge_features(
+        data["protein"].pos, pp_edge_index, num_gaussians=NUM_RBF, cutoff=8.0
+    )
+    data[EDGE_PP].edge_index = pp_edge_index
+    data[EDGE_PP].edge_unit_vectors = edge_unit_vectors
+    data[EDGE_PP].edge_rbf = edge_rbf
+
+    if cached_pw:
+        # Every candidate takes protein atom 0 as its sole source.
+        data[EDGE_PW].edge_index = torch.stack(
+            [
+                torch.zeros(n_wat, dtype=torch.long, device=device),
+                torch.arange(n_wat, dtype=torch.long, device=device),
+            ]
+        )
+    return data
+
+
+@pytest.mark.unit
+class TestConfidenceGVP:
+    def test_forward_output_shape(self, device, gvp_encoder):
+        data = _make_hetero(device, n_prot=10, n_wat=5)
+        model = ConfidenceGVP(
+            encoder=gvp_encoder,
+            hidden_dims=(64, 8),
+            layers=1,
+        ).to(device)
+
+        scores = model(data)
+        assert scores.shape == (5,)
+        assert (scores >= 0).all()
+        assert (scores <= 1.0).all()
+
+    def test_forward_no_water(self, device, gvp_encoder):
+        data = HeteroData()
+        data["protein"].pos = torch.randn(10, 3, device=device)
+        data["protein"].x = torch.randn(10, 16, device=device)
+        data["protein"].batch = torch.zeros(10, dtype=torch.long, device=device)
+        data[EDGE_PP].edge_index = torch.tensor(
+            [[0, 1], [1, 2]], dtype=torch.long, device=device
+        )
+
+        model = ConfidenceGVP(
+            encoder=gvp_encoder,
+            hidden_dims=(64, 8),
+            layers=1,
+        ).to(device)
+        scores = model(data)
+        assert scores.shape == (0,)
+
+    def test_return_logits(self, device, gvp_encoder):
+        data = _make_hetero(device, n_prot=10, n_wat=5)
+        model = ConfidenceGVP(
+            encoder=gvp_encoder,
+            hidden_dims=(64, 8),
+            layers=1,
+        ).to(device)
+        logits = model(data, return_logits=True)
+        assert logits.shape == (5,)
+
+        # Dropout makes two training-mode passes disagree; compare under eval.
+        model.eval()
+        with torch.no_grad():
+            logits2 = model(data, return_logits=True)
+            probs2 = model(data, return_logits=False)
+        assert torch.allclose(probs2, torch.sigmoid(logits2), atol=1e-6)
+
+    def test_no_time_conditioning_in_scalar_encoders(self, device, gvp_encoder):
+        """
+        The scalar encoders take the bare hidden width, not width+1.
+
+        FlowWaterGVP appends a time channel; a confidence model that inherited
+        it would silently consume a feature slot that is never populated.
+        """
+        model = ConfidenceGVP(
+            encoder=gvp_encoder, hidden_dims=(64, 8), layers=1, water_input_dim=16
+        ).to(device)
+        assert model.protein_scalar_encoder[0].in_features == 64
+        assert model.water_scalar_encoder[0].in_features == 16
+
+    def test_only_pw_and_pp_edge_types_are_active(self, device, gvp_encoder):
+        """Candidates are not refined here, so WW and WP carry nothing."""
+        model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to(
+            device
+        )
+        assert set(model.updater.etypes) == {EDGE_PW, EDGE_PP}
+
+    def test_gradients_flow(self, device, gvp_encoder):
+        data = _make_hetero(device, n_prot=10, n_wat=5)
+        model = ConfidenceGVP(
+            encoder=gvp_encoder,
+            hidden_dims=(64, 8),
+            layers=1,
+        ).to(device)
+
+        scores = model(data)
+        target = torch.rand_like(scores)
+        loss = F.mse_loss(scores, target)
+        loss.backward()
+
+        has_grad = any(
+            p.grad is not None and p.grad.abs().sum().item() > 0
+            for p in model.parameters()
+        )
+        assert has_grad
+
+    def test_score_head_receives_gradient(self, device, gvp_encoder):
+        """A head that never learns would leave the backbone doing all the work."""
+        data = _make_hetero(device, n_prot=10, n_wat=5)
+        model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to(
+            device
+        )
+        F.mse_loss(model(data), torch.rand(5, device=device)).backward()
+        assert any(
+            p.grad is not None and p.grad.abs().sum().item() > 0
+            for p in model.score_head.parameters()
+        )
+
+
+@pytest.mark.unit
+class TestConfidenceGVPCachedEdges:
+    def test_cached_edges_skip_dynamic_construction(self, device, gvp_encoder):
+        """
+        With PP and PW both supplied, forward must build no edges at all.
+
+        ConfidenceGVP runs on those two types only, so any call into
+        `build_dynamic_edges` means a cached tensor was ignored and the graph
+        was silently rebuilt.
+        """
+        data = _make_hetero(device, n_prot=8, n_wat=4, cached_pw=True)
+        model = ConfidenceGVP(
+            encoder=gvp_encoder,
+            hidden_dims=(64, 8),
+            layers=1,
+        ).to(device)
+
+        with patch("src.flow.build_dynamic_edges") as build_mock:
+            scores = model(data)
+
+        assert build_mock.call_count == 0, (
+            f"build_dynamic_edges was called {build_mock.call_count} times -- "
+            "cached PW/PP edges should make edge construction unnecessary."
+        )
+        assert scores.shape == (4,)
+        assert (scores >= 0).all() and (scores <= 1.0).all()
+
+    def test_cached_pw_edges_are_used_verbatim(self, device, gvp_encoder):
+        data = _make_hetero(device, n_prot=8, n_wat=4, cached_pw=True)
+        model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to(
+            device
+        )
+        edges = model.updater.build_edges(data)
+        assert torch.equal(edges[EDGE_PW], data[EDGE_PW].edge_index)
+
+    def test_uncached_pw_edges_are_built_by_radius(self, device, gvp_encoder):
+        """Without cached PW, the radius query runs and every candidate is reached."""
+        data = _make_hetero(device, n_prot=8, n_wat=4, cached_pw=False)
+        model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to(
+            device
+        )
+        edges = model.updater.build_edges(data)
+        assert EDGE_PW in edges
+        assert set(edges[EDGE_PW][1].tolist()) == set(range(4))
+
+    def test_candidate_beyond_the_cutoff_still_gets_edges(self, device, gvp_encoder):
+        """Candidates can land in sparse regions; unreached ones would go unscored."""
+        data = _make_hetero(device, n_prot=8, n_wat=4, cached_pw=False)
+        data["water"].pos[3] = torch.tensor([500.0, 500.0, 500.0], device=device)
+        model = ConfidenceGVP(encoder=gvp_encoder, hidden_dims=(64, 8), layers=1).to(
+            device
+        )
+        edges = model.updater.build_edges(data)
+        assert 3 in set(edges[EDGE_PW][1].tolist())