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

ocf_data_sampler/select/geospatial.py

@@ -1,36 +1,45 @@
-"""Geospatial functions"""
+"""Geospatial coordinate transformation functions.
 
-from numbers import Number
-from typing import Union
+Provides utilities for working with different coordinate systems
+commonly used in geospatial applications, particularly for UK-based data.
+
+Supports conversions between:
+- OSGB36 (Ordnance Survey Great Britain, easting/northing in meters)
+- WGS84 (World Geodetic System, latitude/longitude in degrees)
+- Geostationary satellite coordinate systems
+"""
 
 import numpy as np
 import pyproj
+import pyresample
 import xarray as xr
 
-# OSGB is also called "OSGB 1936 / British National Grid -- United
-# Kingdom Ordnance Survey".  OSGB is used in many UK electricity
-# system maps, and is used by the UK Met Office UKV model.  OSGB is a
-# Transverse Mercator projection, using 'easting' and 'northing'
-# coordinates which are in meters.  See https://epsg.io/27700
+# Coordinate Reference System (CRS) identifiers
+# OSGB36: UK Ordnance Survey National Grid (easting/northing in meters)
+# Refer to - https://epsg.io/27700
 OSGB36 = 27700
 
-# WGS84 is short for "World Geodetic System 1984", used in GPS. Uses
-# latitude and longitude.
+# WGS84: World Geodetic System 1984 (latitude/longitude in degrees), used in GPS
 WGS84 = 4326
 
-
+# Pre-init Transformer
 _osgb_to_lon_lat = pyproj.Transformer.from_crs(
-    crs_from=OSGB36, crs_to=WGS84, always_xy=True
+    crs_from=OSGB36,
+    crs_to=WGS84,
+    always_xy=True,
 ).transform
 _lon_lat_to_osgb = pyproj.Transformer.from_crs(
-    crs_from=WGS84, crs_to=OSGB36, always_xy=True
+    crs_from=WGS84,
+    crs_to=OSGB36,
+    always_xy=True,
 ).transform
 
 
 def osgb_to_lon_lat(
-    x: Union[Number, np.ndarray], y: Union[Number, np.ndarray]
-) -> tuple[Union[Number, np.ndarray], Union[Number, np.ndarray]]:
-    """Change OSGB coordinates to lon, lat.
+    x: float | np.ndarray,
+    y: float | np.ndarray,
+) -> tuple[float | np.ndarray, float | np.ndarray]:
+    """Change OSGB coordinates to lon-lat.
 
     Args:
         x: osgb east-west
@@ -41,9 +50,9 @@ def osgb_to_lon_lat(
 
 
 def lon_lat_to_osgb(
-    x: Union[Number, np.ndarray],
-    y: Union[Number, np.ndarray],
-) -> tuple[Union[Number, np.ndarray], Union[Number, np.ndarray]]:
+    x: float | np.ndarray,
+    y: float | np.ndarray,
+) -> tuple[float | np.ndarray, float | np.ndarray]:
     """Change lon-lat coordinates to OSGB.
 
     Args:
@@ -56,11 +65,11 @@ def lon_lat_to_osgb(
 
 
 def lon_lat_to_geostationary_area_coords(
-    longitude: Union[Number, np.ndarray],
-    latitude: Union[Number, np.ndarray],
+    longitude: float | np.ndarray,
+    latitude: float | np.ndarray,
     xr_data: xr.DataArray,
-) -> tuple[Union[Number, np.ndarray], Union[Number, np.ndarray]]:
-    """Loads geostationary area and transformation from lat-lon to geostationary coords
+) -> tuple[float | np.ndarray, float | np.ndarray]:
+    """Loads geostationary area and transformation from lat-lon to geostationary coords.
 
     Args:
         longitude: longitude
@@ -72,12 +81,13 @@ def lon_lat_to_geostationary_area_coords(
     """
     return coordinates_to_geostationary_area_coords(longitude, latitude, xr_data, WGS84)
 
