ocf-data-sampler 0.0.19__py3-none-any.whl → 0.0.43__py3-none-any.whl

ocf_data_sampler/select/location.py

@@ -0,0 +1,62 @@
+"""location"""
+
+from typing import Optional
+
+import numpy as np
+from pydantic import BaseModel, Field, model_validator
+
+
+allowed_coordinate_systems =["osgb", "lon_lat", "geostationary", "idx"]
+
+class Location(BaseModel):
+    """Represent a spatial location."""
+
+    coordinate_system: Optional[str] = "osgb"  # ["osgb", "lon_lat", "geostationary", "idx"]
+    x: float
+    y: float
+    id: Optional[int] = Field(None)
+
+    @model_validator(mode='after')
+    def validate_coordinate_system(self):
+        """Validate 'coordinate_system'"""
+        if self.coordinate_system not in allowed_coordinate_systems:
+            raise ValueError(f"coordinate_system = {self.coordinate_system} is not in {allowed_coordinate_systems}")
+        return self
+
+    @model_validator(mode='after')
+    def validate_x(self):
+        """Validate 'x'"""
+        min_x: float
+        max_x: float
+
+        co = self.coordinate_system
+        if co == "osgb":
+            min_x, max_x = -103976.3, 652897.98
+        if co == "lon_lat":
+            min_x, max_x = -180, 180
+        if co == "geostationary":
+            min_x, max_x = -5568748.275756836, 5567248.074173927
+        if co == "idx":
+            min_x, max_x = 0, np.inf
+        if self.x < min_x or self.x > max_x:
+            raise ValueError(f"x = {self.x} must be within {[min_x, max_x]} for {co} coordinate system")
+        return self
+
+    @model_validator(mode='after')
+    def validate_y(self):
+        """Validate 'y'"""
+        min_y: float
+        max_y: float
+
+        co = self.coordinate_system
+        if co == "osgb":
+            min_y, max_y = -16703.87, 1199851.44
+        if co == "lon_lat":
+            min_y, max_y = -90, 90
+        if co == "geostationary":
+            min_y, max_y = 1393687.2151494026, 5570748.323202133
+        if co == "idx":
+            min_y, max_y = 0, np.inf
+        if self.y < min_y or self.y > max_y:
+            raise ValueError(f"y = {self.y} must be within {[min_y, max_y]} for {co} coordinate system")
+        return self

ocf_data_sampler/select/select_spatial_slice.py

@@ -5,15 +5,15 @@ import logging
 import numpy as np
 import xarray as xr
 
