Address PR review comments

- Fix collate_fn fallback to work with dict batches and return consistent dict
- Gate train/val part building in setup() by stage
- Remove dead `if s.root is None` guard
- Add comment in lightning.py about Dataset_ID for future multi-head use
- Add val_frac comment in configs clarifying it is ignored when split_file is set
- Add tests for AddDatasetID, collate_fn fallback, DatasetSpec, RhoRead init and setup

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: Hananeh Oliaei

src/electrai/configs/MP/config_resnet.yaml

@@ -4,7 +4,7 @@ data:
   datasets:
     - root: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/mp_filelist.txt
       split_file: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/split.json
-      val_frac: 0.005
+      val_frac: 0.005  # ignored when split_file is provided
   precision: f32
   batch_size: 1
   train_workers: 8

src/electrai/configs/MP/config_resunet.yaml

@@ -4,7 +4,7 @@ data:
   datasets:
     - root: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/mp_filelist.txt
       split_file: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/dataset_2/split.json
-      val_frac: 0.005
+      val_frac: 0.005  # ignored when split_file is provided
   precision: f32
   batch_size: 1
   train_workers: 8

src/electrai/dataloader/collate.py

@@ -7,5 +7,4 @@ def collate_fn(batch):
     try:
         return default_collate(batch)
     except RuntimeError:
-        x, y, index, dataset_id = zip(*batch, strict=True)
-        return list(x), list(y), list(index), list(dataset_id)
+        return {k: [d[k] for d in batch] for k in batch[0]}

src/electrai/dataloader/dataset.py

@@ -79,10 +79,6 @@ def __init__(
         ]
         filled: list[DatasetSpec] = []
         for i, s in enumerate(specs):
