ocf-data-sampler 0.0.24__py3-none-any.whl → 0.0.26__py3-none-any.whl

ocf_data_sampler/select/time_slice_for_dataset.py

@@ -0,0 +1,124 @@
+""" 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.time_functions 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),
+                history_duration=minutes(nwp_config.history_minutes),
+                forecast_duration=minutes(nwp_config.forecast_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.history_minutes),
+            interval_end=minutes(-sat_config.live_delay_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(30),
+            interval_end=minutes(gsp_config.forecast_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.history_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.history_minutes),
+            interval_end=minutes(site_config.forecast_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/time_functions.py

@@ -0,0 +1,11 @@
+import pandas as pd
+
+
+def minutes(minutes: int | list[float]) -> pd.Timedelta | pd.TimedeltaIndex:
+    """Timedelta minutes
+
+    Args:
+        minutes: the number of minutes, single value or list
+    """
+    minutes_delta = pd.to_timedelta(minutes, unit="m")
+    return minutes_delta

ocf_data_sampler/torch_datasets/process_and_combine.py

@@ -0,0 +1,153 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+from ocf_data_sampler.config import Configuration
+from ocf_data_sampler.constants import NWP_MEANS, NWP_STDS
+from ocf_data_sampler.numpy_batch import (
+    convert_nwp_to_numpy_batch,
+    convert_satellite_to_numpy_batch,
+    convert_gsp_to_numpy_batch,
+    make_sun_position_numpy_batch,
+    convert_site_to_numpy_batch,
+)
+from ocf_data_sampler.numpy_batch.gsp import GSPBatchKey
+from ocf_data_sampler.numpy_batch.nwp import NWPBatchKey
+from ocf_data_sampler.select.geospatial import osgb_to_lon_lat
+from ocf_data_sampler.select.location import Location
+from ocf_data_sampler.time_functions import minutes
+
+
+def process_and_combine_datasets(
+    dataset_dict: dict,
+    config: Configuration,
+    t0: pd.Timestamp,
+    location: Location,
+    sun_position_key: str = 'gsp'
+) -> dict:
+    """Normalize 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 NumpyBatch
+            nwp_numpy_modalities[nwp_key] = convert_nwp_to_numpy_batch(da_nwp)
+
+        # Combine the NWPs into NumpyBatch
+        numpy_modalities.append({NWPBatchKey.nwp: nwp_numpy_modalities})
+
+    if "sat" in dataset_dict:
+        # Satellite is already in the range [0-1] so no need to standardise
+        da_sat = dataset_dict["sat"]
+
+        # Convert to NumpyBatch
+        numpy_modalities.append(convert_satellite_to_numpy_batch(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_batch(
+                da_gsp, t0_idx=gsp_config.history_minutes // gsp_config.time_resolution_minutes
+            )
+        )
+
+        # Add coordinate data
+        # TODO: Do we need all of these?
+        numpy_modalities.append(
+            {
+                GSPBatchKey.gsp_id: location.id,
+                GSPBatchKey.x_osgb: location.x,
+                GSPBatchKey.y_osgb: location.y,
+            }
+        )
+
+
+    if "site" in dataset_dict:
+        site_config = config.input_data.site
+        da_sites = dataset_dict["site"]
+        da_sites = da_sites / da_sites.capacity_kwp
+
+        numpy_modalities.append(
+            convert_site_to_numpy_batch(
+                da_sites, t0_idx=site_config.history_minutes / site_config.time_resolution_minutes
+            )
+        )
+
+    if sun_position_key == 'gsp':
+        # Make sun coords NumpyBatch
+        datetimes = pd.date_range(
+            t0 - minutes(gsp_config.history_minutes),
+            t0 + minutes(gsp_config.forecast_minutes),
+            freq=minutes(gsp_config.time_resolution_minutes),
+        )
+
+        lon, lat = osgb_to_lon_lat(location.x, location.y)
+        key_prefix = "gsp"
+
+    elif sun_position_key == 'site':
+        # Make sun coords NumpyBatch
+        datetimes = pd.date_range(
+            t0 - minutes(site_config.history_minutes),
+            t0 + minutes(site_config.forecast_minutes),
+            freq=minutes(site_config.time_resolution_minutes),
+        )
+
+        lon, lat = location.x, location.y
+        key_prefix = "site"
+
+    numpy_modalities.append(
+        make_sun_position_numpy_batch(datetimes, lon, lat, key_prefix=key_prefix)
+    )
+
+    # 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(batch: dict) -> dict:
+    """Fills all NaN values in each np.ndarray in the batch dictionary with zeros.
+
+    Operation is performed in-place on the batch.
+    """
+    for k, v in batch.items():
+        if isinstance(v, np.ndarray) and np.issubdtype(v.dtype, np.number):
+            if np.isnan(v).any():
+                batch[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 batch
+
+
+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