Update centroids.select to work properly with mask and indices

CLIMADA-project · Nov 21, 2023 · b22677f · b22677f
1 parent 9ce8e56
commit b22677f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 8 deletions.
diff --git a/climada/entity/exposures/test/test_base.py b/climada/entity/exposures/test/test_base.py
@@ -182,8 +182,8 @@ def test_assign_large_hazard_subset_pass(self):
         haz = Hazard.from_raster([HAZ_DEMO_FL], haz_type='FL')
         exp.assign_centroids(haz)
         assigned_centroids = haz.centroids.select(sel_cen=exp.gdf[INDICATOR_CENTR + 'FL'].values)
-        np.testing.assert_array_equal(assigned_centroids.lat, exp.gdf.latitude)
-        np.testing.assert_array_equal(assigned_centroids.lon, exp.gdf.longitude)
+        np.testing.assert_array_equal(np.unique(assigned_centroids.lat), np.unique(exp.gdf.latitude))
+        np.testing.assert_array_equal(np.unique(assigned_centroids.lon), np.unique(exp.gdf.longitude))
 
     def test_affected_total_value(self):
         haz_type = "RF"

diff --git a/climada/hazard/centroids/centr.py b/climada/hazard/centroids/centr.py
@@ -444,24 +444,30 @@ def select(self, reg_id=None, extent=None, sel_cen=None):
             [lon_max, 180] + [-180, lon_min]
             Borders are inclusive.
         sel_cen : np.array
-            1-dim mask, overrides reg_id and extent
+            1-dim mask or 1-dim centroids indices, complements reg_id and extent
 
         Returns
         -------
         cen : Centroids
             Sub-selection of this object
         """
-        if sel_cen is None:
-            sel_cen = self.select_mask(reg_id=reg_id, extent=extent)
+        if sel_cen.dtype.kind == 'i':  #is integer
+            sel_cen_bool = np.zeros(self.size, dtype=bool)
+            sel_cen_bool[np.unique(sel_cen)] = True
+        else:
+            sel_cen_bool = sel_cen
+        sel_cen_mask = self.select_mask(sel_cen=sel_cen_bool, reg_id=reg_id, extent=extent)
+        return Centroids.from_geodataframe(self.gdf[sel_cen_mask])
 
-        return Centroids.from_geodataframe(self.gdf[sel_cen])
 
-    def select_mask(self, reg_id=None, extent=None):
+    def select_mask(self, sel_cen, reg_id=None, extent=None):
         """
         Make mask of selected centroids
 
         Parameters
         ----------
+        sel_cen: np.array(bool)
+            boolean array mask for centroids
         reg_id : int
             region to filter according to region_id values
         extent : tuple
@@ -476,7 +482,6 @@ def select_mask(self, reg_id=None, extent=None):
             1d mask of selected centroids
 
         """
-        sel_cen = np.ones(self.size, dtype=bool)
         if reg_id is not None:
             sel_cen &= np.isin(self.region_id, reg_id)
         if extent is not None: