Vertical interpolation (pressure, geop) (#26)

docs/api.rst

@@ -139,6 +139,8 @@ methods are almost equivalent):
     DatasetAccessor.roi
     DatasetAccessor.subset
     DatasetAccessor.transform
+    DatasetAccessor.wrf_zlevel
+    DatasetAccessor.wrf_plevel
 
 
 Old-style datasets

docs/index.rst

@@ -35,7 +35,7 @@ Documentation
     faq
     whats-new
     installing
-    xarray
+    xarray_acc
     gis
     plotting
     wrf

docs/whats-new.rst

@@ -1,3 +1,5 @@
+.. currentmodule:: salem
+
 What's New
 ==========
 
@@ -8,6 +10,9 @@ v0.x (unreleased)
 Enhancements
 ~~~~~~~~~~~~
 
+- New :py:func:`~DatasetAccessor.wrf_zlevel` and
+  :py:func:`~DatasetAccessor.wrf_plevel` functions for vertical interpolation
+  (still quite slow, though)
 - Salem can now plot on cartopy's maps thanks to a new
   :py:func:`~salem.proj_to_cartopy` function.
 - Doc improvements

docs/wrf.rst

@@ -9,6 +9,13 @@ Let's open a `WRF model`_ output file with xarray first:
 
 .. _WRF Model: http://www2.mmm.ucar.edu/wrf/users/
 
+
+.. ipython:: python
+   :suppress:
+
+    plt.rcParams['figure.figsize'] = (6, 4)
+    f = plt.figure(figsize=(6, 4))
+
 .. ipython:: python
 
     import xarray as xr
@@ -44,7 +51,8 @@ Salem defines a special parser for these files:
 .. ipython:: python
    :suppress:
 
-    ds.attrs = {'note': 'Global attrs removed.'}
+    ds.attrs = {'note': 'Global attrs removed.',
+                'pyproj_srs': ds.attrs['pyproj_srs']}
 
 This parser greatly simplifies the file structure:
 
@@ -83,6 +91,26 @@ accumulated fields). For a list of diagnostic variables (and TODOs!), refer to
     ds.PRCP.isel(time=-1).salem.quick_map(cmap='Purples', vmax=5)
 
 
+Vertical interpolation
+----------------------
+
+The WRF vertical coordinates are eta-levels, which is not a very practical
+coordinate for analysis or plotting. With the functions
+:py:func:`~DatasetAccessor.wrf_zlevel` and
+:py:func:`~DatasetAccessor.wrf_plevel` it is possible to interpolate the 3d
+data at either altitude or pressure levels:
+
+.. ipython:: python
+
+    ws_h = ds.isel(time=1).salem.wrf_zlevel('WS', levels=10000.)
+    @savefig plot_wrf_zinterp.png width=80%
+    ws_h.salem.quick_map();
+
+Note that currently, the interpolation is quite slow, see :issue:`25`. It's
+probably wise to do it on single time slices or aggregated data, rather than
+huge data cubes.
+
+
 Geogrid simulator
 -----------------
 

docs/xarray.rst → docs/xarray_acc.rst

@@ -1,4 +1,4 @@
-.. _xarray:
+.. _xarray_acc:
 
 xarray accessors
 ================

salem/sio.py

@@ -616,6 +616,41 @@ def quick_map(self, ax=None, interp='nearest', **kwargs):
         """Make a plot of the DataArray."""
         return self._quick_map(self._obj, ax=ax, interp=interp, **kwargs)
 
+    def interpz(self, zcoord, levels, dim_name=''):
+        """Interpolates the array along the vertical dimension
+
+        Parameters
+        ----------
+        zcoord: DataArray
+          the z coordinates of the variable. Must be of same dimensions
+        levels: 1dArray
+          the levels at which to interpolate
+        dim_name: str
+          the name of the new dimension
+
+        Returns
+        -------
+        a new DataArray with the interpolated data
+        """
+
+        if self.z_dim is None:
+            raise RuntimeError('zdimension not recognized')
+
+        data = wrftools.interp3d(self._obj.values, zcoord.values,
+                                 np.atleast_1d(levels))
+
+        dims = list(self._obj.dims)
+        zd = np.nonzero([self.z_dim == d for d in dims])[0][0]
+        dims[zd] = dim_name
+        coords = dict(self._obj.coords)
+        coords.pop(self.z_dim, None)
+        coords[dim_name] = np.atleast_1d(levels)
+        out = xr.DataArray(data, name=self._obj.name, dims=dims, coords=coords)
+        out.attrs['pyproj_srs'] = self.grid.proj.srs
+        if not np.asarray(levels).shape:
+            out = out.isel(**{dim_name:0})
+        return out
+
 
 @xr.register_dataset_accessor('salem')
 class DatasetAccessor(_XarrayAccessorBase):
