Added a stream-geos extension

marquette/__main__.py

@@ -152,10 +152,20 @@ def run_extensions(cfg: DictConfig, edges: zarr.Group) -> None:
         from marquette.merit.extensions import calculate_hf_width
 
         log.info("Adding hf_width to your MERIT River Graph")
-        if "hf_v22_btm_width" in edges:
+        if "hf_v22_ch_slope" in edges:
             log.info("hf_width already exists in zarr format")
         else:
             calculate_hf_width(cfg, edges)
+
+    if "stream_geo_attr" in cfg.extensions:
+        from marquette.merit.extensions import calculate_stream_geo_attr
+
+        log.info("Adding stream_geo_attr to your MERIT River Graph")
+        if "stream_geo_width" in edges:
+            log.info("stream_geo_attr already exists in zarr format")
+        else:
+            calculate_stream_geo_attr(cfg, edges) 
+
 
 
 if __name__ == "__main__":

marquette/conf/config.yaml

@@ -1,6 +1,6 @@
 name: MERIT
 data_path: /projects/mhpi/data/${name}
-zone: 74
+zone: 71
 gpu: 7
 create_edges:
   buffer: 0.3334
@@ -40,6 +40,7 @@ extensions:
   # - incremental_drainage_area
   - q_prime_sum
   - hf_width
+  - stream_geo_attr
   # - upstream_basin_avg_mean_p
   # - q_prime_sum_stats
   # - lstm_stats

marquette/merit/extensions.py

@@ -667,3 +667,44 @@ def calculate_hf_width(cfg: DictConfig, edges: zarr.Group):
         data=ch_slope,
     )
 
+def calculate_stream_geo_attr(cfg: DictConfig, edges: zarr.Group):
+    zone = cfg.zone
+    log.info("reading stream_geo store")
+    file_path = Path(f"/projects/mhpi/yxs275/River_geometry/CONUS_forward/")
+    if file_path.exists() is False:
+        raise FileNotFoundError("global_dhbv_static_inputs data not found")
+    edge_merit_basins: np.ndarray = edges.merit_basin[:]
+    comid_data = []
+    width_data = []
+    depth_data = []
+
+    mapping = np.empty_like(edge_merit_basins, dtype=int)
+    root = zarr.open_group(file_path, mode="r")
+    for key in root.keys():
+        if str(zone) in str(key):
+            log.info(f"Reading: {key}")
+            zone_data = root[key]
+            comids = zone_data.COMID[:].astype(np.int32)
+            width = zone_data.width[:]
+            depth = zone_data.depth[:]
+
+            comid_data.append(comids)
+            width_data.append(width)
+            depth_data.append(depth)
+
+    comid_arr = np.concatenate(comid_data)
+    width_arr = np.concatenate(width_data)
+    depth_arr = np.concatenate(depth_data)
+
+    mask = np.isin(edge_merit_basins, comid_arr)  # note, this is mapping to the edges, not the comids
+    width_mapped = np.zeros(len(edge_merit_basins))
+    depth_mapped = np.zeros(len(edge_merit_basins))
+
+    sorter = np.argsort(comid_arr)
+    indices = sorter[np.searchsorted(comid_arr[sorter], edge_merit_basins[mask])]
+
+    width_mapped[mask] = width_arr[indices]
+    depth_mapped[mask] = depth_arr[indices]
+
+    edges.array(name="stream_geos_width", data=width_mapped)
+    edges.array(name="stream_geos_depth", data=depth_mapped)