Address Confidence Dataset review comments

src/confidence.py

@@ -32,7 +32,7 @@
 # ---------------------------------------------------------------------------
 
 
-def _nearest_gt_distance(candidate_pos: Tensor, gt_pos: Tensor) -> Tensor:
+def nearest_gt_distance(candidate_pos: Tensor, gt_pos: Tensor) -> Tensor:
     """
     Distance from each candidate to its nearest ground-truth water.
 
@@ -128,7 +128,7 @@ def smootherstep_target(
     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
+        nearest_gt_distance(candidate_pos, gt_pos), r_in=r_in, r_out=r_out
     )
 
 

src/confidence_dataset.py

@@ -24,12 +24,11 @@
 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.confidence import nearest_gt_distance, smootherstep_confidence
+from src.constants import ELEM_IDX, NODE_FEATURE_DIM
 from src.dataset import ProteinWaterDataset
 
 
-WATER_FEATURE_DIM = len(ELEMENT_VOCAB) + 1  # matches element_onehot in src/dataset.py
 OXYGEN_INDEX = ELEM_IDX["O"]
 
 
@@ -41,7 +40,7 @@ def _oxygen_features(n: int, device: torch.device | None = None) -> Tensor:
     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()
+    return F.one_hot(idx, num_classes=NODE_FEATURE_DIM).float()
 
 
 class ConfidenceDataset(Dataset):
@@ -106,8 +105,8 @@ def __init__(
             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.")
+        if self.max_candidates is not None and self.max_candidates < 1:
+            raise ValueError("max_candidates must be positive.")
 
         # Map each flow-dataset index to its candidate file (cheap existence
         # checks via the entries list).
@@ -154,34 +153,33 @@ def __len__(self) -> int:
 
     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).
+    ) -> tuple[Tensor, Tensor]:
+        """Return (training target, within-`accept_radius` label).
+
+        The nearest-GT distance (shared with `confidence.nearest_gt_distance`, so
+        both paths agree) drives the regression target (soft smootherstep, or a
+        hard 1[d<=accept_radius] when `hard_label`) and the AUC-PR label
+        (1[d<=accept_radius]).
         """
         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,)
+            return empty, empty
+        d = nearest_gt_distance(candidate_pos, gt_pos)  # (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
+        return target, label
 
     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
+            self._paths[idx], map_location="cpu", weights_only=True
         )["candidate_pos"].float()
         # Optional per-structure cap: random subsample of the candidate cloud.
         # Free for quality (candidates are scored independently -- no
@@ -195,16 +193,13 @@ def __getitem__(self, idx: int) -> HeteroData:
             candidate_pos = candidate_pos[sel]
         n_cand = candidate_pos.size(0)
 
-        target, label_1A, gt_index = self._compute_targets(candidate_pos, gt_pos)
+        target, within_accept_radius = 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
+        data["water"].within_accept_radius = within_accept_radius
 
         return data

tests/test_confidence_dataset.py

@@ -75,7 +75,6 @@ def test_joins_candidates_and_computes_target(self, tmp_path):
         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
@@ -133,7 +132,7 @@ def test_invalid_parameters_raise(self, tmp_path):
         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)
+            ConfidenceDataset(flow, candidate_dir=tmp_path, max_candidates=0)
 
     def test_no_candidates_at_all_raises(self, tmp_path):
         flow = _FakeFlowDataset(["k"])
@@ -145,17 +144,17 @@ def test_flow_dataset_without_entries_raises(self, tmp_path):
         with pytest.raises(TypeError):
             ConfidenceDataset(object(), candidate_dir=tmp_path)
 
-    def test_label_and_gt_index_track_nearest_gt(self, tmp_path):
+    def test_label_tracks_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]))
+        assert torch.equal(
+            sample["water"].within_accept_radius, torch.tensor([1.0, 0.0])
+        )
 
     def test_hard_label_replaces_the_soft_target(self, tmp_path):
         gt = torch.tensor([[0.0, 0.0, 0.0]])
@@ -168,7 +167,9 @@ def test_hard_label_replaces_the_soft_target(self, tmp_path):
 
         # 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, hard["water"].within_accept_radius
+        )
         assert torch.equal(hard["water"].target_confidence, torch.tensor([1.0, 0.0]))
 
     def test_accept_radius_widens_the_label(self, tmp_path):
@@ -178,8 +179,8 @@ def test_accept_radius_widens_the_label(self, tmp_path):
         _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
+        assert narrow["water"].within_accept_radius.item() == 0.0
+        assert wide["water"].within_accept_radius.item() == 1.0
 
     def test_max_candidates_subsamples_the_cloud(self, tmp_path):
         gt = torch.tensor([[0.0, 0.0, 0.0]])
@@ -220,8 +221,32 @@ def test_no_gt_waters_yields_empty_targets(self, tmp_path):
         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,)
+        assert sample["water"].within_accept_radius.shape == (0,)
+
+    def test_items_batch_via_pyg(self, tmp_path):
+        """Different-sized candidate clouds must collate through
+        Batch.from_data_list -- what a DataLoader does to train on many at once."""
+        from torch_geometric.data import Batch
+
+        gt = torch.tensor([[0.0, 0.0, 0.0]])
+        flow = _FakeFlowDataset(
+            ["a", "b"], gt_by_key={"a": gt, "b": gt}, embedding_dim=8
+        )
+        _write_candidate(tmp_path / "a.pt", torch.randn(3, 3))
+        _write_candidate(tmp_path / "b.pt", torch.randn(5, 3))
+        ds = ConfidenceDataset(flow, candidate_dir=tmp_path)
+
+        batch = Batch.from_data_list([ds[0], ds[1]])
+
+        assert batch.num_graphs == 2
+        # candidate water nodes and their per-candidate targets concatenate (3 + 5)
+        assert batch["water"].num_nodes == 8
+        assert batch["water"].target_confidence.shape == (8,)
+        assert batch["water"].within_accept_radius.shape == (8,)
+        assert batch["water"].batch.tolist() == [0, 0, 0, 1, 1, 1, 1, 1]
+        # protein graph and PP edges carried and offset per graph
+        assert batch["protein"].num_nodes == 12
+        assert batch[EDGE_PP].edge_index.shape[1] == 6
 
 
 @pytest.mark.unit