ocf-data-sampler 0.1.10__py3-none-any.whl → 0.1.16__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,45 @@
 """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
 
 
 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`",
@@ -49,31 +49,35 @@ class TimeWindowMixin(Base):
         ...,
         description="Data interval ends at `t0 + interval_end_minutes`",
     )
-    
-    @model_validator(mode='after')
-    def check_interval_range(cls, values):
-        if values.interval_start_minutes > values.interval_end_minutes:
-            raise ValueError('interval_start_minutes must be <= interval_end_minutes')
-        return values
-
-    @field_validator("interval_start_minutes")
-    def interval_start_minutes_divide_by_time_resolution(cls, v: int, info: ValidationInfo) -> int:
-        if v % info.data["time_resolution_minutes"] != 0:
-            raise ValueError("interval_start_minutes must be divisible by time_resolution_minutes")
-        return v
 
-    @field_validator("interval_end_minutes")
-    def interval_end_minutes_divide_by_time_resolution(cls, v: int, info: ValidationInfo) -> int:
-        if v % info.data["time_resolution_minutes"] != 0:
-            raise ValueError("interval_end_minutes must be divisible by time_resolution_minutes")
-        return v
+    @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})",
+            )
+        if start % resolution != 0:
+            raise ValueError(
+                f"interval_start_minutes ({start}) must be divisible "
+                f"by time_resolution_minutes ({resolution})",
+            )
+        if end % resolution != 0:
+            raise ValueError(
+                f"interval_end_minutes ({end}) must be divisible "
+                f"by time_resolution_minutes ({resolution})",
+            )
+        return self
 
 
 class DropoutMixin(Base):
-    """Mixin class, to add dropout minutes"""
+    """Mixin class, to add dropout minutes."""
 
-    dropout_timedeltas_minutes: Optional[List[int]] = Field(
-        default=None,
+    dropout_timedeltas_minutes: list[int] = Field(
+        default=[],
         description="List of possible minutes before t0 where data availability may start. Must be "
         "negative or zero.",
     )
@@ -86,26 +90,27 @@ class DropoutMixin(Base):
     )
 
     @field_validator("dropout_timedeltas_minutes")
-    def dropout_timedeltas_minutes_negative(cls, v: List[int]) -> List[int]:
-        """Validate 'dropout_timedeltas_minutes'"""
-        if v is not None:
-            for m in v:
-                assert m <= 0, "Dropout timedeltas must be negative"
+    def dropout_timedeltas_minutes_negative(cls, v: list[int]) -> list[int]:
+        """Validate 'dropout_timedeltas_minutes'."""
+        for m in v:
+            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 is not None:
+            if self.dropout_timedeltas_minutes != []:
                 raise ValueError("To use dropout timedeltas dropout fraction should be > 0")
         else:
-            if self.dropout_timedeltas_minutes is None:
+            if self.dropout_timedeltas_minutes == []:
                 raise ValueError("To dropout fraction > 0 requires a list of dropout timedeltas")
         return self
 
 
 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(
         ...,
@@ -121,8 +126,8 @@ class SpatialWindowMixin(Base):
 
 
 class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin):
-    """Satellite configuration model"""
-    
+    """Satellite configuration model."""
+
     zarr_path: str | tuple[str] | list[str] = Field(
         ...,
         description="Absolute or relative zarr filepath(s). Prefix with a protocol like s3:// "
@@ -130,82 +135,86 @@ class Satellite(TimeWindowMixin, DropoutMixin, SpatialWindowMixin):
     )
 
     channels: list[str] = Field(
-        ..., description="the satellite channels that are used"
+        ...,
+        description="the satellite channels that are used",
     )
 
 
 class NWP(TimeWindowMixin, DropoutMixin, SpatialWindowMixin):
-    """NWP configuration model"""
-    
+    """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
 
 
 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(
         ...,
@@ -220,17 +229,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/constants.py

@@ -1,33 +1,37 @@
-import xarray as xr
-import numpy as np
+"""Constants for the package."""
 
+import numpy as np
+import xarray as xr
+from typing_extensions import override
 
 NWP_PROVIDERS = [
     "ukv",
     "ecmwf",
-    "gfs"
+    "gfs",
+    "icon_eu",
 ]
-# TODO add ICON
 
 
-def _to_data_array(d):
+def _to_data_array(d: dict) -> xr.DataArray:
+    """Convert a dictionary to a DataArray."""
     return xr.DataArray(
-        [d[k] for k in d.keys()],
-        coords={"channel": [k for k in d.keys()]},
+        [d[k] for k in d],
+        coords={"channel": list(d.keys())},
     ).astype(np.float32)
 
 
 class NWPStatDict(dict):
-    """Custom dictionary class to hold NWP normalization stats"""
+    """Custom dictionary class to hold NWP normalization stats."""
 
-    def __getitem__(self, key):
+    @override
+    def __getitem__(self, key: str) -> xr.DataArray:
         if key not in NWP_PROVIDERS:
             raise KeyError(f"{key} is not a supported NWP provider - {NWP_PROVIDERS}")
         elif key in self.keys():
             return super().__getitem__(key)
         else:
             raise KeyError(
-                f"Values for {key} not yet available in ocf-data-sampler {list(self.keys())}"
+                f"Values for {key} not yet available in ocf-data-sampler {list(self.keys())}",
             )
 
 
@@ -173,16 +177,140 @@ GFS_MEAN = {
 GFS_STD = _to_data_array(GFS_STD)
 GFS_MEAN = _to_data_array(GFS_MEAN)
 
+# ------ ICON-EU
+# Statistics for ICON-EU variables
+ICON_EU_STD = {
+    "alb_rad": 13.7881,
+    "alhfl_s": 73.7198,
+    "ashfl_s": 54.8027,
+    "asob_s": 55.8319,
+    "asob_t": 74.9360,
+    "aswdifd_s": 21.4940,
+    "aswdifu_s": 18.7688,
+    "aswdir_s": 54.4683,
+    "athb_s": 34.8575,
+    "athb_t": 42.9108,
+    "aumfl_s": 0.1460,
+    "avmfl_s": 0.1892,
+    "cape_con": 32.2570,
+    "cape_ml": 106.3998,
+    "clch": 39.9324,
+    "clcl": 36.3961,
+    "clcm": 41.1690,
+    "clct": 34.7696,
+    "clct_mod": 0.4227,
+    "cldepth": 0.1739,
+    "h_snow": 0.9012,
+    "hbas_con": 1306.6632,
+    "htop_con": 1810.5665,
+    "htop_dc": 459.0422,
+    "hzerocl": 1144.6469,
+    "pmsl": 1103.3301,
+    "ps": 4761.3184,
+    "qv_2m": 0.0024,
+    "qv_s": 0.0038,
+    "rain_con": 1.7097,
+    "rain_gsp": 4.2654,
+    "relhum_2m": 15.3779,
+    "rho_snow": 120.2461,
+    "runoff_g": 0.7410,
+    "runoff_s": 2.1930,
+    "snow_con": 1.1432,
+    "snow_gsp": 1.8154,
+    "snowlmt": 656.0699,
+    "synmsg_bt_cl_ir10.8": 17.9438,
+    "t_2m": 7.7973,
+    "t_g": 8.7053,
+    "t_snow": 134.6874,
+    "tch": 0.0052,
+    "tcm": 0.0133,
+    "td_2m": 7.1460,
+    "tmax_2m": 7.8218,
+    "tmin_2m": 7.8346,
+    "tot_prec": 5.6312,
+    "tqc": 0.0976,
+    "tqi": 0.0247,
+    "u_10m": 3.8351,
+    "v_10m": 5.0083,
+    "vmax_10m": 5.5037,
+    "w_snow": 286.1510,
+    "ww": 27.2974,
+    "z0": 0.3901,
+}
+
+ICON_EU_MEAN = {
+    "alb_rad": 15.4437,
+    "alhfl_s": -54.9398,
+    "ashfl_s": -19.4684,
+    "asob_s": 40.9305,
+    "asob_t": 61.9244,
+    "aswdifd_s": 19.7813,
+    "aswdifu_s": 8.8328,
+    "aswdir_s": 29.9820,
+    "athb_s": -53.9873,
+    "athb_t": -212.8088,
+    "aumfl_s": 0.0558,
+    "avmfl_s": 0.0078,
+    "cape_con": 16.7397,
+    "cape_ml": 21.2189,
+    "clch": 26.4262,
+    "clcl": 57.1591,
+    "clcm": 36.1702,
+    "clct": 72.9254,
+    "clct_mod": 0.5561,
+    "cldepth": 0.1356,
+    "h_snow": 0.0494,
+    "hbas_con": 108.4975,
+    "htop_con": 433.0623,
+    "htop_dc": 454.0859,
+    "hzerocl": 1696.6272,
+    "pmsl": 101778.8281,
+    "ps": 99114.4766,
+    "qv_2m": 0.0049,
+    "qv_s": 0.0065,
+    "rain_con": 0.4869,
+    "rain_gsp": 0.9783,
+    "relhum_2m": 78.2258,
+    "rho_snow": 62.5032,
+    "runoff_g": 0.1301,
+    "runoff_s": 0.4119,
+    "snow_con": 0.2188,
+    "snow_gsp": 0.4317,
+    "snowlmt": 1450.3241,
+    "synmsg_bt_cl_ir10.8": 265.0639,
+    "t_2m": 278.8212,
+    "t_g": 279.9216,
+    "t_snow": 162.5582,
+    "tch": 0.0047,
+    "tcm": 0.0091,
+    "td_2m": 274.9544,
+    "tmax_2m": 279.3550,
+    "tmin_2m": 278.2519,
+    "tot_prec": 2.1158,
+    "tqc": 0.0424,
+    "tqi": 0.0108,
+    "u_10m": 1.1902,
+    "v_10m": -0.4733,
+    "vmax_10m": 8.4152,
+    "w_snow": 14.5936,
+    "ww": 15.3570,
+    "z0": 0.2386,
+}
+
+ICON_EU_STD = _to_data_array(ICON_EU_STD)
+ICON_EU_MEAN = _to_data_array(ICON_EU_MEAN)
 
 NWP_STDS = NWPStatDict(
     ukv=UKV_STD,
     ecmwf=ECMWF_STD,
-    gfs=GFS_STD
+    gfs=GFS_STD,
+    icon_eu=ICON_EU_STD,
 )
 NWP_MEANS = NWPStatDict(
     ukv=UKV_MEAN,
     ecmwf=ECMWF_MEAN,
-    gfs=GFS_MEAN
+    gfs=GFS_MEAN,
+    icon_eu=ICON_EU_MEAN,
 )
 
 # ------ Satellite

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)