Assorted enhancements of remap_restarts.py package (EASE, SLES15, NAS, lake/landice log)

CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added EASE grid option for remapping of land restarts in remap_restarts.py package (facilitates use of package in GEOSldas setup script)
+- Added support for SLES15, NAS site and log for remap_lake_landice_saltwater in remap_restarts.py package
+
 ### Changed
 
 ### Fixed

pre/CMakeLists.txt

@@ -2,9 +2,16 @@ add_subdirectory(NSIDC-OSTIA_SST-ICE_blend)
 add_subdirectory(prepare_ocnExtData)
 
 file(GLOB pythonscripts CONFIGURE_DEPENDS ./remap_restart/remap*  ./remap_restart/bin2nc_merra2* ./remap_restart/tests/*.*)
-install(
-   PROGRAMS ${pythonscripts}
-   DESTINATION bin)
+
+list(FILTER pythonscripts EXCLUDE REGEX "remap_utils.py")
+
+foreach (file ${pythonscripts})
+  install(PROGRAMS ${file} DESTINATION bin)
+endforeach()
+
+set(file ./remap_restart/remap_utils.py)
+configure_file(${file} ${file} @ONLY)
+install(PROGRAMS ${CMAKE_CURRENT_BINARY_DIR}/${file} DESTINATION bin)
 
 install(
    PROGRAMS regrid.pl

pre/remap_restart/remap_catchANDcn.py

@@ -18,9 +18,8 @@
 import mimetypes
 import netCDF4 as nc
 from remap_base import remap_base
-from remap_utils import get_landdir
-from remap_utils import get_geomdir
-from remap_utils import get_label
+from remap_params import *
+from remap_utils import *
 
 class catchANDcn(remap_base):
   def __init__(self, **configs):
@@ -60,25 +59,31 @@ def remap(self):
      surflay        = config['output']['surface']['surflay']
      in_tilefile    = config['input']['surface']['catch_tilefile']
 
+     if "gmao_SIteam/ModelData" in out_bc_base: 
+        assert  GEOS_SITE == "NAS", "wrong site to run the package" 
+
      if not in_tilefile :
         agrid        = config['input']['shared']['agrid']
         ogrid        = config['input']['shared']['ogrid']
         omodel       = config['input']['shared']['omodel']
         stretch      = config['input']['shared']['stretch']
-        bc_geomdir   = get_geomdir(in_bc_base, in_bc_version, agrid, ogrid, omodel, stretch)
+        EASE_grid    = config['input']['surface'].get('EASE_grid', None)
+
+        bc_geomdir   = get_geomdir(in_bc_base, in_bc_version, agrid=agrid, ogrid=ogrid, omodel=omodel, stretch=stretch, grid=EASE_grid)
         in_tilefile  = glob.glob(bc_geomdir + '/*.til')[0]
 
      agrid        = config['output']['shared']['agrid']
      ogrid        = config['output']['shared']['ogrid']
      omodel       = config['output']['shared']['omodel']
      stretch      = config['output']['shared']['stretch']
      out_tilefile = config['output']['surface']['catch_tilefile']
+     EASE_grid    = config['output']['surface'].get('EASE_grid', None)
 
      if not out_tilefile :
-        bc_geomdir  = get_geomdir(out_bc_base, out_bc_version, agrid, ogrid, omodel, stretch)
+        bc_geomdir  = get_geomdir(out_bc_base, out_bc_version, agrid=agrid, ogrid=ogrid, omodel=omodel, stretch=stretch, grid=EASE_grid)
         out_tilefile = glob.glob(bc_geomdir+ '/*.til')[0]
 
-     out_bc_landdir = get_landdir(out_bc_base, out_bc_version, agrid, ogrid, omodel, stretch)
+     out_bc_landdir = get_landdir(out_bc_base, out_bc_version, agrid=agrid, ogrid=ogrid, omodel=omodel, stretch=stretch, grid=EASE_grid)
 
  # determine NPE based on *approximate* number of input and output tile
 
@@ -104,14 +109,28 @@ def remap(self):
      else:
         NPE = 160
 
-     QOS = "#SBATCH --qos="+config['slurm']['qos']
-     TIME ="#SBATCH --time=1:00:00"
-     if NPE >= 160:
-       assert config['slurm']['qos'] != 'debug', "qos should be allnccs"
-       TIME = "#SBATCH --time=12:00:00"
-     PARTITION = "#SBATCH --partition=" + config['slurm']['partition']
-
-     account    = config['slurm']['account']
+     PARTITION =''
+     QOS  = config['slurm_pbs']['qos']
+     TIME  = "1:00:00"
+     if QOS != "debug": TIME="12:00:00"
+
+     NNODE = ''
+     job = ''
+     if GEOS_SITE == 'NAS':
+       job = "PBS"
+       CONSTRAINT = 'cas_ait'
+       NNODE = (NPE-1)//40 + 1
+     else:
+       job = "SLURM"
+       partition = config['slurm_pbs']['partition']
+       if (partition != ''):
+         PARTITION = "#SBATCH --partition=" + partition
+
+       CONSTRAINT = '"[cas|sky]"'
+       if BUILT_ON_SLES15:
+         CONSTRAINT = 'mil'
+
+     account    = config['slurm_pbs']['account']
      # even if the (MERRA-2) input restarts are binary, the output restarts will always be nc4 (remap_bin2nc.py)
      label = get_label(config) 
 
@@ -140,18 +159,10 @@ def remap(self):
      in_rstfile = dest
 
      log_name = out_dir+'/'+'mk_catchANDcn_log'
-     mk_catch_j_template = """#!/bin/csh -f
-#SBATCH --account={account}
-#SBATCH --ntasks={NPE}
-#SBATCH --job-name=mk_catchANDcn
-#SBATCH --output={log_name}
-{TIME}
-{QOS}
-{PARTITION}
-#
-
+     job_name = "mk_catchANDcn"
+     mk_catch_j_template = job_directive[job]+ \
+"""
 source {Bin}/g5_modules
-
 limit stacksize unlimited
 
 set esma_mpirun_X = ( {Bin}/esma_mpirun -np {NPE} )
@@ -165,36 +176,45 @@ def remap(self):
 
 """
      catch1script =  mk_catch_j_template.format(Bin = bindir, account = account, out_bcs = out_bc_landdir, \
-                  model = model, out_dir = out_dir, surflay = surflay, log_name = log_name, NPE = NPE,  \
+                  model = model, out_dir = out_dir, surflay = surflay, log_name = log_name, job_name = job_name, \
+                   NPE = NPE, NNODE=NNODE, \
                   in_wemin   = in_wemin, out_wemin = out_wemin, out_tilefile = out_tilefile, in_tilefile = in_tilefile, \
-                  in_rstfile = in_rstfile, out_rstfile = out_rstfile, time = yyyymmddhh_, TIME = TIME, PARTITION = PARTITION, QOS=QOS )
+                  in_rstfile = in_rstfile, out_rstfile = out_rstfile, time = yyyymmddhh_, TIME = TIME, QOS=QOS, CONSTRAINT=CONSTRAINT, PARTITION=PARTITION )
 
      script_name = './mk_catchANDcn.j'
 
      catch_scrpt = open(script_name,'wt')
      catch_scrpt.write(catch1script)
      catch_scrpt.close()
-
-     interactive = os.getenv('SLURM_JOB_ID', default = None)
+
+     interactive = None
+     if GEOS_SITE == 'NAS':
+       interactive = os.getenv('PBS_JOBID', default = None)
+     else:
+       interactive = os.getenv('SLURM_JOB_ID', default = None)
+
      if ( interactive ) :
        print('interactive mode\n')
-       ntasks = os.getenv('SLURM_NTASKS', default = None)
-       if ( not ntasks):
-         nnodes = int(os.getenv('SLURM_NNODES', default = '1'))
-         ncpus  = int(os.getenv('SLURM_CPUS_ON_NODE', default = '40'))
-         ntasks = nnodes * ncpus
-       ntasks = int(ntasks)
+       if GEOS_SITE != "NAS":
+         ntasks = os.getenv('SLURM_NTASKS', default = None)
+         if ( not ntasks):
+           nnodes = int(os.getenv('SLURM_NNODES', default = '1'))
+           ncpus  = int(os.getenv('SLURM_CPUS_ON_NODE', default = '40'))
+           ntasks = nnodes * ncpus
+         ntasks = int(ntasks)
 
-       if (ntasks < NPE):
-         print("\nYou should have at least {NPE} cores. Now you only have {ntasks} cores ".format(NPE=NPE, ntasks=ntasks))
+         if (ntasks < NPE):
+           print("\nYou should have at least {NPE} cores. Now you only have {ntasks} cores ".format(NPE=NPE, ntasks=ntasks))
 
        subprocess.call(['chmod', '755', script_name])
        print(script_name+  '  1>' + log_name  + '  2>&1')
        os.system(script_name + ' 1>' + log_name+ ' 2>&1')
-
+     elif GEOS_SITE == "NAS" :
+       print('qsub -W  block=true '+ script_name +'\n')
+       subprocess.call(['qsub', '-W','block=true', script_name])
      else:
-       print("sbatch -W " + script_name +"\n")
-       subprocess.call(['sbatch','-W', script_name])
+       print('sbatch -W '+ script_name +'\n')
+       subprocess.call(['sbatch', '-W', script_name])
 
      print( "cd " + cwdir)
      os.chdir(cwdir)
@@ -225,6 +245,92 @@ def copy_merra2(self):
     print("Copy file "+f +" to " + rst_dir)
     shutil.copy(f, dest)
 
+def ask_catch_questions():
+   questions =[
+        {
+            "type": "path",
+            "name": "input:shared:rst_dir",
+            "message": "Enter input directory with 'rs' and 'rc_out' as subdirectory:\n",
+        },
+
+        {
+            "type": "text",
+            "name": "input:shared:yyyymmddhh",
+            "message": (message_datetime + ".)\n"),
+            "validate": lambda text: len(text)==10 ,
+        },
+
+        {
+            "type": "path",
+            "name": "output:shared:out_dir",
+            "message": message_out_dir, 
+        },
+
+        {
+            "type": "text",
+            "name": "output:shared:expid",
+            "message": message_expid,
+            "default": "",
+        },  
+
+        {
+            "type": "select",
+            "name": "output:shared:bc_base",
+            "message": ("\nSelect " + message_bc_base_new),
+            "choices": choices_bc_base,
+            "default": get_default_bc_base(),
+        },
+
+        {
+            "type": "select",
+            "name": "output:shared:bc_version",
+            "message": message_bc_ops_new,
+            "choices": choices_bc_ops,
+            "default": "NL3",
+        },
+        {   
+            "type": "select",
+            "name": "output:shared:bc_version",
+            "message": message_bc_other_new,
+            "choices": choices_bc_other,
+            "when": lambda x:  x["output:shared:bc_version"] == 'Other',
+        },
+
+        {   
+            "type": "select",
+            "name": "output:surface:EASE_grid",
+            "message": "Select EASE grid for new restarts",
+            "choices": ['EASEv2_M03', 'EASEv2_M09', 'EASEv2_M25', 'EASEv2_M36']
+        },
+   ]
+
+   answers = questionary.prompt(questions)   
+   yyyy = answers['input:shared:yyyymmddhh'][0:4]
+   mm   = answers['input:shared:yyyymmddhh'][4:6]
+   dd   = answers['input:shared:yyyymmddhh'][6:8]
+   hh   = answers['input:shared:yyyymmddhh'][8:10]
+
+   rst_dir      = answers['input:shared:rst_dir']+'/rs/ens0000/Y'+yyyy +'/M'+mm+'/'
+   rst_file     = glob.glob(rst_dir+'*catch*_internal_rst.'+yyyy+mm+dd+'_'+hh+'00')[0]
+   idx1         = rst_file.find('catch')
+   idx2         = rst_file.find('_internal_rst')
+   catch_model  = rst_file[idx1:idx2]
+
+   tile_dir     = answers['input:shared:rst_dir']+'/rc_out/'
+   in_tilefiles = glob.glob(tile_dir+'*.til')
+   for file in in_tilefiles:
+      answers['input:surface:catch_tilefile'] = file
+      if 'MAPL_' in file:
+         answers['input:surface:catch_tilefile'] = file
+
+   answers['input:shared:rst_dir'] = rst_dir
+   answers['input:surface:catch_model'] =  catch_model
+   answers['input:surface:wemin']  = 13
+   answers['output:surface:wemin'] = 13
+
+   return answers
+
 if __name__ == '__main__' :
+
    catch = catchANDcn(params_file='remap_params.yaml')
    catch.remap()

pre/remap_restart/remap_command_line.py

@@ -72,10 +72,10 @@ def parse_args(program_description):
     p_command.add_argument('-out_bc_base',default="",         help="Boundary conditions base dir (w/o bc_version and resolution info) for new restarts")
     p_command.add_argument('-zoom',                           help= "Zoom parameter (search radius) for input surface restarts")
 
-    p_command.add_argument('-qos',        default="debug",    help="SLURM quality-of-service", choices=['debug', 'allnccs'])
+    p_command.add_argument('-qos',        default="debug",    help="slurm_pbs quality-of-service", choices=['debug', 'allnccs', 'normal'])
     account = get_account()
-    p_command.add_argument('-account',    default=account,    help="SLURM account")
-    p_command.add_argument('-partition',  default='compute',  help="SLURM partition")
+    p_command.add_argument('-account',    default=account,    help="slurm_pbs account")
+    p_command.add_argument('-partition',  default='',         help="slurm_pbs partition")
     p_command.add_argument('-rs',         default='3',        help="Flag indicating which restarts to regrid: 1 (upper air); 2 (surface); 3 (both)", choices=['1','2','3'])
 
     # Parse using parse_known_args so we can pass the rest to the remap scripts
@@ -146,9 +146,9 @@ def get_answers_from_command_line(cml):
    else:
      answers["output:surface:wemin"]   = wemin_default(answers['output:shared:bc_version'])
 
-   answers["slurm:account"]    = cml.account
-   answers["slurm:qos"]        = cml.qos
-   answers["slurm:partition"]  = cml.partition
+   answers["slurm_pbs:account"]    = cml.account
+   answers["slurm_pbs:qos"]        = cml.qos
+   answers["slurm_pbs:partition"]  = cml.partition
 
    return answers