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

ocf_data_sampler/select/select_time_slice.py

@@ -1,57 +1,80 @@
-import xarray as xr
-import pandas as pd
+"""Select a time slice from a Dataset or DataArray."""
+
 import numpy as np
+import pandas as pd
+import xarray as xr
+
 
 def select_time_slice(
-    ds: xr.DataArray,
+    da: xr.DataArray,
     t0: pd.Timestamp,
     interval_start: pd.Timedelta,
     interval_end: pd.Timedelta,
-    sample_period_duration: pd.Timedelta,
-):
-    """Select a time slice from a Dataset or DataArray."""
-    t0_datetime_utc = pd.Timestamp(t0)
-    start_dt = t0_datetime_utc + interval_start
-    end_dt = t0_datetime_utc + interval_end
+    time_resolution: pd.Timedelta,
+) -> xr.DataArray:
+    """Select a time slice from a DataArray.
+
+    Args:
+        da: The DataArray to slice from
+        t0: The init-time
+        interval_start: The start of the interval with respect to t0
+        interval_end: The end of the interval with respect to t0
+        time_resolution: Distance between neighbouring timestamps
+    """
+    start_dt = t0 + interval_start
+    end_dt = t0 + interval_end
 
-    start_dt = start_dt.ceil(sample_period_duration)
-    end_dt = end_dt.ceil(sample_period_duration)
+    start_dt = start_dt.ceil(time_resolution)
+    end_dt = end_dt.ceil(time_resolution)
+
+    return da.sel(time_utc=slice(start_dt, end_dt))
 
-    return ds.sel(time_utc=slice(start_dt, end_dt))
 
 def select_time_slice_nwp(
     da: xr.DataArray,
     t0: pd.Timestamp,
     interval_start: pd.Timedelta,
     interval_end: pd.Timedelta,
-    sample_period_duration: pd.Timedelta,
+    time_resolution: pd.Timedelta,
     dropout_timedeltas: list[pd.Timedelta] | None = None,
     dropout_frac: float | None = 0,
-    accum_channels: list[str] = [],
-    channel_dim_name: str = "channel",
-):
+    accum_channels: list[str] | None = None,
+) -> xr.DataArray:
+    """Select a time slice from an NWP DataArray.
+
+    Args:
+        da: The DataArray to slice from
+        t0: The init-time
+        interval_start: The start of the interval with respect to t0
+        interval_end: The end of the interval with respect to t0
+        time_resolution: Distance between neighbouring timestamps
+        dropout_timedeltas: List of possible timedeltas before t0 where data availability may start
+        dropout_frac: Probability to apply dropout
+        accum_channels: Channels which are accumulated and need to be differenced
+    """
+    if accum_channels is None:
+        accum_channels = []
+
     if dropout_timedeltas is not None:
-        assert all(
-            [t < pd.Timedelta(0) for t in dropout_timedeltas]
-        ), "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
+        if not all(t < pd.Timedelta(0) for t in dropout_timedeltas):
+            raise ValueError("dropout timedeltas must be negative")
+        if len(dropout_timedeltas) < 1:
+            raise ValueError("dropout timedeltas must have at least one element")
+
+    if not (0 <= dropout_frac <= 1):
+        raise ValueError("dropout_frac must be between 0 and 1")
 
-     # The accumatation and non-accumulation channels
-    accum_channels = np.intersect1d(
-        da[channel_dim_name].values, accum_channels
-    )
-    non_accum_channels = np.setdiff1d(
-        da[channel_dim_name].values, accum_channels
-    )
+    consider_dropout = (dropout_timedeltas is not None) and dropout_frac > 0
 
-    start_dt = (t0 + interval_start).ceil(sample_period_duration)
-    end_dt = (t0 + interval_end).ceil(sample_period_duration)
+    # The accumatated and non-accumulated channels
+    accum_channels = np.intersect1d(da.channel.values, accum_channels)
+    non_accum_channels = np.setdiff1d(da.channel.values, accum_channels)
 
-    target_times = pd.date_range(start_dt, end_dt, freq=sample_period_duration)
+    start_dt = (t0 + interval_start).ceil(time_resolution)
+    end_dt = (t0 + interval_end).ceil(time_resolution)
+    target_times = pd.date_range(start_dt, end_dt, freq=time_resolution)
 
-    # Maybe apply NWP dropout
+    # Potentially apply NWP dropout
     if consider_dropout and (np.random.uniform() < dropout_frac):
         dt = np.random.choice(dropout_timedeltas)
         t0_available = t0 + dt
