ocf-data-sampler 0.1.10__py3-none-any.whl → 0.1.16__py3-none-any.whl

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)

ocf_data_sampler/load/nwp/providers/utils.py

@@ -1,8 +1,10 @@
+"""Utility functions for the NWP data processing."""
+
 import xarray as xr
 
 
 def open_zarr_paths(zarr_path: str | list[str], time_dim: str = "init_time") -> xr.Dataset:
-    """Opens the NWP data
+    """Opens the NWP data.
 
     Args:
         zarr_path: Path to the zarr(s) to open

ocf_data_sampler/load/satellite.py

@@ -1,85 +1,76 @@
-"""Satellite loader"""
-
-import subprocess
+"""Satellite loader."""
 
 import xarray as xr
+
 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 _get_single_sat_data(zarr_path: str) -> xr.Dataset:
-    """Helper function to open a Zarr from either a local or GCP path.
+def get_single_sat_data(zarr_path: str) -> xr.Dataset:
+    """Helper function to open a zarr from either a local or GCP path.
 
     Args:
-        zarr_path: Path to a Zarr file. Wildcards (*) are supported **only** for local paths. 
-                   GCS paths (gs://) **do not support** wildcards.
+        zarr_path: path to a zarr file. Wildcards (*) are supported only for local paths
+                   GCS paths (gs://) do not support wildcards
 
     Returns:
-        An xarray Dataset containing satellite data.
+        An xarray Dataset containing satellite data
 
     Raises:
-        ValueError: If a wildcard (*) is used in a GCS (gs://) path.
+        ValueError: If a wildcard (*) is used in a GCS (gs://) path
     """
-
-    # These kwargs are used if the path contains "*"
-    openmf_kwargs = dict(
-        engine="zarr",
-        concat_dim="time",
-        combine="nested",
-        chunks="auto",
-        join="override",
-    )
-
     # Raise an error if a wildcard is used in a GCP path
     if "gs://" in str(zarr_path) and "*" in str(zarr_path):
-        raise ValueError("Wildcard (*) paths are not supported for GCP (gs://) URLs.")
+        raise ValueError("Wildcard (*) paths are not supported for GCP (gs://) URLs")
 
     # Handle multi-file dataset for local paths
     if "*" in str(zarr_path):
-        ds = xr.open_mfdataset(zarr_path, **openmf_kwargs)
+        ds = xr.open_mfdataset(
+            zarr_path,
+            engine="zarr",
+            concat_dim="time",
+            combine="nested",
+            chunks="auto",
+            join="override",
+        )
+        check_time_unique_increasing(ds.time)
     else:
         ds = xr.open_dataset(zarr_path, engine="zarr", chunks="auto")
 
-    # Ensure time is unique and sorted
-    ds = ds.drop_duplicates("time").sortby("time")
-
     return ds
 
 
 def open_sat_data(zarr_path: str | list[str]) -> xr.DataArray:
-    """Lazily opens the Zarr store.
+    """Lazily opens the zarr store.
 
     Args:
-      zarr_path: Cloud URL or local path pattern, or list of these. If GCS URL, it must start with
-          'gs://'.
+      zarr_path: Cloud URL or local path pattern, or list of these. If GCS URL,
+                 it must start with 'gs://'
     """
-
     # Open the data
-    if isinstance(zarr_path, (list, tuple)):
+    if isinstance(zarr_path, list | tuple):
         ds = xr.combine_nested(
-            [_get_single_sat_data(path) for path in zarr_path],
+            [get_single_sat_data(path) for path in zarr_path],
             concat_dim="time",
             combine_attrs="override",
             join="override",
         )
     else:
-        ds = _get_single_sat_data(zarr_path)
+        ds = get_single_sat_data(zarr_path)
 
     ds = ds.rename(
         {
             "variable": "channel",
             "time": "time_utc",
-        }
+        },
     )
 
     check_time_unique_increasing(ds.time_utc)
-
     ds = make_spatial_coords_increasing(ds, x_coord="x_geostationary", y_coord="y_geostationary")
-    
     ds = ds.transpose("time_utc", "channel", "x_geostationary", "y_geostationary")
-    
+
     # TODO: should we control the dtype of the DataArray?
     return get_xr_data_array_from_xr_dataset(ds)

ocf_data_sampler/load/site.py

@@ -1,3 +1,5 @@
+"""Funcitons for loading site data."""
+
 import numpy as np
 import pandas as pd
 import xarray as xr
@@ -5,7 +7,7 @@ import xarray as xr
 
 def open_site(generation_file_path: str, metadata_file_path: str) -> xr.DataArray:
     """Open a site's generation data and metadata.
-    
+
     Args:
         generation_file_path: Path to the site generation netcdf data
         metadata_file_path: Path to the site csv metadata
@@ -13,14 +15,14 @@ def open_site(generation_file_path: str, metadata_file_path: str) -> xr.DataArra
     Returns:
         xr.DataArray: The opened site generation data
     """
-
     generation_ds = xr.open_dataset(generation_file_path)
 
     metadata_df = pd.read_csv(metadata_file_path, index_col="site_id")
 
-    assert metadata_df.index.is_unique
+    if not metadata_df.index.is_unique:
+        raise ValueError("site_id is not unique in metadata")
 
-    # Ensure metadata aligns with the site_id dimension in data_ds
+    # Ensure metadata aligns with the site_id dimension in generation_ds
     metadata_df = metadata_df.reindex(generation_ds.site_id.values)
 
     # Assign coordinates to the Dataset using the aligned metadata
@@ -31,7 +33,9 @@ def open_site(generation_file_path: str, metadata_file_path: str) -> xr.DataArra
     )
 
     # Sanity checks
-    assert np.isfinite(generation_ds.capacity_kwp.values).all()
-    assert (generation_ds.capacity_kwp.values > 0).all()
-    
+    if not np.isfinite(generation_ds.generation_kw.values).all():
+        raise ValueError("generation_kw contains non-finite values")
+    if not (generation_ds.capacity_kwp.values > 0).all():
+        raise ValueError("capacity_kwp contains non-positive values")
+
     return generation_ds.generation_kw

ocf_data_sampler/load/utils.py

@@ -1,43 +1,48 @@
-import xarray as xr
+"""Utility functions for working with xarray objects."""
+
 import pandas as pd
+import xarray as xr
+
 
-def check_time_unique_increasing(datetimes) -> None:
-    """Check that the time dimension is unique and increasing"""
+def check_time_unique_increasing(datetimes: xr.DataArray) -> None:
+    """Check that the time dimension is unique and increasing."""
     time = pd.DatetimeIndex(datetimes)
-    assert time.is_unique
-    assert time.is_monotonic_increasing
+    if not (time.is_unique and time.is_monotonic_increasing):
+        raise ValueError("Time dimension must be unique and monotonically increasing")
 
 
 def make_spatial_coords_increasing(ds: xr.Dataset, x_coord: str, y_coord: str) -> xr.Dataset:
-    """Make sure the spatial coordinates are in increasing order
-    
+    """Make sure the spatial coordinates are in increasing order.
+
     Args:
         ds: Xarray Dataset
         x_coord: Name of the x coordinate
         y_coord: Name of the y coordinate
     """
-
     # Make sure the coords are in increasing order
     if ds[x_coord][0] > ds[x_coord][-1]:
-        ds = ds.isel({x_coord:slice(None, None, -1)})
+        ds = ds.isel({x_coord: slice(None, None, -1)})
     if ds[y_coord][0] > ds[y_coord][-1]:
-       ds = ds.isel({y_coord:slice(None, None, -1)})
+        ds = ds.isel({y_coord: slice(None, None, -1)})
 
     # Check the coords are all increasing now
-    assert (ds[x_coord].diff(dim=x_coord) > 0).all()
-    assert (ds[y_coord].diff(dim=y_coord) > 0).all()
+    if not (ds[x_coord].diff(dim=x_coord) > 0).all():
+        raise ValueError(f"'{x_coord}' coordinate must be increasing")
+    if not (ds[y_coord].diff(dim=y_coord) > 0).all():
+        raise ValueError(f"'{y_coord}' coordinate must be increasing")
 
     return ds
 
 
 def get_xr_data_array_from_xr_dataset(ds: xr.Dataset) -> xr.DataArray:
-    """Return underlying xr.DataArray from passed xr.Dataset. 
+    """Return underlying xr.DataArray from passed xr.Dataset.
+
     Checks only one variable is present and returns it as an xr.DataArray.
 
     Args:
         ds: xr.Dataset to extract xr.DataArray from
     """
-
     datavars = list(ds.var())
-    assert len(datavars) == 1, "Cannot open as xr.DataArray: dataset contains multiple variables"
-    return ds[datavars[0]]
+    if len(datavars) != 1:
+        raise ValueError("Cannot open as xr.DataArray: dataset contains multiple variables")
+    return ds[datavars[0]]

ocf_data_sampler/numpy_sample/collate.py

@@ -1,8 +1,10 @@
+"""Functions for collating samples into batches."""
+
 import numpy as np
 
 
 def stack_np_samples_into_batch(dict_list: list[dict]) -> dict:
-    """Stacks list of dict samples into a dict where all samples are joined along a new axis
+    """Stacks list of dict samples into a dict where all samples are joined along a new axis.
 
     Args:
         dict_list: A list of dict-like samples to stack
@@ -10,7 +12,6 @@ def stack_np_samples_into_batch(dict_list: list[dict]) -> dict:
     Returns:
         Dict of the samples stacked with new batch dimension on axis 0
     """
-    
     batch = {}
 
     keys = list(dict_list[0].keys())
@@ -26,10 +27,10 @@ def stack_np_samples_into_batch(dict_list: list[dict]) -> dict:
             for nwp_provider in nwp_providers:
                 # Keys can be different for different NWPs
                 nwp_keys = list(dict_list[0]["nwp"][nwp_provider].keys())
-                
+
                 # Create dict to store NWP batch for this provider
                 nwp_provider_batch = {}
-                
+
                 for nwp_key in nwp_keys:
                     # Stack values under each NWP key for this provider
                     nwp_provider_batch[nwp_key] = stack_data_list(
@@ -46,16 +47,16 @@ def stack_np_samples_into_batch(dict_list: list[dict]) -> dict:
 
 
 def _key_is_constant(key: str) -> bool:
-    """Check if a key is for value which is constant for all samples"""
+    """Check if a key is for value which is constant for all samples."""
     return key.endswith("t0_idx") or key.endswith("channel_names")
 
 
 def stack_data_list(data_list: list, key: str) -> np.ndarray:
-    """Stack a sequence of data elements along a new axis
+    """Stack a sequence of data elements along a new axis.
 
-     Args:
-        data_list: List of data elements to combine
-        key: string identifying the data type
+    Args:
+       data_list: List of data elements to combine
+       key: string identifying the data type
     """
     if _key_is_constant(key):
         return data_list[0]

ocf_data_sampler/numpy_sample/datetime_features.py

@@ -1,11 +1,11 @@
-"""Functions to create trigonometric date and time inputs"""
+"""Functions to create trigonometric date and time inputs."""
 
 import numpy as np
 import pandas as pd
 
 
 def _get_date_time_in_pi(dt: pd.DatetimeIndex) -> tuple[np.ndarray, np.ndarray]:
-    """Create positional embeddings for the datetimes in radians
+    """Create positional embeddings for the datetimes in radians.
 
     Args:
         dt: DatetimeIndex to create radian embeddings for
@@ -13,7 +13,6 @@ def _get_date_time_in_pi(dt: pd.DatetimeIndex) -> tuple[np.ndarray, np.ndarray]:
     Returns:
         Tuple of numpy arrays containing radian coordinates for date and time
     """
-
     day_of_year = dt.dayofyear
     minute_of_day = dt.minute + dt.hour * 60
 
@@ -24,8 +23,7 @@ def _get_date_time_in_pi(dt: pd.DatetimeIndex) -> tuple[np.ndarray, np.ndarray]:
 
 
 def make_datetime_numpy_dict(datetimes: pd.DatetimeIndex, key_prefix: str = "wind") -> dict:
-    """ Creates dictionary of cyclical datetime features - encoded """
-
+    """Creates dictionary of cyclical datetime features - encoded."""
     date_in_pi, time_in_pi = _get_date_time_in_pi(datetimes)
 
     time_numpy_sample = {}

ocf_data_sampler/numpy_sample/gsp.py

@@ -1,26 +1,28 @@
-"""Convert GSP to Numpy Sample"""
+"""Convert GSP to Numpy Sample."""
 
 import xarray as xr
 
 
 class GSPSampleKey:
+    """Keys for the GSP sample dictionary."""
 
-    gsp = 'gsp'
-    nominal_capacity_mwp = 'gsp_nominal_capacity_mwp'
-    effective_capacity_mwp = 'gsp_effective_capacity_mwp'
-    time_utc = 'gsp_time_utc'
-    t0_idx = 'gsp_t0_idx'
-    solar_azimuth = 'gsp_solar_azimuth'
-    solar_elevation = 'gsp_solar_elevation'
-    gsp_id = 'gsp_id'
-    x_osgb = 'gsp_x_osgb'
-    y_osgb = 'gsp_y_osgb'
+    gsp = "gsp"
+    nominal_capacity_mwp = "gsp_nominal_capacity_mwp"
+    effective_capacity_mwp = "gsp_effective_capacity_mwp"
+    time_utc = "gsp_time_utc"
+    t0_idx = "gsp_t0_idx"
+    gsp_id = "gsp_id"
+    x_osgb = "gsp_x_osgb"
+    y_osgb = "gsp_y_osgb"
 
 
 def convert_gsp_to_numpy_sample(da: xr.DataArray, t0_idx: int | None = None) -> dict:
-    """Convert from Xarray to NumpySample"""
+    """Convert from Xarray to NumpySample.
 
-   # Extract values from the DataArray
+    Args:
+        da: Xarray DataArray containing GSP data
+        t0_idx: Index of the t0 timestamp in the time dimension of the GSP data
+    """
     sample = {
         GSPSampleKey.gsp: da.values,
         GSPSampleKey.nominal_capacity_mwp: da.isel(time_utc=0)["nominal_capacity_mwp"].values,

ocf_data_sampler/numpy_sample/nwp.py

@@ -1,42 +1,36 @@
-"""Convert NWP to NumpySample"""
+"""Convert NWP to NumpySample."""
 
 import pandas as pd
 import xarray as xr
 
 
 class NWPSampleKey:
+    """Keys for NWP NumpySample."""
 
-    nwp = 'nwp'
-    channel_names = 'nwp_channel_names'
-    init_time_utc = 'nwp_init_time_utc'
-    step = 'nwp_step'
-    target_time_utc = 'nwp_target_time_utc'
-    t0_idx = 'nwp_t0_idx'
-    y_osgb = 'nwp_y_osgb'
-    x_osgb = 'nwp_x_osgb'
-
+    nwp = "nwp"
+    channel_names = "nwp_channel_names"
+    init_time_utc = "nwp_init_time_utc"
+    step = "nwp_step"
+    target_time_utc = "nwp_target_time_utc"
+    t0_idx = "nwp_t0_idx"
 
 
 def convert_nwp_to_numpy_sample(da: xr.DataArray, t0_idx: int | None = None) -> dict:
-    """Convert from Xarray to NWP NumpySample"""
-    
-    # Create example and add t if available
+    """Convert from Xarray to NWP NumpySample.
+
+    Args:
+        da: Xarray DataArray containing NWP data
+        t0_idx: Index of the t0 timestamp in the time dimension of the NWP
+    """
     sample = {
         NWPSampleKey.nwp: da.values,
         NWPSampleKey.channel_names: da.channel.values,
         NWPSampleKey.init_time_utc: da.init_time_utc.values.astype(float),
         NWPSampleKey.step: (da.step.values / pd.Timedelta("1h")).astype(int),
+        NWPSampleKey.target_time_utc: da.target_time_utc.values.astype(float),
     }
 
-    if "target_time_utc" in da.coords:
-        sample[NWPSampleKey.target_time_utc] = da.target_time_utc.values.astype(float)
-
-    # TODO: Do we need this at all? Especially since it is only present in UKV data
-    for sample_key, dataset_key in ((NWPSampleKey.y_osgb, "y_osgb"),(NWPSampleKey.x_osgb, "x_osgb"),):
-        if dataset_key in da.coords:
-            sample[sample_key] = da[dataset_key].values
-
     if t0_idx is not None:
         sample[NWPSampleKey.t0_idx] = t0_idx
-        
-    return sample
+
+    return sample

ocf_data_sampler/numpy_sample/satellite.py

@@ -1,30 +1,33 @@
-"""Convert Satellite to NumpySample"""
+"""Convert Satellite to NumpySample."""
+
 import xarray as xr
 
 
 class SatelliteSampleKey:
+    """Keys for the SatelliteSample dictionary."""
 
-    satellite_actual = 'satellite_actual'
-    time_utc = 'satellite_time_utc'
-    x_geostationary = 'satellite_x_geostationary'
-    y_geostationary = 'satellite_y_geostationary'
-    t0_idx = 'satellite_t0_idx'
+    satellite_actual = "satellite_actual"
+    time_utc = "satellite_time_utc"
+    x_geostationary = "satellite_x_geostationary"
+    y_geostationary = "satellite_y_geostationary"
+    t0_idx = "satellite_t0_idx"
 
 
 def convert_satellite_to_numpy_sample(da: xr.DataArray, t0_idx: int | None = None) -> dict:
-    """Convert from Xarray to NumpySample"""
+    """Convert from Xarray to NumpySample.
+
+    Args:
+        da: xarray DataArray containing satellite data
+        t0_idx: Index of the t0 timestamp in the time dimension of the satellite data
+    """
     sample = {
         SatelliteSampleKey.satellite_actual: da.values,
         SatelliteSampleKey.time_utc: da.time_utc.values.astype(float),
+        SatelliteSampleKey.x_geostationary: da.x_geostationary.values,
+        SatelliteSampleKey.y_geostationary: da.y_geostationary.values,
     }
 
-    for sample_key, dataset_key in (
-         (SatelliteSampleKey.x_geostationary, "x_geostationary"),
-        (SatelliteSampleKey.y_geostationary, "y_geostationary"),
-    ):
-        sample[sample_key] = da[dataset_key].values
-
     if t0_idx is not None:
         sample[SatelliteSampleKey.t0_idx] = t0_idx
 
-    return sample
+    return sample