Merge pull request #16 from xylar/detect_esmf_mpi

Run ESMF_RegridWeightGen with mpirun depending on esmf package version

examples/make_mpas_to_lat_lon_mapping.py

@@ -42,7 +42,7 @@
 
 remapper = Remapper(inDescriptor, outDescriptor, mappingFileName)
 
-remapper.build_mapping_file(method='bilinear')
+remapper.build_mapping_file(method='bilinear', mpiTasks=1)
 
 outFileName = 'temp_{}.nc'.format(outGridName)
 ds = xarray.open_dataset(inGridFileName)

pyremap/__init__.py

@@ -6,5 +6,5 @@
 
 from pyremap.polar import get_polar_descriptor_from_file, get_polar_descriptor
 
-__version_info__ = (0, 0, 6)
+__version_info__ = (0, 0, 7)
 __version__ = '.'.join(str(vi) for vi in __version_info__)

pyremap/descriptor.py

@@ -753,8 +753,9 @@ def project_to_lat_lon(self, X, Y):  # {{{
         # -------
         # Xylar Asay-Davis
 
-        Lon, Lat = pyproj.transform(self.projection, self.latLonProjection,
-                                    X, Y)
+        transformer = pyproj.Transformer.from_proj(self.projection,
+                                                   self.latLonProjection)
+        Lon, Lat = transformer.transform(X, Y)
 
         return (Lat, Lon)  # }}}
 

pyremap/polar.py

@@ -125,8 +125,8 @@ def to_polar(points):
     projection = get_antarctic_stereographic_projection()
     latLonProjection = pyproj.Proj(proj='latlong', datum='WGS84')
 
-    x, y = pyproj.transform(latLonProjection, projection, points[:, 0],
-                            points[:, 1], radians=False)
+    transformer = pyproj.Transformer.from_proj(latLonProjection, projection)
+    x, y = transformer.transform(points[:, 0], points[:, 1], radians=False)
     points[:, 0] = x
     points[:, 1] = y
     return points
@@ -137,8 +137,8 @@ def from_polar(points):
     projection = get_antarctic_stereographic_projection()
     latLonProjection = pyproj.Proj(proj='latlong', datum='WGS84')
 
-    lon, lat = pyproj.transform(projection, latLonProjection, points[:, 0],
-                                points[:, 1], radians=False)
+    transformer = pyproj.Transformer.from_proj(projection, latLonProjection)
+    lon, lat = transformer.transform(points[:, 0], points[:, 1], radians=False)
     points[:, 0] = lon
     points[:, 1] = lat
     return points

pyremap/remapper.py

@@ -30,6 +30,9 @@
 from scipy.sparse import csr_matrix
 import xarray as xr
 import sys
+from subprocess import check_output, CalledProcessError
+import json
+import warnings
 
 from pyremap.descriptor import MpasMeshDescriptor, \
     LatLonGridDescriptor, LatLon2DGridDescriptor, ProjectionGridDescriptor, \
@@ -180,12 +183,30 @@ def build_mapping_file(self, method='bilinear', additionalArgs=None,
                 '--netcdf4',
                 '--no_log']
 
-        if mpiTasks > 1:
-            if 'CONDA_PREFIX' in os.environ:
+        if 'CONDA_PREFIX' in os.environ:
+            # this is a conda environment, so we need to find out if esmf
+            # needs mpirun or not
+            conda_args = ['conda', 'list', 'esmf', '--json']
+            output = check_output(conda_args).decode("utf-8")
+            output = json.loads(output)
+            build_string = output[0]['build_string']
+
+            if 'mpi_mpich' in build_string or 'mpi_openmpi' in build_string:
+                # esmf was installed with MPI, so we should use mpirun
                 mpirun_path = '{}/bin/mpirun'.format(os.environ['CONDA_PREFIX'])
             else:
-                mpirun_path = 'mpirun'
+                # esmf was installed without MPI, so we shouldn't try to use it
+                if mpiTasks > 1:
+                    warnings.warn('Requesting {} MPI tasks but the MPI version '
+                                  'of ESMF is not installed'.format(mpiTasks))
+                mpirun_path = None
+
+        elif mpiTasks > 1:
+            mpirun_path = 'mpirun'
+        else:
+            mpirun_path = None
 
+        if mpirun_path is not None:
             args = [mpirun_path, '-np', '{}'.format(mpiTasks)] + args
 
         if self.sourceDescriptor.regional:
@@ -443,10 +464,10 @@ def remap(self, ds, renormalizationThreshold=None):  # {{{
             for var in ds.data_vars:
                 if self._check_drop(ds[var]):
                     drop.append(var)
-            remappedDs = ds.drop(drop)
-            remappedDs = remappedDs.apply(self._remap_data_array,
-                                          keep_attrs=True,
-                                          args=(renormalizationThreshold,))
+            remappedDs = ds.drop_vars(drop)
+            remappedDs = remappedDs.map(self._remap_data_array,
+                                        keep_attrs=True,
+                                        args=(renormalizationThreshold,))
         else:
             raise TypeError('ds not an xarray Dataset or DataArray.')
 
@@ -668,9 +689,4 @@ def _remap_numpy_array(self, inField, remapAxes,
         return outField  # }}}
 
 
-def _get_temp_path():  # {{{
-    '''Returns the name of a temporary NetCDF file'''
-    return '{}/{}.nc'.format(tempfile._get_default_tempdir(),
-                             next(tempfile._get_candidate_names()))  # }}}
-
 # vim: ai ts=4 sts=4 et sw=4 ft=python

pyremap/test/test_interpolate.py

@@ -110,8 +110,8 @@ def check_remap(self, inFileName, outFileName, refFileName, remapper,
             dsRemapped = xarray.open_dataset(outFileName)
             # drop some extra vairables added by ncremap that aren't in the
             # reference data set
-            dsRemapped = dsRemapped.drop(['lat_bnds', 'lon_bnds', 'gw',
-                                          'area'])
+            dsRemapped = dsRemapped.drop_vars(['lat_bnds', 'lon_bnds', 'gw',
+                                               'area'])
             self.assertDatasetApproxEqual(dsRemapped, dsRef)
 
         # now, try in-memory remapping

recipe/meta.yaml

@@ -1,5 +1,5 @@
 {% set name = "pyremap" %}
-{% set version = "0.0.6" %}
+{% set version = "0.0.7" %}
 
 package:
   name: {{ name|lower }}