@@ -59,9 +82,7 @@ def select_time_slice_nwp(
         t0_available = t0
 
     # Forecasts made up to and including t0
-    available_init_times = da.init_time_utc.sel(
-        init_time_utc=slice(None, t0_available)
-    )
+    available_init_times = da.init_time_utc.sel(init_time_utc=slice(None, t0_available))
 
     # Find the most recent available init times for all target times
     selected_init_times = available_init_times.sel(
@@ -74,10 +95,10 @@ def select_time_slice_nwp(
 
     # We want one timestep for each target_time_hourly (obviously!) If we simply do
     # nwp.sel(init_time=init_times, step=steps) then we'll get the *product* of
-    # init_times and steps, which is not what # we want! Instead, we use xarray's
-    # vectorized-indexing mode by using a DataArray indexer.  See the last example here:
+    # init_times and steps, which is not what we want! Instead, we use xarray's
+    # vectorised-indexing mode via using a DataArray indexer.  See the last example here:
     # https://docs.xarray.dev/en/latest/user-guide/indexing.html#more-advanced-indexing
-    
+
     coords = {"target_time_utc": target_times}
     init_time_indexer = xr.DataArray(selected_init_times, coords=coords)
     step_indexer = xr.DataArray(steps, coords=coords)
@@ -90,38 +111,30 @@ def select_time_slice_nwp(
         unique_init_times = np.unique(selected_init_times)
         # - find the min and max steps we slice over. Max is extended due to diff
         min_step = min(steps)
-        max_step = max(steps) + sample_period_duration
+        max_step = max(steps) + time_resolution
 
-        da_min = da.sel(
-            {
-                "init_time_utc": unique_init_times,
-                "step": slice(min_step, max_step),
-            }
-        )
+        da_min = da.sel(init_time_utc=unique_init_times, step=slice(min_step, max_step))
 
         # Slice out the data which does not need to be diffed
-        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
-        )
+        da_non_accum = da_min.sel(channel=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
-        da_accum = da_min.sel({channel_dim_name: accum_channels})
-        
+        da_accum = da_min.sel(channel=accum_channels)
+
         # Take the diff and slice requested data
         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
-        da_sel = xr.concat([da_sel_non_accum, da_sel_accum], dim=channel_dim_name)
-        
+        da_sel = xr.concat([da_sel_non_accum, da_sel_accum], dim="channel")
+
         # Reorder the variable back to the original order
-        da_sel = da_sel.sel({channel_dim_name: da[channel_dim_name].values})
+        da_sel = da_sel.sel(channel=da.channel.values)
 
         # Rename the diffed channels
-        da_sel[channel_dim_name] = [
-            f"diff_{v}" if v in accum_channels else v
-            for v in da_sel[channel_dim_name].values
+        da_sel["channel"] = [
+            f"diff_{v}" if v in accum_channels else v for v in da_sel.channel.values
         ]
 
     return da_sel

ocf_data_sampler/select/spatial_slice_for_dataset.py

@@ -1,4 +1,5 @@
-""" Functions for selecting data around a given location """
+"""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
@@ -9,24 +10,24 @@ def slice_datasets_by_space(
     location: Location,
     config: Configuration,
 ) -> dict:
-    """Slice the dictionary of input data sources around a given location
+    """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"})
+    if not set(datasets_dict.keys()).issubset({"nwp", "sat", "gsp", "site"}):
+        raise ValueError(
+            "'datasets_dict' should only contain keys '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,

ocf_data_sampler/select/time_slice_for_dataset.py

@@ -1,25 +1,26 @@
-""" Slice datasets by time"""
+"""Slice datasets by time."""
+
 import pandas as pd
 import xarray as xr
 
 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.select.dropout import apply_dropout_time, draw_dropout_time
+from ocf_data_sampler.select.select_time_slice import select_time_slice, select_time_slice_nwp
 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
+    """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:
@@ -31,7 +32,7 @@ def slice_datasets_by_time(
             sliced_datasets_dict["nwp"][nwp_key] = select_time_slice_nwp(
                 da_nwp,
                 t0,
-                sample_period_duration=minutes(nwp_config.time_resolution_minutes),
+                time_resolution=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),
@@ -45,7 +46,7 @@ def slice_datasets_by_time(
         sliced_datasets_dict["sat"] = select_time_slice(
             datasets_dict["sat"],
             t0,
-            sample_period_duration=minutes(sat_config.time_resolution_minutes),
+            time_resolution=minutes(sat_config.time_resolution_minutes),
             interval_start=minutes(sat_config.interval_start_minutes),
             interval_end=minutes(sat_config.interval_end_minutes),
         )
@@ -65,11 +66,11 @@ def slice_datasets_by_time(
 
     if "gsp" in datasets_dict:
         gsp_config = config.input_data.gsp
-        
+
         da_gsp_past = select_time_slice(
             datasets_dict["gsp"],
             t0,
-            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            time_resolution=minutes(gsp_config.time_resolution_minutes),
             interval_start=minutes(gsp_config.interval_start_minutes),
             interval_end=minutes(0),
         )
@@ -82,18 +83,18 @@ def slice_datasets_by_time(
         )
 
         da_gsp_past = apply_dropout_time(
-            da_gsp_past, 
-            gsp_dropout_time
+            da_gsp_past,
+            gsp_dropout_time,
         )
-        
+
         da_gsp_future = select_time_slice(
             datasets_dict["gsp"],
             t0,
-            sample_period_duration=minutes(gsp_config.time_resolution_minutes),
+            time_resolution=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"] = xr.concat([da_gsp_past, da_gsp_future], dim="time_utc")
 
     if "site" in datasets_dict:
@@ -102,7 +103,7 @@ def slice_datasets_by_time(
         sliced_datasets_dict["site"] = select_time_slice(
             datasets_dict["site"],
             t0,
-            sample_period_duration=minutes(site_config.time_resolution_minutes),
+            time_resolution=minutes(site_config.time_resolution_minutes),
             interval_start=minutes(site_config.interval_start_minutes),
             interval_end=minutes(site_config.interval_end_minutes),
         )
@@ -120,4 +121,4 @@ def slice_datasets_by_time(
             site_dropout_time,
         )
 
-    return sliced_datasets_dict
+    return sliced_datasets_dict