-            if s.root is None:
-                raise ValueError(
-                    f"`root` is required for dataset {i + 1}. Received: {s.root!r}"
-                )
             dataset_id = s.dataset_id if s.dataset_id is not None else i
             filled.append(
                 DatasetSpec(
@@ -120,10 +116,12 @@ def setup(self, stage: str):
 
             dataset_id = int(spec.dataset_id)
 
-            train_parts.append(AddDatasetID(splits["train"], dataset_id))
-            val_parts.append(AddDatasetID(splits["validation"], dataset_id))
-            if "test" in splits and splits["test"] is not None:
-                test_parts.append(AddDatasetID(splits["test"], dataset_id))
+            if stage == "fit":
+                train_parts.append(AddDatasetID(splits["train"], dataset_id))
+                val_parts.append(AddDatasetID(splits["validation"], dataset_id))
+            elif stage == "test":
+                if "test" in splits and splits["test"] is not None:
+                    test_parts.append(AddDatasetID(splits["test"], dataset_id))
 
         if stage == "fit":
             self.train_set = (

src/electrai/lightning.py

@@ -43,6 +43,7 @@ def validation_step(self, batch):
         return loss
 
     def _loss_calculation(self, batch):
+        # batch["Dataset_ID"] is available for future multi-head model extensions
         x = batch["data"]
         y = batch["label"]
         if isinstance(x, list):

tests/electrai/dataloader/test_dataset.py

@@ -1,19 +1,27 @@
 from __future__ import annotations
 
+import dataclasses
 import json
 from typing import TYPE_CHECKING
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pytest
+import torch
 import zarr
+from torch.utils.data import Dataset
 
-from electrai.dataloader.dataset import RhoData
+from electrai.dataloader.collate import collate_fn
+from electrai.dataloader.dataset import AddDatasetID, DatasetSpec, RhoData, RhoRead
 from electrai.dataloader.utils import load_zarr
 
 if TYPE_CHECKING:
     from pathlib import Path
 
 
+# --- Fixtures ---
+
+
 @pytest.fixture
 def rng():
     return np.random.default_rng(seed=0)
@@ -39,6 +47,24 @@ def filelist(tmp_path: Path, zarr_root: Path) -> Path:
     return filelist_path
 
 
+class SimpleDataset(Dataset):
+    def __init__(self, n: int):
+        self.n = n
+
+    def __len__(self):
+        return self.n
+
+    def __getitem__(self, idx):
+        return {
+            "data": torch.tensor([float(idx)]),
+            "label": torch.tensor([float(idx)]),
+            "index": str(idx),
+        }
+
+
+# --- TestLoadZarr ---
+
+
 class TestLoadZarr:
     def test_returns_arrays_divided_by_volume(self, zarr_root: Path):
         data, label = load_zarr(zarr_root, "mp-1")
@@ -57,6 +83,9 @@ def test_missing_structure_attr_raises(self, tmp_path: Path, rng):
             load_zarr(tmp_path, "bad")
 
 
+# --- TestRhoDataFormatDetection ---
+
+
 class TestRhoDataFormatDetection:
     def test_detects_zarr_format(self, filelist: Path):
         dataset = RhoData(str(filelist), precision="f32", augmentation=False)
@@ -90,3 +119,270 @@ def test_getitem_zarr(self, filelist: Path):
         assert "label" in item
         assert "index" in item
         assert item["data"].shape[0] == 1  # unsqueeze adds channel dim
+
+
+# --- TestAddDatasetID ---
+
+
+class TestAddDatasetID:
+    def test_len_preserved(self):
+        base = SimpleDataset(10)
+        wrapped = AddDatasetID(base, functional_id=3)
+        assert len(wrapped) == 10
+
+    def test_dataset_id_added(self):
+        base = SimpleDataset(5)
+        wrapped = AddDatasetID(base, functional_id=7)
+        sample = wrapped[0]
+        assert sample["Dataset_ID"] == 7
+
+    def test_dataset_id_is_int(self):
+        base = SimpleDataset(3)
+        wrapped = AddDatasetID(base, functional_id="2")
+        assert isinstance(wrapped[0]["Dataset_ID"], int)
+
+    def test_original_keys_preserved(self):
+        base = SimpleDataset(3)
+        wrapped = AddDatasetID(base, functional_id=1)
+        sample = wrapped[0]
+        assert "data" in sample
+        assert "label" in sample
+        assert "index" in sample
+
+    def test_custom_key(self):
+        base = SimpleDataset(3)
+        wrapped = AddDatasetID(base, functional_id=5, key="src_id")
+        assert "src_id" in wrapped[0]
+        assert wrapped[0]["src_id"] == 5
+
+
+# --- TestCollateFn ---
+
+
+class TestCollateFn:
+    def test_uniform_shapes_uses_default_collate(self):
+        batch = [
+            {
+                "data": torch.zeros(4, 4, 4),
+                "label": torch.ones(4, 4, 4),
+                "index": "a",
+                "Dataset_ID": 0,
+            },
+            {
+                "data": torch.zeros(4, 4, 4),
+                "label": torch.ones(4, 4, 4),
+                "index": "b",
+                "Dataset_ID": 0,
+            },
+        ]
+        result = collate_fn(batch)
+        assert isinstance(result, dict)
+        assert result["data"].shape == (2, 4, 4, 4)
+
+    def test_mismatched_shapes_fallback_returns_dict(self):
+        batch = [
+            {
+                "data": torch.zeros(4, 4, 4),
+                "label": torch.ones(4, 4, 4),
+                "index": "a",
+                "Dataset_ID": 0,
+            },
+            {
+                "data": torch.zeros(8, 8, 8),
+                "label": torch.ones(8, 8, 8),
+                "index": "b",
+                "Dataset_ID": 1,
+            },
+        ]
+        result = collate_fn(batch)
+        assert isinstance(result, dict)
+        assert len(result["data"]) == 2
+        assert len(result["label"]) == 2
+
+    def test_fallback_preserves_all_keys(self):
+        batch = [
+            {
+                "data": torch.zeros(4, 4, 4),
+                "label": torch.ones(4, 4, 4),
+                "index": "a",
+                "Dataset_ID": 0,
+            },
+            {
+                "data": torch.zeros(8, 8, 8),
+                "label": torch.ones(8, 8, 8),
+                "index": "b",
+                "Dataset_ID": 1,
+            },
+        ]
+        result = collate_fn(batch)
+        assert set(result.keys()) == {"data", "label", "index", "Dataset_ID"}
+
+    def test_fallback_correct_values(self):
+        batch = [
+            {
+                "data": torch.zeros(2, 2, 2),
+                "label": torch.ones(2, 2, 2),
+                "index": "x",
+                "Dataset_ID": 3,
+            },
+            {
+                "data": torch.zeros(3, 3, 3),
+                "label": torch.ones(3, 3, 3),
+                "index": "y",
+                "Dataset_ID": 5,
+            },
+        ]
+        result = collate_fn(batch)
+        assert result["index"] == ["x", "y"]
+        assert result["Dataset_ID"] == [3, 5]
+
+
+# --- TestDatasetSpec ---
+
+
+class TestDatasetSpec:
+    def test_required_root(self):
+        spec = DatasetSpec(root="/some/path")
+        assert spec.root == "/some/path"
+
+    def test_defaults_are_none(self):
+        spec = DatasetSpec(root="/p")
+        assert spec.split_file is None
+        assert spec.val_frac is None
+        assert spec.dataset_id is None
+
+    def test_frozen(self):
+        spec = DatasetSpec(root="/p")
+        with pytest.raises(dataclasses.FrozenInstanceError):
+            spec.root = "/other"
+
+
+# --- TestRhoReadInit ---
+
+
+class TestRhoReadInit:
+    def test_empty_datasets_raises(self):
+        with pytest.raises(ValueError, match="at least one"):
+            RhoRead(datasets=[])
+
+    def test_none_datasets_raises(self):
+        with pytest.raises(ValueError, match="at least one"):
+            RhoRead(datasets=None)
+
+    def test_auto_assigns_dataset_id(self):
+        reader = RhoRead(
+            datasets=[{"root": "/a/filelist.txt"}, {"root": "/b/filelist.txt"}]
+        )
+        assert reader.specs[0].dataset_id == 0
+        assert reader.specs[1].dataset_id == 1
+
+    def test_explicit_dataset_id_preserved(self):
+        reader = RhoRead(datasets=[{"root": "/a/filelist.txt", "dataset_id": 42}])
+        assert reader.specs[0].dataset_id == 42
+
+    def test_default_val_frac_applied(self):
+        reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}], default_val_frac=0.1)
+        assert reader.specs[0].val_frac == 0.1
+
+    def test_per_spec_val_frac_overrides_default(self):
+        reader = RhoRead(
+            datasets=[{"root": "/a/filelist.txt", "val_frac": 0.2}],
+            default_val_frac=0.1,
+        )
+        assert reader.specs[0].val_frac == 0.2
+
+    def test_default_split_file_applied(self):
+        reader = RhoRead(
+            datasets=[{"root": "/a/filelist.txt"}],
+            default_split_file="/default/split.json",
+        )
+        assert reader.specs[0].split_file == "/default/split.json"
+
+    def test_per_spec_split_file_overrides_default(self):
+        reader = RhoRead(
+            datasets=[{"root": "/a/filelist.txt", "split_file": "/custom/split.json"}],
+            default_split_file="/default/split.json",
+        )
+        assert reader.specs[0].split_file == "/custom/split.json"
+
+
+# --- TestRhoReadSetup ---
+
+
+class TestRhoReadSetup:
+    def _make_fake_splits(self, n: int = 5):
+        ds = SimpleDataset(n)
+        from torch.utils.data import Subset
+
+        train = Subset(ds, list(range(4)))
+        val = Subset(ds, [4])
+        return {"train": train, "validation": val}
+
+    @patch("electrai.dataloader.dataset.split_data")
+    @patch("electrai.dataloader.dataset.RhoData")
+    def test_fit_stage_builds_train_and_val(self, mock_rho, mock_split):
+        mock_rho.return_value = MagicMock(spec=Dataset)
+        mock_split.return_value = self._make_fake_splits()
+
+        reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}])
+        reader.setup("fit")
+
+        assert reader.train_set is not None
+        assert reader.val_set is not None
+        assert reader.test_set is None
+
+    @patch("electrai.dataloader.dataset.split_data")
+    @patch("electrai.dataloader.dataset.RhoData")
+    def test_test_stage_does_not_build_train_or_val(self, mock_rho, mock_split):
+        mock_rho.return_value = MagicMock(spec=Dataset)
+        splits = self._make_fake_splits()
+        splits["test"] = splits["train"]
+        mock_split.return_value = splits
+
+        reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}])
+        reader.setup("test")
+
+        assert reader.test_set is not None
+        assert reader.train_set is None
+        assert reader.val_set is None
+
+    @patch("electrai.dataloader.dataset.split_data")
+    @patch("electrai.dataloader.dataset.RhoData")
+    def test_multi_dataset_concat(self, mock_rho, mock_split):
+        from torch.utils.data import ConcatDataset
+
+        mock_rho.return_value = MagicMock(spec=Dataset)
+        mock_split.return_value = self._make_fake_splits()
+
+        reader = RhoRead(
+            datasets=[{"root": "/a/filelist.txt"}, {"root": "/b/filelist.txt"}]
+        )
+        reader.setup("fit")
+
+        assert isinstance(reader.train_set, ConcatDataset)
+        assert isinstance(reader.val_set, ConcatDataset)
+
+    @patch("electrai.dataloader.dataset.split_data")
+    @patch("electrai.dataloader.dataset.RhoData")
+    def test_single_dataset_no_concat(self, mock_rho, mock_split):
+        from torch.utils.data import ConcatDataset
+
+        mock_rho.return_value = MagicMock(spec=Dataset)
+        mock_split.return_value = self._make_fake_splits()
+
+        reader = RhoRead(datasets=[{"root": "/a/filelist.txt"}])
+        reader.setup("fit")
+
+        assert not isinstance(reader.train_set, ConcatDataset)
+
+    @patch("electrai.dataloader.dataset.split_data")
+    @patch("electrai.dataloader.dataset.RhoData")
+    def test_dataset_id_propagated(self, mock_rho, mock_split):
+        mock_rho.return_value = MagicMock(spec=Dataset)
+        mock_split.return_value = self._make_fake_splits()
+
+        reader = RhoRead(datasets=[{"root": "/a/filelist.txt", "dataset_id": 99}])
+        reader.setup("fit")
+
+        sample = reader.train_set[0]
+        assert sample["Dataset_ID"] == 99