Fix compute_indicatormatrix and compute_intersectionmatrix for non-default index

RELEASE_NOTES.rst

@@ -23,6 +23,10 @@ Upcoming Release
 
 **Bug fixes**
 
+* Fix ``compute_indicatormatrix`` and ``compute_intersectionmatrix`` failing with
+  ``KeyError`` or indexing mismatched shapes when origin or destination geometries
+  are passed as pandas or geopandas Series with a non-default index
+  (https://github.com/PyPSA/atlite/pull/534).
 * Fix ``ExclusionContainer.add_raster`` ignoring ``codes=0``. The code ``0`` was
   treated like a missing ``codes`` argument, so all non-zero cells were excluded
   instead of the cells with value ``0``. Numpy arrays are now accepted as

atlite/gis.py

@@ -185,12 +185,8 @@ def compute_indicatormatrix(
     orig = orig.geometry if isinstance(orig, gpd.GeoDataFrame) else orig
     dest = dest.geometry if isinstance(dest, gpd.GeoDataFrame) else dest
     dest = reproject_shapes(dest, dest_crs, orig_crs)
-    orig_list: list[Any] | pd.Series = (
-        list(orig) if not isinstance(orig, pd.Series) else orig
-    )
-    dest_list: list[Any] | pd.Series = (
-        list(dest) if not isinstance(dest, pd.Series) else dest
-    )
+    orig_list: list[Any] = list(orig)
+    dest_list: list[Any] = list(dest)
     indicator = sp.sparse.lil_matrix((len(dest_list), len(orig_list)), dtype=float)
     tree = STRtree(orig_list)
     idx = {hash(o.wkt): i for i, o in enumerate(orig_list)}
@@ -199,9 +195,11 @@ def compute_indicatormatrix(
         for o in tree.query(d):
             # STRtree query returns a list of indices for shapely >= v2.0
             if isinstance(o, (int | np.integer)):
-                o = orig_list[o]
-            if o.intersects(d):
+                j = int(o)
+                o = orig_list[j]
+            else:
                 j = idx[hash(o.wkt)]
+            if o.intersects(d):
                 area = d.intersection(o).area
                 indicator[i, j] = area / o.area
 
@@ -242,12 +240,8 @@ def compute_intersectionmatrix(
     orig = orig.geometry if isinstance(orig, gpd.GeoDataFrame) else orig
     dest = dest.geometry if isinstance(dest, gpd.GeoDataFrame) else dest
     dest = reproject_shapes(dest, dest_crs, orig_crs)
-    orig_list: list[Any] | pd.Series = (
-        list(orig) if not isinstance(orig, pd.Series) else orig
-    )
-    dest_list: list[Any] | pd.Series = (
-        list(dest) if not isinstance(dest, pd.Series) else dest
-    )
+    orig_list: list[Any] = list(orig)
+    dest_list: list[Any] = list(dest)
     intersection = sp.sparse.lil_matrix((len(dest_list), len(orig_list)), dtype=float)
     tree = STRtree(orig_list)
     idx = {hash(o.wkt): i for i, o in enumerate(orig_list)}
@@ -256,8 +250,10 @@ def compute_intersectionmatrix(
         for o in tree.query(d):
             # STRtree query returns a list of indices for shapely >= v2.0
             if isinstance(o, (int | np.integer)):
-                o = orig_list[o]
-            j = idx[hash(o.wkt)]
+                j = int(o)
+                o = orig_list[j]
+            else:
+                j = idx[hash(o.wkt)]
             intersection[i, j] = o.intersects(d)
 
     return intersection

test/test_gis.py

@@ -29,6 +29,8 @@
 from atlite import Cutout
 from atlite.gis import (
     ExclusionContainer,
+    compute_indicatormatrix,
+    compute_intersectionmatrix,
     pad_extent,
     padded_transform_and_shape,
     regrid,
@@ -721,3 +723,31 @@ def test_invalid_geometry_rejected_at_computation(self, shapes, raster):
         })
         with pytest.raises(ValueError, match="Unsupported buffer geometry"):
             shape_availability(shapes, excluder)
+
+
+def test_compute_indicatormatrix_and_intersectionmatrix_non_default_index():
+    poly1 = box(0, 0, 1, 1)
+    poly2 = box(1, 0, 2, 1)
+    dest_poly = box(0, 0, 1.5, 1)
+
+    orig_default = gpd.GeoSeries([poly1, poly2])
+    dest_default = gpd.GeoSeries([dest_poly])
+    expected_indicator = compute_indicatormatrix(orig_default, dest_default)
+    expected_intersection = compute_intersectionmatrix(orig_default, dest_default)
+
+    orig_custom = gpd.GeoSeries([poly1, poly2], index=[10, 20])
+    dest_custom = gpd.GeoSeries([dest_poly], index=["target"])
+    res_indicator = compute_indicatormatrix(orig_custom, dest_custom)
+    res_intersection = compute_intersectionmatrix(orig_custom, dest_custom)
+
+    assert (res_indicator.toarray() == expected_indicator.toarray()).all()
+    assert (res_intersection.toarray() == expected_intersection.toarray()).all()
+
+    orig_reversed = gpd.GeoSeries([poly2, poly1], index=[1, 0])
+    res_indicator_rev = compute_indicatormatrix(orig_reversed, dest_default)
+    res_intersection_rev = compute_intersectionmatrix(orig_reversed, dest_default)
+
+    assert isclose(res_indicator_rev[0, 0], 0.5)
+    assert isclose(res_indicator_rev[0, 1], 1.0)
+    assert res_intersection_rev[0, 0] == 1.0
+    assert res_intersection_rev[0, 1] == 1.0