pickle free centroids hdf

CHANGELOG.md

@@ -210,6 +210,7 @@ CLIMADA tutorials. [#872](https://github.com/CLIMADA-project/climada_python/pull
 - `Impact.write_hdf5` now throws an error if `event_name` is does not contain strings exclusively [#894](https://github.com/CLIMADA-project/climada_python/pull/894)
 - Split `climada.hazard.trop_cyclone` module into smaller submodules without affecting module usage [#911](https://github.com/CLIMADA-project/climada_python/pull/911)
 - `yearly_steps` parameter of `TropCyclone.apply_climate_scenario_knu` has been made explicit [#991](https://github.com/CLIMADA-project/climada_python/pull/991)
+- `Hazard.write_hdf5` writes centroids as x,y columns (or as wkb in case of polygons) at a compression level of 9, not as pickled `Shapely` objects anymore, which reduces the size of the files significantly.
 
 ### Fixed
 

climada/hazard/centroids/centr.py

@@ -177,10 +177,14 @@
 
         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 @@
         -------
         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."
@@ -906,23 +910,36 @@
             (path and) file name to write to.
         """
         LOGGER.info("Writing %s", file_name)
-        store = pd.HDFStore(file_name, mode=mode)
-        pandas_df = pd.DataFrame(self.gdf)
-        for col in pandas_df.columns:
-            if str(pandas_df[col].dtype) == "geometry":
-                pandas_df[col] = np.asarray(self.gdf[col])
-
-        # Avoid pandas PerformanceWarning when writing HDF5 data
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore", category=pd.errors.PerformanceWarning)
-            # Write dataframe
-            store.put("centroids", pandas_df)
-
-        store.get_storer("centroids").attrs.metadata = {
-            "crs": CRS.from_user_input(self.crs).to_wkt()
-        }
-
-        store.close()
+        xycols = []
+        wkbcols = []
+        store = pd.HDFStore(file_name, mode=mode, complevel=9)
+        try:
+            pandas_df = pd.DataFrame(self.gdf)
+            for col in pandas_df.columns:
+                if str(pandas_df[col].dtype) == "geometry":
+                    if (self.gdf[col].geom_type == "Point").all():
+                        pandas_df[col + ".x"] = self.gdf[col].x
+                        pandas_df[col + ".y"] = self.gdf[col].y
+                        pandas_df.drop(columns=[col], inplace=True)
+                        xycols.append(col)
+                    else:
+                        pandas_df[col] = self.gdf[col].to_wkb()
+                        wkbcols.append(col)
+
+            # Avoid pandas PerformanceWarning when writing HDF5 data
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore", category=pd.errors.PerformanceWarning)
+                # Write dataframe
+                store.put("centroids", pandas_df)
+
+            centroids_metadata = {"crs": CRS.from_user_input(self.crs).to_wkt()}
+            if xycols:
+                centroids_metadata["xy_columns"] = xycols
+            if wkbcols:
+                centroids_metadata["wkb_columns"] = wkbcols
+            store.get_storer("centroids").attrs.metadata = centroids_metadata
+        finally:
+            store.close()
 
     @classmethod
     def from_hdf5(cls, file_name):
@@ -950,7 +967,21 @@
                 # in previous versions of CLIMADA and/or geopandas,
                 # the CRS was stored in '_crs'/'crs'
                 crs = metadata.get("crs")
-                gdf = gpd.GeoDataFrame(store["centroids"], crs=crs)
+                gdf = gpd.GeoDataFrame(store["centroids"])
+                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", inplace=True)
+
         except TypeError:
             with h5py.File(file_name, "r") as data:
                 gdf = cls._gdf_from_legacy_hdf5(data.get("centroids"))

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 = CRS.from_user_input(ALT_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", crs=DEF_CRS)
+            )
+        )
+        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"