Compound flood scripts: minor updates on a few scripts. More work on …

…the STOFS3D-ATL driver script.

src/Utility/Grid_Scripts/Compound_flooding/RiverMapper/RiverMapper/SMS.py

@@ -232,7 +232,7 @@ def merge_maps(mapfile_glob_str, merged_fname):
 
 class SMS_ARC():
     '''class for manipulating arcs in SMS maps'''
-    def __init__(self, points=None, node_idx=None, src_prj=None, dst_prj='epsg:4326'):
+    def __init__(self, points=None, node_idx=None, src_prj=None, dst_prj='epsg:4326', proj_z=True):
         # self.isDummy = (len(points) == 0)
         if node_idx is None:
             node_idx = [0, -1]
@@ -242,7 +242,7 @@ def __init__(self, points=None, node_idx=None, src_prj=None, dst_prj='epsg:4326'
 
         if src_prj == 'cpp' and dst_prj == 'epsg:4326':
             points[:, 0], points[: ,1] = cpp2lonlat(points[:, 0], points[: ,1])
-            if points.shape[1] == 3:
+            if points.shape[1] == 3 and proj_z == True:
                 points[:, 2] = dl_cpp2lonlat(points[:, 2], lat0=points[:, 1])
 
         npoints, ncol = points.shape

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Constant_sinks/set_constant_sink.py

@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+
+import shapefile
+import matplotlib.pyplot as plt
+try:
+    from pylib_utils.utility_functions import inside_polygon
+    from pylib_essentials.schism_file import schism_grid
+    from spp_essentials.Hgrid_extended import read_schism_hgrid_cached  # only for testing purposes
+except:
+    from pylib import inside_polygon, schism_grid
+import numpy as np
+from time import time
+from pylib_essentials.schism_file import source_sink, TimeHistory
+from pathlib import Path
+
+
+
+def find_ele_in_shpfile(shapefile_name, hgrid: str = ""):
+    '''
+    Find element index within polygons of a list of shapefiles
+
+    hgrid:
+        input hgrid to be modified
+    '''
+    # hgrid
+    t = time()
+    if type(hgrid) is str:
+        gd = read_schism_hgrid_cached(hgrid)
+        print(f'Reading hgrid took {time()-t} seconds')
+    elif isinstance(hgrid, schism_grid):
+        gd = hgrid
+    else:
+        raise ValueError('hgrid must be either str or schism_grid')
+
+    gd.compute_ctr()
+
+    # set default value outside polygons
+    t = time()
+    # shapefile # records = sf.records() # sf.shapeType # len(sf) # s = sf.shape(idx)
+    sf = shapefile.Reader(shapefile_name)
+    shapes = sf.shapes()
+
+    ind_list = []
+    for shp in shapes:
+        poly_xy = np.array(shp.points).T
+        ind = inside_polygon(np.c_[gd.xctr, gd.yctr], poly_xy[0], poly_xy[1])  # 1: true; 0: false
+        ind = ind.astype('bool')
+        ind_list.append(ind)
+
+    print(f'Processing {shapefile_name} took {time()-t} seconds')
+
+    return [ind_list, gd]
+
+
+def set_constant_sink(wdir='./', shapefile_name='levee_4_pump_polys.shp', hgrid_utm=None):
+    # copy datafiles
+    shp_basename = Path(shapefile_name).stem
+
+    # set pump capacities
+    [ele_idx_list, gd] = find_ele_in_shpfile(
+        # order matters, latter prevails
+        shapefile_name=f'{wdir}/{shapefile_name}',
+        hgrid=hgrid_utm,
+    )
+    gd.compute_area()
+
+    pump_sinks = np.zeros((gd.ne, ), dtype=float)
+    pump_sink_eles = np.zeros((gd.ne, ), dtype=bool)
+    for k, ele_idx in enumerate(ele_idx_list):
+        # uniform 0.5 inch/h = 0.00000352777 m/s
+        # uni_sinks = np.maximum(-0.00000352777 * gd.area[ele_idx], sinks)
+        pump_sinks[ele_idx] = -0.00000352777 * gd.area[ele_idx]
+        pump_sink_eles[ele_idx] = True
+
+    # set background sink on land
+    land = gd.dpe < -2
+    background_sink = np.zeros((gd.ne, ), dtype=float)
+    background_sink[land] = -0.5/100/3600 * gd.area[land]  # 0.5 cm/h
+    const_sinks = np.minimum(background_sink, pump_sinks)
+
+    total_sink_eles = (land + pump_sink_eles).astype(bool)
+    # better be a list, for instantiating added_ss
+    # index starts from 1
+    total_sink_ele_ids = (np.argwhere(total_sink_eles).reshape(-1, ) + 1).tolist()
+    # index starts from 0
+    total_sink_ele_idx = np.argwhere(total_sink_eles).reshape(-1, )
+
+    # build vsink.th
+    vsink = TimeHistory(
+        data_array=np.r_[np.c_[np.array([0.0]), const_sinks[total_sink_eles].reshape(1, -1)],
+                         np.c_[np.array([100*365*86400]), const_sinks[total_sink_eles].reshape(1, -1)]],
+        columns=total_sink_ele_ids,
+    )
+
+    # save sink ele ids and sink values for operational use
+    np.savetxt(f'{wdir}/pump_sinks.txt', np.c_[total_sink_ele_ids, const_sinks[total_sink_eles]], fmt='%i %10.5f')
+    # save as xyz format
+    gd.compute_ctr()
+    my_ele_xyz = np.c_[gd.xctr[total_sink_ele_idx], gd.yctr[total_sink_ele_idx], const_sinks[total_sink_eles]]
+    np.savetxt(f'{wdir}/sinks.xyz', my_ele_xyz)
+
+    # build a source_sink object
+    const_source_sink = source_sink(vsource=None, vsink=vsink, msource=None)
+    # write source sink files, actually only vsink.th and source_sink.in
+    const_source_sink.writer(wdir=wdir)
+
+    return const_source_sink
+
+if __name__ == "__main__":
+    background_ss = set_constant_sink(wdir='./')
+    original_ss = source_sink(source_dir='../relocated_source_sink/', hgird_utm='./hgrid.utm.gr3')
+    total_ss = original_ss + background_ss
+    total_ss.writer('../')
+    total_ss.nc_writer('../')
+    pass
+
+

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/NWM/gen_sourcesink.py

@@ -15,10 +15,7 @@
 log_level = logging.DEBUG
 logging.getLogger('pyschism').setLevel(log_level)
 
-if __name__ == '__main__':
-
-    startdate = datetime(2022, 1, 1)
-    rnday = 10
+def gen_sourcesink(startdate:datetime, rnday:float):
     hgrid = Hgrid.open("./hgrid.gr3", crs="epsg:4326")
 
     t0 = time()
@@ -49,3 +46,6 @@
 
     nwm.write(output_directory, hgrid, startdate, rnday, overwrite=True)
     print(f'It took {time()-t0} seconds to generate source/sink')
+
+if __name__ == '__main__':
+    gen_sourcesink(datetime(2017, 12, 1), 10)

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Relocate/relocate_source_feeder_lean.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pandas as pd
 
-class SourceSinkIn():
+class SourceSinkIn:
     """ class for *.prop or other similar formats"""
     def __init__(self, filename=None, number_of_groups=2, ele_groups=[]):
         self.n_group = number_of_groups  # 0: source; 1: sink
@@ -38,10 +38,12 @@ def writer(self, filename=None):
                     fout.write(f"{self.ip_group[k][i]}\n")
                 fout.write("\n")  # empty line
 
-def relocate_sources(old_source_sink_in:SourceSinkIn, old_vsource:np.ndarray, outdir=None, relocate_map = None):
+def relocate_sources(old_source_sink_in:SourceSinkIn, old_vsource:np.ndarray, times: np.array, outdir=None, relocate_map = None, output_vsource=False):
 
     # make a dataframe from vsource to facilitate column subsetting and reordering
-    df = pd.DataFrame(data=old_vsource, columns=['time'] + [str(x) for x in old_source_sink_in.ip_group[0]])
+    #df = pd.DataFrame(data=old_vsource, columns=['time'] + [str(x) for x in old_source_sink_in.ip_group[0]])
+    df = pd.DataFrame(data=old_vsource, columns=[str(x) for x in old_source_sink_in.ip_group[0]])
+    df['time'] = times
 
     # read relocate_map
     if relocate_map is not None:
@@ -51,7 +53,7 @@ def relocate_sources(old_source_sink_in:SourceSinkIn, old_vsource:np.ndarray, ou
 
     # Assemble new source_sink.in
     source_sink_in = SourceSinkIn(filename=None, number_of_groups=2, ele_groups=[eleids.tolist(),[]])
-    source_sink_in.writer(f'{outdir}/source_sink.in')
+    source_sink_in.writer(f'{outdir}/source.in')
 
     # Rearrange vsource.th and msource.th
     map_dict = {str(k): str(v) for k, v in np.c_[old_source_sink_in.ip_group[0][new2old_sources], eleids]}
@@ -64,7 +66,8 @@ def relocate_sources(old_source_sink_in:SourceSinkIn, old_vsource:np.ndarray, ou
     df_subset = df[list(map_dict.keys())].rename(columns=map_dict)
 
     # Save the subset dataframe to a new CSV file
-    df_subset.to_csv(f'{outdir}/vsource.th', index=False, header=False, sep=' ')
+    if output_vsource:
+        df_subset.to_csv(f'{outdir}/vsource.th', index=False, header=False, sep=' ')
 
     # msource
     msource = np.c_[
@@ -74,17 +77,23 @@ def relocate_sources(old_source_sink_in:SourceSinkIn, old_vsource:np.ndarray, ou
     ]
     np.savetxt(f'{outdir}/msource.th', msource, fmt='%d', delimiter=' ')
 
-    pass
+    return df_subset
+
 
 if __name__ == "__main__":
     old_source_sink_in = SourceSinkIn(filename='../original_source_sink/source_sink.in')
     old_vsource = np.loadtxt('../original_source_sink/vsource.th')
+    times = old_vsource[:, 0]
+    old_vsource = old_vsource[:, 1:]
     relocate_map = np.loadtxt(f'./relocate_map.txt', dtype=int)
     outdir = './'
 
     relocate_sources(
         old_source_sink_in=old_source_sink_in,
         old_vsource=old_vsource,
+        times=times,
+        relocate_map=relocate_map,
         outdir=outdir,
-        relocate_map=relocate_map
+        output_vsource=True,
     )
+

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample1/msource.th

@@ -0,0 +1,3 @@
+0.0 25 25 26 0.0 0.0 1.0
+3600 20 24 28 0.0 1.0 2.0
+7200 25 25 26 0.0 0.0 3.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample1/source_sink.in

@@ -0,0 +1,9 @@
+3
+100
+200
+300
+
+3
+10
+20
+30

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample1/vsink.th

@@ -0,0 +1,3 @@
+0.0  -1.0 -2.0 -3.0
+3600 -2.0 -4.0 -6.0
+7200 -3.0 -5.0 -7.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample1/vsource.th

@@ -0,0 +1,3 @@
+0.0 1.0 2.0 3.0
+3600 2.0 4.0 6.0
+7200 3.0 5.0 7.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample12/msource.th

@@ -0,0 +1,3 @@
+0.0  25 23.33333333 26 20    0.0 0.0 1.0 0.0
+3600 20 22.66666667 28 20    0.0 0.666667 2.0 0.0
+7200 25 23.125      26 20    0.0 0.0 3.0 0.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample12/source_sink.in

@@ -0,0 +1,10 @@
+4
+100
+200
+300
+400
+
+3
+10
+20
+30

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample12/vsink.th

@@ -0,0 +1,3 @@
+0.0  -1.0 -2.0 -3.0
+3600 -2.0 -4.0 -6.0
+7200 -3.0 -5.0 -7.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample12/vsource.th

@@ -0,0 +1,3 @@
+0.0 1.0 3.0 3.0 2.0
+3600 2.0 6.0 6.0 4.0
+7200 3.0 8.0 7.0 6.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample2/msource.th

@@ -0,0 +1,3 @@
+0.0  20 20 0.0 0.0
+3600 20 20 0.0 0.0
+7200 20 20 0.0 0.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample2/source_sink.in

@@ -0,0 +1,5 @@
+2
+200
+400
+
+0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Test_data/source_sink_sample2/vsource.th

@@ -0,0 +1,2 @@
+0.0 1.0 2.0
+7200 3.0 6.0

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/stofs3d-atl-driver.py

@@ -9,21 +9,26 @@
 import subprocess
 from scipy.spatial import cKDTree
 import copy
+from datetime import datetime
 
 # self-defined modules
 from spp_essentials.Hgrid_extended import read_schism_hgrid_cached as schism_grid  # only for testing purposes
+# from pylib import schism_grid
+
 # support both the experimental and original versions of pylib; the former is imported by default
 try:
     # from pylib_essentials.schism_file import schism_grid
-    from pylib_essentials.schism_file import read_schism_reg
+    from pylib_essentials.schism_file import read_schism_reg, source_sink
     from pylib_utils.utility_functions import inside_polygon
 except:
-    from pylib import inside_polygon, schism_grid, read_schism_reg
+    from pylib import inside_polygon, read_schism_reg
+
 # modules to be moved to pylib
-from schism_py_pre_post.Grid.SourceSinkIn import source_sink
+# from schism_py_pre_post.Grid.SourceSinkIn import source_sink
 
 # import from the sub folders, not the installed package
 from Source_sink.Constant_sinks.set_constant_sink import set_constant_sink
+from Source_sink.NWM.gen_sourcesink import gen_sourcesink
 
 # Global variables:
 # define script path; scripts are available from schism's git repo
@@ -35,12 +40,16 @@ class Config_stofs3d_atlantic():
     i.e., processing the parameters and storing the factory settings.
     '''
     def __init__(self,
+        startdate = datetime(2017, 12, 1),  # start date of the model
+        rnday = 10,  # number of days to run the model
         nudging_zone_width = 1.5,  # in degrees
         nudging_day = 1.0,  # in days
         shapiro_zone_width = 2.5,  # in degrees
         shapiro_tilt = 2.0,  # more abrupt transition in the shapiro zone
     ):
         # see a description of the parameters in the function make_river_map()
+        self.startdate = startdate
+        self.rnday = rnday
         self.nudging_zone_width = nudging_zone_width
         self.nudging_day = nudging_day
         self.shapiro_zone_width = shapiro_zone_width
@@ -156,7 +165,7 @@ def main():
 
     driver_print_prefix = '-----------------STOFS3D-ATL driver:---------------------\n'
     # define the path where the model inputs are generated
-    model_input_path = '/sciclone/schism10/feiye/STOFS3D-v6/Inputs/I16/'
+    model_input_path = '/sciclone/schism10/feiye/STOFS3D-v6/Inputs/I99a/'
     # make the directory if it does not exist
     try_mkdir(model_input_path)
 
@@ -269,28 +278,29 @@ def main():
         print(f'{driver_print_prefix}Generating source_sink.in ...')
         prep_and_change_to_subdir(f'{model_input_path}/{sub_dir}', file_list)
 
-        # generate source_sink files by intersecting NWM river segments with the model land boundary
+        # # generate source_sink files by intersecting NWM river segments with the model land boundary
         # prep_and_change_to_subdir(f'{model_input_path}/{sub_dir}/original_source_sink/', [])
-        pass
+        # os.symlink(f'{model_input_path}/hgrid.gr3', 'hgrid.gr3')
+        # gen_sourcesink(startdate=config.startdate, rnday=config.rnday)
 
-        # relocate source locations to resolved river channels
+        # # relocate source locations to resolved river channels
         # prep_and_change_to_subdir(f'{model_input_path}/{sub_dir}/relocated_source_sink/', [])
         # os.symlink(f'{model_input_path}/hgrid.gr3', 'hgrid.gr3')
         # subprocess.run(
         #     [f'{script_path}/Source_sink/Relocate/relocate_source_feeder.py'],
         #     cwd=f'{model_input_path}/{sub_dir}/relocated_source_sink/'
         # )
-        relocated_ss = source_sink(source_dir=f'{model_input_path}/{sub_dir}/relocated_source_sink/')
+        relocated_ss = source_sink.from_files(source_dir=f'{model_input_path}/{sub_dir}/relocated_source_sink/')
 
         # set constant sinks (pumps and background sinks)
         prep_and_change_to_subdir(f'{model_input_path}/{sub_dir}/constant_sink/', [])
         # copy *.shp to the current directory
         os.system(f'cp {script_path}/Source_sink/Constant_sinks/levee_4_pump_polys.* .')
-
         hgrid_utm = copy.deepcopy(hgrid)
         hgrid_utm.proj(prj0='epsg:4326', prj1='epsg:26918')  # needed because the levee shapefile is in 26918 (lon/lat would lead to problems due to truncation errors)
         background_ss = set_constant_sink(wdir=f'{model_input_path}/{sub_dir}/constant_sink/', hgrid_utm=hgrid_utm)
 
+        # assemble source/sink files and write to model_input_path
         total_ss = relocated_ss + background_ss
         total_ss.writer(f'{model_input_path}/{sub_dir}/')