Refactor channels config

src/ocf_data_sampler/config/model.py

@@ -1,5 +1,6 @@
 """Configuration model for the PVNet dataset."""
 
+from collections.abc import Sequence
 from typing import Annotated, Literal
 
 from pydantic import AfterValidator, BaseModel, Field, field_validator, model_validator
@@ -55,6 +56,13 @@ def _validate_probability_sum(probabilities: dict[int, float]) -> dict[int, floa
 ]
 
 
+def _validate_unique_channel_names(channels: Sequence["Channel"]) -> None:
+    """Reject duplicate source channel names."""
+    names = [channel.name for channel in channels]
+    if len(names) != len(set(names)):
+        raise ValueError("Channel names must be unique")
+
+
 class Base(BaseModel, extra="forbid"):
     """Base configuration model that rejects unknown fields."""
 
@@ -147,8 +155,10 @@ class SpatialWindowMixin(Base):
     )
 
 
-class NormalisationValues(Base):
-    """Normalisation parameters."""
+class Channel(Base):
+    """A source channel and its normalisation parameters."""
+
+    name: str = Field(..., description="Channel name in the source data")
 
     mean: float = Field(..., description="Mean used for normalisation")
     std: float = Field(..., gt=0, description="Standard deviation used for normalisation")
@@ -162,7 +172,7 @@ class NormalisationValues(Base):
     )
 
     @model_validator(mode="after")
-    def validate_clip_range(self) -> "NormalisationValues":
+    def validate_clip_range(self) -> "Channel":
         """Require the minimum clipping bound to be below the maximum when both are set."""
         if (
             self.clip_min is not None
@@ -175,13 +185,23 @@ def validate_clip_range(self) -> "NormalisationValues":
         return self
 
 
-class NormalisationConstantsMixin(Base):
-    """Normalisation constants for multiple channels."""
+class NWPChannel(Channel):
+    """A source channel with normalisation parameters and optional differencing.
+
+    Clipping and normalisation apply to the differences when accumulated is True.
+    """
+
+    accumulated: bool = Field(
+        default=False,
+        description="Whether to difference the channel before normalisation. The normalisation "
+        "values (`mean`, `std`, `min_clip`, and `max_clip`) apply to the differenced values. "
+        "Each difference is labelled with the earlier step's timestamp, and the output channel "
+        "name is prefixed with diff_.",
+    )
 
-    normalisation_constants: dict[str, NormalisationValues]
 
 
-class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin, NormalisationConstantsMixin):
+class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin):
     """Satellite configuration model."""
 
     zarr_path: ZarrSource = Field(
@@ -190,7 +210,7 @@ class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin, Normalisation
         "to read from alternative filesystems.",
     )
 
