ocf-data-sampler 0.0.7__py3-none-any.whl → 0.0.8__py3-none-any.whl

ocf_data_sampler/numpy_batch/__init__.py

@@ -0,0 +1,7 @@
+"""Conversion from Xarray to NumpyBatch"""
+
+from .gsp import convert_gsp_to_numpy_batch
+from .nwp import convert_nwp_to_numpy_batch
+from .satellite import convert_satellite_to_numpy_batch
+from .sun_position import make_sun_position_numpy_batch
+

ocf_data_sampler/numpy_batch/gsp.py

@@ -0,0 +1,23 @@
+"""Convert GSP to Numpy Batch"""
+
+import xarray as xr
+from ocf_datapipes.batch import BatchKey, NumpyBatch
+
+
+def convert_gsp_to_numpy_batch(da: xr.DataArray, t0_idx: int | None = None) -> NumpyBatch:
+    """Convert from Xarray to NumpyBatch"""
+
+    example: NumpyBatch = {
+        BatchKey.gsp: da.values,
+        BatchKey.gsp_id: da.gsp_id.values,
+        BatchKey.gsp_nominal_capacity_mwp: da.isel(time_utc=0)["nominal_capacity_mwp"].values,
+        BatchKey.gsp_effective_capacity_mwp: da.isel(time_utc=0)["effective_capacity_mwp"].values,
+        BatchKey.gsp_time_utc: da["time_utc"].values.astype(float),
+        BatchKey.gsp_x_osgb: da.x_osgb.item(),
+        BatchKey.gsp_y_osgb: da.y_osgb.item(),
+    }
+
+    if t0_idx is not None:
+        example[BatchKey.gsp_t0_idx] = t0_idx
+
+    return example

ocf_data_sampler/numpy_batch/nwp.py

