Merge pull request #3 from xylar/multiple_mpi_tasks

Supports building mapping files with multiple MPI tasks
MPAS-Dev · Nov 18, 2019 · d25ab6f · d25ab6f
2 parents f96b6cf + a595416
commit d25ab6f
Show file tree

Hide file tree

Showing 17 changed files with 27 additions and 9 deletions.
diff --git a/examples/make_mpas_to_Antarctic_stereo_mapping.py b/examples/make_mpas_to_Antarctic_stereo_mapping.py
@@ -39,11 +39,12 @@
                                      projection='antarctic')
 outGridName = outDescriptor.meshName
 
-mappingFileName = 'map_{}_to_{}.nc'.format(inGridName, outGridName)
+mappingFileName = 'map_{}_to_{}_conserve.nc'.format(inGridName, outGridName)
 
 remapper = Remapper(inDescriptor, outDescriptor, mappingFileName)
 
-remapper.build_mapping_file(method='bilinear')
+# conservative remapping with 4 MPI tasks (using mpirun)
+remapper.build_mapping_file(method='conserve', mpiTasks=4)
 
 outFileName = 'temp_{}.nc'.format(outGridName)
 ds = xarray.open_dataset(inGridFileName)

diff --git a/examples/make_mpas_to_lat_lon_mapping.py b/examples/make_mpas_to_lat_lon_mapping.py
@@ -37,7 +37,7 @@
 outDescriptor = get_lat_lon_descriptor(dLon=0.5, dLat=0.5)
 outGridName = outDescriptor.meshName
 
-mappingFileName = 'map_{}_to_{}.nc'.format(inGridName, outGridName)
+mappingFileName = 'map_{}_to_{}_bilinear.nc'.format(inGridName, outGridName)
 
 
 remapper = Remapper(inDescriptor, outDescriptor, mappingFileName)

diff --git a/examples/remap_stereographic.py b/examples/remap_stereographic.py
@@ -25,9 +25,16 @@
                     help="Output file name")
 parser.add_argument('-r', dest='resolution', required=True, type=float,
                     help="Output resolution")
+parser.add_argument('-m', dest='method', required=False, default="bilinear",
+                    help="Method: {'bilinear', 'neareststod', 'conserve'}")
+parser.add_argument('-t', dest='mpiTasks', required=False, type=int, default=1,
+                    help="Number of MPI tasks (default = 1)")
 args = parser.parse_args()
 
 
+if args.method not in ['bilinear', 'neareststod', 'conserve']:
+    raise ValueError('Unexpected method {}'.format(args.method))
+
 dsIn = xarray.open_dataset(args.inFileName)
 
 x = dsIn.x.values
@@ -56,11 +63,12 @@
 outDescriptor = ProjectionGridDescriptor.create(projection, xOut, yOut,
                                                 outMeshName)
 
-mappingFileName = 'map_{}_to_{}.nc'.format(inMeshName, outMeshName)
+mappingFileName = 'map_{}_to_{}_{}.nc'.format(inMeshName, outMeshName,
+                                              args.method)
 
 remapper = Remapper(inDescriptor, outDescriptor, mappingFileName)
 
-remapper.build_mapping_file(method='bilinear')
+remapper.build_mapping_file(method=args.method, mpiTasks=args.mpiTasks)
 
 dsOut = remapper.remap(dsIn, renormalizationThreshold=0.01)
 dsOut.to_netcdf(args.outFileName)
diff --git a/pyremap/__pycache__/__init__.cpython-36.pyc b/pyremap/__pycache__/__init__.cpython-36.pyc
diff --git a/pyremap/__pycache__/__init__.cpython-37.pyc b/pyremap/__pycache__/__init__.cpython-37.pyc
diff --git a/pyremap/__pycache__/descriptor.cpython-36.pyc b/pyremap/__pycache__/descriptor.cpython-36.pyc
diff --git a/pyremap/__pycache__/descriptor.cpython-37.pyc b/pyremap/__pycache__/descriptor.cpython-37.pyc
diff --git a/pyremap/__pycache__/grid.cpython-36.pyc b/pyremap/__pycache__/grid.cpython-36.pyc
diff --git a/pyremap/__pycache__/interp1d.cpython-36.pyc b/pyremap/__pycache__/interp1d.cpython-36.pyc
diff --git a/pyremap/__pycache__/polar.cpython-36.pyc b/pyremap/__pycache__/polar.cpython-36.pyc
diff --git a/pyremap/__pycache__/polar.cpython-37.pyc b/pyremap/__pycache__/polar.cpython-37.pyc
diff --git a/pyremap/__pycache__/remapper.cpython-36.pyc b/pyremap/__pycache__/remapper.cpython-36.pyc
diff --git a/pyremap/__pycache__/remapper.cpython-37.pyc b/pyremap/__pycache__/remapper.cpython-37.pyc
diff --git a/pyremap/__pycache__/res.cpython-36.pyc b/pyremap/__pycache__/res.cpython-36.pyc
diff --git a/pyremap/remapper.py b/pyremap/remapper.py
@@ -97,8 +97,8 @@ def __init__(self, sourceDescriptor, destinationDescriptor,
 
         # }}}
 
-    def build_mapping_file(self, method='bilinear',
-                           additionalArgs=None, logger=None):  # {{{
+    def build_mapping_file(self, method='bilinear', additionalArgs=None,
+                           logger=None, mpiTasks=1):  # {{{
         '''
         Given a source file defining either an MPAS mesh or a lat-lon grid and
         a destination file or set of arrays defining a lat-lon grid, constructs
@@ -117,6 +117,10 @@ def build_mapping_file(self, method='bilinear',
         logger : ``logging.Logger``, optional
             A logger to which ncclimo output should be redirected
 
+        mpiTasks : int, optional
+            The number of MPI tasks (a number > 1 implies that
+            ESMF_RegridWeightGen will be called with ``mpirun``)
+
         Raises
         ------
         OSError
@@ -141,7 +145,9 @@ def build_mapping_file(self, method='bilinear',
             # a valid weight file already exists, so nothing to do
             return
 
-        if find_executable('ESMF_RegridWeightGen') is None:
+        rwgPath = find_executable('ESMF_RegridWeightGen')
+
+        if rwgPath is None:
             raise OSError('ESMF_RegridWeightGen not found. Make sure esmf '
                           'package is installed via\n'
                           'latest nco: \n'
@@ -153,14 +159,17 @@ def build_mapping_file(self, method='bilinear',
         self.sourceDescriptor.to_scrip(_get_temp_path())
         self.destinationDescriptor.to_scrip(_get_temp_path())
 
-        args = ['ESMF_RegridWeightGen',
+        args = [rwgPath,
                 '--source', self.sourceDescriptor.scripFileName,
                 '--destination', self.destinationDescriptor.scripFileName,
                 '--weight', self.mappingFileName,
                 '--method', method,
                 '--netcdf4',
                 '--no_log']
 
+        if mpiTasks > 1:
+            args = ['mpirun', '-np', '{}'.format(mpiTasks)] + args
+
         if self.sourceDescriptor.regional:
             args.append('--src_regional')
 

diff --git a/pyremap/test/__pycache__/__init__.cpython-37.pyc b/pyremap/test/__pycache__/__init__.cpython-37.pyc
diff --git a/pyremap/test/__pycache__/test_interpolate.cpython-37-pytest-5.2.0.pyc b/pyremap/test/__pycache__/test_interpolate.cpython-37-pytest-5.2.0.pyc