@@ -665,6 +700,54 @@ 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):
+        """Interpolates to a specified height above sea level.
+
+        Parameters
+        ----------
+        varname: str
+          the name of the variable to interpolate
+        levels: 1d array
+          levels at which to interpolate (default: some levels I thought of)
+
+        Returns
+        -------
+        an interpolated DataArray
+        """
+        if levels is None:
+            levels = np.array([10, 20, 30, 50, 75, 100, 200, 300, 500, 750,
+                               1000, 2000, 3000, 5000, 7500, 10000])
+
+        zcoord = self._obj['Z']
+        out = self._obj[varname].salem.interpz(zcoord, levels, dim_name='z')
+        out['z'].attrs['description'] = 'height above sea level'
+        out['z'].attrs['units'] = 'm'
+        return out
+
+    def wrf_plevel(self, varname, levels=None):
+        """Interpolates to a specified pressure level (hPa).
+
+        Parameters
+        ----------
+        varname: str
+          the name of the variable to interpolate
+        levels: 1d array
+          levels at which to interpolate (default: some levels I thought of)
+
+        Returns
+        -------
+        an interpolated DataArray
+        """
+        if levels is None:
+            levels = np.array([1000, 975, 950, 925, 900, 850, 800, 750, 700,
+                               650, 600, 550, 500, 450, 400, 300, 200, 100])
+
+        zcoord = self._obj['PRESSURE']
+        out = self._obj[varname].salem.interpz(zcoord, levels, dim_name='p')
+        out['p'].attrs['description'] = 'pressure'
+        out['p'].attrs['units'] = 'hPa'
+        return out
+
 
 class _NetCDF4DataStore(NetCDF4DataStore):
     """Just another way to init xarray's datastore

salem/tests/test_misc.py

@@ -786,6 +786,41 @@ def test_full_wrf_wfile(self):
         assert_allclose(v.mean(), ds.PRCP.mean())
         assert_allclose(v.max(), ds.PRCP.max())
 
+    @requires_xarray
+    def test_3d_interp(self):
+
+        f = get_demo_file('wrf_d01_allvars_cropped.nc')
+        ds = sio.open_wrf_dataset(f)
+
+        out = ds.salem.wrf_zlevel('Z', levels=6000.)
+        ref_2d = out*0. + 6000.
+        assert_allclose(out, ref_2d)
+
+        # this used to raise an error
+        _ = out.isel(time=1)
+
+        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.)
+
+        out = ds.salem.wrf_zlevel('Z')
+        assert_allclose(out.sel(z=7500.), ref_2d*0. + 7500.)
+
+        out = ds.salem.wrf_plevel('PRESSURE', levels=400.)
+        ref_2d = out*0. + 400.
+        assert_allclose(out, ref_2d)
+
+        out = ds.salem.wrf_plevel('PRESSURE', levels=[400., 300.])
+        assert_allclose(out.sel(p=400.), ref_2d)
+        assert_allclose(out.sel(p=300.), ref_2d*0. + 300.)
+
+        out = ds.salem.wrf_plevel('PRESSURE')
+        assert_allclose(out.sel(p=300.), ref_2d*0. + 300.)
+
+        ds = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
+        ws_h = ds.isel(time=1).salem.wrf_zlevel('WS', levels=8000.)
+        assert np.all(np.isfinite(ws_h))
+
 
 class TestGeogridSim(unittest.TestCase):
 

salem/utils.py

@@ -30,7 +30,8 @@
 valid_names['y_dim'] = ['south_north', 'lat', 'latitude', 'latitudes', 'lats',
                         'xlat', 'xlat_m', 'dimlat', 'y','lat_3', 'phony_dim_1',
                         'northings', 'northing']
-valid_names['z_dim'] = ['levelist','level', 'pressure', 'press', 'zlevel', 'z']
+valid_names['z_dim'] = ['levelist','level', 'pressure', 'press', 'zlevel', 'z',
+                        'bottom_top']
 valid_names['t_dim'] = ['time', 'times', 'xtime']
 
 valid_names['lon_var'] = ['lon', 'longitude', 'longitudes', 'lons', 'long']

salem/wrftools.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 import pyproj
+from scipy.interpolate import interp1d
 try:
     from pandas import to_datetime
     from xarray.core import indexing
@@ -344,6 +345,21 @@ def __getitem__(self, item):
         return self.nc.variables['PH'][item] + self.nc.variables['PHB'][item]
 
 
+class Z(FakeVariable):
+    def __init__(self, nc):
+        FakeVariable.__init__(self, nc)
+        self._copy_attrs_from(nc.variables['PH'])
+        self.units = 'm'
+        self.description = 'Full model height'
+
+    @staticmethod
+    def can_do(nc):
+        return np.all([n in nc.variables for n in ['PH', 'PHB']])
+
+    def __getitem__(self, item):
+        return self.nc.variables['GEOPOTENTIAL'][item] / 9.81
+
+
 class SLP(FakeVariable):
     def __init__(self, nc):
         FakeVariable.__init__(self, nc)
@@ -389,6 +405,46 @@ def __getitem__(self, item):
 var_classes.remove('AccumulatedVariable')
 
 
+def interp3d(data, zcoord, levels):
+    """Interpolate on the first dimension of a 3d var
+
+    Useful for WRF pressure or geopotential levels
+
+    Parameters
+    ----------
+    data: ndarrad
+      3d or 4d array of the data to interpolate
+    zcoord: ndarray
+      same dims as data, the z coordinates of the data points
+    levels: 1darray
+      the levels at which to interpolate
+
+    Returns
+    -------
+    a ndarray, with the first dimension now begin of shape nlevels
+    """
+
+    ndims = len(data.shape)
+    if ndims == 4:
+        out = []
+        for d, z in zip(data, zcoord):
+            out.append(np.expand_dims(interp3d(d, z, levels), 0))
+        return np.concatenate(out, axis=0)
+    if ndims != 3:
+        raise ValueError('ndims must be 3')
+
+    # TODO: there got to be a faster way to do this
+    # same problem, no solution: http://stackoverflow.com/questions/27622808/
+    # fast-3d-interpolation-of-atmospheric-data-in-numpy-scipy
+    out = np.zeros((len(levels), data.shape[-2], data.shape[-1]))
+    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)
+            out[:, j, i] = f(levels)
+    return out
+
+
 def _ncl_slp(z, t, p, q):
     """Computes the SLP out of the WRF variables.