Merge pull request #318 from alan-turing-institute/add-dataset-channel-idx-subset

Add channel index support

Author: sgreenbury

src/autocast/data/datamodule.py

@@ -168,9 +168,7 @@ def __init__(
         n_steps_input: int = 1,
         n_steps_output: int = 1,
         stride: int = 1,
-        # TODO: support for passing data from dict
-        input_channel_idxs: tuple[int, ...] | None = None,
-        output_channel_idxs: tuple[int, ...] | None = None,
+        channel_idxs: tuple[int, ...] | None = None,
         batch_size: int = 4,
         dtype: torch.dtype = torch.float32,
         ftype: str = "torch",
@@ -205,8 +203,7 @@ def __init__(
             n_steps_input=n_steps_input,
             n_steps_output=n_steps_output,
             stride=stride,
-            input_channel_idxs=input_channel_idxs,
-            output_channel_idxs=output_channel_idxs,
+            channel_idxs=channel_idxs,
             autoencoder_mode=self.autoencoder_mode,
             full_trajectory_mode=full_trajectory_mode,
             dtype=dtype,
@@ -237,8 +234,7 @@ def __init__(
             n_steps_input=n_steps_input,
             n_steps_output=n_steps_output,
             stride=stride,
-            input_channel_idxs=input_channel_idxs,
-            output_channel_idxs=output_channel_idxs,
+            channel_idxs=channel_idxs,
             autoencoder_mode=self.autoencoder_mode,
             full_trajectory_mode=full_trajectory_mode,
             dtype=dtype,
@@ -254,8 +250,7 @@ def __init__(
             n_steps_input=n_steps_input,
             n_steps_output=n_steps_output,
             stride=stride,
-            input_channel_idxs=input_channel_idxs,
-            output_channel_idxs=output_channel_idxs,
+            channel_idxs=channel_idxs,
             autoencoder_mode=self.autoencoder_mode,
             full_trajectory_mode=full_trajectory_mode,
             dtype=dtype,
@@ -275,8 +270,7 @@ def __init__(
                 n_steps_input=n_steps_input,
                 n_steps_output=n_steps_output,
                 stride=stride,
-                input_channel_idxs=input_channel_idxs,
-                output_channel_idxs=output_channel_idxs,
+                channel_idxs=channel_idxs,
                 full_trajectory_mode=True,
                 dtype=dtype,
                 verbose=self.verbose,
@@ -291,8 +285,7 @@ def __init__(
                 n_steps_input=n_steps_input,
                 n_steps_output=n_steps_output,
                 stride=stride,
-                input_channel_idxs=input_channel_idxs,
-                output_channel_idxs=output_channel_idxs,
+                channel_idxs=channel_idxs,
                 full_trajectory_mode=True,
                 dtype=dtype,
                 verbose=self.verbose,

src/autocast/data/dataset.py

@@ -31,16 +31,14 @@ def to_sample(data: dict) -> Sample:
 class SpatioTemporalDataset(Dataset, BatchMixin):
     """A class for spatio-temporal datasets."""
 
-    def __init__(
+    def __init__(  # noqa: PLR0915
         self,
         data_path: str | None,
         data: dict | None = None,
         n_steps_input: int = 1,
         n_steps_output: int = 1,
         stride: int = 1,
-        # TODO: support for passing data from dict
-        input_channel_idxs: tuple[int, ...] | None = None,
-        output_channel_idxs: tuple[int, ...] | None = None,
+        channel_idxs: tuple[int, ...] | None = None,
         full_trajectory_mode: bool = False,
         autoencoder_mode: bool = False,
         dtype: torch.dtype = torch.float32,
@@ -67,10 +65,9 @@ def __init__(
             Stride for sampling the data.
         data: dict | None
             Preloaded data. Defaults to None.
-        input_channel_idxs: tuple[int, ...] | None
-            Indices of input channels to use. Defaults to None.
-        output_channel_idxs: tuple[int, ...] | None
-            Indices of output channels to use. Defaults to None.
+        channel_idxs: tuple[int, ...] | None
+            Indices of channels to select from the raw data (applied to both
+            input and output). If None, all channels are used. Defaults to None.
         full_trajectory_mode: bool
             If True, use full trajectories without creating subtrajectories.
         autoencoder_mode: bool
@@ -104,8 +101,17 @@ def __init__(
         if data is not None:
             self.parse_data(data)
 
+        if channel_idxs is not None:
+            self.data = self.data[..., list(channel_idxs)]
+
         self.set_up_normalization()
 
+        if channel_idxs is not None and self.norm is not None:
+            self.norm.core_field_names = [
+                self.norm.core_field_names[i] for i in channel_idxs
+            ]
+            self.norm._precompute_flattened_stats()
+
         if autoencoder_mode and full_trajectory_mode:
             msg = "autoencoder_mode and full_trajectory_mode cannot both be True."
             raise ValueError(msg)
@@ -124,8 +130,7 @@ def __init__(
         self.n_steps_input = n_steps_input
         self.n_steps_output = n_steps_output
         self.stride = stride
-        self.input_channel_idxs = input_channel_idxs
-        self.output_channel_idxs = output_channel_idxs
+        self.channel_idxs = channel_idxs
 
         # Destructured here
         (

src/autocast/scripts/eval/encoder_processor_decoder.py

@@ -280,11 +280,13 @@ def _resolve_rollout_channel_names(dataset: Any) -> list[str] | None:
 
     norm = getattr(dataset, "norm", None)
     raw_names = getattr(norm, "core_field_names", None)
+    names_already_subset = raw_names is not None
 
     if not isinstance(raw_names, Sequence) or isinstance(raw_names, str):
         normalization_stats = getattr(dataset, "normalization_stats", None)
         if isinstance(normalization_stats, Mapping):
             raw_names = normalization_stats.get("core_field_names")
+            names_already_subset = False
 
     if not isinstance(raw_names, Sequence) or isinstance(raw_names, str):
         return None
@@ -293,14 +295,14 @@ def _resolve_rollout_channel_names(dataset: Any) -> list[str] | None:
     if not channel_names:
         return None
 
-    output_channel_idxs = getattr(dataset, "output_channel_idxs", None)
-    if output_channel_idxs is not None:
+    channel_idxs = getattr(dataset, "channel_idxs", None)
+    if channel_idxs is not None and not names_already_subset:
         try:
-            channel_names = [channel_names[idx] for idx in output_channel_idxs]
+            channel_names = [channel_names[idx] for idx in channel_idxs]
         except (TypeError, IndexError):
             log.warning(
-                "Could not apply output_channel_idxs=%s to channel names %s.",
-                output_channel_idxs,
+                "Could not apply channel_idxs=%s to channel names %s.",
+                channel_idxs,
                 channel_names,
             )
             return None

tests/data/test_dataset_normalization.py

@@ -137,6 +137,77 @@ def test_normalized_data_is_transformed(deterministic_data, stats_dict):
     )
 
 
+def test_channel_idxs_slices_data_and_subsets_norm(deterministic_data, stats_dict):
+    """`channel_idxs` should slice data channels and align norm field names."""
+    dataset = ReactionDiffusionDataset(
+        data_path=None,
+        data=deterministic_data,
+        n_steps_input=2,
+        n_steps_output=1,
+        channel_idxs=(1,),
+        use_normalization=True,
+        normalization_type=ZScoreNormalization,
+        normalization_stats=stats_dict,
+    )
+
+    # Sliced data keeps only channel 1 (V).
+    assert dataset.data.shape[-1] == 1
+    assert dataset[0].input_fields.shape[-1] == 1
+
+    # Norm field names subset to match sliced channels.
+    assert dataset.norm is not None
+    assert dataset.norm.core_field_names == ["V"]
+
+    # Normalization uses V stats (mean=4.0, std=2.0) against the original V channel.
+    expected = (deterministic_data["data"][0][:2, ..., 1] - 4.0) / 2.0
+    assert torch.allclose(dataset[0].input_fields[..., 0], expected)
+
+
+def test_channel_idxs_none_is_noop(deterministic_data):
+    """`channel_idxs=None` should leave all channels intact."""
+    dataset = ReactionDiffusionDataset(
+        data_path=None,
+        data=deterministic_data,
+        n_steps_input=2,
+        n_steps_output=1,
+        channel_idxs=None,
+        use_normalization=False,
+    )
+    assert dataset.data.shape[-1] == 2
+    assert dataset[0].input_fields.shape[-1] == 2
+
+
+def test_datamodule_threads_channel_idxs(deterministic_data, stats_dict):
+    """DataModule should propagate `channel_idxs` to all sub-datasets."""
+    dm = SpatioTemporalDataModule(
+        data_path=None,
+        data={
+            "train": deterministic_data,
+            "valid": deterministic_data,
+            "test": deterministic_data,
+        },
+        dataset_cls=ReactionDiffusionDataset,
+        n_steps_input=2,
+        n_steps_output=1,
+        batch_size=1,
+        channel_idxs=(0,),
+        use_normalization=True,
+        normalization_type=ZScoreNormalization,
+        normalization_stats=stats_dict,
+    )
+
+    for ds in (
+        dm.train_dataset,
+        dm.val_dataset,
+        dm.test_dataset,
+        dm.rollout_val_dataset,
+        dm.rollout_test_dataset,
+    ):
+        assert ds.data.shape[-1] == 1
+        assert ds.norm is not None
+        assert ds.norm.core_field_names == ["U"]
+
+
 def test_datamodule_with_and_without_normalization(deterministic_data, stats_dict):
     """Test DataModule can be configured with or without normalization."""
 

tests/scripts/test_eval_encoder_processor_decoder.py

@@ -281,10 +281,20 @@ def test_should_skip_metric_variogram_only():
     assert _should_skip_metric("ssr") is False
 
 
-def test_resolve_rollout_channel_names_from_norm_with_output_selection():
+def test_resolve_rollout_channel_names_from_norm_already_subset():
     dataset = SimpleNamespace(
-        norm=SimpleNamespace(core_field_names=["u", "v", "p"]),
-        output_channel_idxs=(2, 0),
+        norm=SimpleNamespace(core_field_names=["p", "u"]),
+        channel_idxs=(2, 0),
+    )
+
+    assert _resolve_rollout_channel_names(dataset) == ["p", "u"]
+
+
+def test_resolve_rollout_channel_names_from_stats_applies_idxs():
+    dataset = SimpleNamespace(
+        norm=None,
+        normalization_stats={"core_field_names": ["u", "v", "p"]},
+        channel_idxs=(2, 0),
     )
 
     assert _resolve_rollout_channel_names(dataset) == ["p", "u"]
@@ -294,16 +304,17 @@ def test_resolve_rollout_channel_names_returns_none_without_norm_names():
     dataset = SimpleNamespace(
         norm=None,
         metadata=SimpleNamespace(field_names={0: ["velocity_x", "velocity_y"]}),
-        output_channel_idxs=None,
+        channel_idxs=None,
     )
 
     assert _resolve_rollout_channel_names(dataset) is None
 
 
 def test_resolve_rollout_channel_names_returns_none_on_invalid_output_indices():
     dataset = SimpleNamespace(
-        norm=SimpleNamespace(core_field_names=["u", "v"]),
-        output_channel_idxs=(0, 3),
+        norm=None,
+        normalization_stats={"core_field_names": ["u", "v"]},
+        channel_idxs=(0, 3),
     )
 
     assert _resolve_rollout_channel_names(dataset) is None