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

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,15 +1,16 @@
-"""Satellite loader"""
+"""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
+    """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
@@ -21,7 +22,6 @@ def get_single_sat_data(zarr_path: str) -> xr.Dataset:
     Raises:
         ValueError: If a wildcard (*) is used in a GCS (gs://) path
     """
-
     # 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")
@@ -44,15 +44,14 @@ def get_single_sat_data(zarr_path: str) -> xr.Dataset:
 
 
 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, 
+      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],
             concat_dim="time",
@@ -66,12 +65,12 @@ def open_sat_data(zarr_path: str | list[str]) -> xr.DataArray:
         {
             "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,12 +15,12 @@ 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 generation_ds
     metadata_df = metadata_df.reindex(generation_ds.site_id.values)
@@ -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,30 +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.
+
     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,27 +1,27 @@
-"""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'
+    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
-    
+    """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,
@@ -32,5 +32,5 @@ def convert_nwp_to_numpy_sample(da: xr.DataArray, t0_idx: int | None = None) ->
 
     if t0_idx is not None:
         sample[NWPSampleKey.t0_idx] = t0_idx
-        
-    return sample
+
+    return sample

ocf_data_sampler/numpy_sample/satellite.py

@@ -1,20 +1,21 @@
-"""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
@@ -30,4 +31,3 @@ def convert_satellite_to_numpy_sample(da: xr.DataArray, t0_idx: int | None = Non
         sample[SatelliteSampleKey.t0_idx] = t0_idx
 
     return sample
-    

ocf_data_sampler/numpy_sample/site.py

@@ -1,37 +1,34 @@
-"""Convert site to Numpy Sample"""
+"""Convert site to Numpy Sample."""
 
 import xarray as xr
 
 
 class SiteSampleKey:
+    """Keys for the site sample dictionary."""
 
     generation = "site"
     capacity_kwp = "site_capacity_kwp"
     time_utc = "site_time_utc"
     t0_idx = "site_t0_idx"
     id = "site_id"
-    solar_azimuth = "site_solar_azimuth"
-    solar_elevation = "site_solar_elevation"
     date_sin = "site_date_sin"
     date_cos = "site_date_cos"
     time_sin = "site_time_sin"
     time_cos = "site_time_cos"
 
+
 def convert_site_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 site data
         t0_idx: Index of the t0 timestamp in the time dimension of the site data
     """
-
     sample = {
         SiteSampleKey.generation: da.values,
         SiteSampleKey.capacity_kwp: da.isel(time_utc=0)["capacity_kwp"].values,
         SiteSampleKey.time_utc: da["time_utc"].values.astype(float),
         SiteSampleKey.id: da["site_id"].values,
-        SiteSampleKey.solar_azimuth: da["solar_azimuth"].values,
-        SiteSampleKey.solar_elevation: da["solar_elevation"].values,
         SiteSampleKey.date_sin: da["date_sin"].values,
         SiteSampleKey.date_cos: da["date_cos"].values,
         SiteSampleKey.time_sin: da["time_sin"].values,

ocf_data_sampler/numpy_sample/sun_position.py

@@ -1,16 +1,17 @@
+"""Module for calculating solar position."""
 
-import pvlib
 import numpy as np
 import pandas as pd
+import pvlib
 
 
 def calculate_azimuth_and_elevation(
-    datetimes: pd.DatetimeIndex, 
-    lon: float, 
-    lat: float
+    datetimes: pd.DatetimeIndex,
+    lon: float,
+    lat: float,
 ) -> tuple[np.ndarray, np.ndarray]:
-    """Calculate the solar coordinates for multiple datetimes at a single location
-    
+    """Calculate the solar coordinates for multiple datetimes at a single location.
+
     Args:
         datetimes: The datetimes to calculate for
         lon: The longitude
@@ -20,45 +21,39 @@ def calculate_azimuth_and_elevation(
         np.ndarray: The azimuth of the datetimes in degrees
         np.ndarray: The elevation of the datetimes in degrees
     """
-
     solpos = pvlib.solarposition.get_solarposition(
         time=datetimes,
         longitude=lon,
         latitude=lat,
-        method='nrel_numpy'
+        method="nrel_numpy",
     )
 
     return solpos["azimuth"].values, solpos["elevation"].values
 
 
 def make_sun_position_numpy_sample(
-    datetimes: pd.DatetimeIndex, 
-    lon: float, 
-    lat: float, 
-    key_prefix: str = "gsp"
+    datetimes: pd.DatetimeIndex,
+    lon: float,
+    lat: float,
 ) -> dict:
-    """Creates NumpySample with standardized solar coordinates
+    """Creates NumpySample with standardized solar coordinates.
 
     Args:
         datetimes: The datetimes to calculate solar angles for
         lon: The longitude
         lat: The latitude
-        key_prefix: The prefix to add to the keys in the NumpySample
     """
-    
     azimuth, elevation = calculate_azimuth_and_elevation(datetimes, lon, lat)
 
     # Normalise
     # Azimuth is in range [0, 360] degrees
     azimuth = azimuth / 360
 
-    # Elevation is in range [-90, 90] degrees
+    # Elevation is in range [-90, 90] degrees
     elevation = elevation / 180 + 0.5
 
     # Make NumpySample
-    sun_numpy_sample = {
-        key_prefix + "_solar_azimuth": azimuth,
-        key_prefix + "_solar_elevation": elevation,
+    return {
+        "solar_azimuth": azimuth,
+        "solar_elevation": elevation,
     }
-
-    return sun_numpy_sample