WIP: Further improvements

oceanmodeling · Oct 24, 2023 · 13d4268 · 13d4268
1 parent a977e4a
commit 13d4268
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 76 deletions.
diff --git a/singularity/ocsmesh/files/hurricane_mesh.py b/singularity/ocsmesh/files/hurricane_mesh.py
@@ -82,7 +82,6 @@ def main(args, clients):
 
     storm_name = str(args.name).lower()
     storm_year = str(args.year).lower()
-    out_dir =  args.out
 
     final_mesh_name = 'hgrid.gr3'
     write_mesh_box = False
@@ -91,15 +90,6 @@ def main(args, clients):
         final_mesh_name = 'final_mesh.2dm'
         write_mesh_box = True
 
-        args.rasters = [
-            i for i in (args.rasters_dir / 'gebco').iterdir() if i.suffix == '.tif'
-        ]
-
-
-        args.out = out_dir
-        args.fine_mesh = args.fine_mesh
-        args.coarse_mesh = args.coarse_mesh
-        args.region_of_interset = args.region_of_interset
     elif cmd == 'hurricane_mesh':
         final_mesh_name = 'mesh_no_bdry.2dm'
 
@@ -110,8 +100,8 @@ def main(args, clients):
 
     #TODO interpolate DEM?
     if write_mesh_box:
-        _write_mesh_box(out_dir, out_dir / final_mesh_name)
-    _generate_mesh_boundary_and_write(out_dir, out_dir / final_mesh_name)
+        _write_mesh_box(args.out, args.out / final_mesh_name)
+    _generate_mesh_boundary_and_write(args.out, args.out / final_mesh_name)
 
 
 class HurricaneMesher:
@@ -124,6 +114,11 @@ def __init__(self, sub_parser):
 
         this_parser = sub_parser.add_parser(self.script_name)
 
+        this_parser.add_argument(
+            "--lo-dem", nargs='+', help="Path to low resolution DEMS")
+        this_parser.add_argument(
+            "--hi-dem", nargs='+', help="Path to high resolution DEMS")
+
         this_parser.add_argument(
             "--nprocs", type=int, help="Number of parallel threads to use when "
             "computing geom and hfun.")
@@ -195,18 +190,19 @@ def run(self, args):
         storm_name = str(args.name).lower()
         storm_year = str(args.year).lower()
 
-        dem_dir = args.rasters_dir
+        lo_res_dems = args.lo_dem
+        hi_res_dems = args.hi_dem
+#        hi_tiles = args.hi_tiles
         shp_dir = args.shapes_dir
         hurr_info = args.windswath
         out_dir = args.out
 
         coarse_geom = shp_dir / 'base_geom'
         fine_geom = shp_dir / 'high_geom'
 
-        gebco_paths = [i for i in (dem_dir / 'gebco').iterdir() if str(i).endswith('.tif')]
-        cudem_paths = [i for i in (dem_dir / 'ncei19').iterdir() if str(i).endswith('.tif')]
-        all_dem_paths = [*gebco_paths, *cudem_paths]
-        tile_idx_path = f'zip://{str(dem_dir)}/tileindex_NCEI_ninth_Topobathy_2014.zip'
+
+        all_dem_paths = [*lo_res_dems, *hi_res_dems]
+        lo_res_paths = lo_res_dems
 
         # Specs
         wind_kt = 34
@@ -261,7 +257,6 @@ def run(self, args):
         logger.info("Reading input shapes...")
         gdf_fine = gpd.read_file(fine_geom)
         gdf_coarse = gpd.read_file(coarse_geom)
-        tile_idx = gpd.read_file(tile_idx_path)
 
         logger.info("Reading hurricane info...")
         gdf = gpd.read_file(hurr_info)
@@ -337,7 +332,6 @@ def run(self, args):
             )
         gdf_dem_box = gdf_dem_box.reset_index()
 
-        lo_res_paths = gebco_paths
 
         # TODO: use sjoin instead?!
         gdf_hi_res_box = gdf_dem_box[gdf_dem_box.geometry.intersects(gdf_refine_super_2.unary_union)].reset_index()
