STOFS-3D-Atl scripts: Minor edits on the tweak_hotstart script. More …

…work on the driver script.

src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Hotstart/tweak_stofs3d_hotstart.py

@@ -1,13 +1,17 @@
 import numpy as np
 import os
-import STOFS3D_scripts
+from pathlib import Path
+
+# pip install git+https://github.com/feiye-vims/schism_py_pre_post.git
 from schism_py_pre_post.Shared_modules.hotstart_proc import Hotstart
 from schism_py_pre_post.Download.download_usgs_with_api import get_usgs_obs_for_stofs3d
 from schism_py_pre_post.Download.download_cbp_with_api import get_cbp_obs_for_stofs3d
 from schism_py_pre_post.Shared_modules.gen_subregion_ic2 import gen_subregion_ic_stofs3d
 from schism_py_pre_post.Grid.Grid_geometry import find_points_in_polyshp
+import schism_py_pre_post  # this is needed for the __file__ attribute
+
+# pip install git+https://github.com/wzhengui/pylibs.git
 from pylib import schism_grid
-from pathlib import Path
 
 
 def gen_elev_ic(hgrid=None, h0=0.1, city_shape_fnames=None):
@@ -29,7 +33,10 @@ def gen_elev_ic(hgrid=None, h0=0.1, city_shape_fnames=None):
 
     return elev_ic
 
-def tweak_stofs3d_hotstart(wdir ='./', hotstart_date_str='2021-05-01', city_shapefile_names=[]):
+def tweak_stofs3d_hotstart(
+    wdir ='./', hotstart_date_str='2021-05-01', city_shapefile_names=[],
+    hycom_TS_file='TS_1.nc', hycom_hot_file='hotstart.nc.hycom'
+):
     '''
     The "wdir" should contain the following files:
         hgrid.gr3 (should be the same as hgrid.ll for STOFS3D), vgrid.in
@@ -41,25 +48,23 @@ def tweak_stofs3d_hotstart(wdir ='./', hotstart_date_str='2021-05-01', city_shap
     only specify the file basename, the script will copy these files to your wdir.
     '''
 
-    # input section
     griddir = wdir
-    output_obs_dir = f'{wdir}/Obs/'
-    hycom_TS_file = f'{wdir}/TS_1.nc'
-    hycom_hot_file = f'{wdir}/hotstart.nc.hycom'
-    my_hot_file = f'{wdir}/hotstart.nc'
-    # end input section
+    output_obs_dir = f'{wdir}/Obs/'  # observations will be downloaded to this directory
+    my_hot_file = f'{wdir}/hotstart.nc'  # this is the final hotstart.nc file
+    hycom_TS_file = f'{wdir}/TS_1.nc'  # this file should already be prepared in the previous step of hotstart.nc generation from HYCOM
+    hycom_hot_file = f'{wdir}/hotstart.nc.hycom'  # this file should already be prepared in the previous step of hotstart.nc generation from HYCOM, renamed with the '.hycom' suffix.
 
     # copy datafiles
-    mydir = os.path.dirname(STOFS3D_scripts.__file__)
+    mydir = os.path.dirname(schism_py_pre_post.__file__)
     for shp in city_shapefile_names:
         shp_basename = Path(shp).stem
-        os.system(f'cp {mydir}/Datafiles/Tweak_hotstart/{shp_basename}.* {wdir}')
+        os.system(f'cp {mydir}/Datafiles/{shp_basename}.* {wdir}')
 
     # download coastal obs from usgs
     get_usgs_obs_for_stofs3d(outdir=output_obs_dir, start_date_str=hotstart_date_str)
 
-    # download coastal obs from CBP
-    get_cbp_obs_for_stofs3d(outdir=output_obs_dir, sample_time=hotstart_date_str)
+    # download coastal obs from CBP; only salinity, no temperature (cbp samples are sparse in time)
+    get_cbp_obs_for_stofs3d(outdir=output_obs_dir, sample_time=hotstart_date_str, varname=['sal'])
 
     # interpolate obs onto model grid
     gen_subregion_ic_stofs3d(wdir=wdir, obsdir=output_obs_dir, hycom_TS_file=hycom_TS_file, date_str=hotstart_date_str)
