Add "extrapolate" option to wrf vertical interpolation methods (#88)

* Add "extrapolate" option to wrf vertical interpolation methods

* Freetype issues

* Freetype issues again

* test must pass without mpl

docs/whats-new.rst

@@ -11,6 +11,9 @@ Enhancements
 ~~~~~~~~~~~~
 
 - :py:func:`~transform_geopandas` can now handle grid to proj transformations.
+- Thw WRF vertical interpolation methods now accept a `fill_value` kwarg
+  which is `np.NaN` per default but can be set to `'extrapolate'` if
+  needed.
 - New :py:func:`~reduce` function, useful to aggregate structured high-res
   grids to lower-res grids (:ref:`sphx_glr_auto_examples_plot_subgrid_mask.py`)
 - new :py:func:`~Grid.to_geometry` method, useful to compute precise

salem/__init__.py

@@ -46,7 +46,7 @@ def _lazy_property(self):
 if not path.exists(download_dir):
     makedirs(download_dir)
 
-sample_data_gh_commit = 'aeb4cff0f61138701ab62a5b2ed2ceac7b809317'
+sample_data_gh_commit = 'b234d1dadc71ed344dada78f29d7b6f08ff26f44'
 sample_data_dir = path.join(cache_dir, 'salem-sample-data-' +
                             sample_data_gh_commit)
 

salem/sio.py

@@ -722,7 +722,8 @@ def deacc(self, as_rate=True):
 
         return out
 
-    def interpz(self, zcoord, levels, dim_name='', use_multiprocessing=True):
+    def interpz(self, zcoord, levels, dim_name='', fill_value=np.NaN,
+                use_multiprocessing=True):
         """Interpolates the array along the vertical dimension
 
         Parameters
@@ -733,6 +734,9 @@ def interpz(self, zcoord, levels, dim_name='', use_multiprocessing=True):
           the levels at which to interpolate
         dim_name: str
           the name of the new dimension
+        fill_value : np.NaN or 'extrapolate', optional
+          how to handle levels below the topography. Default is to mark them
+          as invalid, but you might want the have them extrapolated.
         use_multiprocessing: bool
           set to false if, for some reason, you don't want to use mp
 
@@ -745,8 +749,8 @@ def interpz(self, zcoord, levels, dim_name='', use_multiprocessing=True):
             raise RuntimeError('zdimension not recognized')
 
         data = wrftools.interp3d(self._obj.values, zcoord.values,
-                                 np.atleast_1d(levels), use_multiprocessing=
-                                 use_multiprocessing)
+                                 np.atleast_1d(levels), fill_value=fill_value,
+                                 use_multiprocessing=use_multiprocessing)
 
         dims = list(self._obj.dims)
         zd = np.nonzero([self.z_dim == d for d in dims])[0][0]
@@ -806,7 +810,8 @@ def transform_and_add(self, other, grid=None, interp='nearest', ks=3,
                     new_name = v
                 self._obj[new_name] = out[v]
 
-    def wrf_zlevel(self, varname, levels=None, use_multiprocessing=True):
+    def wrf_zlevel(self, varname, levels=None, fill_value=np.NaN,
+                   use_multiprocessing=True):
         """Interpolates to a specified height above sea level.
 
         Parameters
@@ -815,6 +820,9 @@ def wrf_zlevel(self, varname, levels=None, use_multiprocessing=True):
           the name of the variable to interpolate
         levels: 1d array
           levels at which to interpolate (default: some levels I thought of)
+        fill_value : np.NaN or 'extrapolate', optional
+          how to handle levels below the topography. Default is to mark them
+          as invalid, but you might want the have them extrapolated.
         use_multiprocessing: bool
           set to false if, for some reason, you don't want to use mp
 
@@ -828,13 +836,15 @@ def wrf_zlevel(self, varname, levels=None, use_multiprocessing=True):
 
         zcoord = self._obj['Z']
         out = self._obj[varname].salem.interpz(zcoord, levels, dim_name='z',
+                                               fill_value=fill_value,
                                                use_multiprocessing=
                                                use_multiprocessing)
         out['z'].attrs['description'] = 'height above sea level'
         out['z'].attrs['units'] = 'm'
         return out
 
-    def wrf_plevel(self, varname, levels=None, use_multiprocessing=True):
+    def wrf_plevel(self, varname, levels=None, fill_value=np.NaN,
+                   use_multiprocessing=True):
         """Interpolates to a specified pressure level (hPa).
 
         Parameters
@@ -843,6 +853,9 @@ def wrf_plevel(self, varname, levels=None, use_multiprocessing=True):
           the name of the variable to interpolate
         levels: 1d array
           levels at which to interpolate (default: some levels I thought of)
+        fill_value : np.NaN or 'extrapolate', optional
+          how to handle levels below the topography. Default is to mark them
+          as invalid, but you might want the have them extrapolated.
         use_multiprocessing: bool
           set to false if, for some reason, you don't want to use mp
 
@@ -856,6 +869,7 @@ def wrf_plevel(self, varname, levels=None, use_multiprocessing=True):
 
         zcoord = self._obj['PRESSURE']
         out = self._obj[varname].salem.interpz(zcoord, levels, dim_name='p',
+                                               fill_value=fill_value,
                                                use_multiprocessing=
                                                use_multiprocessing)
         out['p'].attrs['description'] = 'pressure'

salem/tests/test_datasets.py

@@ -446,7 +446,7 @@ def test_center(self):
         # toimage(img).save(get_demo_file('hef_google_roi.png'))
         ref = mpl.image.imread(get_demo_file('hef_google_roi.png'))
         rmsd = np.sqrt(np.mean((ref - img)**2))
-        self.assertTrue(rmsd < 0.1)
+        self.assertTrue(rmsd < 0.2)
         # assert_allclose(ref, img, atol=2e-2)
 
         gm = GoogleCenterMap(center_ll=(10.762660, 46.794221), zoom=13,
@@ -456,7 +456,7 @@ def test_center(self):
         img = gm.get_vardata()
         img[np.nonzero(gm.roi == 0)] /= 2.
         rmsd = np.sqrt(np.mean((ref - img)**2))
-        self.assertTrue(rmsd < 0.1)
+        self.assertTrue(rmsd < 0.2)
         # assert_allclose(ref, img, atol=2e-2)
 
     @requires_internet

salem/tests/test_graphics.py

@@ -15,9 +15,15 @@
     import shapely.geometry as shpg
     import geopandas as gpd
     MPL_VERSION = LooseVersion(mpl.__version__)
+    ftver = LooseVersion(mpl.ft2font.__freetype_version__)
+    if ftver >= LooseVersion('2.8.0'):
+        freetype_subdir = 'freetype_28'
+    else:
+        freetype_subdir = 'freetype_old'
 except ImportError:
     # place holder
     MPL_VERSION = '2.0.0'
+    freetype_subdir = ''
 
 from salem.graphics import ExtendedNorm, DataLevels, Map, get_cmap, shapefiles
 from salem import (Grid, wgs84, mercator_grid, GeoNetcdf,
@@ -33,7 +39,7 @@
 
 baseline_subdir = '2.0.x'
 baseline_dir = os.path.join(sample_data_dir, 'baseline_images',
-                            baseline_subdir)
+                            baseline_subdir, freetype_subdir)
 
 tolpy2 = 5 if python_version == 'py3' else 10
 

salem/tests/test_misc.py

@@ -883,6 +883,7 @@ def test_3d_interp(self):
         out = ds.salem.wrf_zlevel('Z', levels=[6000., 7000.])
         assert_allclose(out.sel(z=6000.), ref_2d)
         assert_allclose(out.sel(z=7000.), ref_2d*0. + 7000.)
+        assert np.all(np.isfinite(out))
 
         out = ds.salem.wrf_zlevel('Z')
         assert_allclose(out.sel(z=7500.), ref_2d*0. + 7500.)
@@ -897,6 +898,11 @@ def test_3d_interp(self):
 
         out = ds.salem.wrf_plevel('PRESSURE')
         assert_allclose(out.sel(p=300.), ref_2d*0. + 300.)
+        assert np.any(~np.isfinite(out))
+
+        out = ds.salem.wrf_plevel('PRESSURE', fill_value='extrapolate')
+        assert_allclose(out.sel(p=300.), ref_2d*0. + 300.)
+        assert np.all(np.isfinite(out))
 
         ds = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
         ws_h = ds.isel(time=1).salem.wrf_zlevel('WS', levels=8000.,

salem/wrftools.py

@@ -435,11 +435,13 @@ def __getitem__(self, item):
 
 
 def _interp1d(args):
-    f = interp1d(args[0], args[1], fill_value=np.NaN, bounds_error=False)
+    f = interp1d(args[0], args[1], fill_value=args[3],
+                 bounds_error=False)
     return f(args[2])
 
 
-def interp3d(data, zcoord, levels, use_multiprocessing=True):
+def interp3d(data, zcoord, levels, fill_value=np.NaN,
+             use_multiprocessing=True):
     """Interpolate on the first dimension of a 3d var
 
     Useful for WRF pressure or geopotential levels
@@ -452,6 +454,9 @@ def interp3d(data, zcoord, levels, use_multiprocessing=True):
       same dims as data, the z coordinates of the data points
     levels: 1darray
       the levels at which to interpolate
+    fill_value : np.NaN or 'extrapolate', optional
+      how to handle levels below the topography. Default is to mark them
+      as invalid, but you might want the have them extrapolated.
     use_multiprocessing: bool
       set to false if, for some reason, you don't want to use mp
 
@@ -464,7 +469,8 @@ def interp3d(data, zcoord, levels, use_multiprocessing=True):
     if ndims == 4:
         out = []
         for d, z in zip(data, zcoord):
-            out.append(np.expand_dims(interp3d(d, z, levels), 0))
+            out.append(np.expand_dims(interp3d(d, z, levels,
+                                               fill_value=fill_value), 0))
         return np.concatenate(out, axis=0)
     if ndims != 3:
         raise ValueError('ndims must be 3')
@@ -473,7 +479,8 @@ def interp3d(data, zcoord, levels, use_multiprocessing=True):
         inp = []
         for j in range(data.shape[-2]):
             for i in range(data.shape[-1]):
-                inp.append((zcoord[:, j, i], data[:, j, i], levels))
+                inp.append((zcoord[:, j, i], data[:, j, i], levels,
+                            fill_value))
         _init_pool()
         out = POOL.map(_interp1d, inp, chunksize=1000)
         out = np.asarray(out).T
@@ -486,7 +493,7 @@ def interp3d(data, zcoord, levels, use_multiprocessing=True):
         for i in range(data.shape[-1]):
             for j in range(data.shape[-2]):
                 f = interp1d(zcoord[:, j, i], data[:, j, i],
-                             fill_value=np.NaN, bounds_error=False)
+                             fill_value=fill_value, bounds_error=False)
                 out[:, j, i] = f(levels)
     return out