[21Pt] PR - Re-work Inundation.py [WIP]

pyproject.toml

@@ -83,3 +83,10 @@ per-file-ignores = """
     tools/tools_shared_functions.py: F821, F841, E711
     data/usgs/rating_curve_get_usgs_curves.py: F841
     """
+
+
+
+[build-system]
+requires = ["maturin>=1.7", 'setuptools>=42']
+#build-backend = "maturin"
+build-backend = 'setuptools.build_meta'

src/run_unit_wb.sh

@@ -1,4 +1,4 @@
-#!/bin/bash -e
+k#!/bin/bash -e
 
 # Do not call this file directly. Call fim_process_unit_wb.sh which calls
 # this file.

src/subdiv_chan_obank_src.py

@@ -373,7 +373,9 @@ def multi_process(variable_mannings_calc, procs_list, log_file, number_of_jobs,
     log_file.writelines(["%s\n" % item for item in map_output])
 
 
-def run_prep(fim_dir, mann_n_table, output_suffix, number_of_jobs, verbose, src_plot_option):
+def run_prep(
+    fim_dir, mann_n_table, output_suffix, number_of_jobs, verbose, src_plot_option, process_huc=None
+):
     procs_list = []
 
     print(f"Writing progress to log file here: {fim_dir}/logs/subdiv_src_{output_suffix}.log")
@@ -407,53 +409,59 @@ def run_prep(fim_dir, mann_n_table, output_suffix, number_of_jobs, verbose, src_
     else:
         print('Running the variable_mannings_calc function...')
 
-        ## Loop through hucs in the fim_dir and create list of variables to feed to multiprocessing
-        huc_list = [d for d in os.listdir(fim_dir) if re.match(r'^\d{8}$', d)]
-        huc_list.sort()  # sort huc_list for helping track progress in future print statments
+        if process_huc is None:
+            ## Loop through hucs in the fim_dir and create list of variables to feed to multiprocessing
+            huc_list = [d for d in os.listdir(fim_dir) if re.match(r'^\d{8}$', d)]
+            huc_list.sort()  # sort huc_list for helping track progress in future print statments
+        else:
+            huc_list = [process_huc]
         for huc in huc_list:
             # if huc != 'logs' and huc[-3:] != 'log' and huc[-4:] != '.csv':
-            if re.match(r'\d{8}', huc):
-                huc_branches_dir = os.path.join(fim_dir, huc, 'branches')
-                for branch_id in os.listdir(huc_branches_dir):
-                    branch_dir = os.path.join(huc_branches_dir, branch_id)
-                    in_src_bankfull_filename = join(branch_dir, 'src_full_crosswalked_' + branch_id + '.csv')
-                    htable_filename = join(branch_dir, 'hydroTable_' + branch_id + '.csv')
-                    huc_plot_output_dir = join(branch_dir, 'src_plots')
-
-                    if isfile(in_src_bankfull_filename) and isfile(htable_filename):
-                        procs_list.append(
-                            [
-                                in_src_bankfull_filename,
-                                df_mann,
-                                huc,
-                                branch_id,
-                                htable_filename,
-                                output_suffix,
-                                src_plot_option,
-                                huc_plot_output_dir,
-                            ]
-                        )
-                    else:
-                        print(
-                            'HUC: '
-                            + str(huc)
-                            + '  branch id: '
-                            + str(branch_id)
-                            + '\nWARNING --> can not find required file (src_full_crosswalked_bankfull_*.csv '
-                            + 'or hydroTable_*.csv) in the fim output dir: '
-                            + str(branch_dir)
-                            + ' - skipping this branch!!!\n'
-                        )
-                        log_file.write(
-                            'HUC: '
-                            + str(huc)
-                            + '  branch id: '
-                            + str(branch_id)
-                            + '\nWARNING --> can not find required file (src_full_crosswalked_bankfull_*.csv '
-                            + 'or hydroTable_*.csv) in the fim output dir: '
-                            + str(branch_dir)
-                            + ' - skipping this branch!!!\n'
+            if process_huc is None or huc in process_huc:
+                if re.match(r'\d{8}', huc):
+                    huc_branches_dir = os.path.join(fim_dir, huc, 'branches')
+                    for branch_id in os.listdir(huc_branches_dir):
+                        branch_dir = os.path.join(huc_branches_dir, branch_id)
+                        in_src_bankfull_filename = join(
+                            branch_dir, 'src_full_crosswalked_' + branch_id + '.csv'
                         )
+                        htable_filename = join(branch_dir, 'hydroTable_' + branch_id + '.csv')
+                        huc_plot_output_dir = join(branch_dir, 'src_plots')
+
+                        if isfile(in_src_bankfull_filename) and isfile(htable_filename):
+                            procs_list.append(
+                                [
+                                    in_src_bankfull_filename,
+                                    df_mann,
+                                    huc,
+                                    branch_id,
+                                    htable_filename,
+                                    output_suffix,
+                                    src_plot_option,
+                                    huc_plot_output_dir,
+                                ]
+                            )
+                        else:
+                            print(
+                                'HUC: '
+                                + str(huc)
+                                + '  branch id: '
+                                + str(branch_id)
+                                + '\nWARNING --> can not find required file (src_full_crosswalked_bankfull_*.csv '
+                                + 'or hydroTable_*.csv) in the fim output dir: '
+                                + str(branch_dir)
+                                + ' - skipping this branch!!!\n'
+                            )
+                            log_file.write(
+                                'HUC: '
+                                + str(huc)
+                                + '  branch id: '
+                                + str(branch_id)
+                                + '\nWARNING --> can not find required file (src_full_crosswalked_bankfull_*.csv '
+                                + 'or hydroTable_*.csv) in the fim output dir: '
+                                + str(branch_dir)
+                                + ' - skipping this branch!!!\n'
+                            )
 
         ## Pass huc procs_list to multiprocessing function
         multi_process(variable_mannings_calc, procs_list, log_file, number_of_jobs, verbose)

src/utils/shared_functions.py

@@ -479,7 +479,7 @@ def print_date_time_duration(start_dt, end_dt):
         total_hours, rem_seconds = divmod(rem_seconds, 60 * 60)
         total_mins, seconds = divmod(rem_seconds, 60)
 
-        time_fmt = f"{total_hours:02d} hours {total_mins:02d} mins {seconds:02d} secs"
+        time_fmt = f"{total_days:02d} days {total_hours:02d} hours {total_mins:02d} mins {seconds:02d} secs"
 
         duration_msg = "Duration: " + time_fmt
         print(duration_msg)

tools/inundate_gms.py

@@ -15,6 +15,7 @@ def Inundate_gms(
     hydrofabric_dir,
     forecast,
     num_workers=1,
+    hydro_table_df=None,
     hucs=None,
     inundation_raster=None,
     inundation_polygon=None,
@@ -67,6 +68,7 @@ def Inundate_gms(
         inundation_polygon,
         depths_raster,
         forecast,
+        hydro_table_df,
         verbose=False,
     )
 
@@ -160,8 +162,16 @@ def __inundate_gms_generator(
     inundation_polygon,
     depths_raster,
     forecast,
+    hydro_table_df=None,
     verbose=False,
 ):
+    """
+    Generator for use in parallelizing inundation
+
+    Parameters
+    ----------
+
+    """
     # Iterate over branches
     for idx, row in hucs_branches.iterrows():
         huc = str(row[0])
@@ -177,35 +187,43 @@ def __inundate_gms_generator(
         catchments_branch = os.path.join(branch_dir, catchments_file_name)
 
         # FIM versions > 4.3.5 use an aggregated hydrotable file rather than individual branch hydrotables
-        hydroTable_huc = os.path.join(huc_dir, "hydrotable.csv")
-        if os.path.isfile(hydroTable_huc):
-            htable_req_cols = [
-                "HUC",
-                "branch_id",
-                "feature_id",
-                "HydroID",
-                "stage",
-                "discharge_cms",
-                "LakeID",
-            ]
-            hydroTable_all = pd.read_csv(
-                hydroTable_huc,
-                dtype={
-                    "HUC": str,
-                    "branch_id": int,
-                    "feature_id": str,
-                    "HydroID": str,
-                    "stage": float,
-                    "discharge_cms": float,
-                    "LakeID": int,
-                },
-                usecols=htable_req_cols,
-            )
-            hydroTable_all.set_index(["HUC", "feature_id", "HydroID"], inplace=True)
-            hydroTable_branch = hydroTable_all.loc[hydroTable_all["branch_id"] == int(branch_id)]
+
+        if hydro_table_df is not None:
+            hydro_table_all = hydro_table_df.set_index(["HUC", "feature_id", "HydroID"], inplace=False)
+            hydro_table_branch = hydro_table_all.loc[hydro_table_all["branch_id"] == int(branch_id)]
         else:
-            # Earlier FIM4 versions only have branch level hydrotables
-            hydroTable_branch = os.path.join(branch_dir, f"hydroTable_{branch_id}.csv")
+            hydro_table_huc = os.path.join(huc_dir, "hydrotable.csv")
+            if os.path.isfile(hydro_table_huc):
+
+                htable_req_cols = [
+                    "HUC",
+                    "branch_id",
+                    "feature_id",
+                    "HydroID",
+                    "stage",
+                    "discharge_cms",
+                    "LakeID",
+                ]
+
+                hydro_table_all = pd.read_csv(
+                    hydro_table_huc,
+                    dtype={
+                        "HUC": str,
+                        "branch_id": int,
+                        "feature_id": str,
+                        "HydroID": str,
+                        "stage": float,
+                        "discharge_cms": float,
+                        "LakeID": int,
+                    },
+                    usecols=htable_req_cols,
+                )
+
+                hydro_table_all.set_index(["HUC", "feature_id", "HydroID"], inplace=True)
+                hydro_table_branch = hydro_table_all.loc[hydro_table_all["branch_id"] == int(branch_id)]
+            else:
+                # Earlier FIM4 versions only have branch level hydrotables
+                hydro_table_branch = os.path.join(branch_dir, f"hydroTable_{branch_id}.csv")
 
         xwalked_file_name = f"gw_catchments_reaches_filtered_addedAttributes_crosswalked_{branch_id}.gpkg"
         catchment_poly = os.path.join(branch_dir, xwalked_file_name)
@@ -237,7 +255,7 @@ def __inundate_gms_generator(
             "rem": rem_branch,
             "catchments": catchments_branch,
             "catchment_poly": catchment_poly,
-            "hydro_table": hydroTable_branch,
+            "hydro_table": hydro_table_branch,
             "forecast": forecast,
             "mask_type": "filter",
             "hucs": None,
@@ -253,10 +271,11 @@ def __inundate_gms_generator(
             "quiet": not verbose,
         }
 
-        yield (inundate_input, identifiers)
+        yield inundate_input, identifiers
 
 
 if __name__ == "__main__":
+
     # parse arguments
     parser = argparse.ArgumentParser(description="Inundate FIM")
     parser.add_argument(