@@ -0,0 +1,33 @@
+"""Convert NWP to NumpyBatch"""
+
+import pandas as pd
+import xarray as xr
+
+from ocf_datapipes.batch import NWPBatchKey, NWPNumpyBatch
+
+
+def convert_nwp_to_numpy_batch(da: xr.DataArray, t0_idx: int | None = None) -> NWPNumpyBatch:
+    """Convert from Xarray to NWP NumpyBatch"""
+
+    example: NWPNumpyBatch = {
+        NWPBatchKey.nwp: da.values,
+        NWPBatchKey.nwp_channel_names: da.channel.values,
+        NWPBatchKey.nwp_init_time_utc: da.init_time_utc.values.astype(float),
+        NWPBatchKey.nwp_step: (da.step.values / pd.Timedelta("1H")).astype(int),
+    }
+
+    if "target_time_utc" in da.coords:
+        example[NWPBatchKey.nwp_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 batch_key, dataset_key in (
+        (NWPBatchKey.nwp_y_osgb, "y_osgb"),
+        (NWPBatchKey.nwp_x_osgb, "x_osgb"),
+    ):
+        if dataset_key in da.coords:
+            example[batch_key] = da[dataset_key].values
+
+    if t0_idx is not None:
+        example[NWPBatchKey.nwp_t0_idx] = t0_idx
+
+    return example

ocf_data_sampler/numpy_batch/satellite.py

@@ -0,0 +1,23 @@
+"""Convert Satellite to NumpyBatch"""
+import xarray as xr
+
+from ocf_datapipes.batch import BatchKey, NumpyBatch
+
+
+def convert_satellite_to_numpy_batch(da: xr.DataArray, t0_idx: int | None = None) -> NumpyBatch:
+    """Convert from Xarray to NumpyBatch"""
+    example: NumpyBatch = {
+        BatchKey.satellite_actual: da.values,
+        BatchKey.satellite_time_utc: da.time_utc.values.astype(float),
+    }
+
+    for batch_key, dataset_key in (
+         (BatchKey.satellite_x_geostationary, "x_geostationary"),
+        (BatchKey.satellite_y_geostationary, "y_geostationary"),
+    ):
+        example[batch_key] = da[dataset_key].values
+
+    if t0_idx is not None:
+        example[BatchKey.satellite_t0_idx] = t0_idx
+
+    return example

ocf_data_sampler/numpy_batch/sun_position.py

@@ -0,0 +1,66 @@
+
+import pvlib
+import numpy as np
+import pandas as pd
+from ocf_datapipes.batch import BatchKey, NumpyBatch
+
+
+def calculate_azimuth_and_elevation(
+    datetimes: pd.DatetimeIndex, 
+    lon: float, 
+    lat: float
+) -> tuple[np.ndarray, np.ndarray]:
+    """Calculate the solar coordinates for multiple datetimes at a single location
+    
+    Args:
+        datetimes: The datetimes to calculate for
+        lon: The longitude
+        lat: The latitude
+
+    Returns:
+        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'
+    )
+    azimuth = solpos["azimuth"].values
+    elevation = solpos["elevation"].values
+    return azimuth, elevation
+
+
+def make_sun_position_numpy_batch(
+        datetimes: pd.DatetimeIndex, 
+        lon: float, 
+        lat: float, 
+        key_preffix: str = "gsp"
+) -> NumpyBatch:
+    """Creates NumpyBatch with standardized solar coordinates
+
+    Args:
+        datetimes: The datetimes to calculate solar angles for
+        lon: The longitude
+        lat: The latitude
+    """
+    
+    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 = elevation / 180 + 0.5
+    
+    # Make NumpyBatch
+    sun_numpy_batch: NumpyBatch = {
+        BatchKey[key_preffix + "_solar_azimuth"]: azimuth,
+        BatchKey[key_preffix + "_solar_elevation"]: elevation,
+    }
+
+    return sun_numpy_batch

ocf_data_sampler/select/__init__.py

@@ -0,0 +1 @@
+

ocf_data_sampler/select/dropout.py

@@ -0,0 +1,38 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+
+
+def draw_dropout_time(
+        t0: pd.Timestamp,
+        dropout_timedeltas: list[pd.Timedelta] | None,
+        dropout_frac: float = 0,
+    ):
+
+    if dropout_timedeltas is not None:
+        assert len(dropout_timedeltas) >= 1, "Must include list of relative dropout timedeltas"
+        assert all(
+            [t < pd.Timedelta("0min") for t in dropout_timedeltas]
+        ), "dropout timedeltas must be negative"
+    assert 0 <= dropout_frac <= 1
+
+    if (dropout_timedeltas is None) or (np.random.uniform() >= dropout_frac):
+        dropout_time = None
+    else:
+        t0_datetime_utc = pd.Timestamp(t0)
+        dt = np.random.choice(dropout_timedeltas)
+        dropout_time = t0_datetime_utc + dt
+
+    return dropout_time
+
+
+def apply_dropout_time(
+        ds: xr.Dataset,
+        dropout_time: pd.Timestamp | None,
+    ):
+
+    if dropout_time is None:
+        return ds
+    else:
+        # This replaces the times after the dropout with NaNs
+        return ds.where(ds.time_utc <= dropout_time)

ocf_data_sampler/select/fill_time_periods.py

@@ -0,0 +1,11 @@
+"""fill time periods"""
+
+import pandas as pd
+import numpy as np
+
+
+def fill_time_periods(time_periods: pd.DataFrame, freq: pd.Timedelta):
+    start_dts = pd.to_datetime(time_periods["start_dt"].values).ceil(freq)
+    end_dts = pd.to_datetime(time_periods["end_dt"].values)
+    date_ranges = [pd.date_range(start_dt, end_dt, freq=freq) for start_dt, end_dt in zip(start_dts, end_dts)]
+    return pd.DatetimeIndex(np.concatenate(date_ranges))

ocf_data_sampler/select/find_contiguous_time_periods.py

@@ -0,0 +1,301 @@
+"""Get contiguous time periods for training"""
+
+import numpy as np
+import pandas as pd
+
+
+
+def find_contiguous_time_periods(
+    datetimes: pd.DatetimeIndex,
+    min_seq_length: int,
+    max_gap_duration: pd.Timedelta,
+) -> pd.DataFrame:
+    """Return a pd.DataFrame where each row records the boundary of a contiguous time period.
+
+    Args:
+      datetimes: pd.DatetimeIndex. Must be sorted.
+      min_seq_length: Sequences of min_seq_length or shorter will be discarded.  Typically, this
+        would be set to the `total_seq_length` of each machine learning example.
+      max_gap_duration: If any pair of consecutive `datetimes` is more than `max_gap_duration`
+        apart, then this pair of `datetimes` will be considered a "gap" between two contiguous
+        sequences. Typically, `max_gap_duration` would be set to the sample period of
+        the timeseries.
+
+    Returns:
+      pd.DataFrame where each row represents a single time period.  The pd.DataFrame
+          has two columns: `start_dt` and `end_dt` (where 'dt' is short for 'datetime').
+    """
+    # Sanity checks.
+    assert len(datetimes) > 0
+    assert min_seq_length > 1
+    assert datetimes.is_monotonic_increasing
+    assert datetimes.is_unique
+
+    # Find indices of gaps larger than max_gap:
+    gap_mask = pd.TimedeltaIndex(np.diff(datetimes)) > max_gap_duration
+    gap_indices = np.argwhere(gap_mask)[:, 0]
+
+    # gap_indicies are the indices into dt_index for the timestep immediately before the gap.
+    # e.g. if the datetimes at 12:00, 12:05, 18:00, 18:05 then gap_indicies will be [1].
+    # So we add 1 to gap_indices to get segment_boundaries, an index into dt_index
+    # which identifies the _start_ of each segment.
+    segment_boundaries = gap_indices + 1
+
+    # Capture the last segment of dt_index.
+    segment_boundaries = np.concatenate((segment_boundaries, [len(datetimes)]))
+
+    periods: list[dict[str, pd.Timestamp]] = []
+    start_i = 0
+    for next_start_i in segment_boundaries:
+        n_timesteps = next_start_i - start_i
+        if n_timesteps > min_seq_length:
+            end_i = next_start_i - 1
+            period = {"start_dt": datetimes[start_i], "end_dt": datetimes[end_i]}
+            periods.append(period)
+        start_i = next_start_i
+
+    assert len(periods) > 0, (
+        f"Did not find an periods from {datetimes}. " f"{min_seq_length=} {max_gap_duration=}"
+    )
+
+    return pd.DataFrame(periods)
+
+
+def trim_contiguous_time_periods(
+    contiguous_time_periods: pd.DataFrame, 
+    history_duration: pd.Timedelta,
+    forecast_duration: pd.Timedelta,
+) -> pd.DataFrame:
+    """Trim the contiguous time periods to allow for history and forecast durations.
+
+    Args:
+        contiguous_time_periods: DataFrame where each row represents a single time period. The 
+            DataFrame must have `start_dt` and `end_dt` columns.
+        history_duration: Length of the historical slice used for a sample
+        forecast_duration: Length of the forecast slice used for a sample
+
+
+    Returns:
+      The contiguous_time_periods DataFrame with the `start_dt` and `end_dt` columns updated.
+    """
+    contiguous_time_periods = contiguous_time_periods.copy()
+
+    contiguous_time_periods["start_dt"] += history_duration
+    contiguous_time_periods["end_dt"] -= forecast_duration
+
+    valid_mask = contiguous_time_periods["start_dt"] <= contiguous_time_periods["end_dt"]
+    contiguous_time_periods = contiguous_time_periods.loc[valid_mask]
+
+    return contiguous_time_periods
+
+
+
+def find_contiguous_t0_periods(
+        datetimes: pd.DatetimeIndex,
+        history_duration: pd.Timedelta,
+        forecast_duration: pd.Timedelta,
+        sample_period_duration: pd.Timedelta,
+    ) -> pd.DataFrame:
+    """Return a pd.DataFrame where each row records the boundary of a contiguous time period.
+
+    Args:
+        datetimes: pd.DatetimeIndex. Must be sorted.
+        history_duration: Length of the historical slice used for each sample
+        forecast_duration: Length of the forecast slice used for each sample
+        sample_period_duration: The sample frequency of the timeseries
+
+
+    Returns:
+        pd.DataFrame where each row represents a single time period.  The pd.DataFrame
+            has two columns: `start_dt` and `end_dt` (where 'dt' is short for 'datetime').
+    """
+    total_duration = history_duration + forecast_duration
+    
+    contiguous_time_periods = find_contiguous_time_periods(
+        datetimes=datetimes,
+        min_seq_length=int(total_duration / sample_period_duration) + 1,
+        max_gap_duration=sample_period_duration,
+    )
+
+    contiguous_t0_periods = trim_contiguous_time_periods(
+        contiguous_time_periods=contiguous_time_periods,
+        history_duration=history_duration,
+        forecast_duration=forecast_duration,
+    )
+
+    assert len(contiguous_t0_periods) > 0
+
+    return contiguous_t0_periods
+
+
+def _find_contiguous_t0_periods_nwp(
+        ds,
+        history_duration: pd.Timedelta,
+        forecast_duration: pd.Timedelta,
+        max_staleness: pd.Timedelta |  None = None,
+        max_dropout: pd.Timedelta = pd.Timedelta(0),
+        time_dim: str = "init_time_utc",
+        end_buffer: pd.Timedelta = pd.Timedelta(0),
+    ):
+
+    assert "step" in ds.coords
+    # It is possible to use up to this amount of max staleness for the dataset and slice
+    # required
+    possible_max_staleness = (
+        pd.Timedelta(ds["step"].max().item())
+        - forecast_duration
+        - end_buffer
+    )
+
+    # If max_staleness is set to None we set it based on the max step ahead of the input
+    # forecast data
+    if max_staleness is None:
+        max_staleness = possible_max_staleness
+    else:
+        # Make sure the max acceptable staleness isn't longer than the max possible
+        assert max_staleness <= possible_max_staleness
+        max_staleness = max_staleness
+
+    contiguous_time_periods = find_contiguous_t0_periods_nwp(
+        datetimes=pd.DatetimeIndex(ds[time_dim]),
+        history_duration=history_duration,
+        max_staleness=max_staleness,
+        max_dropout=max_dropout,
+    )
+    return contiguous_time_periods
+
+
+
+def find_contiguous_t0_periods_nwp(
+    datetimes: pd.DatetimeIndex,
+    history_duration: pd.Timedelta,
+    max_staleness: pd.Timedelta,
+    max_dropout: pd.Timedelta = pd.Timedelta(0),
+) -> pd.DataFrame:
+    """Get all time periods from the NWP init times which are valid as t0 datetimes.
+
+    Args:
+        datetimes: Sorted pd.DatetimeIndex
+        history_duration: Length of the historical slice used for a sample
+        max_staleness: Up to how long after an NWP forecast init_time are we willing to use the
+            forecast. Each init time will only be used up to this t0 time regardless of the forecast
+            valid time.
+        max_dropout: What is the maximum amount of dropout that will be used. This must be <=
+            max_staleness.
+
+    Returns:
+        pd.DataFrame where each row represents a single time period.  The pd.DataFrame
+        has two columns: `start_dt` and `end_dt` (where 'dt' is short for 'datetime').
+    """
+    # Sanity checks.
+    assert len(datetimes) > 0
+    assert datetimes.is_monotonic_increasing
+    assert datetimes.is_unique
+    assert history_duration >= pd.Timedelta(0)
+    assert max_staleness >= pd.Timedelta(0)
+    assert max_dropout <= max_staleness
+
+    hist_drop_buffer = max(history_duration, max_dropout)
+
+    # Store contiguous periods
+    contiguous_periods = []
+
+    # Start first period allowing for history slice and max dropout
+    start_this_period = datetimes[0] + hist_drop_buffer
+
+    # The first forecast is valid up to the max staleness
+    end_this_period = datetimes[0] + max_staleness
+
+    for dt_init in datetimes[1:]:
+        # If the previous init time becomes stale before the next init becomes valid whilst also
+        # considering dropout - then the contiguous period breaks, and new starts with considering
+        # dropout and history duration
+        if end_this_period < dt_init + max_dropout:
+            contiguous_periods.append([start_this_period, end_this_period])
+
+            # And start a new period
+            start_this_period = dt_init + hist_drop_buffer
+        end_this_period = dt_init + max_staleness
+
+    contiguous_periods.append([start_this_period, end_this_period])
+
+    return pd.DataFrame(contiguous_periods, columns=["start_dt", "end_dt"])
+
+
+def intersection_of_multiple_dataframes_of_periods(
+    time_periods: list[pd.DataFrame],
+) -> pd.DataFrame:
+    """Find the intersection of a list of time periods.
+
+    See the docstring of intersection_of_2_dataframes_of_periods() for more details.
+    """
+    assert len(time_periods) > 0
+    intersection = time_periods[0]
+    for time_period in time_periods[1:]:
+        intersection = intersection_of_2_dataframes_of_periods(intersection, time_period)
+    return intersection
+
+
+def intersection_of_2_dataframes_of_periods(a: pd.DataFrame, b: pd.DataFrame) -> pd.DataFrame:
+    """Find the intersection of two pd.DataFrames of time periods.
+
+    Each row of each pd.DataFrame represents a single time period.  Each pd.DataFrame has
+    two columns: `start_dt` and `end_dt` (where 'dt' is short for 'datetime').
+
+    A typical use-case is that each pd.DataFrame represents all the time periods where
+    a `DataSource` has contiguous, valid data.
+
+    Here's a graphical example of two pd.DataFrames of time periods and their intersection:
+
+                 ----------------------> TIME ->---------------------
+               a: |-----|   |----|     |----------|     |-----------|
+               b:    |--------|                       |----|    |---|
+    intersection:    |--|   |-|                         |--|    |---|
+
+    Args:
+        a: pd.DataFrame where each row represents a time period.  The pd.DataFrame has
+        two columns: start_dt and end_dt.
+        b: pd.DataFrame where each row represents a time period.  The pd.DataFrame has
+        two columns: start_dt and end_dt.
+
+    Returns:
+        Sorted list of intersecting time periods represented as a pd.DataFrame with two columns:
+        start_dt and end_dt.
+    """
+    if a.empty or b.empty:
+        return pd.DataFrame(columns=["start_dt", "end_dt"])
+
+    all_intersecting_periods = []
+    for a_period in a.itertuples():
+        # Five ways in which two periods may overlap:
+        # a: |----| or |---|   or  |---| or   |--|   or |-|
+        # b:  |--|       |---|   |---|      |------|    |-|
+        # In all five, `a` must always start before `b` ends,
+        # and `a` must always end after `b` starts:
+
+        # TODO: <= and >= because we should allow overlap time periods of length 1. e.g.
+        # a: |----|      or   |---|   
+        # b:      |--|            |---|
+        # These aren't allowed if we use < and >.
+
+        overlapping_periods = b[(a_period.start_dt < b.end_dt) & (a_period.end_dt > b.start_dt)]
+
+        # There are two ways in which two periods may *not* overlap:
+        # a: |---|        or        |---|
+        # b:       |---|      |---|
+        # `overlapping` will not include periods which do *not* overlap.
+
+        # Now find the intersection of each period in `overlapping_periods` with
+        # the period from `a` that starts at `a_start_dt` and ends at `a_end_dt`.
+        # We do this by clipping each row of `overlapping_periods`
+        # to start no earlier than `a_start_dt`, and end no later than `a_end_dt`.
+
+        # First, make a copy, so we don't clip the underlying data in `b`.
+        intersection = overlapping_periods.copy()
+        intersection["start_dt"] = intersection.start_dt.clip(lower=a_period.start_dt)
+        intersection["end_dt"] = intersection.end_dt.clip(upper=a_period.end_dt)
+
+        all_intersecting_periods.append(intersection)
+
+    all_intersecting_periods = pd.concat(all_intersecting_periods)
+    return all_intersecting_periods.sort_values(by="start_dt").reset_index(drop=True)