+
 def osgb_to_geostationary_area_coords(
-    x: Union[Number, np.ndarray],
-    y: Union[Number, np.ndarray],
+    x: float | np.ndarray,
+    y: float | np.ndarray,
     xr_data: xr.DataArray,
-) -> tuple[Union[Number, np.ndarray], Union[Number, np.ndarray]]:
-    """Loads geostationary area and transformation from OSGB to geostationary coords
+) -> tuple[float | np.ndarray, float | np.ndarray]:
+    """Loads geostationary area and transformation from OSGB to geostationary coords.
 
     Args:
         x: osgb east-west
@@ -87,47 +97,45 @@ def osgb_to_geostationary_area_coords(
     Returns:
         Geostationary coords: x, y
     """
-
     return coordinates_to_geostationary_area_coords(x, y, xr_data, OSGB36)
 
 
-
 def coordinates_to_geostationary_area_coords(
-    x: Union[Number, np.ndarray],
-    y: Union[Number, np.ndarray],
+    x: float | np.ndarray,
+    y: float | np.ndarray,
     xr_data: xr.DataArray,
-    crs_from: int
-) -> tuple[Union[Number, np.ndarray], Union[Number, np.ndarray]]:
-    """Loads geostationary area and transformation from respective coordiates to geostationary coords
-
-        Args:
-            x: osgb east-west, or latitude
-            y: osgb north-south, or longitude
-            xr_data: xarray object with geostationary area
-            crs_from: the cordiates system of x,y
+    crs_from: int,
+) -> tuple[float | np.ndarray, float | np.ndarray]:
+    """Loads geostationary area and transforms to geostationary coords.
 
-        Returns:
-            Geostationary coords: x, y
-        """
-
-    assert crs_from in [OSGB36, WGS84], f"Unrecognized coordinate system: {crs_from}"
+    Args:
+        x: osgb east-west, or latitude
+        y: osgb north-south, or longitude
+        xr_data: xarray object with geostationary area
+        crs_from: the cordiates system of x,y
 
-    # Only load these if using geostationary projection
-    import pyresample
+    Returns:
+        Geostationary coords: x, y
+    """
+    if crs_from not in [OSGB36, WGS84]:
+        raise ValueError(f"Unrecognized coordinate system: {crs_from}")
 
     area_definition_yaml = xr_data.attrs["area"]
 
     geostationary_area_definition = pyresample.area_config.load_area_from_string(
-        area_definition_yaml
+        area_definition_yaml,
     )
     geostationary_crs = geostationary_area_definition.crs
     osgb_to_geostationary = pyproj.Transformer.from_crs(
-        crs_from=crs_from, crs_to=geostationary_crs, always_xy=True
+        crs_from=crs_from,
+        crs_to=geostationary_crs,
+        always_xy=True,
     ).transform
     return osgb_to_geostationary(xx=x, yy=y)
 
 
-def _coord_priority(available_coords):
+def _coord_priority(available_coords: list[str]) -> tuple[str, str, str]:
+    """Determines the coordinate system of spatial coordinates present."""
     if "longitude" in available_coords:
         return "lon_lat", "longitude", "latitude"
     elif "x_geostationary" in available_coords:
@@ -138,7 +146,7 @@ def _coord_priority(available_coords):
         raise ValueError(f"Unrecognized coordinate system: {available_coords}")
 
 
-def spatial_coord_type(ds: xr.DataArray):
+def spatial_coord_type(ds: xr.DataArray) -> tuple[str, str, str]:
     """Searches the data array to determine the kind of spatial coordinates present.
 
     This search has a preference for the dimension coordinates of the xarray object.
@@ -147,9 +155,10 @@ def spatial_coord_type(ds: xr.DataArray):
         ds: Dataset with spatial coords
 
     Returns:
