Tests for weight generation. #52

.gitignore

@@ -7,6 +7,11 @@ access-om2.err
 access-om2.out
 pbs_logs
 
+tools/PET*
+tools/*.nc
+tools/contrib/esmf
+tools/contrib/bin
+
 *.swp
 input_*.tar.gz
 input

test/exp_test_helper.py

@@ -9,6 +9,8 @@
 import glob
 import time
 
+from util import wait_for_qsub, get_git_hash
+
 class ExpTestHelper(object):
 
     def __init__(self, exp_name):
@@ -111,39 +113,12 @@ def do_basic_access_run(self, exp, model='cm'):
                 assert('********** End of MATM **********' in s)
 
 
-    def wait(self, run_id):
-        """
-        Wait for the qsub job to terminate.
-        """
-
-        while True:
-            time.sleep(5)
-            qsub_out = ''
-            try:
-                qsub_out = sp.check_output(['qstat', run_id], stderr=sp.STDOUT)
-            except sp.CalledProcessError as err:
-                qsub_out = err.output
-
-            if 'Job has finished' in qsub_out:
-                break
-
-    def git_hash(self, src_dir):
-        """
-        Get the git hash of src_dir.
-        """
-        mydir = os.getcwd()
-        os.chdir(src_dir)
-        ghash = sp.check_output(['git', 'rev-parse', 'HEAD'])[:8]
-        os.chdir(mydir)
-
-        return ghash
-
     def copy_to_bin(self, src_dir, wildcard):
         exes = glob.glob(wildcard)
         if exes == []:
             return 1
 
-        ghash = self.git_hash(src_dir)
+        ghash = get_git_hash(src_dir)
 
         for e in exes:
             eb = os.path.basename(e)
@@ -225,7 +200,7 @@ def force_run(self):
             print('Error: call to payu-run failed.', file=sys.stderr)
             return 1, None, None, None
 
-        self.wait(run_id)
+        wait_for_qsub(run_id)
         run_id = run_id.split('.')[0]
 
         output_files = []
@@ -261,7 +236,7 @@ def force_run(self):
 
         # Wait for the collate to complete.
         run_id = m.group(1)
-        self.wait(run_id)
+        wait_for_qsub(run_id)
 
         # Return files created by qsub so caller can read or delete.
         collate_files = os.path.join(self.exp_path, '*.[oe]{}'.format(run_id))

test/test_remapping_weights.py

@@ -7,6 +7,8 @@
 import numba
 import netCDF4 as nc
 
+from util import wait_for_qsub
+
 EARTH_RADIUS = 6370997.0
 
 def calc_regridding_err(weights, src, dest):
@@ -101,8 +103,7 @@ def remap(src_data, weights, dest_shape):
                                        n_s, n_b, row, col, s)
     return dest_data
 
-
-class TestBuild():
+class TestRemap():
 
     def test_jra55_to_01deg(self):
 
@@ -128,3 +129,46 @@ def test_jra55_to_1deg(self):
     def test_jra55_to_025deg(self):
         pass
 
+
+class TestCreateWeights():
+    """
+    Create weights and compare to existing.
+    """
+
+    def test_build_esmf(self):
+        """
+        Build ESMF
+        """
+
+        curdir = os.getcwd()
+        contrib_dir = os.path.join(curdir, 'tools', 'contrib')
+        os.chdir(contrib_dir)
+        ret = sp.call('build_esmf_on_raijin.sh')
+        os.chdir(curdir)
+
+        assert ret == 0
+        assert os.path.exists(os.path.join(contrib_dir, 'bin', ESMF_RegridWeightGen)
+
+
+    def test_create_weights(self):
+        """
+        Create weights
+        """
+
+        cmd = os.path.join('tools', 'make_remap_weights.sh'))
+        qsub_id = sp.check_output(['qsub', cmd])
+
+        # Wait for job to complete.
+        wait_for_qsub(qsub_id.strip())
+
+        # Check that weights files have been created.
+        ocn = ['MOM1', 'MOM025', 'MOM01']
+        atm = ['JRA55', 'JRA55_runoff', 'CORE2']
+        method = ['patch', 'conserve2nd']
+
+        for o in ocn:
+            for a in atm:
+                for m in method:
+                    filename = '{}_{}_{}.nc'.format(a, o, m)
+                    assert os.path.exists(os.path.join('tools', filename))
+

test/util.py

@@ -0,0 +1,28 @@
+
+def wait_for_qsub(run_id):
+    """
+    Wait for the qsub job to terminate.
+    """
+
+    while True:
+        time.sleep(5)
+        qsub_out = ''
+        try:
+            qsub_out = sp.check_output(['qstat', run_id], stderr=sp.STDOUT)
+        except sp.CalledProcessError as err:
+            qsub_out = err.output
+
+        if 'Job has finished' in qsub_out:
+            break
+
+def get_git_hash(src_dir):
+    """
+    Get the git hash of src_dir.
+    """
+    mydir = os.getcwd()
+    os.chdir(src_dir)
+    ghash = sp.check_output(['git', 'rev-parse', 'HEAD'])[:8]
+    os.chdir(mydir)
+
+    return ghash
+

tools/contrib/build_esmf_on_raijin.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+git archive --remote=git://git.code.sf.net/p/esmf/esmf --format=tar --prefix=esmf/ ESMF_7_1_0_beta_snapshot_38 | tar xf -
+module load netcdf/4.4.1.1
+module load intel-fc/17.0.1.132
+mkdir -p bin
+cd esmf
+export ESMF_DIR=$(pwd)
+export ESMF_F90COMPILER=ifort
+export ESMF_F90LINKER=ifort
+export ESMF_NETCDF="split"
+export ESMF_NETCDF_INCLUDE=$NETCDF_ROOT/include
+export ESMF_NETCDF_LIBPATH=$NETCDF_ROOT/lib
+export ESMF_NETCDF_LIBS="-lnetcdff -lnetcdf"
+make
+cd src/apps/ESMF_RegridWeightGen
+make
+cd ../../../../
+cp esmf/apps/appsO/*/ESMF_RegridWeightGen bin/
+export PATH=$(pwd)/bin:$PATH

