Merge pull request #25 from GEMScienceTools/subset_source

added subsetting to core and improvements to H3 indexing

Author: cossatot

openquake/hme/core/core.py

@@ -38,6 +38,7 @@
     breakpoint,
     get_eq_df_mag_bins,
     get_mag_bins_from_cfg,
+    subset_source,
 )
 
 from openquake.hme.utils.results_processing import process_results
@@ -254,11 +255,11 @@ def load_ruptures_from_file(cfg: dict):
     logging.info("Reading ruptures from {}".format(rup_file))
     if os.path.exists(rup_file):
         rupture_gdf = read_rupture_file(
-            rup_file, h3_res=h3_res, parallel=parallel
+            rup_file, h3_res=h3_res, parallel=False
         )
 
     else:
-        logging.warn("Rupture file does not exist; reading SSM.")
+        logging.warning("Rupture file does not exist; reading SSM.")
         rupture_gdf = load_ruptures_from_ssm(cfg)
 
     return rupture_gdf
@@ -369,21 +370,23 @@ def load_inputs(cfg: dict) -> dict:
     logging.info("grouping ruptures by cell")
     cell_groups = rupture_gdf.groupby("cell_id")
 
-    logger.info("rupture_gdf shape: {}".format(rupture_gdf.shape))
+    logger.info(f"{len(rupture_gdf):_} ruptures in model")
     logger.debug(
         "rupture_gdf memory: {} GB".format(
             sum(rupture_gdf.memory_usage(index=True, deep=True)) * 1e-9
         )
     )
 
     if cfg["input"]["subset"]["file"] is not None:
