pickle free centroids hdf (#1076)

* hazard.io: avoid pickling geometries and compress hdf5 files

* 'changelog'

* fix column drop

* add unit tests with multiple wkb and xy columns

* avoid crs resetting

* explicitly ask for Points

Co-authored-by: Lukas Riedel <[email protected]>

* fix typo

* fix point condition

* add comment about filtering geometry columns

* remove ducplicated methods and obsolete lines

---------

Co-authored-by: Lukas Riedel <[email protected]>

Author: emanuel-schmid

CHANGELOG.md

@@ -221,6 +221,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 @@ 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."
@@ -906,23 +910,40 @@ def write_hdf5(self, file_name, mode="w"):
             (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)
+            # we replace all columns of type geometry
+            # - with according x and y columns if they are strictly `Point`s
+            # - with wkb values if they have other shapes
+            # this saves a lot of time and disk space
+            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 +971,21 @@ def from_hdf5(cls, file_name):
                 # 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

@@ -556,8 +556,6 @@ def test_write_read_excel(self):
     def test_from_raster_file(self):
         """Test from_raster_file"""
         width, height = 50, 60
-        o_lat, o_lon = (10.42822096697894, -69.33714959699981)
-        res_lat, res_lon = (-0.009000000000000341, 0.009000000000000341)
 
         centr_ras = Centroids.from_raster_file(
             HAZ_DEMO_FL, window=Window(0, 0, width, height)
@@ -725,6 +723,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.assertEqual(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"
@@ -901,41 +944,6 @@ def test_union(self):
             cent.on_land, np.concatenate([on_land, on_land2, [None, None]])
         )
 
-    def test_select_pass(self):
-        """Test Centroids.select method"""
-        region_id = np.zeros(VEC_LAT.size)
-        region_id[[2, 4]] = 10
-        centr = Centroids(lat=VEC_LAT, lon=VEC_LON, region_id=region_id)
-
-        fil_centr = centr.select(reg_id=10)
-        self.assertEqual(fil_centr.size, 2)
-        self.assertEqual(fil_centr.lat[0], VEC_LAT[2])
-        self.assertEqual(fil_centr.lat[1], VEC_LAT[4])
-        self.assertEqual(fil_centr.lon[0], VEC_LON[2])
-        self.assertEqual(fil_centr.lon[1], VEC_LON[4])
-        self.assertTrue(np.array_equal(fil_centr.region_id, np.ones(2) * 10))
-
-    def test_select_extent_pass(self):
-        """Test select extent"""
-        centr = Centroids(
-            lat=np.array([-5, -3, 0, 3, 5]),
-            lon=np.array([-180, -175, -170, 170, 175]),
-            region_id=np.zeros(5),
-        )
-        ext_centr = centr.select(extent=[-175, -170, -5, 5])
-        np.testing.assert_array_equal(ext_centr.lon, np.array([-175, -170]))
-        np.testing.assert_array_equal(ext_centr.lat, np.array([-3, 0]))
-
-        # Cross antimeridian, version 1
-        ext_centr = centr.select(extent=[170, -175, -5, 5])
-        np.testing.assert_array_equal(ext_centr.lon, np.array([-180, -175, 170, 175]))
-        np.testing.assert_array_equal(ext_centr.lat, np.array([-5, -3, 3, 5]))
-
-        # Cross antimeridian, version 2
-        ext_centr = centr.select(extent=[170, 185, -5, 5])
-        np.testing.assert_array_equal(ext_centr.lon, np.array([-180, -175, 170, 175]))
-        np.testing.assert_array_equal(ext_centr.lat, np.array([-5, -3, 3, 5]))
-
     def test_get_meta(self):
         """
         Test that the `get_meta` method correctly generates metadata
@@ -968,7 +976,7 @@ def test_get_meta(self):
         self.assertTrue(u_coord.equal_crs(meta["crs"], expected_meta["crs"]))
         self.assertTrue(meta["transform"].almost_equals(expected_meta["transform"]))
 
-    def test_get_closest_point(self):
+    def test_get_closest_point_1(self):
         """Test get_closest_point"""
         for n, (lat, lon) in enumerate(LATLON):
             x, y, idx = self.centr.get_closest_point(lon * 0.99, lat * 1.01)
@@ -978,7 +986,7 @@ def test_get_closest_point(self):
             self.assertEqual(self.centr.lon[n], x)
             self.assertEqual(self.centr.lat[n], y)
 
-    def test_get_closest_point(self):
+    def test_get_closest_point_2(self):
         """Test get_closest_point"""
         for y_sign in [1, -1]:
             meta = {