tools/contrib/build_esmf_on_ubuntu.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+mkdir bin
+sudo apt-get install libnetcdf-dev libnetcdff-dev
+wget http://s3-ap-southeast-2.amazonaws.com/dp-drop/ocean-regrid/contrib/esmf_7_0_0_src.tar.gz
+tar zxvf esmf_7_0_0_src.tar.gz
+cd esmf
+export ESMF_DIR=$(pwd)
+export ESMF_NETCDF="split"
+export ESMF_NETCDF_INCLUDE=/usr/include/
+export ESMF_NETCDF_LIBPATH=/usr/lib/x86_64-linux-gnu/
+export ESMF_NETCDF_LIBS="-lnetcdff -lnetcdf"
+make
+cd src/apps/ESMF_RegridWeightGen
+make
+cd ../../../../
+cp esmf/apps/appsO/*/ESMF_RegridWeightGen bin/
+export PATH=$(pwd)/bin:$PATH

tools/make_remap_weights.py

@@ -30,7 +30,7 @@
 def convert_to_scrip_output(weights):
 
     _, new_weights = tempfile.mkstemp(suffix='.nc')
-    # FIXME: So that ncrename doesn't prompt for overwrite.
+    # So that ncrename doesn't prompt for overwrite.
     os.remove(new_weights)
 
     cmd = 'ncrename -d n_a,src_grid_size -d n_b,dst_grid_size -d n_s,num_links -d nv_a,src_grid_corners -d nv_b,dst_grid_corners -v yc_a,src_grid_center_lat -v yc_b,dst_grid_center_lat -v xc_a,src_grid_center_lon -v xc_b,dst_grid_center_lon -v yv_a,src_grid_corner_lat -v xv_a,src_grid_corner_lon -v yv_b,dst_grid_corner_lat -v xv_b,dst_grid_corner_lon -v mask_a,src_grid_imask -v mask_b,dst_grid_imask -v area_a,src_grid_area -v area_b,dst_grid_area -v frac_a,src_grid_frac -v frac_b,dst_grid_frac -v col,src_address -v row,dst_address {} {}'.format(weights, new_weights)
@@ -50,44 +50,42 @@ def convert_to_scrip_output(weights):
     return new_weights
 
 
-def create_weights(src_grid, dest_grid, method='conserve',
+def create_weights(src_grid, dest_grid, npes, method,
                    ignore_unmapped=False,
                    unmasked_src=True, unmasked_dest=False):
 
-    _, src_grid_scrip = tempfile.mkstemp(suffix='.nc')
-    _, dest_grid_scrip = tempfile.mkstemp(suffix='.nc')
-    _, regrid_weights = tempfile.mkstemp(suffix='.nc')
+    my_dir = os.path.dirname(os.path.realpath(__file__))
+
+    _, src_grid_scrip = tempfile.mkstemp(suffix='.nc', dir=my_dir)
+    _, dest_grid_scrip = tempfile.mkstemp(suffix='.nc', dir=my_dir)
+    _, regrid_weights = tempfile.mkstemp(suffix='.nc', dir=my_dir)
 
     if unmasked_src:
-        src_grid.write_scrip(src_grid_scrip,
+        src_grid.write_scrip(src_grid_scrip, write_test_scrip=False,
                             mask=np.zeros_like(src_grid.mask_t, dtype=int))
     else:
-        src_grid.write_scrip(src_grid_scrip)
+        src_grid.write_scrip(src_grid_scrip, write_test_scrip=False)
 
     if unmasked_dest:
-        dest_grid.write_scrip(dest_grid_scrip,
+        dest_grid.write_scrip(dest_grid_scrip, write_test_scrip=False,
                               mask=np.zeros_like(dest_grid.mask_t, dtype=int))
     else:
-        dest_grid.write_scrip(dest_grid_scrip)
+        dest_grid.write_scrip(dest_grid_scrip, write_test_scrip=False)
 
     if ignore_unmapped:
         ignore_unmapped = ['--ignore_unmapped']
     else:
         ignore_unmapped = []
 
-    mpirun = ['mpirun', '-np', str(mp.cpu_count() // 2)]
-
-    my_dir = os.path.dirname(os.path.realpath(__file__))
     esmf = os.path.join(my_dir, 'contrib', 'bin', 'ESMF_RegridWeightGen')
     if not os.path.exists(esmf):
         esmf = 'ESMF_RegridWeightGen'
 
     try:
-        cmd = mpirun + [esmf] + [
-                        '-s', src_grid_scrip,
-                        '-d', dest_grid_scrip, '-m', method,
-                        '-w', regrid_weights] + ignore_unmapped
-        print(cmd)
+        cmd = ['mpirun', '-np', npes] + [esmf] + \
+              ['-s', src_grid_scrip,
+               '-d', dest_grid_scrip, '-m', method,
+               '-w', regrid_weights] + ignore_unmapped
         sp.check_output(cmd)
     except sp.CalledProcessError as e:
         print("Error: ESMF_RegridWeightGen failed ret {}".format(e.returncode),
@@ -112,11 +110,11 @@ def find_grid_defs(input_dir, jra55_input):
     """
 
     d = {}
-    d['MOM1'] = (os.path.join(input_dir, 'mom_1deg', 'ocean_hgrid.nc'), 
+    d['MOM1'] = (os.path.join(input_dir, 'mom_1deg', 'ocean_hgrid.nc'),
                  os.path.join(input_dir, 'mom_1deg', 'ocean_mask.nc'))
-    d['MOM025'] = (os.path.join(input_dir, 'mom_025deg', 'ocean_hgrid.nc'), 
+    d['MOM025'] = (os.path.join(input_dir, 'mom_025deg', 'ocean_hgrid.nc'),
                    os.path.join(input_dir, 'mom_025deg', 'ocean_mask.nc'))
-    d['MOM01'] = (os.path.join(input_dir, 'mom_01deg', 'ocean_hgrid.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(input_dir, 'core_nyf', 't_10.0001.nc')
     d['JRA55'] = os.path.join(jra55_input, 'RYF.t_10.1990_1991.nc')
@@ -134,16 +132,18 @@ def main():
                         The JRA55 input directory.""")
     parser.add_argument('--atm', default=None, help="""
                         Atmosphere grid to regrid from, can be one of:
-                        CORE2, JRA55, JRA55_river""")
+                        CORE2, JRA55, JRA55_runoff""")
     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.""")
+    parser.add_argument('--npes', default=1, help="""
+                        The number of PEs to use.""")
 
     args = parser.parse_args()
-    atm_options = ['CORE2', 'JRA55', 'JRA55_runoff']
-    ocean_options = ['MOM1', 'MOM01', 'MOM025']
+    atm_options = ['JRA55', 'JRA55_runoff', 'CORE2']
+    ocean_options = ['MOM1', 'MOM025', 'MOM01']
     method_options = ['patch', 'conserve2nd']
 
     if args.atm is None:
@@ -171,21 +171,21 @@ def main():
 
     grid_file_dict = find_grid_defs(args.input_dir, args.jra55_input)
 
-    for atm in args.atm:
-        for ocean in args.ocean:
+    for ocean in args.ocean:
+        for atm in args.atm:
             for method in args.method:
 
                 if atm == 'CORE2':
                     src_grid = Core2Grid(grid_file_dict[atm])
                 elif atm == 'JRA55':
                     src_grid = Jra55Grid(grid_file_dict[atm])
                 else:
-                    src_grid = Jra552RiverGrid(grid_file_dict[atm])
+                    src_grid = Jra55RiverGrid(grid_file_dict[atm])
 
-                dest_grid = MomGrid.fromfile(grid_file_dict[ocean][0], 
-                                             mask_file=grid_file_dict[ocean][1]) 
+                dest_grid = MomGrid.fromfile(grid_file_dict[ocean][0],
+                                             mask_file=grid_file_dict[ocean][1])
 
-                weights = create_weights(src_grid, dest_grid, method=method)
+                weights = create_weights(src_grid, dest_grid, args.npes, method)
                 weights = convert_to_scrip_output(weights)
 
                 shutil.move(weights, '{}_{}_{}.nc'.format(atm, ocean, method))

tools/make_remap_weights.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+#PBS -P x77
+#PBS -q normal
+#PBS -l ncpus=128,mem=496GB,walltime=02:00:00,jobfs=100GB
+#PBS -l wd
+
+time ./make_remap_weights.py /short/x77/nah599/access-om2/input/ /g/data1/ua8/JRA55-do/RYF/v1-3/  --npes 128
+