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

ocf_data_sampler/select/select_spatial_slice.py

@@ -0,0 +1,358 @@
+"""Select spatial slices"""
+
+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,
+    lon_lat_to_osgb,
+    osgb_to_geostationary_area_coords,
+    osgb_to_lon_lat,
+    spatial_coord_type,
+)
+from ocf_datapipes.utils.utils import searchsorted
+
+logger = logging.getLogger(__name__)
+
+
+# -------------------------------- utility functions --------------------------------
+
+
+def convert_coords_to_match_xarray(
+        x: float | np.ndarray, 
+        y: float | np.ndarray, 
+        from_coords: str, 
+        da: xr.DataArray
+    ):
+    """Convert x and y coords to cooridnate system matching xarray data
+
+    Args:
+        x: Float or array-like
+        y: Float or array-like
+        from_coords: String describing coordinate system of x and y
+        da: DataArray to which coordinates should be matched
+    """
+
+    target_coords, *_ = spatial_coord_type(da)
+
+    assert from_coords in ["osgb", "lon_lat"]
+    assert target_coords in ["geostationary", "osgb", "lon_lat"]
+
+    if target_coords == "geostationary":
+        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)
+
+        # else the from_coords=="lon_lat" and we don't need to convert
+
+    elif target_coords == "osgb":
+        if from_coords == "lon_lat":
+            x, y = lon_lat_to_osgb(x, y)
+
+        # else the from_coords=="osgb" and we don't need to convert
+
+    return x, y
+
+# TODO: This function and _get_idx_of_pixel_closest_to_poi_geostationary() should not be separate
+# We should combine them, and consider making a Coord class to help with this
+def _get_idx_of_pixel_closest_to_poi(
+    da: xr.DataArray,
+    location: Location,
+) -> Location:
+    """
+    Return x and y index location of pixel at center of region of interest.
+
+    Args:
+        da: xarray DataArray
+        location: Location to find index of
+    Returns:
+        The Location for the center pixel
+    """
+    xr_coords, x_dim, y_dim = spatial_coord_type(da)
+
+    if xr_coords not in ["osgb", "lon_lat"]:
+        raise NotImplementedError(f"Only 'osgb' and 'lon_lat' are supported - not '{xr_coords}'")
+
+    # Convert location coords to match xarray data
+    x, y = convert_coords_to_match_xarray(
+        location.x,
+        location.y,
+        from_coords=location.coordinate_system,
+        da=da,
+    )
+
+    # Check that the requested point lies within the data
+    assert da[x_dim].min() < x < da[x_dim].max()
+    assert da[y_dim].min() < y < da[y_dim].max()
+
+    x_index = da.get_index(x_dim)
+    y_index = da.get_index(y_dim)
+
+    closest_x = x_index.get_indexer([x], method="nearest")[0]
+    closest_y = y_index.get_indexer([y], method="nearest")[0]
+
+    return Location(x=closest_x, y=closest_y, coordinate_system="idx")
+
+
+def _get_idx_of_pixel_closest_to_poi_geostationary(
+    da: xr.DataArray,
+    center_osgb: Location,
+) -> Location:
+    """
+    Return x and y index location of pixel at center of region of interest.
+
+    Args:
+        da: xarray DataArray
+        center_osgb: Center in OSGB coordinates
+
+    Returns:
+        Location for the center pixel in geostationary coordinates
+    """
+
+    _, 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)
+    center_geostationary = Location(x=x, y=y, coordinate_system="geostationary")
+
+    # Check that the requested point lies within the data
+    assert da[x_dim].min() < x < da[x_dim].max()
+    assert da[y_dim].min() < y < da[y_dim].max()
+
+    # 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
+    )
+
+    return Location(x=x_index_at_center, y=y_index_at_center, coordinate_system="idx")
+
+
+# ---------------------------- sub-functions for slicing ----------------------------
+
+
+def _select_partial_spatial_slice_pixels(
+    da,
+    left_idx,
+    right_idx,
+    bottom_idx,
+    top_idx,
+    left_pad_pixels,
+    right_pad_pixels,
+    bottom_pad_pixels,
+    top_pad_pixels,
+    x_dim,
+    y_dim,
+):
+    """Return spatial window of given pixel size when window partially overlaps input data"""
+
+    # We should never be padding on both sides of a window. This would mean our desired window is 
+    # larger than the size of the input data
+    assert left_pad_pixels==0 or right_pad_pixels==0
+    assert bottom_pad_pixels==0 or top_pad_pixels==0
+
+    dx = np.median(np.diff(da[x_dim].values))
+    dy = np.median(np.diff(da[y_dim].values))
+
+    # Pad the left of the window
+    if left_pad_pixels > 0:
+        x_sel = np.concatenate(
+            [
+                da[x_dim].values[0] + np.arange(-left_pad_pixels, 0) * dx,
+                da[x_dim].values[0:right_idx],
+            ]
+        )
+        da = da.isel({x_dim: slice(0, right_idx)}).reindex({x_dim: x_sel})
+
+    # Pad the right of the window
+    elif right_pad_pixels > 0:
+        x_sel = np.concatenate(
+            [
+                da[x_dim].values[left_idx:],
+                da[x_dim].values[-1] + np.arange(1, right_pad_pixels + 1) * dx,
+            ]
+        )
+        da = da.isel({x_dim: slice(left_idx, None)}).reindex({x_dim: x_sel})
+
+    # No left-right padding required
+    else:
+        da = da.isel({x_dim: slice(left_idx, right_idx)})
+
+    # Pad the bottom of the window
+    if bottom_pad_pixels > 0:
+        y_sel = np.concatenate(
+            [
+                da[y_dim].values[0] + np.arange(-bottom_pad_pixels, 0) * dy,
+                da[y_dim].values[0:top_idx],
+            ]
+        )
+        da = da.isel({y_dim: slice(0, top_idx)}).reindex({y_dim: y_sel})
+
+    # Pad the top of the window
+    elif top_pad_pixels > 0:
+        y_sel = np.concatenate(
+            [
+                da[y_dim].values[bottom_idx:],
+                da[y_dim].values[-1] + np.arange(1, top_pad_pixels + 1) * dy,
+            ]
+        )
+        da = da.isel({y_dim: slice(left_idx, None)}).reindex({y_dim: y_sel})
+
+    # No bottom-top padding required
+    else:
+        da = da.isel({y_dim: slice(bottom_idx, top_idx)})
+
+    return da
+
+
+def _select_spatial_slice_pixels(
+    da: xr.DataArray, 
+    center_idx: Location, 
+    width_pixels: int, 
+    height_pixels: int, 
+    x_dim: str, 
+    y_dim: str, 
+    allow_partial_slice: bool,
+):
+    """Select a spatial slice from an xarray object
+
+    Args:
+        da: xarray DataArray to slice from
+        center_idx: Location object describing the centre of the window with index coordinates
+        width_pixels: Window with in pixels
+        height_pixels: Window height in pixels
+        x_dim: Name of the x-dimension in `da`
+        y_dim: Name of the y-dimension in `da`
+        allow_partial_slice: Whether to allow a partially filled window
+    """
+
+    assert center_idx.coordinate_system == "idx"
+    # TODO: It shouldn't take much effort to allow height and width to be odd
+    assert (width_pixels % 2)==0, "Width must be an even number"
+    assert (height_pixels % 2)==0, "Height must be an even number"
+
+    half_width = width_pixels // 2
+    half_height = height_pixels // 2
+
+    left_idx = int(center_idx.x - half_width)
+    right_idx = int(center_idx.x + half_width)
+    bottom_idx = int(center_idx.y - half_height)
+    top_idx = int(center_idx.y + half_height)
+
+    data_width_pixels = len(da[x_dim])
+    data_height_pixels = len(da[y_dim])
+
+    left_pad_required = left_idx < 0
+    right_pad_required = right_idx > data_width_pixels
+    bottom_pad_required = bottom_idx < 0
+    top_pad_required = top_idx > data_height_pixels
+
+    pad_required = left_pad_required | right_pad_required | bottom_pad_required | top_pad_required
+
+    if pad_required:
+        if allow_partial_slice:
+
+            left_pad_pixels = (-left_idx) if left_pad_required else 0
+            right_pad_pixels = (right_idx - data_width_pixels) if right_pad_required else 0
+
+            bottom_pad_pixels = (-bottom_idx) if bottom_pad_required else 0
+            top_pad_pixels = (top_idx - data_height_pixels) if top_pad_required else 0
+
+
+            da = _select_partial_spatial_slice_pixels(
+                da,
+                left_idx,
+                right_idx,
+                bottom_idx,
+                top_idx,
+                left_pad_pixels,
+                right_pad_pixels,
+                bottom_pad_pixels,
+                top_pad_pixels,
+                x_dim,
+                y_dim,
+            )
+        else:
+            raise ValueError(
+                f"Window for location {center_idx} not available. Missing (left, right, bottom, "
+                f"top) pixels  = ({left_pad_required}, {right_pad_required}, "
+                f"{bottom_pad_required}, {top_pad_required}). "
+                f"You may wish to set `allow_partial_slice=True`"
+            )
+
+    else:
+        da = da.isel(
+            {
+                x_dim: slice(left_idx, right_idx),
+                y_dim: slice(bottom_idx, top_idx),
+            }
+        )
+
+    assert len(da[x_dim]) == width_pixels, (
+        f"Expected x-dim len {width_pixels} got {len(da[x_dim])} "
+        f"for location {center_idx} for slice {left_idx}:{right_idx}"
+    )
+    assert len(da[y_dim]) == height_pixels, (
+        f"Expected y-dim len {height_pixels} got {len(da[y_dim])} "
+        f"for location {center_idx} for slice {bottom_idx}:{top_idx}"
+    )
+
+    return da
+
+
+# ---------------------------- main functions for slicing ---------------------------
+
+
+def select_spatial_slice_pixels(
+    da: xr.DataArray,
+    location: Location,
+    width_pixels: int,
+    height_pixels: int,
+    allow_partial_slice: bool = False,
+):
+    """
+    Select spatial slice based off pixels from location point of interest
+
+    If `allow_partial_slice` is set to True, then slices may be made which intersect the border
+    of the input data. The additional x and y cordinates that would be required for this slice
+    are extrapolated based on the average spacing of these coordinates in the input data.
+    However, currently slices cannot be made where the centre of the window is outside of the
+    input data.
+
+    Args:
+        da: xarray DataArray to slice from
+        location: Location of interest
+        height_pixels: Height of the slice in pixels
+        width_pixels: Width of the slice in pixels
+        allow_partial_slice: Whether to allow a partial slice.
+    """
+
+    xr_coords, x_dim, y_dim = spatial_coord_type(da)
+
+    if xr_coords == "geostationary":
+        center_idx: Location = _get_idx_of_pixel_closest_to_poi_geostationary(da, location)
+    else:
+        center_idx: Location = _get_idx_of_pixel_closest_to_poi(da, location)
+
+    selected = _select_spatial_slice_pixels(
+        da,
+        center_idx,
+        width_pixels,
+        height_pixels,
+        x_dim,
+        y_dim,
+        allow_partial_slice=allow_partial_slice,
+    )
+
+    return selected

