raster geolocation subsetting and square bracket index subsetting

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rasters"
-version = "1.14.0"
+version = "1.15.0"
 description = "raster processing toolkit"
 readme = "README.md"
 authors = [

rasters/raster_geolocation.py

@@ -13,9 +13,11 @@
 
 if TYPE_CHECKING:
     from .CRS import CRS
+    from .bbox import BBox
     from .point import Point
     from .polygon import Polygon
     from .raster_grid import RasterGrid
+    from .spatial_geometry import SpatialGeometry
 
 
 class RasterGeolocation(RasterGeometry):
@@ -340,10 +342,19 @@ def to_dict(self, output_dict: Dict = None, write_geolocation_arrays: bool = Fal
 
         return output_dict
 
-    def window(self, geometry) -> Window:
+    def window(
+            self,
+            geometry: Union[SpatialGeometry, Tuple[float, float, float, float]],
+            buffer: int = None) -> Window:
         """
         Returns a rasterio.windows.Window covering the target geometry.
-        Equivalent to RasterGrid.window, but for geolocation arrays.
+
+        Args:
+            geometry: The geometry to create a window for
+            buffer: Optional buffer in pixels to add around the geometry
+
+        Returns:
+            Window: A rasterio Window object covering the geometry
         """
 
         mask = self.index(geometry)
@@ -357,4 +368,29 @@ def window(self, geometry) -> Window:
         height = int(rows.max() - rows.min() + 1)
         width = int(cols.max() - cols.min() + 1)
 
-        return Window(col_off=col_off, row_off=row_off, width=width, height=height)
+        # Apply buffer if specified
+        if buffer is not None and buffer > 0:
+            row_off = max(0, row_off - buffer)
+            col_off = max(0, col_off - buffer)
+            height = min(self.rows - row_off, height + 2 * buffer)
+            width = min(self.cols - col_off, width + 2 * buffer)
+
+        return Window(col_off=col_off, row_off=row_off, width=width, height=height)
+
+    def subset(self, target: Union[Window, Point, Polygon, BBox, RasterGeometry]) -> 'RasterGeolocation':
+        """
+        Subset the raster geolocation using a Window or other geometry.
+
+        Args:
+            target: Window object or geometry to subset with
+
+        Returns:
+            RasterGeolocation: A new RasterGeolocation object representing the subset
+        """
+        if not isinstance(target, Window):
+            target = self.window(target)
+
+        row_slice = slice(target.row_off, target.row_off + target.height)
+        col_slice = slice(target.col_off, target.col_off + target.width)
+
+        return self._slice(row_slice, col_slice)

rasters/raster_geometry.py

@@ -10,6 +10,7 @@
 import numpy as np
 import shapely
 from pyproj import Transformer
+from rasterio.windows import Window
 from scipy.ndimage import shift
 from shapely.geometry.base import BaseGeometry
 from shapely.ops import transform as shapely_transform
@@ -91,10 +92,16 @@ def __repr__(self):
     def __eq__(self, other: RasterGeometry) -> bool:
         pass
 
-    def __getitem__(self, key: Union[slice, int, tuple]):
-        if isinstance(key, (slice, int, tuple)):
+    def __getitem__(self, key):
+        from .point import Point
+        from .polygon import Polygon
+        from .bbox import BBox
+
+        # Check if key is a type accepted by subset method
+        if isinstance(key, (Window, Point, Polygon, BBox, RasterGeometry)):
+            return self.subset(key)
+        elif isinstance(key, (slice, int, tuple)):
             return self._key(key)
-
         else:
             raise KeyError('unrecognized key')
 
@@ -271,6 +278,36 @@ def index_point(self, point: Point) -> Tuple[int, int]:
     def index(self, geometry: Union[RasterGeometry, Point, Polygon, Tuple[float, float, float, float]]):
         pass
 
+    @abstractmethod
+    def window(
+            self,
+            geometry: Union[SpatialGeometry, Tuple[float, float, float, float]],
+            buffer: int = None) -> Window:
+        """
+        Returns a rasterio.windows.Window covering the target geometry.
+
+        Args:
+            geometry: The geometry to create a window for
+            buffer: Optional buffer in pixels to add around the geometry
+
+        Returns:
+            Window: A rasterio Window object covering the geometry
+        """
+        pass
+
+    @abstractmethod
+    def subset(self, target: Union[Window, Point, Polygon, BBox, RasterGeometry]) -> RasterGeometry:
+        """
+        Subset the raster geometry using a Window or other geometry.
+
+        Args:
+            target: Window object or geometry to subset with
+
+        Returns:
+            RasterGeometry: A new raster geometry object representing the subset
+        """
+        pass
+
     @property
     @abstractmethod
     def x(self) -> np.ndarray: