ocf-data-sampler 0.1.11__py3-none-any.whl → 0.1.17__py3-none-any.whl

ocf_data_sampler/config/load.py

@@ -1,13 +1,13 @@
-"""Load configuration from a yaml file"""
+"""Load configuration from a yaml file."""
 
 import fsspec
 from pyaml_env import parse_config
+
 from ocf_data_sampler.config import Configuration
 
 
 def load_yaml_configuration(filename: str) -> Configuration:
-    """
-    Load a yaml file which has a configuration in it
+    """Load a yaml file which has a configuration in it.
 
     Args:
         filename: the yaml file name that you want to load.  Will load from local, AWS, or GCP

ocf_data_sampler/config/model.py

@@ -1,45 +1,50 @@
 """Configuration model for the dataset.
 
-
-Absolute or relative zarr filepath(s). Prefix with a protocol like s3:// to read from alternative filesystems. 
-
+Absolute or relative zarr filepath(s).
+Prefix with a protocol like s3:// to read from alternative filesystems.
 """
 
-from typing import Dict, List, Optional
-from typing_extensions import Self
+from collections.abc import Iterator
 
-from pydantic import BaseModel, Field, RootModel, field_validator, ValidationInfo, model_validator
+from pydantic import BaseModel, Field, RootModel, field_validator, model_validator
+from typing_extensions import override
 
-from ocf_data_sampler.constants import NWP_PROVIDERS
+NWP_PROVIDERS = [
+    "ukv",
+    "ecmwf",
+    "gfs",
+    "icon_eu",
+]
 
 
 class Base(BaseModel):
-    """Pydantic Base model where no extras can be added"""
+    """Pydantic Base model where no extras can be added."""
 
     class Config:
-        """config class"""
+        """Config class."""
 
         extra = "forbid"  # forbid use of extra kwargs
 
 
 class General(Base):
-    """General pydantic model"""
+    """General pydantic model."""
 
     name: str = Field("example", description="The name of this configuration file")
     description: str = Field(
-        "example configuration", description="Description of this configuration file"
+        "example configuration",
+        description="Description of this configuration file",
     )
 
 
 class TimeWindowMixin(Base):
-    """Mixin class, to add interval start, end and resolution minutes"""
+    """Mixin class, to add interval start, end and resolution minutes."""
 
     time_resolution_minutes: int = Field(
         ...,
         gt=0,
         description="The temporal resolution of the data in minutes",
     )
-    
+
     interval_start_minutes: int = Field(
         ...,
         description="Data interval starts at `t0 + interval_start_minutes`",
@@ -50,32 +55,33 @@ class TimeWindowMixin(Base):
         description="Data interval ends at `t0 + interval_end_minutes`",
     )
 
-    @model_validator(mode='after')
-    def validate_intervals(cls, values):
-        start = values.interval_start_minutes
-        end = values.interval_end_minutes
-        resolution = values.time_resolution_minutes
+    @model_validator(mode="after")
+    def validate_intervals(self) -> "TimeWindowMixin":
+        """Validator for time interval fields."""
+        start = self.interval_start_minutes
+        end = self.interval_end_minutes
+        resolution = self.time_resolution_minutes
         if start > end:
             raise ValueError(
-                f"interval_start_minutes ({start}) must be <= interval_end_minutes ({end})"
+                f"interval_start_minutes ({start}) must be <= interval_end_minutes ({end})",
             )
-        if (start % resolution != 0):
+        if start % resolution != 0:
             raise ValueError(
                 f"interval_start_minutes ({start}) must be divisible "
-                f"by time_resolution_minutes ({resolution})"
+                f"by time_resolution_minutes ({resolution})",
             )
-        if (end % resolution != 0):
+        if end % resolution != 0:
             raise ValueError(
                 f"interval_end_minutes ({end}) must be divisible "
-                f"by time_resolution_minutes ({resolution})"
+                f"by time_resolution_minutes ({resolution})",
             )