ocf_data_sampler/select/select_time_slice.py

@@ -0,0 +1,184 @@
+import xarray as xr
+import pandas as pd
+import numpy as np
+
+
+
+def _sel_fillnan(
+        da: xr.DataArray, 
+        start_dt: pd.Timestamp, 
+        end_dt: pd.Timestamp, 
+        sample_period_duration: pd.Timedelta,
+    ) -> xr.DataArray:
+    """Select a time slice from a DataArray, filling missing times with NaNs."""
+    requested_times = pd.date_range(start_dt, end_dt, freq=sample_period_duration)
+    return da.reindex(time_utc=requested_times)
+
+
+def _sel_default(
+        da: xr.DataArray, 
+        start_dt: pd.Timestamp, 
+        end_dt: pd.Timestamp, 
+        sample_period_duration: pd.Timedelta,
+    ) -> xr.DataArray:
+    """Select a time slice from a DataArray, without filling missing times."""
+    return da.sel(time_utc=slice(start_dt, end_dt))
+
+
+# TODO either implement this or remove it, which would tidy up the code
+def _sel_fillinterp(
+        da: xr.DataArray, 
+        start_dt: pd.Timestamp, 
+        end_dt: pd.Timestamp, 
+        sample_period_duration: pd.Timedelta,
+    ) -> xr.DataArray:
+    """Select a time slice from a DataArray, filling missing times with linear interpolation."""
+    return NotImplemented
+
+
+def select_time_slice(
+    ds: xr.Dataset | xr.DataArray,
+    t0: pd.Timestamp,
+    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
+    elif fill_selection and max_steps_gap > 0:
+        _sel = _sel_fillinterp
+    else:
+        _sel = _sel_default
+
+    t0_datetime_utc = pd.Timestamp(t0)
+    start_dt = t0_datetime_utc + interval_start
+    end_dt = t0_datetime_utc + interval_end
+
+    start_dt = start_dt.ceil(sample_period_duration)
+    end_dt = end_dt.ceil(sample_period_duration)
+
+    return _sel(ds, start_dt, end_dt, sample_period_duration)
+
+
+def select_time_slice_nwp(
+    ds: xr.Dataset | xr.DataArray,
+    t0: pd.Timestamp,
+    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] = [],
+    channel_dim_name: str = "channel",
+):
+
+    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
+
+
+    # The accumatation and non-accumulation channels
+    accum_channels = np.intersect1d(
+        ds[channel_dim_name].values, accum_channels
+    )
+    non_accum_channels = np.setdiff1d(
+        ds[channel_dim_name].values, accum_channels
+    )
+
+    start_dt = (t0 - history_duration).ceil(sample_period_duration)
+    end_dt = (t0 + forecast_duration).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):
+        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(
+        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(
+        init_time_utc=target_times,
+        method="ffill",  # forward fill from init times to target times
+    ).values
+
+    # Find the required steps for all target times
+    steps = target_times - selected_init_times
+    
+    # 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:
+    # 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)
+
+    if len(accum_channels) == 0:
+        xr_sel = ds.sel(step=step_indexer, init_time_utc=init_time_indexer)
+
+    else:
+        # First minimise the size of the dataset we are diffing
+        # - find the init times we are slicing from
+        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
+
+        xr_min = ds.sel(
+            {
+                "init_time_utc": unique_init_times,
+                "step": slice(min_step, max_step),
+            }
+        )
+
+        # 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(
+            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})
+
+        # 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)
+
+        # Join diffed and non-diffed variables
+        xr_sel = xr.concat([xr_sel_non_accum, xr_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})
+
+        # Rename the diffed channels
+        xr_sel[channel_dim_name] = [
+            f"diff_{v}" if v in accum_channels else v
+            for v in xr_sel[channel_dim_name].values
+        ]
+
+    return xr_sel

ocf_data_sampler/torch_datasets/__init__.py

@@ -0,0 +1 @@
+