Fix indexing bug; add tests (#21)

Author: dcherian

src/rasterix/raster_index.py

@@ -209,10 +209,13 @@ def sel(self, labels, method=None, tolerance=None):
         label = labels[coord_name]
 
         if isinstance(label, slice):
+            if label.start is None:
+                label = slice(0, label.stop, label.step)
             if label.step is None:
                 # continuous interval slice indexing (preserves the index)
                 pos = self.transform.reverse({coord_name: np.array([label.start, label.stop])})
-                pos = np.round(pos[self.dim]).astype("int")
+                # np.round rounds to even, this way we round upwards
+                pos = np.floor(pos[self.dim] + 0.5).astype("int")
                 new_start = max(pos[0], 0)
                 new_stop = min(pos[1], self.axis_transform.size)
                 return IndexSelResult({self.dim: slice(new_start, new_stop)})
@@ -368,7 +371,7 @@ def sel(self, labels: dict[Any, Any], method=None, tolerance=None) -> IndexSelRe
         for coord_names, index in self._wrapped_indexes.items():
             if not isinstance(coord_names, tuple):
                 coord_names = (coord_names,)
-            index_labels = {k: v for k, v in labels if k in coord_names}
+            index_labels = {k: v for k, v in labels.items() if k in coord_names}
             if index_labels:
                 results.append(index.sel(index_labels, method=method, tolerance=tolerance))
 
@@ -403,3 +406,14 @@ def __repr__(self):
             items += [repr(coord_names) + ":", textwrap.indent(repr(index), "    ")]
 
         return "RasterIndex\n" + "\n".join(items)
+
+    def transform(self) -> Affine:
+        """Returns Affine transform for top-left corners."""
+        if len(self._wrapped_indexes) > 1:
+            x = self._wrapped_indexes["x"].axis_transform.affine
+            y = self._wrapped_indexes["y"].axis_transform.affine
+            aff = Affine(x.a, x.b, x.c, y.d, y.e, y.f)
+        else:
+            index = next(iter(self._wrapped_indexes.values()))
+            aff = index.affine
+        return aff * Affine.translation(-0.5, -0.5)

tests/test_raster_index.py

@@ -1,21 +1,44 @@
+import numpy as np
+import rioxarray  # noqa
 import xarray as xr
+from affine import Affine
 
 from rasterix import RasterIndex
 
 
-def test_rectilinear():
-    source = "/vsicurl/https://noaadata.apps.nsidc.org/NOAA/G02135/south/daily/geotiff/2024/01_Jan/S_20240101_concentration_v3.0.tif"
-    da_no_raster_index = xr.open_dataarray(source, engine="rasterio")
-    x_dim = da_no_raster_index.rio.x_dim
-    y_dim = da_no_raster_index.rio.y_dim
+def set_raster_index(obj):
+    x_dim = obj.rio.x_dim
+    y_dim = obj.rio.y_dim
 
     index = RasterIndex.from_transform(
-        da_no_raster_index.rio.transform(),
-        da_no_raster_index.sizes[x_dim],
-        da_no_raster_index.sizes[y_dim],
-        x_dim=x_dim,
-        y_dim=y_dim,
+        obj.rio.transform(), obj.sizes[x_dim], obj.sizes[y_dim], x_dim=x_dim, y_dim=y_dim
     )
     coords = xr.Coordinates.from_xindex(index)
-    da_raster_index = da_no_raster_index.assign_coords(coords)
+    return obj.assign_coords(coords)
+
+
+def test_rectilinear():
+    source = "/vsicurl/https://noaadata.apps.nsidc.org/NOAA/G02135/south/daily/geotiff/2024/01_Jan/S_20240101_concentration_v3.0.tif"
+    da_no_raster_index = xr.open_dataarray(source, engine="rasterio")
+    da_raster_index = set_raster_index(da_no_raster_index)
     assert da_raster_index.equals(da_no_raster_index)
+
+
+# TODO: parameterize over
+# 1. y points up;
+# 2. y points down
+def test_sel_slice():
+    ds = xr.Dataset({"foo": (("y", "x"), np.ones((10, 12)))})
+    transform = Affine.identity()
+    ds = ds.rio.write_transform(transform)
+    ds = set_raster_index(ds)
+
+    assert ds.xindexes["x"].transform() == transform
+
+    actual = ds.sel(x=slice(4), y=slice(3, 5))
+    assert isinstance(actual.xindexes["x"], RasterIndex)
+    assert isinstance(actual.xindexes["y"], RasterIndex)
+    actual_transform = actual.xindexes["x"].transform()
+
+    assert actual_transform == actual.rio.transform()
+    assert actual_transform == (transform * Affine.translation(0, 3))