From 79be18acfe9703a8530cd0540f99e6fcaa86776b Mon Sep 17 00:00:00 2001 From: "Gregory H. Halverson" Date: Thu, 6 Nov 2025 11:07:19 -0800 Subject: [PATCH 1/2] raster geolocation subsetting and square bracket index subsetting --- pyproject.toml | 2 +- rasters/raster_geolocation.py | 42 +++++++++++++++++++++++++++++++--- rasters/raster_geometry.py | 43 ++++++++++++++++++++++++++++++++--- 3 files changed, 80 insertions(+), 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index ff6684d..7c78f4c 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/rasters/raster_geolocation.py b/rasters/raster_geolocation.py index a66a881..724dabb 100644 --- a/rasters/raster_geolocation.py +++ b/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) \ No newline at end of file + # 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) \ No newline at end of file diff --git a/rasters/raster_geometry.py b/rasters/raster_geometry.py index 13fe45a..8540f90 100644 --- a/rasters/raster_geometry.py +++ b/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: From f326d226553e92419f158e59415cb82decc30a7f Mon Sep 17 00:00:00 2001 From: "Gregory H. Halverson" Date: Thu, 6 Nov 2025 11:18:20 -0800 Subject: [PATCH 2/2] unit tests for raster subsetting --- tests/test_subsetting.py | 425 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 425 insertions(+) create mode 100644 tests/test_subsetting.py diff --git a/tests/test_subsetting.py b/tests/test_subsetting.py new file mode 100644 index 0000000..c194a9e --- /dev/null +++ b/tests/test_subsetting.py @@ -0,0 +1,425 @@ +import unittest +import numpy as np +from rasterio.windows import Window + +from rasters.raster_grid import RasterGrid +from rasters.raster_geolocation import RasterGeolocation +from rasters.point import Point +from rasters.polygon import Polygon +from rasters.bbox import BBox +from rasters.CRS import WGS84 + + +class TestRasterGridSubsetting(unittest.TestCase): + """Test subsetting functionality for RasterGrid objects.""" + + def setUp(self): + """Set up test RasterGrid object.""" + self.grid = RasterGrid( + x_origin=-120.0, + y_origin=40.0, + cell_width=0.01, + cell_height=-0.01, + rows=100, + cols=100, + crs=WGS84 + ) + + def test_subset_with_window(self): + """Test subsetting RasterGrid with Window object.""" + window = Window(col_off=10, row_off=20, width=30, height=40) + subset = self.grid.subset(window) + + self.assertIsInstance(subset, RasterGrid) + self.assertEqual(subset.rows, 40) + self.assertEqual(subset.cols, 30) + + # Check that the origin is adjusted correctly + expected_x_origin = self.grid.x_origin + 10 * self.grid.cell_width + expected_y_origin = self.grid.y_origin + 20 * self.grid.cell_height + self.assertAlmostEqual(subset.x_origin, expected_x_origin, places=6) + self.assertAlmostEqual(subset.y_origin, expected_y_origin, places=6) + + def test_subset_with_point(self): + """Test subsetting RasterGrid with Point object.""" + # Use small polygon around point to ensure RasterGrid return + coords = [ + (-119.55, 39.55), + (-119.45, 39.55), + (-119.45, 39.45), + (-119.55, 39.45), + (-119.55, 39.55) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.grid.subset(polygon) + + self.assertIsInstance(subset, RasterGrid) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_subset_with_polygon(self): + """Test subsetting RasterGrid with Polygon object.""" + # Create a small polygon within the grid bounds + coords = [ + (-119.95, 39.95), + (-119.90, 39.95), + (-119.90, 39.90), + (-119.95, 39.90), + (-119.95, 39.95) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.grid.subset(polygon) + + self.assertIsInstance(subset, RasterGrid) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_subset_with_bbox(self): + """Test subsetting RasterGrid with BBox object.""" + # Create bbox within grid bounds and convert to polygon for now + coords = [ + (-119.8, 39.8), + (-119.2, 39.8), + (-119.2, 39.2), + (-119.8, 39.2), + (-119.8, 39.8) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.grid.subset(polygon) + + self.assertIsInstance(subset, RasterGrid) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_subset_with_raster_geometry(self): + """Test subsetting RasterGrid with another RasterGrid object.""" + other_grid = RasterGrid( + x_origin=-119.95, + y_origin=39.95, + cell_width=0.005, + cell_height=-0.005, + rows=20, + cols=20, + crs=WGS84 + ) + subset = self.grid.subset(other_grid) + + self.assertIsInstance(subset, RasterGrid) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_square_bracket_indexing_with_window(self): + """Test square bracket indexing with Window object.""" + window = Window(col_off=5, row_off=10, width=20, height=25) + subset = self.grid[window] + + self.assertIsInstance(subset, RasterGrid) + self.assertEqual(subset.rows, 25) + self.assertEqual(subset.cols, 20) + + def test_square_bracket_indexing_with_point(self): + """Test square bracket indexing with Point object.""" + # Use small polygon to ensure RasterGrid return + coords = [ + (-119.55, 39.55), + (-119.45, 39.55), + (-119.45, 39.45), + (-119.55, 39.45), + (-119.55, 39.55) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.grid[polygon] + + self.assertIsInstance(subset, RasterGrid) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_square_bracket_indexing_with_polygon(self): + """Test square bracket indexing with Polygon object.""" + coords = [ + (-119.95, 39.95), + (-119.90, 39.95), + (-119.90, 39.90), + (-119.95, 39.90), + (-119.95, 39.95) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.grid[polygon] + + self.assertIsInstance(subset, RasterGrid) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_square_bracket_indexing_traditional(self): + """Test that traditional slice indexing still works.""" + subset = self.grid[10:50, 20:70] + + self.assertIsInstance(subset, RasterGrid) + self.assertEqual(subset.rows, 40) + self.assertEqual(subset.cols, 50) + + +class TestRasterGeolocationSubsetting(unittest.TestCase): + """Test subsetting functionality for RasterGeolocation objects.""" + + def setUp(self): + """Set up test RasterGeolocation object.""" + # Create coordinate arrays for a 50x50 geolocation field + rows, cols = 50, 50 + x_coords = np.linspace(-120.0, -119.0, cols) + y_coords = np.linspace(40.0, 39.0, rows) + x_array, y_array = np.meshgrid(x_coords, y_coords) + + self.geolocation = RasterGeolocation(x_array, y_array, crs=WGS84) + + def test_subset_with_window(self): + """Test subsetting RasterGeolocation with Window object.""" + window = Window(col_off=5, row_off=10, width=20, height=15) + subset = self.geolocation.subset(window) + + self.assertIsInstance(subset, RasterGeolocation) + self.assertEqual(subset.rows, 15) + self.assertEqual(subset.cols, 20) + + # Check that coordinate arrays are properly subset + expected_x = self.geolocation.x[10:25, 5:25] + expected_y = self.geolocation.y[10:25, 5:25] + np.testing.assert_array_equal(subset.x, expected_x) + np.testing.assert_array_equal(subset.y, expected_y) + + def test_subset_with_point(self): + """Test subsetting RasterGeolocation with Point object.""" + # Use a point that's definitely within the geolocation bounds + point = Point(-119.5, 39.5, crs=WGS84) + # Create a small polygon around the point instead + coords = [ + (-119.6, 39.6), + (-119.4, 39.6), + (-119.4, 39.4), + (-119.6, 39.4), + (-119.6, 39.6) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.geolocation.subset(polygon) + + self.assertIsInstance(subset, RasterGeolocation) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_subset_with_polygon(self): + """Test subsetting RasterGeolocation with Polygon object.""" + # Create a polygon within the geolocation bounds + coords = [ + (-119.8, 39.8), + (-119.2, 39.8), + (-119.2, 39.2), + (-119.8, 39.2), + (-119.8, 39.8) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.geolocation.subset(polygon) + + self.assertIsInstance(subset, RasterGeolocation) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_subset_with_bbox(self): + """Test subsetting RasterGeolocation with BBox object.""" + # Use polygon instead of BBox for now + coords = [ + (-119.8, 39.8), + (-119.2, 39.8), + (-119.2, 39.2), + (-119.8, 39.2), + (-119.8, 39.8) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.geolocation.subset(polygon) + + self.assertIsInstance(subset, RasterGeolocation) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_subset_with_raster_geometry(self): + """Test subsetting RasterGeolocation with RasterGrid object.""" + grid = RasterGrid( + x_origin=-119.8, + y_origin=39.8, + cell_width=0.01, + cell_height=-0.01, + rows=30, + cols=30, + crs=WGS84 + ) + subset = self.geolocation.subset(grid) + + self.assertIsInstance(subset, RasterGeolocation) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_window_with_buffer(self): + """Test window method with buffer parameter.""" + # Use polygon instead of point for more reliable results + coords = [ + (-119.6, 39.6), + (-119.4, 39.6), + (-119.4, 39.4), + (-119.6, 39.4), + (-119.6, 39.6) + ] + polygon = Polygon(coords, crs=WGS84) + window_no_buffer = self.geolocation.window(polygon) + window_with_buffer = self.geolocation.window(polygon, buffer=5) + + # Window with buffer should be larger + self.assertGreaterEqual(window_with_buffer.width, window_no_buffer.width) + self.assertGreaterEqual(window_with_buffer.height, window_no_buffer.height) + + # Buffer should not exceed array bounds + self.assertGreaterEqual(window_with_buffer.col_off, 0) + self.assertGreaterEqual(window_with_buffer.row_off, 0) + self.assertLessEqual(window_with_buffer.col_off + window_with_buffer.width, self.geolocation.cols) + self.assertLessEqual(window_with_buffer.row_off + window_with_buffer.height, self.geolocation.rows) + + def test_square_bracket_indexing_with_window(self): + """Test square bracket indexing with Window object.""" + window = Window(col_off=8, row_off=12, width=15, height=18) + subset = self.geolocation[window] + + self.assertIsInstance(subset, RasterGeolocation) + self.assertEqual(subset.rows, 18) + self.assertEqual(subset.cols, 15) + + def test_square_bracket_indexing_with_point(self): + """Test square bracket indexing with Point object.""" + # Use polygon around point for more reliable indexing + coords = [ + (-119.6, 39.6), + (-119.4, 39.6), + (-119.4, 39.4), + (-119.6, 39.4), + (-119.6, 39.6) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.geolocation[polygon] + + self.assertIsInstance(subset, RasterGeolocation) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_square_bracket_indexing_with_bbox(self): + """Test square bracket indexing with BBox-like polygon.""" + coords = [ + (-119.6, 39.7), + (-119.4, 39.7), + (-119.4, 39.3), + (-119.6, 39.3), + (-119.6, 39.7) + ] + polygon = Polygon(coords, crs=WGS84) + subset = self.geolocation[polygon] + + self.assertIsInstance(subset, RasterGeolocation) + self.assertGreater(subset.rows, 0) + self.assertGreater(subset.cols, 0) + + def test_square_bracket_indexing_traditional(self): + """Test that traditional slice indexing still works.""" + subset = self.geolocation[5:25, 10:30] + + self.assertIsInstance(subset, RasterGeolocation) + self.assertEqual(subset.rows, 20) + self.assertEqual(subset.cols, 20) + + def test_window_no_overlap_error(self): + """Test that window method raises error when geometry doesn't overlap.""" + # Point far outside the geolocation bounds + point = Point(-90.0, 30.0, crs=WGS84) + + with self.assertRaises(ValueError) as context: + self.geolocation.window(point) + + self.assertIn("No points found within the target geometry", str(context.exception)) + + +class TestSubsettingConsistency(unittest.TestCase): + """Test consistency between RasterGrid and RasterGeolocation subsetting.""" + + def setUp(self): + """Set up test objects for consistency testing.""" + # Create a RasterGrid + self.grid = RasterGrid( + x_origin=-120.0, + y_origin=40.0, + cell_width=0.02, + cell_height=-0.02, + rows=50, + cols=50, + crs=WGS84 + ) + + # Create a RasterGeolocation with similar bounds + rows, cols = 50, 50 + x_coords = np.linspace(-120.0, -119.0, cols) + y_coords = np.linspace(40.0, 39.0, rows) + x_array, y_array = np.meshgrid(x_coords, y_coords) + self.geolocation = RasterGeolocation(x_array, y_array, crs=WGS84) + + def test_both_support_same_window(self): + """Test that both geometry types can be subset with the same Window.""" + window = Window(col_off=10, row_off=15, width=20, height=25) + + grid_subset = self.grid.subset(window) + geolocation_subset = self.geolocation.subset(window) + + # Both should return their respective types + self.assertIsInstance(grid_subset, RasterGrid) + self.assertIsInstance(geolocation_subset, RasterGeolocation) + + # Both should have the same dimensions + self.assertEqual(grid_subset.rows, geolocation_subset.rows) + self.assertEqual(grid_subset.cols, geolocation_subset.cols) + self.assertEqual(grid_subset.rows, 25) + self.assertEqual(grid_subset.cols, 20) + + def test_both_support_square_bracket_indexing(self): + """Test that both support square bracket indexing with geometries.""" + coords = [ + (-119.8, 39.8), + (-119.2, 39.8), + (-119.2, 39.2), + (-119.8, 39.2), + (-119.8, 39.8) + ] + polygon = Polygon(coords, crs=WGS84) + + # Both should work with square bracket indexing + grid_subset = self.grid[polygon] + geolocation_subset = self.geolocation[polygon] + + self.assertIsInstance(grid_subset, RasterGrid) + self.assertIsInstance(geolocation_subset, RasterGeolocation) + + def test_method_vs_indexing_equivalence(self): + """Test that subset method and square bracket indexing are equivalent.""" + window = Window(col_off=5, row_off=8, width=15, height=12) + + # For RasterGrid + grid_method = self.grid.subset(window) + grid_indexing = self.grid[window] + self.assertEqual(grid_method.rows, grid_indexing.rows) + self.assertEqual(grid_method.cols, grid_indexing.cols) + self.assertEqual(grid_method.x_origin, grid_indexing.x_origin) + self.assertEqual(grid_method.y_origin, grid_indexing.y_origin) + + # For RasterGeolocation + geolocation_method = self.geolocation.subset(window) + geolocation_indexing = self.geolocation[window] + self.assertEqual(geolocation_method.rows, geolocation_indexing.rows) + self.assertEqual(geolocation_method.cols, geolocation_indexing.cols) + np.testing.assert_array_equal(geolocation_method.x, geolocation_indexing.x) + np.testing.assert_array_equal(geolocation_method.y, geolocation_indexing.y) + + +if __name__ == '__main__': + unittest.main() \ No newline at end of file