@@ -423,7 +417,7 @@ def run(self, args):
         # For interpolation after meshing and use GEBCO for mesh size calculation in refinement area.
         hfun_hi_rast_paths = hi_res_paths
         if len(hi_res_paths) > max_n_hires_dem:
-            hfun_hi_rast_paths = gebco_paths
+            hfun_hi_rast_paths = lo_res_paths
 
         logger.info("Create high-res size function...")
         hfun_hi = Hfun(
@@ -505,7 +499,7 @@ def run(self, args):
 
         dst_crs = "EPSG:4326"
         interp_rast_list = [
-            *get_rasters(gebco_paths, dst_crs),
+            *get_rasters(lo_res_paths, dst_crs),
             *get_rasters(gdf_hi_res_box.path.values, dst_crs)]
 
         # TODO: Fix the deadlock issue with multiple cores when interpolating
@@ -533,8 +527,6 @@ def run(self, args):
         "name", help="name of the storm", type=str)
     parser.add_argument(
         "year", help="year of the storm", type=int)
-    parser.add_argument(
-        "--rasters-dir", help="top-level directory that contains rasters")
 
     subparsers = parser.add_subparsers(dest='cmd')
     subset_client = SubsetAndCombine(subparsers)

diff --git a/singularity/ocsmesh/ocsmesh.def b/singularity/ocsmesh/ocsmesh.def
@@ -46,7 +46,7 @@ From: continuumio/miniconda3:23.3.1-0-alpine
 
 
 %runscript
-    conda run -n ocsmesh --no-capture-output python -m $*
+    conda run -n ocsmesh --no-capture-output python -m hurricane_mesh $*
 
 
 %labels

diff --git a/singularity/scripts/mesh.sbatch b/singularity/scripts/mesh.sbatch
@@ -3,4 +3,4 @@
 #SBATCH --exclusive
 #SBATCH --mem=0
 
-singularity run --bind /lustre ${SCRIPT_DIR}/ocsmesh ${STORM} ${YEAR} ${KWDS} 
+singularity run --bind /lustre ${SCRIPT_DIR}/ocsmesh.sif ${STORM} ${YEAR} ${KWDS} 
diff --git a/singularity/scripts/workflow.sh b/singularity/scripts/workflow.sh
@@ -1,6 +1,9 @@
 #!/bin/bash
 set -e
 
+# TODO: Avoid hardcoding dirs
+# TODO: Make this an input
+subset_mesh=1
 storm=$1
 year=$2
 uuid=$(uuidgen)
@@ -43,65 +46,51 @@ singularity run --bind /lustre $SCRIPT_DIR/info.sif \
     --hours-before-landfall $hr_prelandfall \
     $storm $year
 
-exit 0
-
 # To redirect all the temp file creations in OCSMesh to luster file sys
-# TODO: Avoid hardcoding temporary dir
 export TMPDIR=/lustre/.tmp
 mkdir -p $TMPDIR
 
