From f8596162f7b8af2c0987951e94a02003f79358e0 Mon Sep 17 00:00:00 2001 From: Fei Ye Date: Tue, 29 Oct 2024 14:38:47 -0400 Subject: [PATCH] STOFS-3D scripts: Minor updates on the driver script and tested on Hercules. --- .../Pre_processing/Grid/clip_autoarcs.py | 80 +-- .../Pre_processing/Prop/gen_tvd.py | 72 ++- .../Relocate/relocate_source_feeder.py | 13 +- .../Pre_processing/stofs3d_atl_driver.py | 462 ++++++++++++------ 4 files changed, 424 insertions(+), 203 deletions(-) diff --git a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Grid/clip_autoarcs.py b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Grid/clip_autoarcs.py index d1f50ee1a..586225e47 100755 --- a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Grid/clip_autoarcs.py +++ b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Grid/clip_autoarcs.py @@ -13,73 +13,85 @@ import geopandas as gpd # ------------------------- inputs --------------------------- -wdir = '/sciclone/schism10/Hgrid_projects/STOFS3D-v8/v43s2_RiverMapper/v44/Clip/' -crs = 'esri:102008' +WDIR = '/sciclone/schism10/Hgrid_projects/STOFS3D-v8/v22/Clip/' +CRS = 'esri:102008' -# manual polygons defined in the coastal coverage -coastal_shpfile = f'{wdir}/inputs/coastal.shp' # esri:102008 +# manual polygons defined in the coastal coverage, +# may need to delete some polygons so that they are treated as watershed +coastal_shpfile = f'{WDIR}/inputs/coastal.shp' # esri:102008 -# land boundary + coastal, i.e., watershed region between the coastline and the land boundary + manual polygons in the coastal coverage -lbnd_coastal_shpfile = f'{wdir}/inputs/lbnd_coastal.shp' # esri:102008 +# land boundary + coastal, i.e., watershed region between the coastline and the land boundary +# as well as manual polygons in the coastal coverage +lbnd_coastal_shpfile = f'{WDIR}/inputs/lbnd_coastal.shp' # esri:102008 # use linestring as the levee shapefile -levee_shpfile = f'{wdir}/inputs/levee.shp' # esri:102008 -levee_buf_distance = 5 +levee_shpfile = f'{WDIR}/inputs/levee.shp' # esri:102008 +LEVEE_BUF_DISTANCE = 5 # used to select manual polygons that are supposedly better refined than the auto arcs -select_nearshore_shpfile = f'{wdir}/inputs/select_nearshore_v3.shp' # esri:102008 +select_nearshore_shpfile = f'{WDIR}/inputs/select_nearshore_v3.shp' # esri:102008 -auto_arcs_file = f'{wdir}/inputs/total_arcs.shp' # lonlat -auto_polys_file = f'{wdir}/inputs/total_river_polys.shp' # lonlat +auto_arcs_file = f'{WDIR}/inputs/total_arcs.shp' # lonlat +auto_polys_file = f'{WDIR}/inputs/total_river_polys.shp' # lonlat # ------------------------- end inputs --------------------------- # make outputs directory -output_dir = f'{wdir}/outputs' +output_dir = f'{WDIR}/outputs' os.makedirs(output_dir, exist_ok=True) +# save a copy of this script to the working directory +os.system(f'cp {__file__} {WDIR}') coastal = gpd.read_file(coastal_shpfile) -coastal.set_crs(crs, inplace=True) +coastal.set_crs(CRS, inplace=True) lbnd_coastal = gpd.read_file(lbnd_coastal_shpfile) -lbnd_coastal.set_crs(crs, inplace=True) +lbnd_coastal.set_crs(CRS, inplace=True) -# raw_watershed is a rough estimate of watershed region, which excludes manual polygons but may include levee +# raw_watershed is a rough estimate of watershed region, +# which excludes manual polygons but may include levee # This includes most of the watershed region EXCEPT for the manual polygons in the coastal coverage raw_watershed = lbnd_coastal.overlay(coastal, how='difference').dissolve() -# extract the refined coastal polygons, i.e., coastal minus those intersecting the select_nearshore polygons, -# which are the manual polygons that are supposedly better refined than the auto arcs -# (However this is not always the case because some manualy polygons are not accurate), -# where no intersections with auto arcs are allowed, hence the buffer +# Extract the refined coastal polygons, i.e., +# coastal minus the un-refined polyogns (those intersecting the select_nearshore), +# which are the manual polygons that are supposedly better configured than the auto arcs. +# It also implies that refined coastal polygons accomodate the connectivity to watershed +# rivers, e.g., a lake shoreline should have denser nodes where it connects to a river. +# However, this is not always the case because some manualy polygons are not accurate, +# No intersections between refined polygons and auto arcs are allowed, hence the buffer. +# Unrefined polygons are allowed to intersect with auto arcs to ensure channel connectivity select_nearshore = gpd.read_file(select_nearshore_shpfile) nearshore_coastal = gpd.sjoin(coastal, select_nearshore, how="inner", predicate='intersects') refined_coastal = coastal.overlay(nearshore_coastal, how='difference').dissolve() refined_coastal_buf = gpd.GeoDataFrame(geometry=refined_coastal.buffer(50)) -refined_coastal_buf.to_file(f'{output_dir}/coastal_refined.shp', crs=crs) +refined_coastal_buf.to_file(f'{output_dir}/coastal_refined.shp') +print(f'coastal_refined.shp saved to {output_dir}') -# get final watershed by subtracting polygons that cannot be intersected by auto arcs from the raw watershed +# get final watershed +print('subtracting refined coastal from watershed') watershed = raw_watershed.overlay(refined_coastal_buf, how='difference').dissolve() -# add additional watershed region (manually specified) to the watershed -pass - -# subtract levee from watershed -levee_buf = gpd.GeoDataFrame(geometry=gpd.read_file(levee_shpfile).buffer(levee_buf_distance)) +print('subtracting levee from watershed') +levee_buf = gpd.GeoDataFrame(geometry=gpd.read_file(levee_shpfile).buffer(LEVEE_BUF_DISTANCE)) +levee_buf.to_file(f'{output_dir}/levee_buf.shp') watershed = watershed.overlay(levee_buf, how='difference').dissolve() -watershed.to_file(f'{output_dir}/watershed.shp', crs=crs) +watershed.to_file(f'{output_dir}/watershed.shp') + +# add additional watershed region (manually specified) to the watershed +print('break here if using qgis for clipping, which is faster') watershed = gpd.read_file(f'{output_dir}/watershed.shp') # clip the auto arcs to the watershed boundary -total_arcs = gpd.read_file(auto_arcs_file).to_crs(crs) +total_arcs = gpd.read_file(auto_arcs_file).to_crs(CRS) total_arcs_clipped = total_arcs.clip(watershed) -total_arcs_clipped.to_file(f'{output_dir}/total_arcs_clipped.shp', crs=crs) +total_arcs_clipped.to_file(f'{output_dir}/total_arcs_clipped.shp', crs=CRS) -# clip the polygons formed by auto arcs to the watershed boundary; this step may fail, in that case do it manually in qgis -total_arcs_polygons = gpd.read_file(auto_polys_file).to_crs(crs) +# clip the polygons formed by auto arcs to the watershed boundary; +# this step may fail, in that case do it manually in qgis +total_arcs_polygons = gpd.read_file(auto_polys_file).to_crs(CRS) total_arcs_polygons_clipped = total_arcs_polygons.buffer(0).clip(watershed) # put total_arcs_polygons_clipped back to gdf and dissolve total_arcs_polygons_clipped = gpd.GeoDataFrame(geometry=total_arcs_polygons_clipped).dissolve() -total_arcs_polygons_clipped.to_file(f'{output_dir}/total_river_polys_clipped_test.shp', crs=crs) - -pass \ No newline at end of file +total_arcs_polygons_clipped.to_file(f'{output_dir}/total_river_polys_clipped_test.shp', crs=CRS) +print('done') diff --git a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Prop/gen_tvd.py b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Prop/gen_tvd.py index 8dd57c66c..9afb3f786 100644 --- a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Prop/gen_tvd.py +++ b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Prop/gen_tvd.py @@ -1,37 +1,59 @@ -import os +''' +Generate tvd.prop file for the high-order transport schemes in SCHISM +''' import numpy as np -from pylib import schism_grid, read_schism_bpfile, inside_polygon +import socket +import pylib +from pylib import read_schism_bpfile, inside_polygon +if 'sciclone' in socket.gethostname(): + from pylib_experimental.schism_file import cread_schism_hgrid as schism_read +else: + from pylib import schism_grid as schism_read -if __name__ == '__main__': - - #Read hgrid - hgrid_name='hgrid.gr3' - gr3_pkl = 'hgrid.pkl' - if os.path.exists(gr3_pkl): - print('Reading hgrid.pkl') - gd = schism_grid(gr3_pkl) - else: - print('Reading hgrid.gr3') - gd = schism_grid(hgrid_name) - gd.save(gr3_pkl) +def gen_tvd_prop(hg: pylib.schism_grid, regions: list): + ''' + Generate tvd.prop file for the high-order transport schemes in SCHISM - #gd = read_schism_hgrid('hgrid.gr3') - gd.compute_ctr() + Parameters + ---------- + hg : schism_grid + SCHISM grid object + regions : list + list of region files, + in the final output tvd.prop: 1 for inside region, 0 for outside region - regions = ['tvd0_1.reg', 'tvd0_2.reg', 'tvd0_3.reg', 'tvd0_4.reg', \ - 'tvd0_5.reg', 'tvd0_6.reg', 'tvd0_7.reg', \ - 'upwind_east_Carribbean.rgn', 'upwind_west_Carribbean.rgn'] + Returns + ------- + None - #write constant (=1) *prop - ab = np.ones(gd.ne) + ''' + hg.compute_ctr() + tvd = np.zeros(hg.ne) - #reset values (=0) within region for region in regions: bp = read_schism_bpfile(region, fmt=1) - sind = inside_polygon(np.c_[gd.xctr,gd.yctr], bp.x, bp.y) - ab[sind==1] = 0 + idx = inside_polygon(np.c_[hg.xctr, hg.yctr], bp.x, bp.y) + tvd[idx == 1] = 1 + + hg.write_prop('tvd.prop', value=tvd, fmt='{:d}') + - gd.write_prop('tvd.prop',value=ab, fmt='{:d}') +def sample_usage(): + ''' + Sample usage of gen_tvd_prop + ''' + gd = schism_read('hgrid.gr3') + regions = [ + 'tvd0_1.reg', 'tvd0_2.reg', 'tvd0_3.reg', 'tvd0_4.reg', + 'tvd0_5.reg', 'tvd0_6.reg', 'tvd0_7.reg', + 'upwind_east_Carribbean.rgn', 'upwind_west_Carribbean.rgn' + ] + + gen_tvd_prop(gd, regions) + + +if __name__ == '__main__': + sample_usage() diff --git a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Relocate/relocate_source_feeder.py b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Relocate/relocate_source_feeder.py index a712b283d..a5ce781a5 100755 --- a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Relocate/relocate_source_feeder.py +++ b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/Source_sink/Relocate/relocate_source_feeder.py @@ -13,7 +13,8 @@ import numpy as np from scipy import spatial -from pylib_essentials.schism_file import source_sink, TimeHistory, read_schism_hgrid_cached, schism_bpfile +from pylib import schism_grid +from pylib_essentials.schism_file import source_sink, TimeHistory, schism_bpfile from pylib_essentials.utility_functions import inside_polygon # Global Var @@ -161,7 +162,7 @@ def relocate_sources2( # -------------------------------------info from old source_sink------------------------------------- # Read original source/sink generated by pyschism based on NWM old_source_sink = source_sink.from_files(source_dir=old_ss_dir) - old_gd = read_schism_hgrid_cached(f'{old_ss_dir}/hgrid.gr3', overwrite_cache=False) + old_gd = schism_grid(f'{old_ss_dir}/hgrid.gr3') old_gd.compute_ctr() # computer element center coordinates # read source to fid mapping from json @@ -197,7 +198,7 @@ def relocate_sources2( # -------------------------------------new hgrid------------------------------------- # get new hgrid (with feeders) - new_gd = read_schism_hgrid_cached(hgrid_fname, overwrite_cache=False) + new_gd = schism_grid(hgrid_fname) new_gd.compute_ctr() # -------------------- process nan values in mandatory_sources_coor ------------------ @@ -399,7 +400,7 @@ def relocate_sources( # -------------------------------------info from old source_sink------------------------------------- # get old hgrid (without feeders); hgrid with feeders may also be used - old_gd = read_schism_hgrid_cached(f'{old_ss_dir}/hgrid.gr3', overwrite_cache=False) + old_gd = schism_grid(f'{old_ss_dir}/hgrid.gr3') # get old source coordinates old_gd.compute_ctr() @@ -424,7 +425,7 @@ def relocate_sources( # -------------------------------------new hgrid------------------------------------- # get new hgrid (with feeders) - new_gd = read_schism_hgrid_cached(hgrid_fname, overwrite_cache=False) + new_gd = schism_grid(hgrid_fname) new_gd.compute_ctr() # -------------------------------------manual intervetion as pre-requisits ------------------------------------- @@ -609,7 +610,7 @@ def samples(): # copy data files from git to the working directory for record if needed # read hgrid - hgrid = read_schism_hgrid_cached(hgrid_fname, overwrite_cache=False) + hgrid = schism_grid(hgrid_fname) # read region region = schism_bpfile() diff --git a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/stofs3d_atl_driver.py b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/stofs3d_atl_driver.py index ef290d4b0..0a3e023f1 100755 --- a/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/stofs3d_atl_driver.py +++ b/src/Utility/Pre-Processing/STOFS-3D-Atl-shadow-VIMS/Pre_processing/stofs3d_atl_driver.py @@ -1,26 +1,37 @@ #!/usr/bin/env python3 ''' -This is the driver script to run the STOFS3D-ATL model. +This is the driver script to set up the STOFS-3D-ATL model. Simpler tasks such as generating uniform *.gr3 are included in this script. More complex tasks such as generating source_sink are imported as modules from subfolders. Usage: - python stofs3d-atl-driver.py in any directory + Under the current schism git directory, -Set run parameters at the beginning of the main function. + python stofs3d-atl-driver.py + + You can discard any changes when you git pull again, + because a copy of the script and imported modules will be automatically saved in the model + input folder to keep a record of the model setup. + + Alternatively, you can copy the current directory (Pre_processing/) to your preferred + location and run the script there. + +Default settings for the STOFS-3D-ATL model are set in the ConfigStofs3dAtlantic class. +Read the docstring of the main function for how to use a configuration preset. +Also, set paths at the beginning of the main function for each new setup. The script will prepare the model folders and keep a record of itself in the model input folder. ''' -# Import modules +# ------------------------- Import modules --------------------------- import os +import subprocess +import socket from pathlib import Path import shutil from datetime import datetime -import logging -import subprocess import copy import numpy as np @@ -28,16 +39,23 @@ import shapefile # from pyshp # self-defined modules -from pylib import schism_grid, inside_polygon, read_schism_reg -from pylib_experimental.schism_file import cread_schism_hgrid, source_sink +import pylib +from pylib import inside_polygon, read_schism_reg +from pylib_experimental.schism_file import source_sink +if 'sciclone' in socket.gethostname(): + from pylib_experimental.schism_file import cread_schism_hgrid as schism_read + print('Using c++ function to accelerate hgrid reading') +else: + from pylib import schism_grid as schism_read + print('Using python function to read hgrid') from pyschism.mesh import Hgrid -# import from the sub folders, not from installed packages -from logging_config import setup_logging +# Import from the sub folders. These are not from installed packages. from Source_sink.NWM.gen_sourcesink_nwm import gen_sourcesink_nwm from Source_sink.Constant_sinks.set_constant_sink import set_constant_sink from Source_sink.Relocate.relocate_source_feeder import relocate_sources2 from Source_sink.Relocate.relocate_source_feeder import v19p2_mandatory_sources_coor +from Prop.gen_tvd import gen_tvd_prop from Bctides.bctides.bctides import Bctides # temporary, bctides.py will be merged into pyschism # from pyschism.forcing.bctides import Bctides @@ -45,26 +63,35 @@ # 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() +script_path = os.path.dirname(os.path.realpath(__file__)) +print(f"script_path: {script_path}") +DRIVER_PRINT_PREFIX = '\n-----------------STOFS3D-ATL driver:---------------------\n' -# Classes and functions: + +# --------------------------------------------------------------------- +# Classes +# --------------------------------------------------------------------- class ConfigStofs3dAtlantic(): '''A class to handle the configuration of STOFS-3D-ATL model, i.e., processing the parameters and storing the factory settings. ''' - def __init__(self, - startdate=datetime(2017, 12, 1), # start date of the model - rnday=60, # 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 - nwm_cache_folder=None, - relocate_source=True, - feeder_info_file=None, # file containing feeder info, made by make_feeder_channel.py in RiverMapper - gr3_values=None - ): + def __init__( + self, + startdate=datetime(2017, 12, 1), # start date of the model + rnday=60, # 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 + bctides_flags=None, # a list of lists, each sublist is a set of flags for an open boundary + nwm_cache_folder=None, + relocate_source=True, + feeder_info_file=None, # file containing feeder info, + # made by make_feeder_channel.py in RiverMapper + gr3_values=None, + tvd_regions=None + ): self.startdate = startdate self.rnday = rnday @@ -75,6 +102,17 @@ def __init__(self, self.relocate_source = relocate_source self.nwm_cache_folder = nwm_cache_folder self.feeder_info_file = feeder_info_file + if bctides_flags is None: + self.bctides_flags = [ + [3, 3, 0, 0], # tides for elev and vel + [3, 3, 0, 0], # tides for elev and vel + [3, 3, 0, 0], # tides for elev and vel + [3, 3, 0, 0], # tides for elev and vel + [3, 3, 0, 0], # tides for elev and vel + ] + else: + self.bctides_flags = bctides_flags + if gr3_values is None: self.gr3_values = { # uniform gr3 values 'albedo': 0.1, @@ -86,6 +124,11 @@ def __init__(self, else: self.gr3_values = gr3_values + if tvd_regions is None: + self.tvd_regions = [] + else: + self.tvd_regions = tvd_regions + @classmethod def v6(cls): '''Factory method to create a configuration for STOFS3D-v6''' @@ -93,28 +136,81 @@ def v6(cls): nudging_zone_width=.3, # very wide nudging zone shapiro_zone_width=11.5, # very wide shapiro zone shapiro_tilt=3.5, # very abrupt transition in the shapiro zone - feeder_info_file='/sciclone/schism10/feiye/STOFS3D-v5/Inputs/v14/Parallel/SMS_proj/feeder/feeder.pkl') + feeder_info_file=( + '/sciclone/schism10/feiye/STOFS3D-v5/Inputs/v14/Parallel/' + 'SMS_proj/feeder/feeder.pkl'), + bctides_flags=[[5, 5, 4, 4]], + tvd_regions=[ + 'tvd0_1.reg', 'tvd0_2.reg', 'tvd0_3.reg', 'tvd0_4.reg', + 'tvd0_5.reg', 'tvd0_6.reg', 'tvd0_7.reg', + 'upwind_east_Carribbean.rgn', 'upwind_west_Carribbean.rgn' + ] + ) @classmethod def v7(cls): '''Factory method to create a configuration for STOFS3D-v7''' - return cls(nudging_zone_width=7.3, # default nudging zone - shapiro_zone_width=11.5, # default shapiro zone - shapiro_tilt=3.5, # default abrupt transition in the shapiro zone - feeder_info_file='/sciclone/schism10/Hgrid_projects/STOFS3D-v7/v20.0/Feeder/feeder_heads_bases.xy', - nwm_cache_folder=Path('/sciclone/schism10/whuang07/schism20/NWM_v2.1/')) + return cls( + nudging_zone_width=7.3, # default nudging zone + shapiro_zone_width=11.5, # default shapiro zone + shapiro_tilt=3.5, # default abrupt transition in the shapiro zone + feeder_info_file=( + '/sciclone/schism10/Hgrid_projects/STOFS3D-v7/v20.0/Feeder/' + 'feeder_heads_bases.xy'), + nwm_cache_folder=Path('/sciclone/schism10/whuang07/schism20/NWM_v2.1/'), + bctides_flags=[ + [3, 3, 0, 0], # Atlantic Ocean + [3, 3, 0, 0], # Gulf of St. Lawrence + [0, 1, 0, 0], # St. Lawrence River + ], + tvd_regions=[ + 'tvd0_1.reg', 'tvd0_2.reg', 'tvd0_3.reg', 'tvd0_4.reg', + 'tvd0_5.reg', 'tvd0_6.reg', 'tvd0_7.reg', + 'upwind_east_Carribbean.rgn', 'upwind_west_Carribbean.rgn' + ] + ) + + @classmethod + def v7_hercules_test(cls): + '''Factory method to create a configuration for STOFS3D-v7''' + return cls( + nudging_zone_width=7.3, # default nudging zone + shapiro_zone_width=11.5, # default shapiro zone + shapiro_tilt=3.5, # default abrupt transition in the shapiro zone + relocate_source=True, # need the feeder info file + feeder_info_file=( + '/work/noaa/nosofs/feiye/STOFS-3D-Atl-Example-Setup/DATA/' + 'Feeder_channels/feeder_heads_bases.xy'), + nwm_cache_folder=None, + bctides_flags=[ + [3, 3, 0, 0], # Atlantic Ocean + [3, 3, 0, 0], # Gulf of St. Lawrence + [0, 1, 0, 0], # St. Lawrence River + ], + tvd_regions=[ + 'tvd0_1.reg', 'tvd0_2.reg', 'tvd0_3.reg', 'tvd0_4.reg', + 'tvd0_5.reg', 'tvd0_6.reg', 'tvd0_7.reg', + 'upwind_east_Carribbean.rgn', 'upwind_west_Carribbean.rgn' + ] + ) @classmethod def v8_louisianna(cls): '''Factory method to create a configuration for STOFS3D-v8's local test in Louisianna''' - return cls(nudging_zone_width=0, # default nudging zone - shapiro_zone_width=0, # default shapiro zone - shapiro_tilt=0, # default abrupt transition in the shapiro zone - feeder_info_file='', - relocate_source=False, - nwm_cache_folder=Path('/sciclone/schism10/whuang07/schism20/NWM_v2.1/')) + return cls( + nudging_zone_width=0, # default nudging zone + shapiro_zone_width=0, # default shapiro zone + shapiro_tilt=0, # default abrupt transition in the shapiro zone + feeder_info_file='', + relocate_source=False, + nwm_cache_folder=Path('/sciclone/schism10/whuang07/schism20/NWM_v2.1/'), + bctides_flags=[[5, 5, 4, 4]] + ) +# --------------------------------------------------------------------- +# Utility functions +# --------------------------------------------------------------------- def mkcd_new_dir(path): '''Make a new directory and change to it''' remove_folder(path) @@ -131,6 +227,7 @@ def remove_folder(path): print(f"{path} does not exist, no need to remove.") except OSError as error: print(f"Error removing {path}: {error}") + raise def try_remove(file): @@ -138,10 +235,10 @@ def try_remove(file): try: os.remove(file) except FileNotFoundError: - pass - except OSError as e: - print("Error removing path:", file) - print("Error message:", str(e)) + print(f"{file} does not exist, no need to remove.") + except OSError: + print(f"Error removing: {file}") + raise def find_points_in_polyshp(pt_xy, shapefile_names): @@ -161,14 +258,60 @@ def find_points_in_polyshp(pt_xy, shapefile_names): return ind -def gen_nudge_coef(hgrid: schism_grid, rlmax=1.5, rnu_day=0.25, open_bnd_list=None): +def prep_run_dir(parent_dir, runid, scr_dir=None): + ''' + Prepare the run directory and return the path to the model input folder + - parent_dir: the parent directory where the run directory is created + - runid: the run directory name + - scr_dir: the parent directory where outputs are stored, e.g.: + scr_dir/R{runid}/outputs/, + if None, use the run directory + parent_dir/R{runid}/outputs/ + + Model inputs and the scripts are saved in the Input dir, e.g., I01; + The run directory is saved in the Run dir, e.g., R01; + The outputs are saved on the scratch disk and linked under parent dir. e.g., + parent_dir/O01/outputs/. + Other post-processed outputs can be put under parent_dir/O01/ + ''' + if scr_dir is None: + scr_dir = parent_dir + + model_input_path = f'{parent_dir}/I{runid}' + rundir = f'{parent_dir}/R{runid}' + output_dir = f'{scr_dir}/R{runid}/outputs' + + os.makedirs(rundir, exist_ok=True) + os.makedirs(model_input_path, exist_ok=True) + os.makedirs(output_dir, exist_ok=True) + os.makedirs(f'{parent_dir}/O{runid}', exist_ok=True) + + os.chdir(rundir) + subprocess.run( + f'ln -sf {output_dir} .', + shell=True, stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, check=False) # ignore errors, existing folder okay + + os.chdir(f'{parent_dir}/O{runid}') + os.system(f'ln -sf {output_dir} .') + os.system(f'ln -sf {model_input_path}/hgrid.gr3 .') + + return model_input_path + + +# --------------------------------------------------------------------- +# Simple pre-processing functions that do not need subfolders +# --------------------------------------------------------------------- + +def gen_nudge_coef(hgrid: pylib.schism_grid, rlmax=1.5, rnu_day=0.25, open_bnd_list=None): """ Set up nudge zone within rlmax distance from the ocean boundary: - - schism_grid (from pylib) must be in lon/lat coordinates + - hgrid must be in lon/lat coordinates - rlmax can be a uniform value, e.g., rl_max = 1.5, or a 2D array of the same size as the hgrid's number of nodes, e.g., rl_max = np.zeros(NP) with further tuning of the nudging zone width. - - If there are more than one ocean boundary, specify the boundary index in open_bnd_list as a list + - If there are more than one ocean boundary, + specify the boundary index in open_bnd_list as a list """ if open_bnd_list is None: open_bnd_list = [0, 1] @@ -178,11 +321,14 @@ def gen_nudge_coef(hgrid: schism_grid, rlmax=1.5, rnu_day=0.25, open_bnd_list=No nudge_coeff = np.zeros((hgrid.np, ), dtype=float) for open_bnd_idx in open_bnd_list: - print(f'boundary {open_bnd_idx}: {len(hgrid.iobn[open_bnd_idx])} open boundary nodes') + print( + f'boundary {open_bnd_idx}: {len(hgrid.iobn[open_bnd_idx])} open boundary nodes') bnd_node_idx = hgrid.iobn[open_bnd_idx] # distance between each node to its nearest open boundary node - distance, _ = cKDTree(np.c_[hgrid.x[bnd_node_idx], hgrid.y[bnd_node_idx]]).query(np.c_[hgrid.x, hgrid.y]) + distance, _ = cKDTree(np.c_[ + hgrid.x[bnd_node_idx], hgrid.y[bnd_node_idx] + ]).query(np.c_[hgrid.x, hgrid.y]) # nudge coefficient based on the current open boundary this_nudge_coeff = np.clip((1.0-distance/rlmax)*rnu_max, 0.0, rnu_max) @@ -191,7 +337,7 @@ def gen_nudge_coef(hgrid: schism_grid, rlmax=1.5, rnu_day=0.25, open_bnd_list=No return nudge_coeff -def gen_drag(hgrid: schism_grid): +def gen_drag(hgrid: pylib.schism_grid): '''generate drag coefficient based on the depth and regions''' # 1) overall: depth based @@ -215,7 +361,9 @@ def gen_drag(hgrid: schism_grid): 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]) + drag[idx] = np.interp( + hgrid.dp[idx], grid_depths, drag_coef, + left=drag_coef[0], right=drag_coef[-1]) # 5) tweak: limit some areas to be no more than 0.0025 for region in [ @@ -251,7 +399,7 @@ def gen_drag(hgrid: schism_grid): return drag -def gen_shapiro_strength(hgrid: schism_grid, init_shapiro_dist: np.ndarray = None, tilt=2.5): +def gen_shapiro_strength(hgrid: pylib.schism_grid, init_shapiro_dist: np.ndarray = None, tilt=2.5): '''generate shapiro strength based on the depth and regions''' shapiro_min = 100 @@ -309,12 +457,13 @@ def gen_elev_ic(hgrid=None, h0=0.1, city_shape_fnames=None, base_elev_ic=None): 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) + in_city = find_points_in_polyshp( + pt_xy=np.c_[hgrid.x, hgrid.y], shapefile_names=city_shape_fnames) high_ground = hgrid.dp < 0 if base_elev_ic is not None: - elev_ic = cread_schism_hgrid(base_elev_ic).dp + elev_ic = schism_read(base_elev_ic).dp else: elev_ic = np.zeros(hgrid.dp.shape, dtype=float) @@ -327,29 +476,9 @@ def gen_elev_ic(hgrid=None, h0=0.1, city_shape_fnames=None, base_elev_ic=None): return elev_ic -def prep_run_dir(parent_dir, runid): - '''Prepare the run directory and return the path to the model input folder''' - - model_input_path = f'{parent_dir}/I{runid}' - os.makedirs(model_input_path, exist_ok=True) - os.makedirs(f'/sciclone/scr10/feiye/R{runid}/outputs', exist_ok=True) - - rundir = f'{parent_dir}/R{runid}' - output_dir = f'/sciclone/scr10/feiye/R{runid}/outputs' - - os.makedirs(rundir, exist_ok=True) - os.chdir(rundir) - os.system(f'ln -sf {output_dir} .') - - os.makedirs(f'{parent_dir}/O{runid}', exist_ok=True) - os.chdir(f'{parent_dir}/O{runid}') - os.system(f'ln -sf {output_dir} .') - os.system(f'ln -sf {model_input_path}/hgrid.gr3 .') - os.system(f'ln -sf {model_input_path}/vgrid.in .') - - return model_input_path - - +# --------------------------------------------------------------------- +# Main function to generate inputs for STOFS-3D-ATL +# --------------------------------------------------------------------- def main(): ''' Main function to generate inputs for STOFS3D-ATL. @@ -357,65 +486,70 @@ def main(): Set the configuration parameters at the beginning of this function, i.e., between "---------input---------" and "---------end input---------". - hgrid.gr3 (in lon/lat) must be prepared before running this script + hgrid.gr3 (in lon/lat, with boundaries) must be prepared before running this script ''' - setup_logging(level=logging.DEBUG) # set to INFO to reduce the amount of output - # -----------------input--------------------- # hgrid generated by SMS, pre-processed, and converted to *.gr3 - hgrid_path = '/sciclone/schism10/feiye/STOFS3D-v8/I09/hgrid.ll' - - # get a default configuration and make changes if necessary - config = ConfigStofs3dAtlantic.v8_louisianna() - config.nwm_cache_folder = None - config.rnday = 35 + hgrid_path = ('/work/noaa/nosofs/feiye/STOFS-3D-Atl-Example-Setup/' + '2D_Setup/hgrid_levee_xGEOID20b.gr3') + + # get a configuration preset and make changes if necessary + # alternatively, you can set the parameters directly on an + # new instance of ConfigStofs3dAtlantic + config = ConfigStofs3dAtlantic.v7_hercules_test() + config.rnday = 3 config.startdate = datetime(2024, 3, 5) - # define the path where the model inputs are generated - project_dir = '/sciclone/schism10/feiye/STOFS3D-v8/' - runid = '09' + # define the project dir, where the run folders are located + project_dir = '/work/noaa/nosofs/feiye/STOFS-3D-Atl-Example-Setup/2D_Setup/' + + # run ID. e.g, 00, 01a, 10b, etc.; the run folders are named using this ID as follows: + # I{runid}: input directory for holding model input files and scripts; + # R{runid}: run directory, where the run will be submitted to queue; + # O{runid}: output directory for holding raw outputs and post-processing. + # under project_dir + runid = '00' - # swithes to generate different input files + # swithes to generate differ input_files = { - 'bctides': False, + 'bctides': True, 'vgrid': False, - 'gr3': False, - 'nudge_gr3': False, - 'shapiro': False, - 'drag': False, + 'gr3': True, + 'nudge_gr3': True, + 'shapiro': True, + 'drag': True, 'elev_ic': True, 'source_sink': True, - # 'hotstart.nc`` - # '*D.th.nc' - # '*nu.nc' - # '*.prop' + 'hotstart.nc': False, + '*D.th.nc': False, + '*nu.nc': False, + '*.prop': True, } # -----------------end input--------------------- - hgrid = cread_schism_hgrid(hgrid_path) - # hgrid_pyschism = Hgrid.open(hgrid_path, crs='epsg:4326') + print(f'{DRIVER_PRINT_PREFIX}reading hgrid from {hgrid_path} ...') + hgrid = schism_read(hgrid_path) - # ------------------------------------------------------------------- # -----------------begin generating model inputs--------------------- - # ------------------------------------------------------------------- - driver_print_prefix = '-----------------STOFS3D-ATL driver:---------------------\n' # define and make the model_input_path, the run_dir and the output dir model_input_path = prep_run_dir(project_dir, runid) # make a copy of the script itself to the model_input_path os.system(f'cp -r {script_path} {model_input_path}/') - os.system(f'cp {hgrid_path} {model_input_path}/') + # make a copy of the hgrid to the model_input_path + os.system(f'cp {hgrid_path} {model_input_path}/hgrid.gr3') # -----------------bctides--------------------- if input_files['bctides']: + hgrid_pyschism = Hgrid.open(hgrid_path, crs='epsg:4326') sub_dir = 'Bctides' - print(f'{driver_print_prefix}Generating bctides.in ...') + print(f'{DRIVER_PRINT_PREFIX}Generating bctides.in ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') Bctides( hgrid=hgrid_pyschism, # needs hgrid in pyschism's Hgrid class - flags=[[3, 3, 0, 0]], + flags=config.bctides_flags, database='fes2014' ).write( output_directory=f'{model_input_path}/{sub_dir}', @@ -426,7 +560,7 @@ def main(): # -----------------Vgrid--------------------- if input_files['vgrid']: sub_dir = 'Vgrid' - print(f'{driver_print_prefix}Generating vgrid.in ...') + print(f'{DRIVER_PRINT_PREFIX}Generating vgrid.in ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') os.system('ln -sf ../hgrid.gr3 .') @@ -441,7 +575,7 @@ def main(): # the command line argument 0 means no outputs of a sample vgrid along a given transect fortran_process.communicate(input=str(0).encode()) - print(f'{driver_print_prefix}converting the format of vgrid.in ...') + print(f'{DRIVER_PRINT_PREFIX}converting the format of vgrid.in ...') # rename vgrid.in to vgrid.in.old try_remove(f'{model_input_path}/{sub_dir}/vgrid.in.old') os.rename('vgrid.in', 'vgrid.in.old') @@ -461,25 +595,28 @@ def main(): # -----------------spatially uniform Gr3--------------------- if input_files['gr3']: sub_dir = 'Gr3' - print(f'{driver_print_prefix}Generating spatially uniform gr3 files ...') + print(f'{DRIVER_PRINT_PREFIX}Generating spatially uniform gr3 files ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') for name, value in config.gr3_values.items(): - print(f'{driver_print_prefix}Generating {name}.gr3 ...') + print(f'{DRIVER_PRINT_PREFIX}Generating {name}.gr3 ...') try_remove(f'{model_input_path}/{sub_dir}/{name}.gr3') hgrid.write(f'{model_input_path}/{sub_dir}/{name}.gr3', value=value) os.chdir(model_input_path) - # >>> spatially varying Gr3 >>>>>>>>>>>>>>>>>>>>>>>>>>>>> + # ------------------------------------------------- + # ---------begin spatially varying Gr3 ------------ + # -----------------nudge.gr3--------------------- if input_files['nudge_gr3']: sub_dir = 'Nudge_gr3' - print(f'{driver_print_prefix}Generating nudge.gr3 ...') + print(f'{DRIVER_PRINT_PREFIX}Generating nudge.gr3 ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') # generate nudging coefficient based on the proximity to the open boundaries - nudge_coef = gen_nudge_coef(hgrid, rlmax=config.nudging_zone_width, rnu_day=config.nudging_day) + nudge_coef = gen_nudge_coef( + hgrid, rlmax=config.nudging_zone_width, rnu_day=config.nudging_day) # write nudge.gr3 hgrid.save(f'{model_input_path}/{sub_dir}/nudge.gr3', value=nudge_coef) @@ -495,7 +632,7 @@ def main(): # -----------------shapiro.gr3--------------------- if input_files['shapiro']: sub_dir = 'Shapiro' - print(f'{driver_print_prefix}Generating shapiro.gr3 ...') + print(f'{DRIVER_PRINT_PREFIX}Generating shapiro.gr3 ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') if config.shapiro_zone_width > 0: @@ -504,7 +641,8 @@ def main(): else: distribute_coef = None # using a distribution coefficient simliar to the nudging coefficient - shapiro = gen_shapiro_strength(hgrid, init_shapiro_dist=distribute_coef, tilt=config.shapiro_tilt) + shapiro = gen_shapiro_strength( + hgrid, init_shapiro_dist=distribute_coef, tilt=config.shapiro_tilt) hgrid.save(f'{model_input_path}/{sub_dir}/shapiro.gr3', value=shapiro) @@ -513,7 +651,7 @@ def main(): # -----------------drag.gr3--------------------- if input_files['drag']: sub_dir = 'Drag' - print(f'{driver_print_prefix}Generating drag.gr3 ...') + print(f'{DRIVER_PRINT_PREFIX}Generating drag.gr3 ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') drag = gen_drag(hgrid) @@ -524,7 +662,7 @@ def main(): # -----------------elev_ic--------------------- if input_files['elev_ic']: sub_dir = 'Elev_ic' - print(f'{driver_print_prefix}Generating elev_ic.gr3 ...') + print(f'{DRIVER_PRINT_PREFIX}Generating elev_ic.gr3 ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') elev_ic = gen_elev_ic( hgrid, h0=0.1, @@ -534,23 +672,29 @@ def main(): ) 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}/{sub_dir}/elev.ic') + os.symlink( + f'{model_input_path}/{sub_dir}/elev_ic.gr3', + f'{model_input_path}/{sub_dir}/elev.ic') os.chdir(model_input_path) - # <<< end spatially varying Gr3 <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<< + # ----- end spatially varying Gr3 ----------------- + # ------------------------------------------------- # -----------------source_sink--------------------- if input_files['source_sink']: sub_dir = 'Source_sink' - print(f'{driver_print_prefix}Generating source_sink.in ...') + print(f'{DRIVER_PRINT_PREFIX}Generating source_sink.in ...') - mkcd_new_dir(f'{model_input_path}/{sub_dir}') + # mkcd_new_dir(f'{model_input_path}/{sub_dir}') - # 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 mkcd_new_dir(f'{model_input_path}/{sub_dir}/original_source_sink/') os.symlink(f'{model_input_path}/hgrid.gr3', 'hgrid.gr3') - gen_sourcesink_nwm(startdate=config.startdate, rnday=config.rnday, cache_folder=config.nwm_cache_folder) + gen_sourcesink_nwm( + startdate=config.startdate, rnday=config.rnday, + cache_folder=config.nwm_cache_folder) # relocate source locations to resolved river channels, the result is the "base" source/sink if config.relocate_source: @@ -565,14 +709,17 @@ def main(): max_search_radius=2100, mandatory_sources_coor=v19p2_mandatory_sources_coor, allow_neglection=False ) - # regenerated source/sink based on relocated sources.json and sinks.json + # regenerate source/sink based on relocated sources.json and sinks.json os.symlink('../original_source_sink/sinks.json', 'sinks.json') # dummy gen_sourcesink_nwm( - startdate=config.startdate, rnday=config.rnday, cache_folder=config.nwm_cache_folder) - relocated_ss = source_sink.from_files(f'{model_input_path}/{sub_dir}/relocated_source_sink/') + startdate=config.startdate, rnday=config.rnday, + cache_folder=config.nwm_cache_folder) + relocated_ss = source_sink.from_files( + f'{model_input_path}/{sub_dir}/relocated_source_sink/') # remove sinks os.unlink('vsink.th') - base_ss = source_sink(vsource=relocated_ss.vsource, vsink=None, msource=relocated_ss.msource) + base_ss = source_sink( + vsource=relocated_ss.vsource, vsink=None, msource=relocated_ss.msource) base_ss.writer(f'{model_input_path}/{sub_dir}/relocated_source_sink/') else: base_ss = source_sink.from_files(f'{model_input_path}/{sub_dir}/original_source_sink/') @@ -608,37 +755,67 @@ def main(): os.chdir(model_input_path) # -----------------*prop--------------------- - if input_files['*prop']: + if input_files['*.prop']: sub_dir = 'Prop' - print(f'{driver_print_prefix}Generating *prop files ...') + print(f'{DRIVER_PRINT_PREFIX}Generating *prop files ...') mkcd_new_dir(f'{model_input_path}/{sub_dir}') - # generate tvd.prop - os.chdir(model_input_path) + os.system(f'cp {script_path}/Prop/* .') + gen_tvd_prop(hgrid, config.tvd_regions) + + # -----------------hotstart.nc--------------------- + if input_files['hotstart.nc']: + sub_dir = 'Hotstart' + print(f'{DRIVER_PRINT_PREFIX}Generating hotstart.nc ...') + mkcd_new_dir(f'{model_input_path}/{sub_dir}') + + # generate hotstart.nc + raise NotImplementedError('hotstart.nc is not implemented yet') + + # -----------------*D.th.nc--------------------- + if input_files['*D.th.nc']: + sub_dir = 'D.th' + print(f'{DRIVER_PRINT_PREFIX}Generating *D.th.nc files ...') + mkcd_new_dir(f'{model_input_path}/{sub_dir}') + + # generate *D.th.nc + raise NotImplementedError('*D.th.nc is not implemented yet') + + # -----------------*nu.nc--------------------- + if input_files['*nu.nc']: + sub_dir = 'Nu' + print(f'{DRIVER_PRINT_PREFIX}Generating *nu.nc files ...') + mkcd_new_dir(f'{model_input_path}/{sub_dir}') + + # generate *nu.nc + raise NotImplementedError('*nu.nc is not implemented yet') def test(): '''Temporary tests''' print("--------------------- Temporary tests ---------------------") - setup_logging(level=logging.DEBUG) - # inputs model_input_path = '/sciclone/schism10/feiye/STOFS3D-v7/Inputs/I12z/' config = ConfigStofs3dAtlantic.v7() config.startdate = datetime(2024, 3, 5) config.rnday = 5 config.nwm_cache_folder = Path( - '/sciclone/schism10/feiye/STOFS3D-v7/Inputs/I12z/Source_sink/original_source_sink/20240305/') - hgrid = cread_schism_hgrid('/sciclone/schism10/feiye/STOFS3D-v7/Inputs/I12y/hgrid.gr3') + '/sciclone/schism10/feiye/STOFS3D-v7/Inputs/I12z/' + 'Source_sink/original_source_sink/20240305/') + hgrid = schism_read('/sciclone/schism10/feiye/STOFS3D-v7/Inputs/I12y/hgrid.gr3') # end inputs sub_dir = 'Source_sink' - # 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 os.chdir(f'{model_input_path}/{sub_dir}/relocated_source_sink/') - gen_sourcesink_nwm(startdate=config.startdate, rnday=config.rnday, cache_folder=config.nwm_cache_folder) - relocated_ss = source_sink.from_files(f'{model_input_path}/{sub_dir}/relocated_source_sink/') + gen_sourcesink_nwm( + startdate=config.startdate, rnday=config.rnday, + cache_folder=config.nwm_cache_folder) + relocated_ss = source_sink.from_files( + f'{model_input_path}/{sub_dir}/relocated_source_sink/') # set constant sinks (pumps and background sinks) mkcd_new_dir(f'{model_input_path}/{sub_dir}/constant_sink/') @@ -646,7 +823,8 @@ def test(): 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') - background_ss = set_constant_sink(wdir=f'{model_input_path}/{sub_dir}/constant_sink/', hgrid_utm=hgrid_utm) + 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 @@ -656,11 +834,19 @@ def test(): hgrid.compute_ctr() np.savetxt( f'{model_input_path}/{sub_dir}/vsource.xyz', - np.c_[hgrid.xctr[total_ss.source_eles-1], hgrid.yctr[total_ss.source_eles-1], total_ss.vsource.df.mean().values] + np.c_[ + hgrid.xctr[total_ss.source_eles-1], + hgrid.yctr[total_ss.source_eles-1], + total_ss.vsource.df.mean().values + ] ) np.savetxt( f'{model_input_path}/{sub_dir}/vsink.xyz', - np.c_[hgrid.xctr[total_ss.sink_eles-1], hgrid.yctr[total_ss.sink_eles-1], total_ss.vsink.df.mean().values] + np.c_[ + hgrid.xctr[total_ss.sink_eles-1], + hgrid.yctr[total_ss.sink_eles-1], + total_ss.vsink.df.mean().values + ] ) os.chdir(model_input_path)