-from ocf_datapipes.utils import Location
-from ocf_datapipes.utils.geospatial import (
-    lon_lat_to_geostationary_area_coords,
+from ocf_data_sampler.select.location import Location
+from ocf_data_sampler.select.geospatial import (
     lon_lat_to_osgb,
+    lon_lat_to_geostationary_area_coords,
     osgb_to_geostationary_area_coords,
     osgb_to_lon_lat,
     spatial_coord_type,
 )
-from ocf_datapipes.utils.utils import searchsorted
+
 
 logger = logging.getLogger(__name__)
 
@@ -45,9 +45,6 @@ def convert_coords_to_match_xarray(
         if from_coords == "osgb":
             x, y = osgb_to_geostationary_area_coords(x, y, da)
 
-        elif from_coords == "lon_lat":
-            x, y = lon_lat_to_geostationary_area_coords(x, y, da)
-
     elif target_coords == "lon_lat":
         if from_coords == "osgb":
             x, y = osgb_to_lon_lat(x, y)
@@ -105,7 +102,7 @@ def _get_idx_of_pixel_closest_to_poi(
 
 def _get_idx_of_pixel_closest_to_poi_geostationary(
     da: xr.DataArray,
-    center_osgb: Location,
+    center: Location,
 ) -> Location:
     """
     Return x and y index location of pixel at center of region of interest.
@@ -120,7 +117,12 @@ def _get_idx_of_pixel_closest_to_poi_geostationary(
 
     _, x_dim, y_dim = spatial_coord_type(da)
 
-    x, y = osgb_to_geostationary_area_coords(x=center_osgb.x, y=center_osgb.y, xr_data=da)
+    if center.coordinate_system == 'osgb':
+        x, y = osgb_to_geostationary_area_coords(x=center.x, y=center.y, xr_data=da)
+    elif center.coordinate_system == 'lon_lat':
+        x, y = lon_lat_to_geostationary_area_coords(longitude=center.x, latitude=center.y, xr_data=da)
+    else:
+        x,y = center.x, center.y
     center_geostationary = Location(x=x, y=y, coordinate_system="geostationary")
 
     # Check that the requested point lies within the data
@@ -130,13 +132,8 @@ def _get_idx_of_pixel_closest_to_poi_geostationary(
         f"{y} is not in the interval {da[y_dim].min().values}: {da[y_dim].max().values}"
 
     # Get the index into x and y nearest to x_center_geostationary and y_center_geostationary:
-    x_index_at_center = searchsorted(
-        da[x_dim].values, center_geostationary.x, assume_ascending=True
-    )
-
-    y_index_at_center = searchsorted(
-        da[y_dim].values, center_geostationary.y, assume_ascending=True
-    )
+    x_index_at_center = np.searchsorted(da[x_dim].values, center_geostationary.x)
+    y_index_at_center = np.searchsorted(da[y_dim].values, center_geostationary.y)
 
     return Location(x=x_index_at_center, y=y_index_at_center, coordinate_system="idx")
 

ocf_data_sampler/select/select_time_slice.py

@@ -39,23 +39,14 @@ def _sel_fillinterp(
 def select_time_slice(
     ds: xr.DataArray,
     t0: pd.Timestamp,
+    interval_start: pd.Timedelta,
+    interval_end: pd.Timedelta,
     sample_period_duration: pd.Timedelta,
-    history_duration: pd.Timedelta | None = None,
-    forecast_duration: pd.Timedelta | None = None,
-    interval_start: pd.Timedelta | None = None,
-    interval_end: pd.Timedelta | None = None,
     fill_selection: bool = False,
     max_steps_gap: int = 0,
 ):
     """Select a time slice from a Dataset or DataArray."""
-    used_duration = history_duration is not None and forecast_duration is not None
-    used_intervals = interval_start is not None and interval_end is not None
-    assert used_duration ^ used_intervals, "Either durations, or intervals must be supplied"
     assert max_steps_gap >= 0, "max_steps_gap must be >= 0 "
-
-    if used_duration:
-        interval_start = - history_duration
-        interval_end = forecast_duration
     
     if fill_selection and max_steps_gap == 0:
         _sel = _sel_fillnan
@@ -75,11 +66,11 @@ def select_time_slice(
 
 
 def select_time_slice_nwp(
-    ds: xr.DataArray,
+    da: xr.DataArray,
     t0: pd.Timestamp,
+    interval_start: pd.Timedelta,
+    interval_end: pd.Timedelta,
     sample_period_duration: pd.Timedelta,
-    history_duration: pd.Timedelta,
-    forecast_duration: pd.Timedelta,
     dropout_timedeltas: list[pd.Timedelta] | None = None,
     dropout_frac: float | None = 0,
     accum_channels: list[str] = [],
@@ -92,31 +83,31 @@ def select_time_slice_nwp(
         ), "dropout timedeltas must be negative"
         assert len(dropout_timedeltas) >= 1
     assert 0 <= dropout_frac <= 1
-    _consider_dropout = (dropout_timedeltas is not None) and dropout_frac > 0
+    consider_dropout = (dropout_timedeltas is not None) and dropout_frac > 0
 
 
     # The accumatation and non-accumulation channels
     accum_channels = np.intersect1d(
-        ds[channel_dim_name].values, accum_channels
+        da[channel_dim_name].values, accum_channels
     )
     non_accum_channels = np.setdiff1d(
-        ds[channel_dim_name].values, accum_channels
+        da[channel_dim_name].values, accum_channels
     )
 
-    start_dt = (t0 - history_duration).ceil(sample_period_duration)
-    end_dt = (t0 + forecast_duration).ceil(sample_period_duration)
+    start_dt = (t0 + interval_start).ceil(sample_period_duration)
+    end_dt = (t0 + interval_end).ceil(sample_period_duration)
 
     target_times = pd.date_range(start_dt, end_dt, freq=sample_period_duration)
 
     # Maybe apply NWP dropout
-    if _consider_dropout and (np.random.uniform() < dropout_frac):
+    if consider_dropout and (np.random.uniform() < dropout_frac):
         dt = np.random.choice(dropout_timedeltas)
         t0_available = t0 + dt
     else:
         t0_available = t0
 
     # Forecasts made up to and including t0
-    available_init_times = ds.init_time_utc.sel(
+    available_init_times = da.init_time_utc.sel(
         init_time_utc=slice(None, t0_available)
     )
 
@@ -139,7 +130,7 @@ def select_time_slice_nwp(
     step_indexer = xr.DataArray(steps, coords=coords)
 
     if len(accum_channels) == 0:
-        xr_sel = ds.sel(step=step_indexer, init_time_utc=init_time_indexer)
+        da_sel = da.sel(step=step_indexer, init_time_utc=init_time_indexer)
 
     else:
         # First minimise the size of the dataset we are diffing
@@ -149,7 +140,7 @@ def select_time_slice_nwp(
         min_step = min(steps)
         max_step = max(steps) + sample_period_duration
 
-        xr_min = ds.sel(
+        da_min = da.sel(
             {
                 "init_time_utc": unique_init_times,
                 "step": slice(min_step, max_step),
@@ -157,28 +148,28 @@ def select_time_slice_nwp(
         )
 
         # Slice out the data which does not need to be diffed
-        xr_non_accum = xr_min.sel({channel_dim_name: non_accum_channels})
-        xr_sel_non_accum = xr_non_accum.sel(
+        da_non_accum = da_min.sel({channel_dim_name: non_accum_channels})
+        da_sel_non_accum = da_non_accum.sel(
             step=step_indexer, init_time_utc=init_time_indexer
         )
 
         # Slice out the channels which need to be diffed
-        xr_accum = xr_min.sel({channel_dim_name: accum_channels})
+        da_accum = da_min.sel({channel_dim_name: accum_channels})
 
         # Take the diff and slice requested data
-        xr_accum = xr_accum.diff(dim="step", label="lower")
-        xr_sel_accum = xr_accum.sel(step=step_indexer, init_time_utc=init_time_indexer)
+        da_accum = da_accum.diff(dim="step", label="lower")
+        da_sel_accum = da_accum.sel(step=step_indexer, init_time_utc=init_time_indexer)
 
         # Join diffed and non-diffed variables
-        xr_sel = xr.concat([xr_sel_non_accum, xr_sel_accum], dim=channel_dim_name)
+        da_sel = xr.concat([da_sel_non_accum, da_sel_accum], dim=channel_dim_name)
 
         # Reorder the variable back to the original order
-        xr_sel = xr_sel.sel({channel_dim_name: ds[channel_dim_name].values})
+        da_sel = da_sel.sel({channel_dim_name: da[channel_dim_name].values})
 
         # Rename the diffed channels
-        xr_sel[channel_dim_name] = [
+        da_sel[channel_dim_name] = [
             f"diff_{v}" if v in accum_channels else v
-            for v in xr_sel[channel_dim_name].values
+            for v in da_sel[channel_dim_name].values
         ]
 
-    return xr_sel
+    return da_sel

ocf_data_sampler/select/spatial_slice_for_dataset.py

@@ -0,0 +1,53 @@
+""" Functions for selecting data around a given location """
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.select.location import Location
+from ocf_data_sampler.select.select_spatial_slice import select_spatial_slice_pixels
+
+
+def slice_datasets_by_space(
+    datasets_dict: dict,
+    location: Location,
+    config: Configuration,
+) -> dict:
+    """Slice the dictionary of input data sources around a given location
+
+    Args:
+        datasets_dict: Dictionary of the input data sources
+        location: The location to sample around
+        config: Configuration object.
+    """
+
+    assert set(datasets_dict.keys()).issubset({"nwp", "sat", "gsp", "site"})
+
+    sliced_datasets_dict = {}
+
+    if "nwp" in datasets_dict:
+
+        sliced_datasets_dict["nwp"] = {}
+
+        for nwp_key, nwp_config in config.input_data.nwp.items():
+
+            sliced_datasets_dict["nwp"][nwp_key] = select_spatial_slice_pixels(
+                datasets_dict["nwp"][nwp_key],
+                location,
+                height_pixels=nwp_config.image_size_pixels_height,
+                width_pixels=nwp_config.image_size_pixels_width,
+            )
+
+    if "sat" in datasets_dict:
+        sat_config = config.input_data.satellite
+
+        sliced_datasets_dict["sat"] = select_spatial_slice_pixels(
+            datasets_dict["sat"],
+            location,
+            height_pixels=sat_config.image_size_pixels_height,
+            width_pixels=sat_config.image_size_pixels_width,
+        )
+
+    if "gsp" in datasets_dict:
+        sliced_datasets_dict["gsp"] = datasets_dict["gsp"].sel(gsp_id=location.id)
+
+    if "site" in datasets_dict:
+        sliced_datasets_dict["site"] = datasets_dict["site"].sel(site_id=location.id)
+
+    return sliced_datasets_dict

ocf_data_sampler/select/time_slice_for_dataset.py

@@ -0,0 +1,125 @@
+""" Slice datasets by time"""
+import pandas as pd
+
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.select.dropout import draw_dropout_time, apply_dropout_time
+from ocf_data_sampler.select.select_time_slice import select_time_slice_nwp, select_time_slice
+from ocf_data_sampler.utils import minutes
+
+
+def slice_datasets_by_time(
+    datasets_dict: dict,
+    t0: pd.Timestamp,
+    config: Configuration,
+) -> dict:
+    """Slice the dictionary of input data sources around a given t0 time
+
+    Args:
+        datasets_dict: Dictionary of the input data sources
+        t0: The init-time
+        config: Configuration object.
+    """
+
+    sliced_datasets_dict = {}
+
+    if "nwp" in datasets_dict:
+        
+        sliced_datasets_dict["nwp"] = {}
+        
+        for nwp_key, da_nwp in datasets_dict["nwp"].items():
+                            
+            nwp_config = config.input_data.nwp[nwp_key]
+
+            sliced_datasets_dict["nwp"][nwp_key] = select_time_slice_nwp(
+                da_nwp,
+                t0,
+                sample_period_duration=minutes(nwp_config.time_resolution_minutes),
+                interval_start=minutes(nwp_config.interval_start_minutes),
+                interval_end=minutes(nwp_config.interval_end_minutes),
+                dropout_timedeltas=minutes(nwp_config.dropout_timedeltas_minutes),
+                dropout_frac=nwp_config.dropout_fraction,
+                accum_channels=nwp_config.accum_channels,
+            )
+
+    if "sat" in datasets_dict:
+
+        sat_config = config.input_data.satellite
+
+        sliced_datasets_dict["sat"] = select_time_slice(
+            datasets_dict["sat"],
+            t0,
+            sample_period_duration=minutes(sat_config.time_resolution_minutes),
+            interval_start=minutes(sat_config.interval_start_minutes),
+            interval_end=minutes(sat_config.interval_end_minutes),
+            max_steps_gap=2,
+        )
+
+        # Randomly sample dropout
+        sat_dropout_time = draw_dropout_time(
+            t0,
+            dropout_timedeltas=minutes(sat_config.dropout_timedeltas_minutes),
+            dropout_frac=sat_config.dropout_fraction,
+        )
+
+        # Apply the dropout
+        sliced_datasets_dict["sat"] = apply_dropout_time(
+            sliced_datasets_dict["sat"],
+            sat_dropout_time,
+        )
+
+    if "gsp" in datasets_dict:
+        gsp_config = config.input_data.gsp
+
+        sliced_datasets_dict["gsp_future"] = select_time_slice(
+            datasets_dict["gsp"],
+            t0,
+            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            interval_start=minutes(gsp_config.time_resolution_minutes),
+            interval_end=minutes(gsp_config.interval_end_minutes),
+        )
+    
+        sliced_datasets_dict["gsp"] = select_time_slice(
+            datasets_dict["gsp"],
+            t0,
+            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            interval_start=minutes(gsp_config.interval_start_minutes),
+            interval_end=minutes(0),
+        )
+
+        # Dropout on the GSP, but not the future GSP
+        gsp_dropout_time = draw_dropout_time(
+            t0,
+            dropout_timedeltas=minutes(gsp_config.dropout_timedeltas_minutes),
+            dropout_frac=gsp_config.dropout_fraction,
+        )
+
+        sliced_datasets_dict["gsp"] = apply_dropout_time(
+            sliced_datasets_dict["gsp"], 
+            gsp_dropout_time
+        )
+    
+    if "site" in datasets_dict:
+        site_config = config.input_data.site
+
+        sliced_datasets_dict["site"] = select_time_slice(
+            datasets_dict["site"],
+            t0,
+            sample_period_duration=minutes(site_config.time_resolution_minutes),
+            interval_start=minutes(site_config.interval_start_minutes),
+            interval_end=minutes(site_config.interval_end_minutes),
+        )
+
+        # Randomly sample dropout
+        site_dropout_time = draw_dropout_time(
+            t0,
+            dropout_timedeltas=minutes(site_config.dropout_timedeltas_minutes),
+            dropout_frac=site_config.dropout_fraction,
+        )
+
+        # Apply the dropout
+        sliced_datasets_dict["site"] = apply_dropout_time(
+            sliced_datasets_dict["site"],
+            site_dropout_time,
+        )
+
+    return sliced_datasets_dict

ocf_data_sampler/torch_datasets/__init__.py

@@ -1 +1,2 @@
-
+from .pvnet_uk_regional import PVNetUKRegionalDataset
+from .site import SitesDataset

ocf_data_sampler/torch_datasets/process_and_combine.py

@@ -0,0 +1,131 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+from typing import Optional
+
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.constants import NWP_MEANS, NWP_STDS,RSS_MEAN,RSS_STD
+from ocf_data_sampler.numpy_sample import (
+    convert_nwp_to_numpy_sample,
+    convert_satellite_to_numpy_sample,
+    convert_gsp_to_numpy_sample,
+    make_sun_position_numpy_sample,
+)
+from ocf_data_sampler.numpy_sample.gsp import GSPSampleKey
+from ocf_data_sampler.numpy_sample.nwp import NWPSampleKey
+from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
+from ocf_data_sampler.select.location import Location
+from ocf_data_sampler.utils import minutes
+
+
+def process_and_combine_datasets(
+    dataset_dict: dict,
+    config: Configuration,
+    t0: Optional[pd.Timestamp] = None,
+    location: Optional[Location] = None,
+    target_key: str = 'gsp'
+) -> dict:
+
+    """Normalise and convert data to numpy arrays"""
+    numpy_modalities = []
+
+    if "nwp" in dataset_dict:
+
+        nwp_numpy_modalities = dict()
+
+        for nwp_key, da_nwp in dataset_dict["nwp"].items():
+            # Standardise
+            provider = config.input_data.nwp[nwp_key].provider
+            da_nwp = (da_nwp - NWP_MEANS[provider]) / NWP_STDS[provider]
+            # Convert to NumpySample
+            nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_sample(da_nwp)
+
+        # Combine the NWPs into NumpySample
+        numpy_modalities.append({NWPSampleKey.nwp: nwp_numpy_modalities})
+
+
+    if "sat" in dataset_dict:
+        # Standardise
+        da_sat = dataset_dict["sat"]
+        da_sat = (da_sat - RSS_MEAN) / RSS_STD
+
+        # Convert to NumpySample
+        numpy_modalities.append(convert_satellite_to_numpy_sample(da_sat))
+
+
+    gsp_config = config.input_data.gsp
+
+    if "gsp" in dataset_dict:
+        da_gsp = xr.concat([dataset_dict["gsp"], dataset_dict["gsp_future"]], dim="time_utc")
+        da_gsp = da_gsp / da_gsp.effective_capacity_mwp
+
+        numpy_modalities.append(
+            convert_gsp_to_numpy_sample(
+                da_gsp, 
+                t0_idx=-gsp_config.interval_start_minutes / gsp_config.time_resolution_minutes
+            )
+        )
+
+        # Add coordinate data
+        # TODO: Do we need all of these?
+        numpy_modalities.append(
+            {
+                GSPSampleKey.gsp_id: location.id,
+                GSPSampleKey.x_osgb: location.x,
+                GSPSampleKey.y_osgb: location.y,
+            }
+        )
+
+    if target_key == 'gsp':
+        # Make sun coords NumpySample
+        datetimes = pd.date_range(
+            t0+minutes(gsp_config.interval_start_minutes),
+            t0+minutes(gsp_config.interval_end_minutes),
+            freq=minutes(gsp_config.time_resolution_minutes),
+        )
+
+        lon, lat = osgb_to_lon_lat(location.x, location.y)
+
+    numpy_modalities.append(
+        make_sun_position_numpy_sample(datetimes, lon, lat, key_prefix=target_key)
+    )
+
+    # Combine all the modalities and fill NaNs
+    combined_sample = merge_dicts(numpy_modalities)
+    combined_sample = fill_nans_in_arrays(combined_sample)
+
+    return combined_sample
+
+def merge_dicts(list_of_dicts: list[dict]) -> dict:
+    """Merge a list of dictionaries into a single dictionary"""
+    # TODO: This doesn't account for duplicate keys, which will be overwritten
+    combined_dict = {}
+    for d in list_of_dicts:
+        combined_dict.update(d)
+    return combined_dict
+
+def fill_nans_in_arrays(sample: dict) -> dict:
+    """Fills all NaN values in each np.ndarray in the sample dictionary with zeros.
+
+    Operation is performed in-place on the sample.
+    """
+    for k, v in sample.items():
+        if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.number):
+            if np.isnan(v).any():
+                sample[k] = np.nan_to_num(v, copy=False, nan=0.0)
+
+        # Recursion is included to reach NWP arrays in subdict
+        elif isinstance(v, dict):
+            fill_nans_in_arrays(v)
+
+    return sample
+
+
+def compute(xarray_dict: dict) -> dict:
+    """Eagerly load a nested dictionary of xarray DataArrays"""
+    for k, v in xarray_dict.items():
+        if isinstance(v, dict):
+            xarray_dict[k] = compute(v)
+        else:
+            xarray_dict[k] = v.compute(scheduler="single-threaded")
+    return xarray_dict