-    channels: list[str] = Field(
+    channels: list[Channel] = Field(
         ...,
         min_length=1,
         description="Channels to include, in output order",
@@ -207,20 +227,14 @@ def validate_dropout_interval_end(self) -> "Satellite":
         _validate_dropout_interval_end(self.dropout_minute_probabilities, self.interval_end_minutes)
         return self
 
-    @model_validator(mode="after")
-    def check_all_channels_have_normalisation_constants(self) -> "Satellite":
-        """Check that all the channels have normalisation constants."""
-        normalisation_channels = set(self.normalisation_constants)
-        missing_channels = set(self.channels) - normalisation_channels
-        if missing_channels:
-            raise ValueError(
-                "Normalisation constants must be provided for all channels. Missing values for "
-                f"channels: {missing_channels}",
-            )
-        return self
+    @field_validator("channels")
+    def validate_channel_names(cls, channels: list[Channel]) -> list[Channel]:
+        """Reject duplicate source channel names."""
+        _validate_unique_channel_names(channels)
+        return channels
 
 
-class NWP(TimeWindowMixin, FillValueMixin, SpatialWindowMixin, NormalisationConstantsMixin):
+class NWP(TimeWindowMixin, FillValueMixin, SpatialWindowMixin):
     """NWP configuration model."""
 
     zarr_path: ZarrSource = Field(
@@ -229,7 +243,7 @@ class NWP(TimeWindowMixin, FillValueMixin, SpatialWindowMixin, NormalisationCons
         "to read from alternative filesystems.",
     )
 
-    channels: list[str] = Field(
+    channels: list[NWPChannel] = Field(
         ...,
         min_length=1,
         description="Channels to include, in output order",
@@ -244,8 +258,6 @@ class NWP(TimeWindowMixin, FillValueMixin, SpatialWindowMixin, NormalisationCons
 
     provider: str = Field(..., description="The provider of the NWP data")
 
-    accum_channels: list[str] = Field([], description="The channels which need to be diffed")
-
     max_staleness_minutes: int | None = Field(
         None,
         description="Sets a limit on how stale an NWP init-time is allowed to be whilst still "
@@ -258,17 +270,6 @@ def validate_resolution(cls, resolution: int) -> int:
         """Validate the supported temporal resolution."""
         return _validate_grid_resolution_minutes(resolution)
 
-    @model_validator(mode="after")
-    def validate_accum_channels_subset(self) -> "NWP":
-        """Validate accum_channels is subset of channels."""
-        invalid_channels = set(self.accum_channels) - set(self.channels)
-        if invalid_channels:
-            raise ValueError(
-                f"NWP provider '{self.provider}': all values in 'accum_channels' should "
-                f"be present in 'channels'. Extra values found: {invalid_channels}",
-            )
-        return self
-
     @field_validator("provider")
     def validate_provider(cls, provider: str) -> str:
         """Require a registered provider name."""
@@ -277,29 +278,11 @@ def validate_provider(cls, provider: str) -> str:
             raise ValueError(f"Unknown NWP provider {provider!r}. Supported: {supported}")
         return provider
 
-    @model_validator(mode="after")
-    def check_all_channels_have_normalisation_constants(self) -> "NWP":
-        """Check that all the channels have normalisation constants."""
-        normalisation_channels = set(self.normalisation_constants)
-        non_accum_channels = [c for c in self.channels if c not in self.accum_channels]
-        accum_channel_names = [f"diff_{c}" for c in self.accum_channels]
-
-        missing_channels = set(non_accum_channels) - normalisation_channels
-        if missing_channels:
-            raise ValueError(
-                "Normalisation constants must be provided for all channels. Missing values for "
-                f"channels: {missing_channels}",
-            )
-
-        missing_channels = set(accum_channel_names) - normalisation_channels
-        if missing_channels:
-            raise ValueError(
-                "Normalisation constants must be provided for all channels. Accumulated "
-                "channels which will be diffed require normalisation constant names which "
-                "start with the prefix 'diff_'. The following channels were missing: "
-                f"{missing_channels}.",
-            )
-        return self
+    @field_validator("channels")
+    def validate_channel_names(cls, channels: list[NWPChannel]) -> list[NWPChannel]:
+        """Reject duplicate source channel names."""
+        _validate_unique_channel_names(channels)
+        return channels
 
 
 class GenerationWindow(Base):

src/ocf_data_sampler/datasets/pvnet/availability.py

@@ -70,7 +70,7 @@ def _find_non_generation_t0_periods(
                 max_delay = minutes(max(nwp_config.delay_minute_probabilities))
 
             # The last step of the forecast is lost if we have to diff channels
-            if len(nwp_config.accum_channels) > 0:
+            if any(channel.accumulated for channel in nwp_config.channels):
                 end_buffer = minutes(nwp_config.time_resolution_minutes)
             else:
                 end_buffer = minutes(0)

src/ocf_data_sampler/datasets/pvnet/loading.py

@@ -78,7 +78,7 @@ def get_dataset_dict(
                 source=f"nwp/{nwp_source}",
             )
 
-            da_nwp = da_nwp.sel(channel=list(nwp_config.channels))
+            da_nwp = da_nwp.sel(channel=[channel.name for channel in nwp_config.channels])
 
             datasets_dict["nwp"][nwp_source] = da_nwp
 
@@ -94,7 +94,7 @@ def get_dataset_dict(
             source="satellite",
         )
 
-        da_sat = da_sat.sel(channel=list(config.satellite.channels))
+        da_sat = da_sat.sel(channel=[channel.name for channel in config.satellite.channels])
 
         datasets_dict["sat"] = da_sat
 

src/ocf_data_sampler/datasets/pvnet/preprocess.py

@@ -5,7 +5,7 @@
 import numpy as np
 
 from ocf_data_sampler.common.types import TArray
-from ocf_data_sampler.config.model import NormalisationValues, PVNetDataConfig
+from ocf_data_sampler.config.model import Channel, PVNetDataConfig
 from ocf_data_sampler.datasets.pvnet.types import (
     NormalisationArrays,
     SourceDict,
@@ -15,9 +15,7 @@
 from ocf_data_sampler.select.dropout import apply_dropout, choose_delay
 
 
-def _build_source_normalisation_arrays(
-    values: Sequence[NormalisationValues]
-) -> NormalisationArrays:
+def _build_source_normalisation_arrays(values: Sequence[Channel]) -> NormalisationArrays:
     """Build parameters broadcastable over (time, channel, x, y)."""
     def channel_array(channel_values: list[float]) -> np.ndarray:
         return np.array(channel_values, dtype=np.float32)[None, :, None, None]
@@ -36,18 +34,11 @@ def build_normalisation_arrays(config: PVNetDataConfig) -> SourceDict[Normalisat
     if config.nwp is not None:
         normalisation_arrays["nwp"] = {}
         for nwp_source, nwp_config in config.nwp.items():
-            channels = [
-                f"diff_{channel}" if channel in nwp_config.accum_channels else channel
-                for channel in nwp_config.channels
-            ]
             normalisation_arrays["nwp"][nwp_source] = _build_source_normalisation_arrays(
-                [nwp_config.normalisation_constants[channel] for channel in channels]
+                nwp_config.channels,
             )
     if config.satellite is not None:
-        sat_config = config.satellite
-        normalisation_arrays["sat"] = _build_source_normalisation_arrays(
-            [sat_config.normalisation_constants[channel] for channel in sat_config.channels]
-        )
+        normalisation_arrays["sat"] = _build_source_normalisation_arrays(config.satellite.channels)
     return normalisation_arrays
 
 
@@ -127,7 +118,9 @@ def diff_nwp_data(dataset_dict: SourceDict, config: PVNetDataConfig) -> SourceDi
     """
     if "nwp" in dataset_dict:
         for nwp_key, da_nwp in dataset_dict["nwp"].items():
-            accum_channels = config.nwp[nwp_key].accum_channels
+            accum_channels = [
+                channel.name for channel in config.nwp[nwp_key].channels if channel.accumulated
+            ]
             if len(accum_channels)>0:
                 # diff_channels() is an in-place operation and modifies the input
                 dataset_dict["nwp"][nwp_key] = diff_channels(da_nwp, accum_channels)

src/ocf_data_sampler/datasets/pvnet/slicing.py

@@ -171,7 +171,7 @@ def slice_datasets_by_time(
             nwp_config = config.nwp[nwp_key]
 
             # Add a buffer if we need to diff some of the channels in time
-            if len(nwp_config.accum_channels)>0:
+            if any(channel.accumulated for channel in nwp_config.channels):
                 interval_end_mins = (
                     nwp_config.interval_end_minutes
                     + nwp_config.time_resolution_minutes

tests/config/test_config.py

@@ -149,24 +149,15 @@ def test_dropout_rejects_cutoffs_at_or_after_interval_end(
     _validate_configuration(config)
 
 
-def test_accum_channels_validation(config_filename):
-    """Test accum_channels validation with required normalization constants."""
-    config, nwp_name = _load_config_and_provider(config_filename)
-
-    # Test invalid channel scenario
-    invalid_config = config.model_copy(deep=True)
-    invalid_nwp = invalid_config.nwp[nwp_name]
-    invalid_nwp.accum_channels = ["invalid_channel"]
-
-    # Verify exact error message
-    expected_error = (
-        rf"nwp.{nwp_name}\n"
-        fr"  Value error, NWP provider '{nwp_name}': all values in 'accum_channels' "
-        r"should be present in 'channels'\. "
-        r"Extra values found: {'invalid_channel'}.*"
-    )
-    with pytest.raises(ValidationError, match=expected_error):
-        _validate_configuration(invalid_config)
+@pytest.mark.parametrize("source_name", ["satellite", "ukv"])
+def test_duplicate_channel_names(config_filename, source_name):
+    config = load_yaml_configuration(config_filename)
+    sources = {"satellite": config.satellite, "ukv": config.nwp["ukv"]}
+    source = sources[source_name]
+    source.channels.append(source.channels[0].model_copy())
+
+    with pytest.raises(ValidationError, match="Channel names must be unique"):
+        _validate_configuration(config)
 
 
 def test_generation_interval_divisibility_raises(config_filename):

tests/config/test_save.py

@@ -14,10 +14,8 @@ def test_save_yaml_configuration_basic(tmp_path, config_filename):
 def test_save_load_yaml_configuration(tmp_path, config_filename):
     """Make sure a saved configuration is the same after loading"""
 
-    # Start with this config
     initial_config = load_yaml_configuration(config_filename)
 
-    # Save it - then load and check it is identical
     filepath = tmp_path / "config.yaml"
     save_yaml_configuration(initial_config, filepath)
     assert load_yaml_configuration(filepath) == initial_config

tests/datasets/pvnet/test_loading.py

@@ -4,7 +4,9 @@
 import pyproj
 import xarray as xr
 
-from ocf_data_sampler.datasets.pvnet.loading import add_source_coordinates
+from ocf_data_sampler.config import load_yaml_configuration
+from ocf_data_sampler.config.model import Channel, NWPChannel
+from ocf_data_sampler.datasets.pvnet.loading import add_source_coordinates, get_dataset_dict
 from ocf_data_sampler.datasets.pvnet.types import Location
 from ocf_data_sampler.spatial import WGS84
 from tests.conftest import UK_SAT_AREA
@@ -49,3 +51,20 @@ def test_add_source_coordinates_different_geostationary_sources():
         )
         assert (coordinate.x_dim, coordinate.y_dim) == ("x_geostationary", "y_geostationary")
     assert sat_coord.x != cloud_coord.x
+
+
+def test_get_dataset_dict_preserves_channel_order(pvnet_config_filename):
+    config = load_yaml_configuration(pvnet_config_filename)
+    config.generation = None
+    config.satellite.channels = [
+        Channel(name="IR_039", mean=0, std=1), Channel(name="IR_016", mean=0, std=1),
+    ]
+    config.nwp["ukv"].channels = [
+        NWPChannel(name="t", mean=0, std=1), NWPChannel(name="si10", mean=0, std=1),
+    ]
+
+    datasets = get_dataset_dict(config, location_ids=[])
+
+    # Both configured orders differ from their source coordinate order.
+    np.testing.assert_array_equal(datasets["sat"]["channel"].values, ["IR_039", "IR_016"])
+    np.testing.assert_array_equal(datasets["nwp"]["ukv"]["channel"].values, ["t", "si10"])

tests/datasets/pvnet/test_preprocess.py

@@ -4,7 +4,7 @@
 import xarray as xr
 
 from ocf_data_sampler.config import load_yaml_configuration
-from ocf_data_sampler.config.model import NormalisationValues
+from ocf_data_sampler.config.model import Channel, NWPChannel
 from ocf_data_sampler.datasets.pvnet.preprocess import (
     apply_dropout_to_datasets,
     build_normalisation_arrays,
@@ -145,17 +145,14 @@ def test_apply_dropout_to_datasets(config_filename):
 
 def test_normalise_channels(config_filename):
     config = load_yaml_configuration(config_filename)
-    config.satellite.channels = ["b", "a"]
-    config.satellite.normalisation_constants = {
-        "a": NormalisationValues(mean=10, std=2),
-        "b": NormalisationValues(mean=2, std=2, clip_min=0, clip_max=4),
-    }
-    config.nwp["ukv"].channels = ["b", "a"]
-    config.nwp["ukv"].accum_channels = ["b"]
-    config.nwp["ukv"].normalisation_constants = {
-        "a": NormalisationValues(mean=10, std=2),
-        "diff_b": NormalisationValues(mean=2, std=2, clip_min=0, clip_max=4),
-    }
+    config.satellite.channels = [
+        Channel(name="b", mean=2, std=2, clip_min=0, clip_max=4),
+        Channel(name="a", mean=10, std=2),
+    ]
+    config.nwp["ukv"].channels = [
+        NWPChannel(name="b", accumulated=True, mean=2, std=2, clip_min=0, clip_max=4),
+        NWPChannel(name="a", mean=10, std=2),
+    ]
     sat = xr.DataArray(np.array([[-2, 2, 8], [6, 10, np.nan]], dtype=np.float32)[None, :, None, :])
     nwp = sat.copy(deep=True)
     datasets = {"sat": sat, "nwp": {"ukv": nwp}}
@@ -167,3 +164,4 @@ def test_normalise_channels(config_filename):
     for da in (sat, nwp):
         assert da.dtype == np.float32
         np.testing.assert_allclose(da.values, expected)
+