linting

taddyb · Apr 1, 2024 · 3a13315 · 3a13315
1 parent b9f4251
commit 3a13315
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 19 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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__)