diff --git a/src/ocf_data_sampler/config/model.py b/src/ocf_data_sampler/config/model.py index 480bb6fd..295e9682 100644 --- a/src/ocf_data_sampler/config/model.py +++ b/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): diff --git a/src/ocf_data_sampler/datasets/pvnet/availability.py b/src/ocf_data_sampler/datasets/pvnet/availability.py index b75a4ba0..33d7c5d2 100644 --- a/src/ocf_data_sampler/datasets/pvnet/availability.py +++ b/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) diff --git a/src/ocf_data_sampler/datasets/pvnet/loading.py b/src/ocf_data_sampler/datasets/pvnet/loading.py index e99fe506..bbdeadaa 100644 --- a/src/ocf_data_sampler/datasets/pvnet/loading.py +++ b/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 diff --git a/src/ocf_data_sampler/datasets/pvnet/preprocess.py b/src/ocf_data_sampler/datasets/pvnet/preprocess.py index d49c8503..513a4ebd 100644 --- a/src/ocf_data_sampler/datasets/pvnet/preprocess.py +++ b/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) diff --git a/src/ocf_data_sampler/datasets/pvnet/slicing.py b/src/ocf_data_sampler/datasets/pvnet/slicing.py index bede4cc8..c0d50196 100644 --- a/src/ocf_data_sampler/datasets/pvnet/slicing.py +++ b/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 diff --git a/tests/config/test_config.py b/tests/config/test_config.py index b4caaab1..b6784791 100644 --- a/tests/config/test_config.py +++ b/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): diff --git a/tests/config/test_save.py b/tests/config/test_save.py index f91408fa..27b00ed7 100644 --- a/tests/config/test_save.py +++ b/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 diff --git a/tests/datasets/pvnet/test_loading.py b/tests/datasets/pvnet/test_loading.py index c4fafe58..e1b6d017 100644 --- a/tests/datasets/pvnet/test_loading.py +++ b/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"]) diff --git a/tests/datasets/pvnet/test_preprocess.py b/tests/datasets/pvnet/test_preprocess.py index bb1f402a..b34b8fee 100644 --- a/tests/datasets/pvnet/test_preprocess.py +++ b/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) + diff --git a/tests/fixtures/configs/pvnet_test_config.yaml b/tests/fixtures/configs/pvnet_test_config.yaml index 2cb48b7e..2f69b495 100644 --- a/tests/fixtures/configs/pvnet_test_config.yaml +++ b/tests/fixtures/configs/pvnet_test_config.yaml @@ -21,18 +21,16 @@ nwp: interval_end_minutes: 120 time_resolution_minutes: 60 channels: - - t # 2-metre temperature + - name: t # 2-metre temperature + mean: 283.64913206 + std: 4.38818501 + clip_min: 270 + clip_max: 310 image_size_pixels_height: 2 image_size_pixels_width: 2 delay_minute_probabilities: 180: 1.0 max_staleness_minutes: null - normalisation_constants: - t: - mean: 283.64913206 - std: 4.38818501 - clip_min: 270 - clip_max: 310 satellite: zarr_path: set_in_temp_file @@ -40,13 +38,11 @@ satellite: interval_end_minutes: 0 time_resolution_minutes: 5 channels: - - IR_016 - image_size_pixels_height: 2 - image_size_pixels_width: 2 - normalisation_constants: - IR_016: + - name: IR_016 mean: 0.17594202 std: 0.21462157 + image_size_pixels_height: 2 + image_size_pixels_width: 2 solar_position: interval_start_minutes: -15