-        str: The kind of the coordinate system
-        x_coord: Name of the x-coordinate
-        y_coord: Name of the y-coordinate
+        Three strings with:
+            1. The kind of the coordinate system
+            2. Name of the x-coordinate
+            3. Name of the y-coordinate
     """
     if isinstance(ds, xr.DataArray):
         # Search dimension coords of dataarray

ocf_data_sampler/select/location.py

@@ -1,62 +1,27 @@
-"""location"""
+"""Location model with coordinate system validation."""
 
-from typing import Optional
-
-import numpy as np
 from pydantic import BaseModel, Field, model_validator
 
+allowed_coordinate_systems = ["osgb", "lon_lat", "geostationary", "idx"]
 
-allowed_coordinate_systems =["osgb", "lon_lat", "geostationary", "idx"]
 
 class Location(BaseModel):
     """Represent a spatial location."""
 
-    coordinate_system: Optional[str] = "osgb"  # ["osgb", "lon_lat", "geostationary", "idx"]
-    x: float
-    y: float
-    id: Optional[int] = Field(None)
-
-    @model_validator(mode='after')
-    def validate_coordinate_system(self):
-        """Validate 'coordinate_system'"""
-        if self.coordinate_system not in allowed_coordinate_systems:
-            raise ValueError(f"coordinate_system = {self.coordinate_system} is not in {allowed_coordinate_systems}")
-        return self
-
-    @model_validator(mode='after')
-    def validate_x(self):
-        """Validate 'x'"""
-        min_x: float
-        max_x: float
-
-        co = self.coordinate_system
-        if co == "osgb":
-            min_x, max_x = -103976.3, 652897.98
-        if co == "lon_lat":
-            min_x, max_x = -180, 180
-        if co == "geostationary":
-            min_x, max_x = -5568748.275756836, 5567248.074173927
-        if co == "idx":
-            min_x, max_x = 0, np.inf
-        if self.x < min_x or self.x > max_x:
-            raise ValueError(f"x = {self.x} must be within {[min_x, max_x]} for {co} coordinate system")
-        return self
+    coordinate_system: str = Field(...,
+        description="Coordinate system for the location must be lon_lat, osgb, or geostationary",
+    )
 
-    @model_validator(mode='after')
-    def validate_y(self):
-        """Validate 'y'"""
-        min_y: float
-        max_y: float
+    x: float = Field(..., description="x coordinate - i.e. east-west position")
+    y: float = Field(..., description="y coordinate - i.e. north-south position")
+    id: int | None = Field(None, description="ID of the location - e.g. GSP ID")
 
-        co = self.coordinate_system
-        if co == "osgb":
-            min_y, max_y = -16703.87, 1199851.44
-        if co == "lon_lat":
-            min_y, max_y = -90, 90
-        if co == "geostationary":
-            min_y, max_y = 1393687.2151494026, 5570748.323202133
-        if co == "idx":
-            min_y, max_y = 0, np.inf
-        if self.y < min_y or self.y > max_y:
-            raise ValueError(f"y = {self.y} must be within {[min_y, max_y]} for {co} coordinate system")
+    @model_validator(mode="after")
+    def validate_coordinate_system(self) -> "Location":
+        """Validate 'coordinate_system'."""
+        if self.coordinate_system not in allowed_coordinate_systems:
+            raise ValueError(
+                f"coordinate_system = {self.coordinate_system} "
+                f"is not in {allowed_coordinate_systems}",
+            )
         return self

ocf_data_sampler/select/select_spatial_slice.py

@@ -1,19 +1,18 @@
-"""Select spatial slices"""
+"""Select spatial slices."""
 
 import logging
 
 import numpy as np
 import xarray as xr
 
-from ocf_data_sampler.select.location import Location
 from ocf_data_sampler.select.geospatial import (
-    lon_lat_to_osgb,
     lon_lat_to_geostationary_area_coords,
+    lon_lat_to_osgb,
     osgb_to_geostationary_area_coords,
     osgb_to_lon_lat,
     spatial_coord_type,
 )
-
+from ocf_data_sampler.select.location import Location
 
 logger = logging.getLogger(__name__)
 
@@ -22,12 +21,12 @@ logger = logging.getLogger(__name__)
 
 
 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
+    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.
 
     Args:
         x: Float or array-like
@@ -35,38 +34,42 @@ def convert_coords_to_match_xarray(
         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":
+    match (from_coords, target_coords):
+        case ("osgb", "geostationary"):
             x, y = osgb_to_geostationary_area_coords(x, y, da)
 
-    elif target_coords == "lon_lat":
-        if from_coords == "osgb":
+        case ("osgb", "lon_lat"):
             x, y = osgb_to_lon_lat(x, y)
 
-        # else the from_coords=="lon_lat" and we don't need to convert
+        case ("osgb", "osgb"):
+            pass
 
-    elif target_coords == "osgb":
-        if from_coords == "lon_lat":
+        case ("lon_lat", "osgb"):
             x, y = lon_lat_to_osgb(x, y)
 
-        # else the from_coords=="osgb" and we don't need to convert
+        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",
+            )
 
     return x, y
 
-# TODO: This function and _get_idx_of_pixel_closest_to_poi_geostationary() should not be separate
+
+# 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.
+    """Return x and y index location of pixel at center of region of interest.
 
     Args:
         da: xarray DataArray
