Spatial slice select update

ocf_data_sampler/select/select_spatial_slice.py

@@ -17,80 +17,64 @@
 logger = logging.getLogger(__name__)
 
 
-# -------------------------------- utility functions --------------------------------
-
-
-def convert_coords_to_match_xarray(
+def convert_coordinates(
+    from_coords: str,
     x: float | np.ndarray,
     y: float | np.ndarray,
-    from_coords: str,
     da: xr.DataArray,
 ) -> tuple[float | np.ndarray, float | np.ndarray]:
-    """Convert x and y coords to cooridnate system matching xarray data.
+    """Convert x and y coordinates to coordinate 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
+        from_coords: The coordinate system to convert from.
+        x: The x-coordinate to convert.
+        y: The y-coordinate to convert.
+        da: The xarray DataArray used for context (e.g., for geostationary conversion).
+
+    Returns:
+        The converted (x, y) coordinates.
     """
     target_coords, *_ = spatial_coord_type(da)
 
     match (from_coords, target_coords):
         case ("osgb", "geostationary"):
             x, y = osgb_to_geostationary_area_coords(x, y, da)
-
         case ("osgb", "lon_lat"):
             x, y = osgb_to_lon_lat(x, y)
-
         case ("osgb", "osgb"):
             pass
-
         case ("lon_lat", "osgb"):
             x, y = lon_lat_to_osgb(x, y)
-
         case ("lon_lat", "geostationary"):
             x, y = lon_lat_to_geostationary_area_coords(x, y, da)
-
         case ("lon_lat", "lon_lat"):
             pass
-
         case (_, _):
             raise NotImplementedError(
-                f"Conversion from {from_coords} to {target_coords} is not supported",
+                f"Conversion from {from_coords} to "
+                f"{target_coords} is not supported",
             )
-
     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.
+def _get_pixel_index_location(da: xr.DataArray, location: Location) -> Location:
+    """Find pixel index location closest to given Location.
 
     Args:
-        da: xarray DataArray
-        location: Location to find index of
+        da: The xarray DataArray.
+        location: The Location object representing the point of interest.
+
     Returns:
-        The Location for the center pixel
+        A Location object with x and y attributes representing the pixel indices.
+
+    Raises:
+        ValueError: If the location is outside the bounds of the DataArray.
     """
     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}'")
+    x, y = convert_coordinates(location.coordinate_system, location.x, location.y, da)
 
-    # 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
+    # Check that requested point lies within the data
     if not (da[x_dim].min() < x < da[x_dim].max()):
         raise ValueError(
             f"{x} is not in the interval {da[x_dim].min().values}: {da[x_dim].max().values}",
@@ -102,84 +86,53 @@ def _get_idx_of_pixel_closest_to_poi(
 
     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: Location,
-) -> Location:
-    """Return x and y index location of pixel at center of region of interest.
-
-    Args:
-        da: xarray DataArray
-        center: Center in OSGB coordinates
-
-    Returns:
-        Location for the center pixel in geostationary coordinates
-    """
-    _, x_dim, y_dim = spatial_coord_type(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
-    if not (da[x_dim].min() < x < da[x_dim].max()):
-        raise ValueError(
-            f"{x} is not in the interval {da[x_dim].min().values}: {da[x_dim].max().values}",
-        )
-    if not (da[y_dim].min() < y < da[y_dim].max()):
-        raise ValueError(
-            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 = 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")
-
-
-# ---------------------------- sub-functions for slicing ----------------------------
-
-
-def _select_partial_spatial_slice_pixels(
+def _select_padded_slice(
     da: xr.DataArray,
     left_idx: int,
     right_idx: int,
     bottom_idx: int,
     top_idx: int,
-    left_pad_pixels: int,
-    right_pad_pixels: int,
-    bottom_pad_pixels: int,
-    top_pad_pixels: int,
     x_dim: str,
     y_dim: str,
 ) -> xr.DataArray:
-    """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
-    if (left_pad_pixels != 0 and right_pad_pixels != 0) or (
-        bottom_pad_pixels != 0 and top_pad_pixels != 0
+    """Selects spatial slice - padding where necessary if indices are out of bounds.
+
+    Args:
+        da: xarray DataArray.
+        left_idx: The leftmost index of the slice.
+        right_idx: The rightmost index of the slice.
+        bottom_idx: The bottommost index of the slice.
+        top_idx: The topmost index of the slice.
+        x_dim: Name of the x dimension.
+        y_dim: Name of the y dimension.
+
+    Returns:
+        An xarray DataArray with padding, if necessary.
+    """
+    data_width_pixels = len(da[x_dim])
+    data_height_pixels = len(da[y_dim])
+
+    left_pad_pixels = max(0, -left_idx)
+    right_pad_pixels = max(0, right_idx - data_width_pixels)
+    bottom_pad_pixels = max(0, -bottom_idx)
+    top_pad_pixels = max(0, top_idx - data_height_pixels)
+
+    if (left_pad_pixels > 0 and right_pad_pixels > 0) or (
+        bottom_pad_pixels > 0 and top_pad_pixels > 0
     ):
         raise ValueError("Cannot pad both sides of the window")
 
     dx = np.median(np.diff(da[x_dim].values))
     dy = np.median(np.diff(da[y_dim].values))
 
+    # Create a new DataArray which has indices which go outside
+    # the original DataArray
     # Pad the left of the window
     if left_pad_pixels > 0:
         x_sel = np.concatenate(
@@ -222,7 +175,7 @@ def _select_partial_spatial_slice_pixels(
                 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})
+        da = da.isel({y_dim: slice(bottom_idx, None)}).reindex({y_dim: y_sel})
 
     # No bottom-top padding required
     else:
@@ -231,34 +184,38 @@ def _select_partial_spatial_slice_pixels(
     return da
 
 
-def _select_spatial_slice_pixels(
+def select_spatial_slice_pixels(
     da: xr.DataArray,
-    center_idx: Location,
+    location: Location,
     width_pixels: int,
     height_pixels: int,
-    x_dim: str,
-    y_dim: str,
-    allow_partial_slice: bool,
+    allow_partial_slice: bool = False,
 ) -> xr.DataArray:
-    """Select a spatial slice from an xarray object.
+    """Select spatial slice based off pixels from location point of interest.
 
     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
+        location: Location of interest that will be the center of the returned slice
+        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.
+
+    Returns:
+        The selected DataArray slice.
+
+    Raises:
+        ValueError: If the dimensions are not even or the slice is not allowed
+                    when padding is required.
+
     """
-    if center_idx.coordinate_system != "idx":
-        raise ValueError(f"Expected center_idx to be in 'idx' coordinates, got '{center_idx}'")
-    # TODO: It shouldn't take much effort to allow height and width to be odd
     if (width_pixels % 2) != 0:
         raise ValueError("Width must be an even number")
     if (height_pixels % 2) != 0:
         raise ValueError("Height must be an even number")
 
+    _, x_dim, y_dim = spatial_coord_type(da)
+    center_idx = _get_pixel_index_location(da, location)
+
     half_width = width_pixels // 2
     half_height = height_pixels // 2
 
@@ -270,104 +227,29 @@ def _select_spatial_slice_pixels(
     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
+    # Padding checks
+    pad_required = (
+        left_idx < 0
+        or right_idx > data_width_pixels
+        or bottom_idx < 0
+        or top_idx > data_height_pixels
+    )
 
     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,
-            )
+            da = _select_padded_slice(da, left_idx, right_idx, bottom_idx, top_idx, 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`",
+                f"Window for location {location} not available.  Padding required. "
+                "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),
-            },
-        )
+        # Standard selection - without padding
+        da = da.isel({x_dim: slice(left_idx, right_idx), y_dim: slice(bottom_idx, top_idx)})
 
     if len(da[x_dim]) != width_pixels:
-        raise ValueError(
-            f"Expected x-dim len {width_pixels} got {len(da[x_dim])} "
-            f"for location {center_idx} for slice {left_idx}:{right_idx}",
-        )
+        raise ValueError(f"x-dim has size {len(da[x_dim])}, expected {width_pixels}")
     if len(da[y_dim]) != height_pixels:
-        raise ValueError(
-            f"Expected y-dim len {height_pixels} got {len(da[y_dim])} "
-            f"for location {center_idx} for slice {bottom_idx}:{top_idx}",
-        )
+        raise ValueError(f"y-dim has size {len(da[y_dim])}, expected {height_pixels}")
 
     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,
-) -> xr.DataArray:
-    """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