@@ -95,11 +100,15 @@ def tweak_stofs3d_hotstart(wdir ='./', hotstart_date_str='2021-05-01', city_shap
 if __name__ == "__main__":
     '''
     Modify hycom-based hotstart.nc.hycom with coastal observation values
-    See instructions in "def tweak_stofs3d_hotstart()"
     '''
+
+    # Sample usage:
+    # The "wdir" should be the dir where the original hotstart.nc from hycom is generated.
+    # It should include hgrid.gr3 and vgrid.in, TS_1.nc;
+    # in addition, rename the original hotstart.nc from hycom to hotstart.nc.hycom
     tweak_stofs3d_hotstart(
         wdir='/sciclone/schism10/feiye/STOFS3D-v5/Inputs/I24/Hot_test/',
         hotstart_date_str='2021-05-01',
-        city_shapefile_names = ["city_polys_from_v10_lonlat.shp"]
+        city_shapefile_names = ["city_polys_from_v10_lonlat.shp"],  # polygon shapefile specifying cities
     )
 

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

@@ -2,9 +2,6 @@
 This is the driver script to run the STOFS3D-ATL model.
 '''
 
-# Global variables:
-# define script path; scripts are available from schism's git repo
-script_path = '/sciclone/data10/feiye/SCHISM_REPOSITORY/schism/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/'
 
 # Import modules
 import os
@@ -15,17 +12,27 @@
 import copy
 from datetime import datetime
 from pyproj import Transformer
+from pathlib import Path
+import shutil
+import shapefile
 
 # self-defined modules
-iTest = True
 from pylib_essentials.schism_file import read_schism_reg, source_sink, schism_grid, read_schism_hgrid_cached
 from pylib_essentials.utility_functions import inside_polygon
+from pyschism.mesh import Hgrid
 
-# import from the sub folders, not from the installed packages
-from Source_sink.Constant_sinks.set_constant_sink import set_constant_sink
+# import from the sub folders, not from installed packages
 from Source_sink.NWM.gen_sourcesink import gen_sourcesink
+from Source_sink.Constant_sinks.set_constant_sink import set_constant_sink
 from Source_sink.Relocate.relocate_source_feeder import relocate_sources, v16_mandatory_sources_coor
+from Bctides.bctides.bctides import Bctides  # temporary, bctides.py will be merged into pyschism
+# from pyschism.forcing.bctides import Bctides
 
+# Global variables:
+# use the full path of the Pre-Processing dir inside your schism repo
+# e.g, script_path = '/my_dir/schism/src/Utility/Pre-Processing/'
+# If you are running this script in the schism repo, you can also use the following line:
+script_path = Path('./').absolute()
 
 # Classes and functions:
 class Config_stofs3d_atlantic():
@@ -91,6 +98,21 @@ def try_remove(file):
         print("Error removing path:", file)
         print("Error message:", str(e))
 
+def find_points_in_polyshp(pt_xy, shapefile_names):
+    ind = np.zeros(pt_xy[:, 0].shape)
+    for shapefile_name in shapefile_names:
+        # shapefile # records = sf.records() # sf.shapeType # len(sf) # s = sf.shape(idx)
+        sf = shapefile.Reader(shapefile_name)
+        shapes = sf.shapes()
+
+        for i, shp in enumerate(shapes):
+            poly_xy = np.array(shp.points).T
+            print(f'shape {i+1} of {len(shapes)}, {poly_xy[:, 0]}')
+            ind += inside_polygon(pt_xy, poly_xy[0], poly_xy[1])  # 1: true; 0: false
+
+    ind = ind.astype('bool')
+    return ind
+
 def gen_nudge_coef(hgrid:schism_grid, rlmax = 1.5, rnu_day=0.25, open_bnd_list=[0, 1]):
         """
         Set up nudge zone within rlmax distance from the ocean boundary:
@@ -119,20 +141,27 @@ def gen_nudge_coef(hgrid:schism_grid, rlmax = 1.5, rnu_day=0.25, open_bnd_list=[
         return nudge_coeff
 
 def gen_drag(hgrid:schism_grid):
-    # depth based
-    grid_depths = [-9e9,  -3,    -1,     9e9]
-    drag_coef = [0.025, 0.025, 0.0025, 0.0025]
-
-    drag = np.interp(hgrid.dp, grid_depths, drag_coef)
-
-    # region based
-    region_files = [f'{script_path}/Gr3/Drag/GoME+0.001.reg', f'{script_path}/Gr3/Drag/Lake_Charles_0.reg']
-    region_tweaks = [0.001, 0.0]
+    # 1) overall: depth based
+    grid_depths = [-3, -1]
+    drag_coef = [0.025, 0.0025]
+    # linear interpolation with constant extrapolation of nearest end values
+    drag = np.interp(hgrid.dp, grid_depths, drag_coef, left=drag_coef[0], right=drag_coef[-1])
+
+    # 2) tweak: regions with constant drag
+    region_files = [f'{script_path}/Gr3/Drag/Lake_Charles_0.reg']
+    region_tweaks = [0.0]
     for region_tweak, region_file in zip(region_tweaks, region_files):
         reg = read_schism_reg(region_file)
         idx = inside_polygon(np.c_[hgrid.x, hgrid.y], reg.x, reg.y).astype(bool)
         drag[idx] = region_tweak
-
+
+    # 3) tweak: for nodes inside GoME and dp > -1 m, set drag between 0.01 and 0.02
+    grid_depths = [5, 20]
+    drag_coef = [0.02, 0.01]
+    reg = read_schism_reg(f'{script_path}/Gr3/Drag/GoME2.reg')
+    idx = inside_polygon(np.c_[hgrid.x, hgrid.y], reg.x, reg.y).astype(bool) & (hgrid.dp > -1)
+    drag[idx] = np.interp(hgrid.dp[idx], grid_depths, drag_coef, left=drag_coef[0], right=drag_coef[-1])
+
     return drag
 
 def gen_shapiro_strength(hgrid:schism_grid, init_shapiro_dist:np.ndarray=None, tilt=2.5):
@@ -169,29 +198,68 @@ def gen_shapiro_strength(hgrid:schism_grid, init_shapiro_dist:np.ndarray=None, t
 
     return shapiro
 
+def gen_elev_ic(hgrid=None, h0=0.1, city_shape_fnames=None):
+    '''
+    set initial elevation: 0 in the ocean, h0 below ground on higher grounds and in cities
+    city_shape_fname: the shapefile must be in the same projection as the hgrid
+    '''
+    if city_shape_fnames is None:
+        in_city = np.ones(hgrid.dp.shape, dtype=bool)
+    else:
+        in_city = find_points_in_polyshp(pt_xy=np.c_[hgrid.x, hgrid.y], shapefile_names=city_shape_fnames)
+
+    above_NAVD88_0 = hgrid.dp < 0
+
+    elev_ic = np.zeros(hgrid.dp.shape, dtype=float)
+
+    ind = np.logical_or(above_NAVD88_0, in_city)
+    elev_ic[ind] = - hgrid.dp[ind] - h0
+
+    return elev_ic
+
 def main():
     # prepare a configuration
     config = Config_stofs3d_atlantic.v6()
+    config.rnday = 40
+    config.startdate = datetime(2021, 8, 1)
 
     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/I99a/'
+    model_input_path = '/sciclone/schism10/feiye/STOFS3D-v7/Inputs/I00/'
     # make the directory if it does not exist
     try_mkdir(model_input_path)
 
     # hgrid must be prepared before running this script
-    hgrid = read_schism_hgrid_cached(f'{model_input_path}/hgrid.gr3')
+    hgrid = read_schism_hgrid_cached(f'{model_input_path}/hgrid.gr3', overwrite_cache=False)
+    hgrid_pyschism = Hgrid.open(f'{model_input_path}/hgrid.gr3', crs='epsg:4326')
 
     # set a dict to indicate input files to be generated
     input_files = {
+        'bctides': True,
         'vgrid': False,
         'gr3': False,
         'nudge_gr3': False,
         'shapiro': False,
         'drag': False,
-        'source_sink': True,
+        'source_sink': False,
+        'elev_ic': False,
     }
 
+    # -----------------bctides---------------------
+    if input_files['bctides']:
+        sub_dir = 'Bctides'
+        print(f'{driver_print_prefix}Generating bctides.in ...')
+        prep_and_change_to_subdir(f'{model_input_path}/{sub_dir}', 'bctides.in')
+        Bctides(
+            hgrid=hgrid_pyschism,  # needs hgrid in pyschism's Hgrid class
+            flags=[[5,3,0,0],[0,1,0,0],[0,1,0,0]],
+            database='fes2014'
+        ).write(
+            output_directory=f'{model_input_path}/{sub_dir}',
+            start_date=config.startdate,
+            rnday=config.rnday,
+        )
+
     # -----------------Vgrid---------------------
     if input_files['vgrid']:
         sub_dir = 'Vgrid'
@@ -286,6 +354,18 @@ def main():
         hgrid.save(f'{model_input_path}/{sub_dir}/drag.gr3', value=drag)
 
         os.chdir(model_input_path)
+
+    # -----------------elev_ic---------------------
+    if input_files['elev_ic']:
+        sub_dir = 'Elev_ic'
+        print(f'{driver_print_prefix}Generating elev_ic.gr3 ...')
+        prep_and_change_to_subdir(f'{model_input_path}/{sub_dir}', 'elev_ic.gr3')
+        elev_ic = gen_elev_ic(hgrid, city_shape_fnames=[f'{script_path}/Hotstart/city_polys_from_v10_lonlat.shp'])
+
+        hgrid.save(f'{model_input_path}/{sub_dir}/elev_ic.gr3', value=elev_ic)
+        os.symlink(f'{model_input_path}/{sub_dir}/elev_ic.gr3', f'{model_input_path}/elev.ic')
+
+        os.chdir(model_input_path)
 
     # <<< end spatially varying Gr3 <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<