add unit tests with multiple wkb and xy columns

Author: emanuel-schmid

climada/hazard/centroids/centr.py

@@ -177,10 +177,14 @@ def __eq__(self, other):
 
         try:
             pd.testing.assert_frame_equal(self.gdf, other.gdf, check_like=True)
-            return True
         except AssertionError:
             return False
 
+        if not (self.gdf.geometry == other.gdf.geometry).all():
+            return False
+
+        return True
+
     def to_default_crs(self, inplace=True):
         """Project the current centroids to the default CRS (epsg4326)
 
@@ -483,11 +487,11 @@ def plot(self, *, axis=None, figsize=(9, 13), **kwargs):
         -------
         ax : cartopy.mpl.geoaxes.GeoAxes instance
         """
-        if axis == None:
+        if axis is None:
             fig, axis = plt.subplots(
                 figsize=figsize, subplot_kw={"projection": ccrs.PlateCarree()}
             )
-        if type(axis) != cartopy.mpl.geoaxes.GeoAxes:
+        if type(axis) is not cartopy.mpl.geoaxes.GeoAxes:
             raise AttributeError(
                 f"The axis provided is of type: {type(axis)} "
                 "The function requires a cartopy.mpl.geoaxes.GeoAxes."
@@ -964,13 +968,19 @@ def from_hdf5(cls, file_name):
                 # the CRS was stored in '_crs'/'crs'
                 crs = metadata.get("crs")
                 gdf = gpd.GeoDataFrame(store["centroids"])
-                for xycol in metadata.get("xy_columns", []):
-                    gdf[xycol] = gpd.points_from_xy(
-                        x=gdf[xycol + ".x"], y=gdf[xycol + ".y"], crs=crs
-                    )
-                    gdf.drop(columns=[xycol + ".x", xycol + ".y"], inplace=True)
-                for wkbcol in metadata.get("wkb_columns", []):
-                    gdf[wkbcol] = gpd.GeoSeries.from_wkb(gdf[wkbcol], crs=crs)
+                with warnings.catch_warnings():
+                    # setting a column named 'geometry' triggers a future warning
+                    # with geopandas 0.14
+                    warnings.simplefilter(action="ignore", category=FutureWarning)
+
+                    for xycol in metadata.get("xy_columns", []):
+                        gdf[xycol] = gpd.points_from_xy(
+                            x=gdf[xycol + ".x"], y=gdf[xycol + ".y"], crs=crs
+                        )
+                        gdf.drop(columns=[xycol + ".x", xycol + ".y"], inplace=True)
+                    for wkbcol in metadata.get("wkb_columns", []):
+                        gdf[wkbcol] = gpd.GeoSeries.from_wkb(gdf[wkbcol], crs=crs)
+                gdf.set_geometry("geometry", crs=crs, inplace=True)
 
         except TypeError:
             with h5py.File(file_name, "r") as data:

climada/hazard/centroids/test/test_centr.py

@@ -725,6 +725,51 @@ def test_read_write_hdf5(self):
         self.assertTrue(centroids_w == centroids_r)
         tmpfile.unlink()
 
+    def test_read_write_hdf5_with_additional_columns(self):
+        tmpfile = Path("test_write_hdf5.out.hdf5")
+        crs = DEF_CRS
+        centroids_w = Centroids(
+            lat=VEC_LAT,
+            lon=VEC_LON,
+            crs=crs,
+            region_id=REGION_ID,
+            on_land=ON_LAND,
+        )
+        centroids_w.gdf = (
+            centroids_w.gdf.join(
+                gpd.GeoDataFrame(
+                    {"more_points": [shapely.Point(i, i) for i in range(8)]}
+                ).set_geometry("more_points")
+            )
+            .join(
+                gpd.GeoDataFrame(
+                    {
+                        "some_shapes": [
+                            shapely.Point((2, 2)),
+                            shapely.Point((3, 3)),
+                            shapely.Polygon([(0, 0), (1, 1), (1, 0), (0, 0)]),
+                            shapely.LineString([(0, 1), (1, 0)]),
+                        ]
+                        * 2
+                    }
+                ).set_geometry("some_shapes")
+            )
+            .join(
+                gpd.GeoDataFrame(
+                    {
+                        "more_shapes": [
+                            shapely.LineString([(0, 1), (1, 2)]),
+                        ]
+                        * 8
+                    }
+                ).set_geometry("more_shapes")
+            )
+        )
+        centroids_w.write_hdf5(tmpfile)
+        centroids_r = Centroids.from_hdf5(tmpfile)
+        self.assertTrue(centroids_w == centroids_r)
+        tmpfile.unlink()
+
     def test_from_hdf5_nonexistent_file(self):
         """Test raising FileNotFoundError when creating Centroids object from a nonexistent HDF5 file"""
         file_name = "/path/to/nonexistentfile.h5"