-        return values
+        return self
 
 
 class DropoutMixin(Base):
-    """Mixin class, to add dropout minutes"""
+    """Mixin class, to add dropout minutes."""
 
-    dropout_timedeltas_minutes: List[int] = Field(
+    dropout_timedeltas_minutes: list[int] = Field(
         default=[],
         description="List of possible minutes before t0 where data availability may start. Must be "
         "negative or zero.",
@@ -89,14 +95,16 @@ class DropoutMixin(Base):
     )
 
     @field_validator("dropout_timedeltas_minutes")
-    def dropout_timedeltas_minutes_negative(cls, v: List[int]) -> List[int]:
-        """Validate 'dropout_timedeltas_minutes'"""
+    def dropout_timedeltas_minutes_negative(cls, v: list[int]) -> list[int]:
+        """Validate 'dropout_timedeltas_minutes'."""
         for m in v:
-            assert m <= 0, "Dropout timedeltas must be negative"
+            if m > 0:
+                raise ValueError("Dropout timedeltas must be negative")
         return v
 
     @model_validator(mode="after")
-    def dropout_instructions_consistent(self) -> Self:
+    def dropout_instructions_consistent(self) -> "DropoutMixin":
+        """Validator for dropout instructions."""
         if self.dropout_fraction == 0:
             if self.dropout_timedeltas_minutes != []:
                 raise ValueError("To use dropout timedeltas dropout fraction should be > 0")
@@ -107,7 +115,7 @@ class DropoutMixin(Base):
 
 
 class SpatialWindowMixin(Base):
-    """Mixin class, to add path and image size"""
+    """Mixin class, to add path and image size."""
 
     image_size_pixels_height: int = Field(
         ...,
@@ -122,9 +130,37 @@ class SpatialWindowMixin(Base):
     )
 
 
-class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin):
-    """Satellite configuration model"""
-    
+class NormalisationValues(Base):
+    """Normalisation mean and standard deviation."""
+    mean: float = Field(..., description="Mean value for normalization")
+    std: float = Field(..., gt=0, description="Standard deviation (must be positive)")
+
+
+class NormalisationConstantsMixin(Base):
+    """Normalisation constants for multiple channels."""
+    normalisation_constants: dict[str, NormalisationValues]
+
+    @property
+    def channel_means(self) -> dict[str, float]:
+        """Return the channel means."""
+        return {
+            channel: norm_values.mean
+            for channel, norm_values in self.normalisation_constants.items()
+        }
+
+
+    @property
+    def channel_stds(self) -> dict[str, float]:
+        """Return the channel standard deviations."""
+        return {
+            channel: norm_values.std
+            for channel, norm_values in self.normalisation_constants.items()
+        }
+
+
+class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin, NormalisationConstantsMixin):
+    """Satellite configuration model."""
+
     zarr_path: str | tuple[str] | list[str] = Field(
         ...,
         description="Absolute or relative zarr filepath(s). Prefix with a protocol like s3:// "
@@ -132,82 +168,123 @@ class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin):
     )
 
     channels: list[str] = Field(
-        ..., description="the satellite channels that are used"
+        ...,
+        description="the satellite channels that are used",
     )
 
+    @model_validator(mode="after")
+    def check_all_channel_have_normalisation_constants(self) -> "Satellite":
+        """Check that all the channels have normalisation constants."""
+        normalisation_channels = set(self.normalisation_constants.keys())
+        missing_norm_values = set(self.channels) - set(normalisation_channels)
+        if len(missing_norm_values)>0:
+            raise ValueError(
+                "Normalsation constants must be provided for all channels. Missing values for "
+                f"channels: {missing_norm_values}",
+            )
+        return self
+
+
+class NWP(TimeWindowMixin, DropoutMixin, SpatialWindowMixin, NormalisationConstantsMixin):
+    """NWP configuration model."""
 