@@ -88,8 +91,14 @@ def _get_idx_of_pixel_closest_to_poi(
     )
 
     # 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()
+    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}",
+        )
 
     x_index = da.get_index(x_dim)
     y_index = da.get_index(y_dim)
@@ -104,32 +113,38 @@ 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.
+    """Return x and y index location of pixel at center of region of interest.
 
     Args:
         da: xarray DataArray
-        center_osgb: Center in OSGB coordinates
+        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':
+    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)
+    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
+        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
-    assert da[x_dim].min() < x < da[x_dim].max(), \
-        f"{x} is not in the interval {da[x_dim].min().values}: {da[x_dim].max().values}"
-    assert da[y_dim].min() < y < da[y_dim].max(), \
-        f"{y} is not in the interval {da[y_dim].min().values}: {da[y_dim].max().values}"
+    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)
@@ -142,24 +157,25 @@ def _get_idx_of_pixel_closest_to_poi_geostationary(
 
 
 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 
+    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
-    assert left_pad_pixels==0 or right_pad_pixels==0
-    assert bottom_pad_pixels==0 or top_pad_pixels==0
+    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))
@@ -170,7 +186,7 @@ def _select_partial_spatial_slice_pixels(
             [
                 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})
 
@@ -180,7 +196,7 @@ def _select_partial_spatial_slice_pixels(
             [
                 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})
 
@@ -194,7 +210,7 @@ def _select_partial_spatial_slice_pixels(
             [
                 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})
 
@@ -204,7 +220,7 @@ def _select_partial_spatial_slice_pixels(
             [
                 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})
 
@@ -216,15 +232,15 @@ def _select_partial_spatial_slice_pixels(
 
 
 def _select_spatial_slice_pixels(
-    da: xr.DataArray, 
-    center_idx: Location, 
-    width_pixels: int, 
-    height_pixels: int, 
-    x_dim: str, 
-    y_dim: str, 
+    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
+) -> xr.DataArray:
+    """Select a spatial slice from an xarray object.
 
     Args:
         da: xarray DataArray to slice from
@@ -235,11 +251,13 @@ def _select_spatial_slice_pixels(
         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"
+    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
-    assert (width_pixels % 2)==0, "Width must be an even number"
-    assert (height_pixels % 2)==0, "Height must be an even number"
+    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")
 
     half_width = width_pixels // 2
     half_height = height_pixels // 2
@@ -261,14 +279,12 @@ def _select_spatial_slice_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,
@@ -287,7 +303,7 @@ def _select_spatial_slice_pixels(
                 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"You may wish to set `allow_partial_slice=True`",
             )
 
     else:
@@ -295,17 +311,19 @@ def _select_spatial_slice_pixels(
             {
                 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}"
-    )
+    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}",
+        )
+    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}",
+        )
 
     return da
 
@@ -319,9 +337,8 @@ def select_spatial_slice_pixels(
     width_pixels: int,
     height_pixels: int,
     allow_partial_slice: bool = False,
-):
-    """
-    Select spatial slice based off pixels from location point of interest
+) -> 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
@@ -336,7 +353,6 @@ def select_spatial_slice_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":
@@ -354,4 +370,4 @@ def select_spatial_slice_pixels(
         allow_partial_slice=allow_partial_slice,
     )
 
-    return selected
+    return selected