Add GEOS-IT option for remap restarts

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Added remapping for GEOS-IT restarts 
 - Added new res C1120 
 - NOTE: If running on SLES15 remap tests will not be zero diff for GOCART RST but are zero diff for all other
 

pre/remap_restart/remap_analysis.py

@@ -23,7 +23,10 @@
 class analysis(remap_base):
   def __init__(self, **configs):
      super().__init__(**configs)
-     self.copy_merra2()
+     if self.config['input']['shared']['MERRA-2']:
+        self.copy_merra2()
+     if self.config['input']['shared']['GEOS-IT']:
+        self.copy_geosit()
 
   def remap(self):
      config = self.config
@@ -125,7 +128,8 @@ def remap(self):
      print( "cd " + cwdir)
      os.chdir(cwdir)
 
-     self.remove_merra2()
+     if self.config['input']['shared']['MERRA-2']:
+        self.remove_merra2()
 
   def get_grid_kind(this, grid):
      hgrd = {}
@@ -243,3 +247,4 @@ def copy_merra2(self):
 if __name__ == '__main__' :
    ana = analysis(params_file='remap_params.yaml')
    ana.remap()
+   ana.remove_geosit()

pre/remap_restart/remap_base.py

@@ -33,6 +33,10 @@ def remove_merra2(self):
     if self.config['input']['shared']['MERRA-2']:
       print(" remove temporary folder that contains MERRA-2 archived files ... \n")
       subprocess.call(['/bin/rm', '-rf', self.config['input']['shared']['rst_dir']])
+  def remove_geosit(self):
+    if self.config['input']['shared']['GEOS-IT']:
+      print(" remove temporary folder that contains GEOS-IT archived files ... \n")
+      subprocess.call(['/bin/rm', '-rf', self.config['input']['shared']['rst_dir']])
 
   def copy_without_remap(self, restarts_in, compared_file_in, compared_file_out, suffix, catch=False):
 #