-# TODO: Add mesh from scratch option?
-# NOTE: Paths are hardcoded in the mesh Python script
-#param_storm_name, param_storm_year,
-#"--rasters-dir", 'dem',
-## If subsetting flag is False
-#_use_if(param_subset_mesh, False, "hurricane_mesh"),
-#_use_if(param_subset_mesh, False, "--hmax"),
-#_use_if(param_subset_mesh, False, param_mesh_hmax),
-#_use_if(param_subset_mesh, False, "--hmin-low"),
-#_use_if(param_subset_mesh, False, param_mesh_hmin_low),
-#_use_if(param_subset_mesh, False, "--rate-low"),
-#_use_if(param_subset_mesh, False, param_mesh_rate_low),
-#_use_if(param_subset_mesh, False, "--transition-elev"),
-#_use_if(param_subset_mesh, False, param_mesh_trans_elev),
-#_use_if(param_subset_mesh, False, "--hmin-high"),
-#_use_if(param_subset_mesh, False, param_mesh_hmin_high),
-#_use_if(param_subset_mesh, False, "--rate-high"),
-#_use_if(param_subset_mesh, False, param_mesh_rate_high),
-#_use_if(param_subset_mesh, False, "--shapes-dir"),
-#_use_if(param_subset_mesh, False, 'shape'),
-#_use_if(param_subset_mesh, False, "--windswath"),
-#_tag_n_use_if(
-#    param_subset_mesh, False, 'hurricanes/{tag}/windswath'
-#),
-## If subsetting flag is True
-#_use_if(param_subset_mesh, True, "subset_n_combine"),
-#_use_if(param_subset_mesh, True, 'grid/HSOFS_250m_v1.0_fixed.14'),
-#_use_if(param_subset_mesh, True, 'grid/WNAT_1km.14'),
-#_tag_n_use_if(
-#    param_subset_mesh, True, 'hurricanes/{tag}/windswath'
-#),
-## Other shared options
-#"--out", _tag('hurricanes/{tag}/mesh'),
-sbatch --wait \
-    --export=ALL,KWDS="--tag $tag subset_n_combine _1 _2 _3",STORM=$storm,YEAR=$year \
-    $SCRIPT_DIR/mesh.sbatch
+KWDS=""
+if [ $subset_mesh == 1 ]; then
+    KWDS+="subset_n_combine"
+    KWDS+=" /lustre/grid/HSOFS_250m_v1.0_fixed.14"
+    KWDS+=" /lustre/grid/WNAT_1km.14"
+    KWDS+=" /lustre/hurricanes/${tag}/windswath"
+else
+    # TODO: Get param_* values from somewhere
+    KWDS+="hurricane_mesh"
+    KWDS+=" --hmax $param_mesh_hmax"
+    KWDS+=" --hmin-low $param_mesh_hmin_low"
+    KWDS+=" --rate-low $param_mesh_rate_low"
+    KWDS+=" --transition-elev $param_mesh_trans_elev"
+    KWDS+=" --hmin-high $param_mesh_hmin_high"
+    KWDS+=" --rate-high $param_mesh_rate_high"
+    KWDS+=" --shapes-dir /lustre/shape"
+    KWDS+=" --windswath hurricanes/${tag}/windswath"
+    KWDS+=" --lo-dem /lustre/dem/gebco/*.tif"
+    KWDS+=" --hi-dem /lustre/dem/ncei19/*.tif"
+fi
+KWDS+=" --out /lustre/hurricanes/${tag}/mesh"
+export KWDS
+sbatch --wait --export=ALL,KWDS,STORM=$storm,YEAR=$year $SCRIPT_DIR/mesh.sbatch
 
-singularity run --bind /lustre $SCRIPT_DIR/info.sif setup_ensemble.py \
-            --track-file $run_dir/nhc_track/hurricane-track.dat \
-            --output-directory $run_dir/setup/ensemble.dir/ \
-            --num-perturbations $num_perturb \
-            --mesh-directory $run_dir/mesh/ \
-            --sample-from-distribution \
-            --sample-rule $sample_rule \
-            --hours-before-landfall $hr_prelandfall \
-            --nwm-file /lustre/nwm/NWM_v2.0_channel_hydrofabric/nwm_v2_0_hydrofabric.gdb \
-            --date-range-file $run_dir/setup/dates.csv \
-            --tpxo-dir /lustre/tpxo \
-            $storm $year
+singularity run --bind /lustre $SCRIPT_DIR/info.sif setup_ensemble \
+        --track-file $run_dir/nhc_track/hurricane-track.dat \
+        --output-directory $run_dir/setup/ensemble.dir/ \
+        --num-perturbations $num_perturb \
+        --mesh-directory $run_dir/mesh/ \
+        --sample-from-distribution \
+        --sample-rule $sample_rule \
+        --hours-before-landfall $hr_prelandfall \
+        --nwm-file /lustre/nwm/NWM_v2.0_channel_hydrofabric/nwm_v2_0_hydrofabric.gdb \
+        --date-range-file $run_dir/setup/dates.csv \
+        --tpxo-dir /lustre/tpxo \
+        $storm $year
 #            _use_if(param_wind_coupling, True, "--use-wwm"),
 
 
+exit 0
+
 spinup_id=$(sbatch --parsable --export=ALL,STORM_PATH="$run_dir/setup/ensemble.dir/spinup",SCHISM_EXEC="$spinup_exec" $SCRIPT_DIR/schism.sbatch)
 
 jobarr=()