-class NWP(TimeWindowMixin, DropoutMixin, SpatialWindowMixin):
-    """NWP configuration model"""
-    
     zarr_path: str | tuple[str] | list[str] = Field(
         ...,
         description="Absolute or relative zarr filepath(s). Prefix with a protocol like s3:// "
         "to read from alternative filesystems.",
     )
-    
+
     channels: list[str] = Field(
-        ..., description="the channels used in the nwp data"
+        ...,
+        description="the channels used in the nwp data",
     )
 
     provider: str = Field(..., description="The provider of the NWP data")
 
     accum_channels: list[str] = Field([], description="the nwp channels which need to be diffed")
 
-    max_staleness_minutes: Optional[int] = Field(
+    max_staleness_minutes: int | None = Field(
         None,
         description="Sets a limit on how stale an NWP init time is allowed to be whilst still being"
         " used to construct an example. If set to None, then the max staleness is set according to"
         " the maximum forecast horizon of the NWP and the requested forecast length.",
     )
 
-
     @field_validator("provider")
     def validate_provider(cls, v: str) -> str:
-        """Validate 'provider'"""
+        """Validator for 'provider'."""
         if v.lower() not in NWP_PROVIDERS:
-            message = f"NWP provider {v} is not in {NWP_PROVIDERS}"
-            raise Exception(message)
+            raise OSError(f"NWP provider {v} is not in {NWP_PROVIDERS}")
         return v
 
 
+    @model_validator(mode="after")
+    def check_all_channel_have_normalisation_constants(self) -> "NWP":
+        """Check that all the channels have normalisation constants."""
+        normalisation_channels = set(self.normalisation_constants.keys())
+        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_norm_values = set(non_accum_channels) - set(normalisation_channels)
+        if len(missing_norm_values)>0:
+            raise ValueError(
+                "Normalsation constants must be provided for all channels. Missing values for "
+                f"channels: {missing_norm_values}",
+            )
+
+        missing_norm_values = set(accum_channel_names) - set(normalisation_channels)
+        if len(missing_norm_values)>0:
+            raise ValueError(
+                "Normalsation 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_norm_values}.",
+            )
+        return self
+
+
 class MultiNWP(RootModel):
-    """Configuration for multiple NWPs"""
+    """Configuration for multiple NWPs."""
 
-    root: Dict[str, NWP]
+    root: dict[str, NWP]
 
-    def __getattr__(self, item):
+    @override
+    def __getattr__(self, item: str) -> NWP:
         return self.root[item]
 
-    def __getitem__(self, item):
+    @override
+    def __getitem__(self, item: str) -> NWP:
         return self.root[item]
 
-    def __len__(self):
+    @override
+    def __len__(self) -> int:
         return len(self.root)
 
-    def __iter__(self):
+    @override
+    def __iter__(self) -> Iterator:
         return iter(self.root)
 
-    def keys(self):
-        """Returns dictionary-like keys"""
+    def keys(self) -> Iterator[str]:
+        """Returns dictionary-like keys."""
         return self.root.keys()
 
-    def items(self):
-        """Returns dictionary-like items"""
+    def items(self) -> Iterator[tuple[str, NWP]]:
+        """Returns dictionary-like items."""
         return self.root.items()
 
 
 class GSP(TimeWindowMixin, DropoutMixin):
-    """GSP configuration model"""
+    """GSP configuration model."""
 
     zarr_path: str = Field(
-        ..., 
+        ...,
         description="Absolute or relative zarr filepath. Prefix with a protocol like s3:// "
         "to read from alternative filesystems.",
     )
 
 
 class Site(TimeWindowMixin, DropoutMixin):
-    """Site configuration model"""
+    """Site configuration model."""
 
     file_path: str = Field(
         ...,
@@ -222,17 +299,22 @@ class Site(TimeWindowMixin, DropoutMixin):
     # TODO validate the csv for metadata
 
 
+class SolarPosition(TimeWindowMixin):
+    """Solar position configuration model."""
+
+
 class InputData(Base):
-    """Input data model"""
+    """Input data model."""
 
-    satellite: Optional[Satellite] = None
-    nwp: Optional[MultiNWP] = None
-    gsp: Optional[GSP] = None
-    site: Optional[Site] = None
+    satellite: Satellite | None = None
+    nwp: MultiNWP | None = None
+    gsp: GSP | None = None
+    site: Site | None = None
+    solar_position: SolarPosition | None = None
 
 
 class Configuration(Base):
-    """Configuration model for the dataset"""
+    """Configuration model for the dataset."""
 
     general: General = General()
     input_data: InputData = InputData()

ocf_data_sampler/config/save.py

@@ -5,12 +5,14 @@ supporting local and cloud storage locations.
 """
 
 import json
+import os
+
 import fsspec
 import yaml
-import os
 
 from ocf_data_sampler.config import Configuration
 
+
 def save_yaml_configuration(configuration: Configuration, filename: str) -> None:
     """Save a configuration object to a YAML file.
 
@@ -20,12 +22,11 @@ def save_yaml_configuration(configuration: Configuration, filename: str) -> None
                  cloud storage URL (e.g., 'gs://', 's3://'). For local paths,
                  absolute paths are recommended.
     """
-    
     if os.path.exists(filename):
         raise FileExistsError(f"File already exists: {filename}")
 
     # Serialize configuration to JSON-compatible dictionary
     config_dict = json.loads(configuration.model_dump_json())
 
-    with fsspec.open(filename, mode='w') as yaml_file:
-        yaml.safe_dump(config_dict, yaml_file, default_flow_style=False)
+    with fsspec.open(filename, mode="w") as yaml_file:
+        yaml.safe_dump(config_dict, yaml_file, default_flow_style=False)

ocf_data_sampler/load/gsp.py

@@ -1,26 +1,27 @@
-import pkg_resources
+"""Functions for loading GSP data."""
+
+from importlib.resources import files
 
 import pandas as pd
 import xarray as xr
 
 
 def open_gsp(zarr_path: str) -> xr.DataArray:
-    """Open the GSP data
-    
+    """Open the GSP data.
+
     Args:
         zarr_path: Path to the GSP zarr data
 
     Returns:
         xr.DataArray: The opened GSP data
     """
-
     ds = xr.open_zarr(zarr_path)
 
     ds = ds.rename({"datetime_gmt": "time_utc"})
 
     # Load UK GSP locations
     df_gsp_loc = pd.read_csv(
-        pkg_resources.resource_filename(__name__, "../data/uk_gsp_locations.csv"),
+        files("ocf_data_sampler.data").joinpath("uk_gsp_locations.csv"),
         index_col="gsp_id",
     )
 

ocf_data_sampler/load/load_dataset.py

@@ -1,17 +1,17 @@
-""" Loads all data sources """
+"""Loads all data sources."""
+
 import xarray as xr
 
 from ocf_data_sampler.config import InputData
-from ocf_data_sampler.load import open_nwp, open_gsp, open_sat_data, open_site
+from ocf_data_sampler.load import open_gsp, open_nwp, open_sat_data, open_site
 
 
 def get_dataset_dict(input_config: InputData) -> dict[str, dict[xr.DataArray] | xr.DataArray]:
-    """Construct dictionary of all of the input data sources
+    """Construct dictionary of all of the input data sources.
 
     Args:
         input_config: InputData configuration object
     """
-
     datasets_dict = {}
 
     # Load GSP data unless the path is None
@@ -23,10 +23,8 @@ def get_dataset_dict(input_config: InputData) -> dict[str, dict[xr.DataArray] |
 
     # Load NWP data if in config
     if input_config.nwp:
-
         datasets_dict["nwp"] = {}
         for nwp_source, nwp_config in input_config.nwp.items():
-
             da_nwp = open_nwp(nwp_config.zarr_path, provider=nwp_config.provider)
 
             da_nwp = da_nwp.sel(channel=list(nwp_config.channels))
@@ -48,6 +46,7 @@ def get_dataset_dict(input_config: InputData) -> dict[str, dict[xr.DataArray] |
             generation_file_path=input_config.site.file_path,
             metadata_file_path=input_config.site.metadata_file_path,
         )
+
         datasets_dict["site"] = da_sites
 
     return datasets_dict

ocf_data_sampler/load/nwp/nwp.py

@@ -1,22 +1,34 @@
+"""Module for opening NWP data."""
+
 import xarray as xr
 
-from ocf_data_sampler.load.nwp.providers.ukv import open_ukv
 from ocf_data_sampler.load.nwp.providers.ecmwf import open_ifs
+from ocf_data_sampler.load.nwp.providers.gfs import open_gfs
+from ocf_data_sampler.load.nwp.providers.icon import open_icon_eu
+from ocf_data_sampler.load.nwp.providers.ukv import open_ukv
 
 
 def open_nwp(zarr_path: str | list[str], provider: str) -> xr.DataArray:
-    """Opens NWP zarr
+    """Opens NWP zarr.
 
     Args:
         zarr_path: path to the zarr file
         provider: NWP provider
+
+    Returns:
+        Xarray DataArray of the NWP data
     """
+    provider = provider.lower()
 
-    if provider.lower() == "ukv":
+    if provider == "ukv":
         _open_nwp = open_ukv
-    elif provider.lower() == "ecmwf":
+    elif provider == "ecmwf":
         _open_nwp = open_ifs
+    elif provider == "icon-eu":
+        _open_nwp = open_icon_eu
+    elif provider == "gfs":
+        _open_nwp = open_gfs
     else:
         raise ValueError(f"Unknown provider: {provider}")
-    return _open_nwp(zarr_path)
 
+    return _open_nwp(zarr_path)

ocf_data_sampler/load/nwp/providers/ecmwf.py

@@ -1,17 +1,17 @@
-"""ECMWF provider loaders"""
+"""ECMWF provider loaders."""
 
 import xarray as xr
+
 from ocf_data_sampler.load.nwp.providers.utils import open_zarr_paths
 from ocf_data_sampler.load.utils import (
     check_time_unique_increasing,
+    get_xr_data_array_from_xr_dataset,
     make_spatial_coords_increasing,
-    get_xr_data_array_from_xr_dataset
 )
 
 
 def open_ifs(zarr_path: str | list[str]) -> xr.DataArray:
-    """
-    Opens the ECMWF IFS NWP data
+    """Opens the ECMWF IFS NWP data.
 
     Args:
         zarr_path: Path to the zarr to open
@@ -19,9 +19,8 @@ def open_ifs(zarr_path: str | list[str]) -> xr.DataArray:
     Returns:
         Xarray DataArray of the NWP data
     """
-
     ds = open_zarr_paths(zarr_path)
-    
+
     # LEGACY SUPPORT - rename variable to channel if it exists
     ds = ds.rename({"init_time": "init_time_utc", "variable": "channel"})
 
@@ -30,6 +29,6 @@ def open_ifs(zarr_path: str | list[str]) -> xr.DataArray:
     ds = make_spatial_coords_increasing(ds, x_coord="longitude", y_coord="latitude")
 
     ds = ds.transpose("init_time_utc", "step", "channel", "longitude", "latitude")
-    
+
     # TODO: should we control the dtype of the DataArray?
     return get_xr_data_array_from_xr_dataset(ds)

ocf_data_sampler/load/nwp/providers/gfs.py

@@ -0,0 +1,36 @@
+"""Open GFS Forecast data."""
+
+import logging
+
+import xarray as xr
+
+from ocf_data_sampler.load.nwp.providers.utils import open_zarr_paths
+from ocf_data_sampler.load.utils import check_time_unique_increasing, make_spatial_coords_increasing
+
+_log = logging.getLogger(__name__)
+
+
+def open_gfs(zarr_path: str | list[str]) -> xr.DataArray:
+    """Opens the GFS data.
+
+    Args:
+        zarr_path: Path to the zarr to open
+
+    Returns:
+        Xarray DataArray of the NWP data
+    """
+    _log.info("Loading NWP GFS data")
+
+    # Open data
+    gfs: xr.Dataset = open_zarr_paths(zarr_path, time_dim="init_time_utc")
+    nwp: xr.DataArray = gfs.to_array()
+
+    del gfs
+
+    nwp = nwp.rename({"variable": "channel","init_time": "init_time_utc"})
+    check_time_unique_increasing(nwp.init_time_utc)
+    nwp = make_spatial_coords_increasing(nwp, x_coord="longitude", y_coord="latitude")
+
+    nwp = nwp.transpose("init_time_utc", "step", "channel", "latitude", "longitude")
+
+    return nwp

ocf_data_sampler/load/nwp/providers/icon.py

@@ -0,0 +1,46 @@
+"""DWD ICON Loading."""
+
+import xarray as xr
+
+from ocf_data_sampler.load.nwp.providers.utils import open_zarr_paths
+from ocf_data_sampler.load.utils import check_time_unique_increasing, make_spatial_coords_increasing
+
+
+def remove_isobaric_lelvels_from_coords(nwp: xr.Dataset) -> xr.Dataset:
+    """Removes the isobaric levels from the coordinates of the NWP data.
+
+    Args:
+        nwp: NWP data
+
+    Returns:
+        NWP data without isobaric levels in the coordinates
+    """
+    variables_to_drop = [var for var in nwp.data_vars if "isobaricInhPa" in nwp[var].dims]
+    return nwp.drop_vars(["isobaricInhPa", *variables_to_drop])
+
+
+def open_icon_eu(zarr_path: str) -> xr.Dataset:
+    """Opens the ICON data.
+
+    ICON EU Data is on a regular lat/lon grid
+    It has data on multiple pressure levels, as well as the surface
+    Each of the variables is its own data variable
+
+    Args:
+        zarr_path: Path to the zarr to open
+
+    Returns:
+        Xarray DataArray of the NWP data
+    """
+    # Open the data
+    nwp = open_zarr_paths(zarr_path, time_dim="time")
+    nwp = nwp.rename({"time": "init_time_utc"})
+    # Sanity checks.
+    check_time_unique_increasing(nwp.init_time_utc)
+    # 0-78 one hour steps, rest 3 hour steps
+    nwp = nwp.isel(step=slice(0, 78))
+    nwp = remove_isobaric_lelvels_from_coords(nwp)
+    nwp = nwp.to_array().rename({"variable": "channel"})
+    nwp = nwp.transpose("init_time_utc", "step", "channel", "latitude", "longitude")
+    nwp = make_spatial_coords_increasing(nwp, x_coord="longitude", y_coord="latitude")
+    return nwp

ocf_data_sampler/load/nwp/providers/ukv.py

@@ -1,18 +1,17 @@
-"""UKV provider loaders"""
+"""UKV provider loaders."""
 
 import xarray as xr
 
 from ocf_data_sampler.load.nwp.providers.utils import open_zarr_paths
 from ocf_data_sampler.load.utils import (
     check_time_unique_increasing,
+    get_xr_data_array_from_xr_dataset,
     make_spatial_coords_increasing,
-    get_xr_data_array_from_xr_dataset
 )
 
 
 def open_ukv(zarr_path: str | list[str]) -> xr.DataArray:
-    """
-    Opens the NWP data
+    """Opens the NWP data.
 
     Args:
         zarr_path: Path to the zarr to open
@@ -28,7 +27,7 @@ def open_ukv(zarr_path: str | list[str]) -> xr.DataArray:
             "variable": "channel",
             "x": "x_osgb",
             "y": "y_osgb",
-        }
+        },
     )
 
     check_time_unique_increasing(ds.init_time_utc)