@@ -89,3 +93,60 @@ def copy_without_remap(self, restarts_in, compared_file_in, compared_file_out, s
        return True
 
      return False
+
+  def copy_geosit(self):
+    if not self.config['input']['shared']['GEOS-IT']:
+        return 
+
+    expid = self.config['input']['shared']['expid']
+    yyyymmddhh_ = str(self.config['input']['shared']['yyyymmddhh'])
+    yyyy_ = yyyymmddhh_[0:4]
+    mm_   = yyyymmddhh_[4:6]
+    day_  = yyyymmddhh_[6:8]  # Extract the day from yyyymmddhh_
+
+    time_suffix = '_21z' 
+    time_suffix_nc4 = '_2100z' 
+
+    geos_it_rst_dir = '/discover/nobackup/projects/gmao/geos-it/dao_ops/archive/' + expid + '/rs/Y' + yyyy_ + '/M' + mm_ + '/'
+    rst_dir = self.config['input']['shared']['rst_dir'] + '/'
+    os.makedirs(rst_dir, exist_ok=True)
+
+    print('Stage GEOS-IT restarts \n from \n    ' + geos_it_rst_dir + '\n to\n    ' + rst_dir + '\n')
+
+    # Only use the specific day from yyyymmddhh_
+    filename = f'{expid}.rst.{yyyy_}{mm_}{day_}{time_suffix}.tar'
+    src_file = os.path.join(geos_it_rst_dir, filename)
+    dest_file = os.path.join(rst_dir, filename)
+    if os.path.exists(dest_file):
+       print('tar file is copied and untar, no need to copy')
+       return
+    if os.path.exists(src_file):
+       print(f"Copying file {src_file} to {dest_file}")
+       shutil.copy(src_file, dest_file)
+
+       # Untar the .tar file using the tar command
+       if os.path.exists(dest_file):
+           print(f"Untarring {dest_file} to {rst_dir}")
+           try:
+               subprocess.run(['tar', '-xf', dest_file, '-C', rst_dir], check=True)
+               print(f"Untarred {dest_file} successfully.")
+           except subprocess.CalledProcessError as e:
+               print(f"Error untarring {dest_file}: {e}")
+
+        # Optionally remove the tar file after extraction
+        #os.remove(dest_file)
+    else:
+       print(f"Tar file {src_file} does not exist.")
+
+    # Now handle the additional .nc4 file
+    cp_agcm_import_rst = self.config['output']['air']['agcm_import_rst']
+    if (cp_agcm_import_rst) :
+       nc4_filename = f'{expid}.agcm_import_rst.{yyyy_}{mm_}{day_}{time_suffix_nc4}.nc4'
+       src_nc4_file = os.path.join(geos_it_rst_dir, nc4_filename)
+       dest_nc4_file = os.path.join(rst_dir, nc4_filename)
+       # Copy the .nc4 file if it exists
+       if os.path.exists(src_nc4_file):
+          print(f"Copying .nc4 file {src_nc4_file} to {dest_nc4_file}")
+          shutil.copy(src_nc4_file, dest_nc4_file)
+       else:
+          print(f" agcm_import_rst file {src_nc4_file} does not exist.")

pre/remap_restart/remap_catchANDcn.py

@@ -23,7 +23,10 @@
 class catchANDcn(remap_base):
   def __init__(self, **configs):
      super().__init__(**configs)
-     self.copy_merra2()
+     if self.config['input']['shared']['MERRA-2']:
+        self.copy_merra2()
+     if self.config['input']['shared']['GEOS-IT']:
+        self.copy_geosit()
 
   def remap(self):
      if not self.config['output']['surface']['remap_catch']:
@@ -222,7 +225,8 @@ def remap(self):
      print( "cd " + cwdir)
      os.chdir(cwdir)
 
-     self.remove_merra2()
+     if self.config['input']['shared']['MERRA-2']:
+        self.remove_merra2()
 
   def copy_merra2(self):
     if not self.config['input']['shared']['MERRA-2']:
@@ -342,7 +346,7 @@ def has_catch_rst(text):
         {
             "type": "select",
             "name": "output:shared:ogrid",
-            "message": message_ogrid_in,
+            "message": message_ogrid_new,
             "choices": choices_ogrid_data,
             "default": lambda x: data_ocean_default(x.get('output:shared:agrid')),
             "when": lambda x : x['output:surface:EASE_grid'] == 'Cubed-Sphere',
@@ -436,8 +440,7 @@ def remap_land_only():
 
    catch = catchANDcn(params_file=config_yaml)
    catch.remap()
-
-
+   catch.remove_geosit()
 
 if __name__ == '__main__' :
 

pre/remap_restart/remap_command_line.py

@@ -39,6 +39,7 @@ def parse_args(program_description):
      help = "Use command line as input",
     )
     p_command.add_argument('-merra2', action='store_true', default= False, help='use merra2 restarts')
+    p_command.add_argument('-geosit', action='store_true', default= False, help='use GEOSIT restarts')
 
     p_command.add_argument('-ymdh',              help='yyyymmddhh year month date hour of input and new restarts')
     p_command.add_argument('-grout',             help='Grid ID/resolution of new restarts, format C[xxx] (cubed-sphere only for now)')
@@ -91,20 +92,25 @@ def get_answers_from_command_line(cml):
 
    answers = {}
    answers["input:shared:MERRA-2"]     = cml.merra2
+   answers["input:shared:GEOS-IT"]     = cml.geosit
    answers["input:shared:yyyymmddhh"]  = cml.ymdh
-   answers["input:shared:omodel"]      = cml.ocnmdlin
    answers["output:shared:out_dir"]    = os.path.abspath(cml.out_dir + '/')
    if  cml.merra2:
       init_merra2(answers)
+   elif  cml.geosit:
+      init_geosit(answers)
    else:   
-      answers["input:shared:bc_version"]   = cml.bcvin
-      answers["input:surface:catch_model"] = cml.catch_model
-      answers["input:shared:rst_dir"]      = os.path.abspath(cml.rst_dir + '/')
+      answers["input:shared:bc_base"]     =  cml.in_bc_base
+      answers["input:shared:omodel"]      = cml.ocnmdlin
+      answers["input:shared:bc_version"]  = cml.bcvin
+      answers["input:surface:catch_model"]= cml.catch_model
+      answers["input:shared:stretch"]     = cml.in_stretch
+      answers["input:shared:rst_dir"]     = os.path.abspath(cml.rst_dir + '/')
       fvcore_info(answers) 
-      ogrid                            = cml.oceanin
+      ogrid                               = cml.oceanin
       if ogrid == "CS":
          ogrid = answers["input:shared:agrid"]
-      answers["input:shared:ogrid"]    = ogrid
+      answers["input:shared:ogrid"]       = ogrid
 
    answers["output:shared:agrid"]      = cml.grout
    answers["output:air:nlevel"]        = cml.levsout
@@ -117,11 +123,9 @@ def get_answers_from_command_line(cml):
       ogrid = answers["output:shared:agrid"]
    answers["output:shared:ogrid"]      = ogrid
 
-   answers["input:shared:bc_base"]  =  cml.in_bc_base
    answers["output:shared:bc_base"] =  cml.out_bc_base
 
    answers["output:shared:stretch"]    = cml.out_stretch
-   answers["input:shared:stretch"]     = cml.in_stretch
    answers["output:analysis:bkg"]      = not cml.nobkg
    answers["output:analysis:lcv"]      = not cml.nolcv
    if cml.rs == '1':

pre/remap_restart/remap_lake_landice_saltwater.py

@@ -12,6 +12,7 @@
 import os
 import subprocess as sp
 import shutil
+import subprocess
 import glob
 import ruamel.yaml
 import shlex
@@ -23,7 +24,10 @@
 class lake_landice_saltwater(remap_base):
   def __init__(self, **configs):
      super().__init__(**configs)
-     self.copy_merra2()
+     if self.config['input']['shared']['MERRA-2']:
+        self.copy_merra2()
+     if self.config['input']['shared']['GEOS-IT']:
+        self.copy_geosit()
 
   def remap(self):
      if not self.config['output']['surface']['remap_water']:
@@ -170,7 +174,8 @@ def remap(self):
      print('cd ' + cwdir)
      os.chdir(cwdir)
 
-     self.remove_merra2()
+     if self.config['input']['shared']['MERRA-2']:
+        self.remove_merra2()
 
   def run_and_log(self, cmd, log_name):
      print('\n'+cmd)
@@ -247,3 +252,4 @@ def copy_merra2(self):
 if __name__ == '__main__' :
    lls = lake_landice_saltwater(params_file='remap_params.yaml')
    lls.remap()
+   lls.remove_geosit()

pre/remap_restart/remap_params.tpl

@@ -10,6 +10,7 @@ input:
     hydrostatic: 0
   shared:
     MERRA-2: false
+    GEOS-IT: false
     stretch: false
     # (coupled) ocean model: data, MOM5, MOM6
     omodel: data
@@ -44,6 +45,7 @@ output:
     # remap upper air or not
     remap: true
     nlevel:
+    agcm_import_rst: false
   surface:
     split_saltwater: false
     surflay: 50.