diff --git a/marquette/__main__.py b/marquette/__main__.py index b7c95f6..02059a8 100644 --- a/marquette/__main__.py +++ b/marquette/__main__.py @@ -23,8 +23,12 @@ def main(cfg: DictConfig) -> None: if cfg.name.lower() == "hydrofabric": raise ImportError("Hydrofabric functionality not yet supported") elif cfg.name.lower() == "merit": - from marquette.merit.create import (create_edges, create_N, create_TMs, - write_streamflow) + from marquette.merit.create import ( + create_edges, + create_N, + create_TMs, + write_streamflow, + ) start = time.perf_counter() log.info(f"Creating MERIT {cfg.zone} River Graph") diff --git a/marquette/merit/create.py b/marquette/merit/create.py index 446c126..8fe5095 100644 --- a/marquette/merit/create.py +++ b/marquette/merit/create.py @@ -10,19 +10,29 @@ from omegaconf import DictConfig from tqdm import tqdm -from marquette.merit._connectivity_matrix import (create_gage_connectivity, - map_gages_to_zone, - new_zone_connectivity) +from marquette.merit._connectivity_matrix import ( + create_gage_connectivity, + map_gages_to_zone, + new_zone_connectivity, +) from marquette.merit._edge_calculations import ( - calculate_num_edges, create_segment, find_flowlines, - many_segment_to_edge_partition, singular_segment_to_edge_partition, - sort_xarray_dataarray) + calculate_num_edges, + create_segment, + find_flowlines, + many_segment_to_edge_partition, + singular_segment_to_edge_partition, + sort_xarray_dataarray, +) from marquette.merit._streamflow_conversion_functions import ( - calculate_huc10_flow_from_individual_files, calculate_merit_flow, - separate_basins) -from marquette.merit._TM_calculations import (create_HUC_MERIT_TM, - create_MERIT_FLOW_TM, - join_geospatial_data) + calculate_huc10_flow_from_individual_files, + calculate_merit_flow, + separate_basins, +) +from marquette.merit._TM_calculations import ( + create_HUC_MERIT_TM, + create_MERIT_FLOW_TM, + join_geospatial_data, +) log = logging.getLogger(__name__) diff --git a/marquette/merit/extensions.py b/marquette/merit/extensions.py index fa71e28..b383e0c 100644 --- a/marquette/merit/extensions.py +++ b/marquette/merit/extensions.py @@ -115,16 +115,22 @@ def soils_data(cfg: DictConfig, edges: zarr.Group) -> None: attr, attr_values, mapping, polyline_gdf ) root.array(name=names[i], data=_attr_nan_filter[mapping]) - - + + def pet_forcing(cfg: DictConfig, edges: zarr.Group) -> None: pet_zarr_data_path = Path(cfg.data_path) / f"extensions/pet_forcing/{cfg.zone}" if pet_zarr_data_path.exists(): log.info("PET forcing data already exists in zarr format") else: root = zarr.group(store=pet_zarr_data_path) - num_timesteps = pd.date_range(start=cfg.create_streamflow.start_date, end=cfg.create_streamflow.end_date, freq='d').shape[0] - pet_file_path = Path(f"/projects/mhpi/hjj5218/data/global/zarr_sub_zone/{cfg.zone}") + pet_file_path = Path( + f"/projects/mhpi/hjj5218/data/global/zarr_sub_zone/{cfg.zone}" + ) + num_timesteps = pd.date_range( + start=cfg.create_streamflow.start_date, + end=cfg.create_streamflow.end_date, + freq="d", + ).shape[0] if pet_file_path.exists() is False: raise FileNotFoundError("PET forcing data not found") edge_merit_basins: np.ndarray = edges.merit_basin[:] @@ -140,6 +146,16 @@ def pet_forcing(cfg: DictConfig, edges: zarr.Group) -> None: pet_edge_data.append(pet) pet_comid_arr = np.concatenate(pet_comid_data) pet_arr = np.concatenate(pet_edge_data) + + if pet_arr.shape[0] != len(edge_merit_basins): + raise ValueError( + "PET forcing data is not consistent. Check the number of comids in the data and the edge_merit_basins array." + ) + if pet_arr.shape[1] != num_timesteps: + raise ValueError( + "PET forcing data is not consistent. Check the number of timesteps in the data and the num_timesteps variable." + ) + for i, id in enumerate(tqdm(pet_comid_arr, desc="\rProcessing PET data")): idx = np.where(edge_merit_basins == id)[0] mapping[idx] = i diff --git a/marquette/merit/map.py b/marquette/merit/map.py index ac89a99..600570d 100644 --- a/marquette/merit/map.py +++ b/marquette/merit/map.py @@ -12,8 +12,12 @@ from scipy.sparse import csr_matrix from tqdm import tqdm -from marquette.merit._graph import (Segment, data_to_csv, get_edge_counts, - segments_to_edges) +from marquette.merit._graph import ( + Segment, + data_to_csv, + get_edge_counts, + segments_to_edges, +) log = logging.getLogger(__name__)