diff --git a/src/confidence.py b/src/confidence.py index 761c30a..0cf4292 100644 --- a/src/confidence.py +++ b/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 ) diff --git a/src/confidence_dataset.py b/src/confidence_dataset.py index 27d55a1..b2e7f7a 100644 --- a/src/confidence_dataset.py +++ b/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,26 +153,25 @@ 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] @@ -181,7 +179,7 @@ def __getitem__(self, idx: int) -> HeteroData: 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 diff --git a/tests/test_confidence_dataset.py b/tests/test_confidence_dataset.py index 9cc6033..56b9767 100644 --- a/tests/test_confidence_dataset.py +++ b/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,7 +144,7 @@ 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]]) @@ -153,9 +152,9 @@ def test_label_and_gt_index_track_nearest_gt(self, tmp_path): _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