diff --git a/src/confidence_dataset.py b/src/confidence_dataset.py
new file mode 100644
index 0000000..27d55a1
--- /dev/null
+++ b/src/confidence_dataset.py
@@ -0,0 +1,210 @@
+"""
+Dataset for confidence-model training.
+
+`ConfidenceDataset` rides on the flow dataset/cache layout: it composes a
+`ProteinWaterDataset` (protein graph + embeddings + GT waters from
+`{processed_dir}/geometry[_mates]` + `esm/`) with per-structure candidate
+files (`<candidate_dir>/<pdb_id>.pt = {"candidate_pos": (Nc, 3)}`). Each item
+swaps the GT water nodes for the sampled candidates and computes the target on
+the fly, so no bespoke confidence-cache format is needed.
+
+Targets are computed per item rather than cached: benchmarking showed this is
+~80x cheaper than the disk load that already happens, and caching would bake
+the sharpness hyperparameters into the files.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from loguru import logger
+from torch import Tensor
+from torch.utils.data import Dataset
+from torch_geometric.data import HeteroData
+
+from src.confidence import smootherstep_confidence
+from src.constants import ELEM_IDX, ELEMENT_VOCAB
+from src.dataset import ProteinWaterDataset
+
+
+WATER_FEATURE_DIM = len(ELEMENT_VOCAB) + 1  # matches element_onehot in src/dataset.py
+OXYGEN_INDEX = ELEM_IDX["O"]
+
+
+def _oxygen_features(n: int, device: torch.device | None = None) -> Tensor:
+    """Oxygen one-hot feature tensor, equal to `element_onehot(['O'] * n)`.
+
+    Candidates are all oxygen by construction, so this skips the per-atom
+    string lookups of the general encoder (~6x faster in `__getitem__`, which
+    runs per structure per epoch); equivalence is pinned by a unit test.
+    """
+    idx = torch.full((n,), OXYGEN_INDEX, dtype=torch.long, device=device)
+    return F.one_hot(idx, num_classes=WATER_FEATURE_DIM).float()
+
+
+class ConfidenceDataset(Dataset):
+    """Confidence dataset over the flow dataset/cache layout + candidate files.
+
+    Composes a `ProteinWaterDataset` (the flow model's own dataset -- protein
+    graph, embeddings, GT waters, PP edges) with a directory of per-structure
+    candidate files. For each structure it swaps the GT water nodes for the
+    flow-sampled candidates and computes the target on the fly from each
+    candidate's distance to its nearest GT water.
+
+    PW edges are NOT cached, even though the candidate positions are fixed:
+    at scoring time candidates come straight from the flow sampler with no
+    cache in sight, so `ConfidenceGVP` builds its PW edges dynamically
+    (`dynamic_edge_policy="knn_if_isolated"`) and training must exercise that
+    same path to avoid a train/inference skew. `max_candidates` also redraws
+    the candidate subset each epoch, which would invalidate cached edge
+    indices anyway -- and rebuilding is cheap next to storing edge indices +
+    RBF features per structure. The candidate file therefore stores only
+    `candidate_pos`.
+
+    Args:
+        flow_dataset: A constructed `ProteinWaterDataset` (or compatible) whose
+            items are flow `HeteroData` with `data["water"].pos` = GT waters and
+            `data.pdb_id` = the cache key. Exposes `.entries[i]["cache_key"]`.
+        candidate_dir: Directory holding one candidate file per structure,
+            `<pdb_id>.pt = {"candidate_pos": (Nc, 3)}`.
+        r_in, r_out: smootherstep plateau / floor radii (Å) -- target is 1
+            within `r_in`, decays C2-smoothly, and is 0 past `r_out`.
+        hard_label: Train on `1[d <= accept_radius]` instead of the soft target.
+        accept_radius: Radius (Å) defining the binary AUC-PR label.
+        max_candidates: Optional per-structure cap on the candidate cloud.
+        strict: If True, every structure must have a candidate file. If False,
+            structures without one are dropped (logged).
+    """
+
+    def __init__(
+        self,
+        flow_dataset: ProteinWaterDataset,
+        candidate_dir: str | Path,
+        *,
+        r_in: float = 0.5,
+        r_out: float = 1.5,
+        hard_label: bool = False,
+        accept_radius: float = 1.0,
+        max_candidates: int | None = None,
+        strict: bool = True,
+    ):
+        self.flow_dataset = flow_dataset
+        self.candidate_dir = Path(candidate_dir)
+        if not self.candidate_dir.exists():
+            raise FileNotFoundError(
+                f"Candidate directory not found: {self.candidate_dir}. "
+                "Build it with scripts/cache_candidates.py first."
+            )
+        self.r_in = float(r_in)
+        self.r_out = float(r_out)
+        self.hard_label = bool(hard_label)
+        self.accept_radius = float(accept_radius)
+        self.max_candidates = max_candidates
+        if self.r_out <= self.r_in:
+            raise ValueError("r_out must exceed r_in.")
+        if self.accept_radius < 0:
+            raise ValueError("accept_radius must be non-negative.")
+        if self.max_candidates is not None and self.max_candidates < 0:
+            raise ValueError("max_candidates must be non-negative.")
+
+        # Map each flow-dataset index to its candidate file (cheap existence
+        # checks via the entries list).
+        entries = getattr(flow_dataset, "entries", None)
+        if entries is None:
+            raise TypeError(
+                "flow_dataset must expose `.entries` (a ProteinWaterDataset)."
+            )
+        self._indices: list[int] = []
+        self._paths: list[Path] = []  # per kept index: its candidate file
+        missing: list[str] = []
+        for i, entry in enumerate(entries):
+            key = entry["cache_key"]
+            path = self.candidate_dir / f"{key}.pt"
+            if path.exists():
+                self._indices.append(i)
+                self._paths.append(path)
+            else:
+                missing.append(key)
+
+        if missing and strict:
+            raise FileNotFoundError(
+                f"{len(missing)} structures lack a candidate file under "
+                f"{self.candidate_dir} (first: {missing[0]}). Generate them "
+                "with scripts/cache_candidates.py or pass strict=False."
+            )
+        if missing:
+            logger.warning(
+                f"ConfidenceDataset: skipping {len(missing)} structures with no "
+                f"candidate file (first: {missing[0]})."
+            )
+        if not self._indices:
+            raise RuntimeError(
+                f"ConfidenceDataset: no candidate files matched under "
+                f"{self.candidate_dir} for the {len(entries)} requested entries."
+            )
+        logger.info(
+            f"ConfidenceDataset: {len(self._indices)} structures; smootherstep "
+            f"target (r_in={self.r_in}, r_out={self.r_out})."
+        )
+
+    def __len__(self) -> int:
+        return len(self._indices)
+
+    def _compute_targets(
+        self, candidate_pos: Tensor, gt_pos: Tensor
+    ) -> tuple[Tensor, Tensor, Tensor]:
+        """Return (training target, within-`accept_radius` label, nearest-GT index).
+
+        The nearest-GT distance / argmin is computed once and reused for: the
+        regression target (soft smootherstep, or a hard 1[d<=accept_radius] when
+        `hard_label`), the AUC-PR label (1[d<=accept_radius]), and the
+        per-candidate nearest-GT index (for the optional coverage loss, which
+        groups candidates by the GT site they could cover).
+        """
+        if candidate_pos.numel() == 0 or gt_pos.numel() == 0:
+            empty = candidate_pos.new_empty(0)
+            return empty, empty, candidate_pos.new_empty(0, dtype=torch.long)
+        d, gt_index = torch.cdist(candidate_pos, gt_pos).min(dim=1)  # (Nc,), (Nc,)
+        label = (d <= self.accept_radius).float()
+        target = (
+            label
+            if self.hard_label
+            else smootherstep_confidence(d, r_in=self.r_in, r_out=self.r_out)
+        )
+        return target, label, gt_index
+
+    def __getitem__(self, idx: int) -> HeteroData:
+        flow_idx = self._indices[idx]
+        data = self.flow_dataset[flow_idx]
+
+        gt_pos: Tensor = data["water"].pos.float().clone()
+        candidate_pos: Tensor = torch.load(
+            self._paths[idx], map_location="cpu", weights_only=False
+        )["candidate_pos"].float()
+        # Optional per-structure cap: random subsample of the candidate cloud.
+        # Free for quality (candidates are scored independently -- no
+        # water-water edges) but bounds per-step memory. A fresh draw each
+        # epoch covers all candidates over training.
+        if (
+            self.max_candidates is not None
+            and candidate_pos.size(0) > self.max_candidates
+        ):
+            sel = torch.randperm(candidate_pos.size(0))[: self.max_candidates]
+            candidate_pos = candidate_pos[sel]
+        n_cand = candidate_pos.size(0)
+
+        target, label_1A, gt_index = self._compute_targets(candidate_pos, gt_pos)
+
+        # Swap GT water nodes for the candidates to be scored.
+        data["water"].pos = candidate_pos
+        data["water"].x = _oxygen_features(n_cand, device=candidate_pos.device)
+        data["water"].num_nodes = n_cand
+        data["water"].target_confidence = target
+        data["water"].label_1A = label_1A
+        data["water"].gt_index = gt_index
+        data.n_gt = torch.tensor([gt_pos.size(0)], dtype=torch.long)
+        data["water"].gt_pos = gt_pos
+
+        return data
diff --git a/tests/test_confidence_dataset.py b/tests/test_confidence_dataset.py
new file mode 100644
index 0000000..9cc6033
--- /dev/null
+++ b/tests/test_confidence_dataset.py
@@ -0,0 +1,235 @@
+"""Unit tests for src/confidence_dataset.py -- candidate join and target computation."""
+
+import pytest
+import torch
+import torch.nn.functional as F
+from torch_geometric.data import HeteroData
+
+from src.confidence import smootherstep_target
+from src.confidence_dataset import _oxygen_features, ConfidenceDataset
+from src.constants import EDGE_PP, ELEM_IDX, NUM_RBF
+from src.dataset import element_onehot
+
+
+class _FakeFlowDataset:
+    """Minimal stand-in for ProteinWaterDataset: exposes `.entries` and yields
+    flow HeteroData with GT waters at `data["water"].pos` and `data.pdb_id`."""
+
+    def __init__(self, keys, gt_by_key=None, n_prot=6, n_gt=3, embedding_dim=None):
+        self.entries = [{"cache_key": k} for k in keys]
+        self._gt_by_key = gt_by_key or {}
+        self._n_prot = n_prot
+        self._n_gt = n_gt
+        self._embedding_dim = embedding_dim
+
+    def __len__(self):
+        return len(self.entries)
+
+    def __getitem__(self, idx):
+        key = self.entries[idx]["cache_key"]
+        data = HeteroData()
+        data["protein"].x = F.one_hot(
+            torch.randint(0, 16, (self._n_prot,)), num_classes=16
+        ).float()
+        data["protein"].pos = torch.randn(self._n_prot, 3)
+        data["protein"].residue_index = torch.arange(self._n_prot, dtype=torch.long)
+        data["protein"].is_ligand = torch.zeros(self._n_prot, dtype=torch.bool)
+        data["protein"].num_nodes = self._n_prot
+        if self._embedding_dim is not None:
+            data["protein"].embedding = torch.randn(self._n_prot, self._embedding_dim)
+            data["protein"].embedding_type = "esm"
+
+        gt = self._gt_by_key.get(key)
+        if gt is None:
+            gt = torch.randn(self._n_gt, 3)
+        data["water"].pos = gt
+        data["water"].x = F.one_hot(
+            torch.full((gt.size(0),), ELEM_IDX["O"]), num_classes=16
+        ).float()
+        data["water"].num_nodes = gt.size(0)
+
+        data[EDGE_PP].edge_index = torch.tensor([[0, 1, 2], [1, 2, 3]])
+        data[EDGE_PP].edge_unit_vectors = torch.randn(3, 3)
+        data[EDGE_PP].edge_rbf = torch.randn(3, NUM_RBF)
+        data.pdb_id = key
+        return data
+
+
+def _write_candidate(path, cand):
+    torch.save({"candidate_pos": cand}, path)
+
+
+@pytest.mark.unit
+class TestConfidenceDataset:
+    def test_joins_candidates_and_computes_target(self, tmp_path):
+        cand = torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [10.0, 0.0, 0.0]])
+        gt = torch.tensor([[0.0, 0.0, 0.0]])
+        flow = _FakeFlowDataset(["pdb_a"], gt_by_key={"pdb_a": gt}, embedding_dim=8)
+        _write_candidate(tmp_path / "pdb_a.pt", cand)
+
+        ds = ConfidenceDataset(flow, candidate_dir=tmp_path)
+        assert len(ds) == 1
+        sample = ds[0]
+
+        # candidates became the scored water nodes
+        assert sample["water"].num_nodes == 3
+        assert torch.equal(sample["water"].pos, cand)
+        assert (sample["water"].x[:, ELEM_IDX["O"]] == 1.0).all()  # oxygen one-hot
+        assert torch.equal(sample["water"].gt_pos, gt)
+
+        # target on the fly: candidate on GT -> ~1; far candidate -> ~0
+        tc = sample["water"].target_confidence
+        assert tc.shape == (3,)
+        assert tc[0].item() == pytest.approx(1.0, abs=1e-5)
+        assert tc[2].item() < 0.05
+
+        # protein graph + embedding + PP edges carried from the flow item
+        assert sample["protein"].num_nodes == 6
+        assert sample["protein"].embedding.shape == (6, 8)
+        assert sample[EDGE_PP].edge_index.shape == (2, 3)
+        assert sample.pdb_id == "pdb_a"
+
+    def test_default_target_is_smootherstep(self, tmp_path):
+        torch.manual_seed(1)
+        cand = torch.randn(9, 3)
+        gt = torch.randn(3, 3)
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": gt})
+        _write_candidate(tmp_path / "k.pt", cand)
+        ds = ConfidenceDataset(flow, candidate_dir=tmp_path, r_in=0.4, r_out=2.0)
+        expected = smootherstep_target(cand, gt, r_in=0.4, r_out=2.0)
+        assert torch.allclose(ds[0]["water"].target_confidence, expected, atol=1e-5)
+
+    def test_empty_candidates(self, tmp_path):
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": torch.randn(2, 3)})
+        _write_candidate(tmp_path / "k.pt", torch.empty(0, 3))
+        ds = ConfidenceDataset(flow, candidate_dir=tmp_path)
+        sample = ds[0]
+        assert sample["water"].num_nodes == 0
+        assert sample["water"].target_confidence.shape == (0,)
+
+    def test_strict_missing_raises(self, tmp_path):
+        flow = _FakeFlowDataset(["have", "gone"])
+        _write_candidate(tmp_path / "have.pt", torch.randn(3, 3))
+        with pytest.raises(FileNotFoundError):
+            ConfidenceDataset(flow, candidate_dir=tmp_path, strict=True)
+
+    def test_non_strict_filters_missing(self, tmp_path):
+        flow = _FakeFlowDataset(["have", "gone"])
+        _write_candidate(tmp_path / "have.pt", torch.randn(3, 3))
+        ds = ConfidenceDataset(flow, candidate_dir=tmp_path, strict=False)
+        assert len(ds) == 1
+        assert ds[0].pdb_id == "have"
+
+    def test_missing_candidate_dir_raises(self, tmp_path):
+        flow = _FakeFlowDataset(["k"])
+        with pytest.raises(FileNotFoundError):
+            ConfidenceDataset(flow, candidate_dir=tmp_path / "missing")
+
+    def test_invalid_parameters_raise(self, tmp_path):
+        flow = _FakeFlowDataset(["k"])
+        _write_candidate(tmp_path / "k.pt", torch.randn(2, 3))
+        with pytest.raises(ValueError):
+            ConfidenceDataset(flow, candidate_dir=tmp_path, r_in=1.5, r_out=0.5)
+        with pytest.raises(ValueError):
+            ConfidenceDataset(flow, candidate_dir=tmp_path, accept_radius=-1.0)
+        with pytest.raises(ValueError):
+            ConfidenceDataset(flow, candidate_dir=tmp_path, max_candidates=-3)
+
+    def test_no_candidates_at_all_raises(self, tmp_path):
+        flow = _FakeFlowDataset(["k"])
+        with pytest.raises(RuntimeError):
+            ConfidenceDataset(flow, candidate_dir=tmp_path, strict=False)
+
+    def test_flow_dataset_without_entries_raises(self, tmp_path):
+        _write_candidate(tmp_path / "k.pt", torch.randn(2, 3))
+        with pytest.raises(TypeError):
+            ConfidenceDataset(object(), candidate_dir=tmp_path)
+
+    def test_label_and_gt_index_track_nearest_gt(self, tmp_path):
+        # candidates at 0.5 A and 3 A from GT site 1; accept_radius default 1.0
+        gt = torch.tensor([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]])
+        cand = torch.tensor([[10.5, 0.0, 0.0], [13.0, 0.0, 0.0]])
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": gt})
+        _write_candidate(tmp_path / "k.pt", cand)
+        sample = ConfidenceDataset(flow, candidate_dir=tmp_path)[0]
+
+        assert torch.equal(sample["water"].label_1A, torch.tensor([1.0, 0.0]))
+        assert torch.equal(sample["water"].gt_index, torch.tensor([1, 1]))
+        assert torch.equal(sample.n_gt, torch.tensor([2]))
+
+    def test_hard_label_replaces_the_soft_target(self, tmp_path):
+        gt = torch.tensor([[0.0, 0.0, 0.0]])
+        cand = torch.tensor([[0.8, 0.0, 0.0], [1.2, 0.0, 0.0]])
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": gt})
+        _write_candidate(tmp_path / "k.pt", cand)
+
+        soft = ConfidenceDataset(flow, candidate_dir=tmp_path)[0]
+        hard = ConfidenceDataset(flow, candidate_dir=tmp_path, hard_label=True)[0]
+
+        # soft target is strictly between the plateau and the floor at 0.8 A
+        assert 0.0 < soft["water"].target_confidence[0].item() < 1.0
+        assert torch.equal(hard["water"].target_confidence, hard["water"].label_1A)
+        assert torch.equal(hard["water"].target_confidence, torch.tensor([1.0, 0.0]))
+
+    def test_accept_radius_widens_the_label(self, tmp_path):
+        gt = torch.tensor([[0.0, 0.0, 0.0]])
+        cand = torch.tensor([[1.5, 0.0, 0.0]])
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": gt})
+        _write_candidate(tmp_path / "k.pt", cand)
+        narrow = ConfidenceDataset(flow, candidate_dir=tmp_path)[0]
+        wide = ConfidenceDataset(flow, candidate_dir=tmp_path, accept_radius=2.0)[0]
+        assert narrow["water"].label_1A.item() == 0.0
+        assert wide["water"].label_1A.item() == 1.0
+
+    def test_max_candidates_subsamples_the_cloud(self, tmp_path):
+        gt = torch.tensor([[0.0, 0.0, 0.0]])
+        cand = torch.arange(20, dtype=torch.float32).reshape(20, 1).repeat(1, 3)
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": gt})
+        _write_candidate(tmp_path / "k.pt", cand)
+
+        ds = ConfidenceDataset(flow, candidate_dir=tmp_path, max_candidates=5)
+        sample = ds[0]
+        assert sample["water"].num_nodes == 5
+        assert sample["water"].target_confidence.shape == (5,)
+        # a subset of the candidate cloud, not a truncation
+        rows = {tuple(r.tolist()) for r in sample["water"].pos}
+        assert rows.issubset({tuple(r.tolist()) for r in cand})
+
+        # A fresh draw each epoch: pin the RNG so the two reads are known to
+        # subsample differently (an unseeded pair could collide, rarely).
+        torch.manual_seed(0)
+        first = ds[0]["water"].pos
+        torch.manual_seed(1)
+        second = ds[0]["water"].pos
+        assert not torch.equal(first, second)
+
+    def test_max_candidates_above_the_cloud_is_a_no_op(self, tmp_path):
+        gt = torch.tensor([[0.0, 0.0, 0.0]])
+        cand = torch.randn(4, 3)
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": gt})
+        _write_candidate(tmp_path / "k.pt", cand)
+        sample = ConfidenceDataset(flow, candidate_dir=tmp_path, max_candidates=10)[0]
+        assert torch.equal(sample["water"].pos, cand)
+
+    def test_no_gt_waters_yields_empty_targets(self, tmp_path):
+        # A structure whose waters were all filtered out. Targets come back
+        # empty while num_nodes stays at the candidate count -- such an item
+        # cannot be collated, so the trainer must exclude these structures.
+        flow = _FakeFlowDataset(["k"], gt_by_key={"k": torch.empty(0, 3)})
+        _write_candidate(tmp_path / "k.pt", torch.randn(3, 3))
+        sample = ConfidenceDataset(flow, candidate_dir=tmp_path)[0]
+        assert sample["water"].num_nodes == 3
+        assert sample["water"].target_confidence.shape == (0,)
+        assert sample["water"].label_1A.shape == (0,)
+        assert sample["water"].gt_index.shape == (0,)
+
+
+@pytest.mark.unit
+class TestOxygenFeatures:
+    def test_matches_element_onehot(self):
+        assert torch.equal(_oxygen_features(4), element_onehot(["O"] * 4))
+
+    def test_empty(self):
+        feats = _oxygen_features(0)
+        assert feats.shape == (0, 16)
+        assert feats.dtype == torch.float32