-        # logger.info("   Subsetting bin_gdf")
-        # bin_gdf = subset_source(
-        #    bin_gdf,
-        #    subset_file=cfg["input"]["subset"]["file"],
-        #    buffer=cfg["input"]["subset"]["buffer"],
-        # )
-        logger.warning("CANNOT SUBSET SOURCE YET!!!")
+        logger.info("Subsetting rupture_gdf")
+        rupture_gdf = subset_source(
+            rupture_gdf,
+            cfg["input"]["subset"]["file"],
+            buffer=cfg["input"]["subset"].get("buffer"),
+            fid=cfg["input"]["subset"].get("fid"),
+            feature_number=cfg["input"]["subset"].get("feature_number"),
+        )
+        logger.info(f"{len(rupture_gdf):_} ruptures after subset")
 
     if len(rupture_gdf) != len(rupture_gdf.index.unique()):
         logging.warning(

openquake/hme/utils/io/io.py

@@ -20,7 +20,6 @@
 from ..simple_rupture import SimpleRupture, rup_to_dict
 from openquake.hme.utils.io.source_processing import (
     _get_h3_cell_for_rupture_df,
-    _get_h3_cell_for_rupture_df_parallel,
 )
 
 
@@ -116,6 +115,7 @@ def oq_rupture_to_json(
 def read_rupture_file(
     rupture_file, h3_res: int = 3, parallel=False
 ) -> pd.DataFrame:
+    # parallel not currently used but leaving for future possibilities
     rup_file_type = rupture_file.split(".")[-1]
 
     if rup_file_type == "hdf5":
@@ -127,10 +127,7 @@ def read_rupture_file(
     else:
         raise ValueError("Cannot read filetype {}".format(rup_file_type))
 
-    if parallel is False:
-        _get_h3_cell_for_rupture_df(rupture_df, h3_res)
-    else:
-        _get_h3_cell_for_rupture_df_parallel(rupture_df, h3_res)
+    _get_h3_cell_for_rupture_df(rupture_df, h3_res)
 
     return rupture_df
 

openquake/hme/utils/io/source_processing.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 from time import sleep
 from typing import Union, Optional, Tuple
 from multiprocessing import Pool
@@ -10,6 +11,12 @@
 from shapely.geometry import Point, Polygon
 from tqdm.autonotebook import tqdm
 
+warnings.filterwarnings(
+    "ignore",
+    message="Modules under `h3.unstable` are experimental, and may change at any time.",
+)
+
+from h3.unstable import vect as h3_vect  # name in 3.x
 
 from openquake.hazardlib.source.rupture import (
     NonParametricProbabilisticRupture,
@@ -24,7 +31,7 @@
     flatten_list,
     get_nonparametric_rupture_occurrence_rate,
     _get_class_name,
-    breakpoint
+    breakpoint,
 )
 
 
@@ -401,7 +408,8 @@ def rupture_dict_from_logic_tree_dict(
             logic_tree_dict.items()
         ):
             logging.info(
-                f"processing {branch_name} ({i+1}/{len(logic_tree_dict.keys())})"
+                f"processing {branch_name} ({i+1}/"
+                + f"{len(logic_tree_dict.keys())})"
             )
             rup_dict[branch_name] = rupture_list_from_source_list_parallel(
                 source_list,
@@ -417,7 +425,8 @@ def rupture_dict_from_logic_tree_dict(
             logic_tree_dict.items()
         ):
             logging.info(
-                f"processing {branch_name} ({i+1}/{len(logic_tree_dict.keys())})"
+                f"processing {branch_name} ({i+1}/"
+                + f"{len(logic_tree_dict.keys())})"
             )
             rup_dict[branch_name] = rupture_df_from_source_list(
                 source_list,
@@ -452,10 +461,11 @@ def rupture_dict_to_gdf(
         if df.occurrence_rate.min() <= 0.0:
             logging.info("trimming zero-rate ruptures")
             df_len = len(df)
-            df = df[df.occurrence_rate > 0.]
+            df = df[df.occurrence_rate > 0.0]
             df_trim_len = len(df)
             logging.info(
-                f"{df_trim_len} ruptures remaining, {df_len - df_trim_len} removed"
+                f"{df_trim_len} ruptures remaining, "
+                + f"{df_len - df_trim_len} removed"
             )
 
     if return_gdf:
@@ -472,35 +482,17 @@ def parse_geometry(row, x="longitude", y="latitude", z="depth"):
     return df
 
 
-def _get_h3_cell_for_rupture_df(rupture_df, h3_res):
-    logging.info("getting H3 cells")
-    cell_ids = list(
-        tqdm(
-            (
-                h3.geo_to_h3(row.latitude, row.longitude, h3_res)
-                for i, row in rupture_df.iterrows()
-            ),
-            total=len(rupture_df),
-        )
-    )
-    rupture_df["cell_id"] = cell_ids
-
-
 def _get_h3_cell(args):
     return h3.geo_to_h3(*args)
 
 
-def _get_h3_cell_for_rupture_df_parallel(rupture_df, h3_res):
-    logging.info("getting H3 cells in parallel")
-
-    lons = rupture_df.longitude.values
-    lats = rupture_df.latitude.values
-
-    args = ((lat, lons[i], h3_res) for i, lat in enumerate(lats))
-
-    with Pool(_n_procs) as pool:
-        cell_ids = pool.map(_get_h3_cell, args)
+def _get_h3_cell_for_rupture_df(rupture_df, h3_res):
+    lats = rupture_df["latitude"].to_numpy()
+    lons = rupture_df["longitude"].to_numpy()
 
+    cell_ints = h3_vect.geo_to_h3(lats, lons, h3_res)
+    h3_to_string = np.frompyfunc(h3.h3_to_string, 1, 1)
+    cell_ids = h3_to_string(cell_ints)
     rupture_df["cell_id"] = cell_ids
 
 

openquake/hme/utils/tests/test_utils.py

@@ -53,7 +53,7 @@ def test_get_mag_duration_from_comp_table(self):
         )
 
 
-def test_subset_source():
+def test_subset_source_no_buffer():
     rups = pd.read_csv(
         os.path.join(test_data_dir, "sm1_rups.csv"), index_col=0
     )
@@ -65,3 +65,19 @@ def test_subset_source():
     assert "88.0_313_0" not in rups_inside.index
 
     assert len(rups_inside) == 4814
+
+
+def test_subset_source_buffer():
+    rups = pd.read_csv(
+        os.path.join(test_data_dir, "sm1_rups.csv"), index_col=0
+    )
+    subset_gj_file = os.path.join(test_data_dir, "data", "phl_subset.geojson")
+
+    rups_inside = subset_source(
+        rups, subset_gj_file, feature_number=0, buffer=0.001
+    )  # small but checks if fns work
+
+    assert "88.3_20_0" in rups_inside.index
+    assert "88.0_313_0" not in rups_inside.index
+
+    assert len(rups_inside) == 4814

openquake/hme/utils/utils.py

@@ -21,7 +21,8 @@
 from h3 import h3
 
 from tqdm.autonotebook import tqdm
-from shapely.geometry import Point
+
+from shapely.geometry import Point, shape, mapping
 
 from openquake.hazardlib.geo.point import Point as OQPoint
 
@@ -992,8 +993,49 @@ def datetime_to_string(dt):
     return timestamp.strftime("%Y-%m-%d %H:%M:%S")
 
 
+def buffer_geojson_polygon(geojson_feature, buffer_distance):
+    """
+    Buffer a GeoJSON polygon and return a new GeoJSON dict.
+
+    Parameters:
+    -----------
+    geojson_feature : dict
+        GeoJSON Feature or Geometry dict with a Polygon
+    buffer_distance : float
+        Buffer distance in degrees
+
+    Returns:
+    --------
+    dict
+        GeoJSON geometry dict with buffered polygon
+    """
+    is_feature = "geometry" in geojson_feature
+
+    if is_feature:
+        geometry_dict = geojson_feature["geometry"]
+    else:
+        geometry_dict = geojson_feature
+
+    polygon = shape(geometry_dict)
+    buffered_polygon = polygon.buffer(buffer_distance)
+    buffered_geometry = mapping(buffered_polygon)
+
+    if is_feature:
+        return {
+            "type": "Feature",
+            "geometry": buffered_geometry,
+            "properties": geojson_feature.get("properties", {}),
+        }
+    else:
+        return buffered_geometry
+
+
 def subset_source(
-    rupture_gdf, subset_geojson_file, fid=None, feature_number=None
+    rupture_gdf,
+    subset_geojson_file,
+    buffer=None,
+    fid=None,
+    feature_number=None,
 ):
 
     with open(subset_geojson_file) as f:
@@ -1017,13 +1059,16 @@ def subset_source(
 
     if subset_geojson["geometry"]["type"] not in [
         "Polygon",
-        "MultiPolygon",
     ]:
         raise ValueError(
-            "`subset_geojson_file` not Polygon or MultiPolygon: "
+            "`subset_geojson_file` not Polygon: "
             + f"is {subset_geojson['geometry']['type']}"
         )
 
+    if buffer is not None:
+        logging.info(f"  Buffering subset polygon at {buffer} degrees.")
+        subset_geojson = buffer_geojson_polygon(subset_geojson, buffer)
+
     h3_res = h3.h3_get_resolution(rupture_gdf.iloc[0]["cell_id"])
 
     subset_h3_cells = h3.polyfill_geojson(