diff --git a/marquette/__main__.py b/marquette/__main__.py index 4909031..907d251 100644 --- a/marquette/__main__.py +++ b/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__": diff --git a/marquette/conf/config.yaml b/marquette/conf/config.yaml index 0f0b9ce..182c09f 100644 --- a/marquette/conf/config.yaml +++ b/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 diff --git a/marquette/merit/extensions.py b/marquette/merit/extensions.py index 1e4053e..c0dd29c 100644 --- a/marquette/merit/extensions.py +++ b/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)