update to Nic's /g/data/v45/nah599/access-om2/tools/make_remap_weights.* to support ERA5; COSIMA/esmgrids#4 ; #242

Author: aekiss

tools/make_remap_weights.py

@@ -19,6 +19,7 @@
 from esmgrids.mom_grid import MomGrid  # noqa
 from esmgrids.core2_grid import Core2Grid  # noqa
 from esmgrids.jra55_grid import Jra55Grid  # noqa
+from esmgrids.era5_grid import Era5Grid  # noqa
 from esmgrids.jra55_river_grid import Jra55RiverGrid  # noqa
 from esmgrids.daitren_runoff_grid import DaitrenRunoffGrid  # noqa
 
@@ -121,9 +122,9 @@ def create_weights(src_grid, dest_grid, npes, method,
     return regrid_weights
 
 
-def find_grid_defs(input_dir, jra55_input, core_input):
+def find_ocean_grid_defs(input_dir):
     """
-    Return a dictionary containing the grid definition files.
+    Return a dictionary containing ocean grid definition files.
     """
 
     d = {}
@@ -133,51 +134,39 @@ def find_grid_defs(input_dir, jra55_input, core_input):
                    os.path.join(input_dir, 'mom_025deg', 'ocean_mask.nc'))
     d['MOM01'] = (os.path.join(input_dir, 'mom_01deg', 'ocean_hgrid.nc'),
                   os.path.join(input_dir, 'mom_01deg', 'ocean_mask.nc'))
-    d['CORE2'] = os.path.join(core_input, 't_10.0001.nc')
-    d['JRA55'] = os.path.join(jra55_input, 'RYF.t_10.1990_1991.nc')
-    d['JRA55_runoff'] = os.path.join(jra55_input,
-                                     'RYF.runoff_all.1990_1991.nc')
-    d['Daitren_runoff'] = os.path.join(core_input, 'runoff.daitren.clim.10FEB2011.nc')
-
     return d
 
 
 def main():
 
     parser = argparse.ArgumentParser()
-    parser.add_argument('input_dir', help="""
+    parser.add_argument('--accessom2_input_dir', required=True, help="""
                         The ACCESS-OM2 input directory.""")
-    parser.add_argument('jra55_input', help="""
-                        The JRA55 input directory.""")
-    parser.add_argument('core_input', help="""
-                        The CORE input directory.""")
-    parser.add_argument('--atm', default=None, help="""
+    parser.add_argument('--atm_forcing_file', required=True, help="""
+                        An atmospheric or runoff forcing field.""")
+    parser.add_argument('--atm', required=True, default=None, help="""
                         Atmosphere grid to regrid from, can be one of:
-                        CORE2, JRA55, JRA55_runoff, Daitren_runoff""")
+                        CORE2, JRA55, JRA55_runoff, Daitren_runoff, ERA5""")
     parser.add_argument('--ocean', default=None, help="""
                         Ocean grid to regrid to, can be one of:
                         MOM1, MOM01, MOM025""")
     parser.add_argument('--method', default=None, help="""
                         The interpolation method to use, can be patch, conserve or conserve2nd""")
     parser.add_argument('--npes', default=None, help="""
                         The number of PEs to use.""")
-    parser.add_argument('--unmask_dest', 
+    parser.add_argument('--unmask_dest',
                         action='store_true',
                         help='Ignore destination grid mask')
 
     args = parser.parse_args()
-    atm_options = ['JRA55', 'JRA55_runoff', 'CORE2', 'Daitren_runoff']
+    atm_options = ['JRA55', 'JRA55_runoff', 'CORE2', 'Daitren_runoff', 'ERA5']
     ocean_options = ['MOM1', 'MOM025', 'MOM01']
     method_options = ['patch', 'conserve', 'conserve2nd']
 
-    if args.atm is None:
-        args.atm = atm_options
-    else:
-        if args.atm not in atm_options:
-            print("Error: bad atm grid.", file=sys.stderr)
-            parser.print_help()
-            return 1
-        args.atm = [args.atm]
+    if args.atm not in atm_options:
+        print("Error: bad atm grid.", file=sys.stderr)
+        parser.print_help()
+        return 1
 
     if args.ocean is None:
         args.ocean = ocean_options
@@ -197,37 +186,38 @@ def main():
         import multiprocessing as mp
         args.npes = mp.cpu_count() // 2
 
-    grid_file_dict = find_grid_defs(args.input_dir, args.jra55_input, args.core_input)
+    ocean_grid_file_dict = find_ocean_grid_defs(args.accessom2_input_dir)
 
     for ocean in args.ocean:
-        umask_file = grid_file_dict[ocean][1]
-        dest_grid = MomGrid.fromfile(grid_file_dict[ocean][0],
+        umask_file = ocean_grid_file_dict[ocean][1]
+        dest_grid = MomGrid.fromfile(ocean_grid_file_dict[ocean][0],
                                      mask_file=umask_file)
-        for atm in args.atm:
-
-            if atm == 'CORE2':
-                src_grid = Core2Grid(grid_file_dict[atm])
-            elif atm == 'Daitren_runoff':
-                src_grid = DaitrenRunoffGrid(grid_file_dict[atm])
-            elif atm == 'JRA55':
-                src_grid = Jra55Grid(grid_file_dict[atm])
-            elif atm == 'JRA55_runoff':
-                src_grid = Jra55RiverGrid(grid_file_dict[atm], calc_areas=False)
-            else:
-                print('Unrecognised atmosphere grid: {}'.format(atm))
-                return 1
 
-            for method in args.method:
+        if args.atm == 'CORE2':
+            src_grid = Core2Grid(args.atm_forcing_file)
+        elif args.atm == 'Daitren_runoff':
+            src_grid = DaitrenRunoffGrid(args.atm_forcing_file)
+        elif args.atm == 'JRA55':
+            src_grid = Jra55Grid(args.atm_forcing_file)
+        elif args.atm == 'JRA55_runoff':
+            src_grid = Jra55RiverGrid(args.atm_forcing_file, calc_areas=False)
+        elif args.atm == 'ERA5':
+            src_grid = Era5Grid(args.atm_forcing_file)
+        else:
+            print('Unrecognised atmosphere grid: {}'.format(atm))
+            return 1
+
+        for method in args.method:
 
-                weights = create_weights(src_grid, dest_grid, args.npes,
-                                         method, unmasked_dest=args.unmask_dest)
-                if not weights:
-                    return 1
-                weights = convert_to_scrip_output(weights)
-                if not weights:
-                    return 1
+            weights = create_weights(src_grid, dest_grid, args.npes,
+                                     method, unmasked_dest=args.unmask_dest)
+            if not weights:
+                return 1
+            weights = convert_to_scrip_output(weights)
+            if not weights:
+                return 1
 
-                shutil.move(weights, '{}_{}_{}.nc'.format(atm, ocean, method))
+            shutil.move(weights, '{}_{}_{}.nc'.format(args.atm, ocean, method))
 
     return 0
 

tools/make_remap_weights.sh

@@ -7,10 +7,12 @@
 module purge
 module load openmpi
 module load nco
-module load esmf/7.1.0r-intel
+module load esmf
 module use /g/data/hh5/public/modules
 module load conda/analysis3
 
+ulimit -s unlimited
+
 # Make all 1 deg weights.
 time ./make_remap_weights.py /short/x77/nah599/access-om2/input/ /g/data/ua8/JRA55-do/RYF/v1-3/ /short/x77/nah599/access-om2/input/yatm_1deg/ --ocean MOM1 --npes 256