Improved horizontal interpolation

- Can now specify single level, multiple levels, or interpolate
  to a 2D surface like PBLH.
- Should perform better when interpolating multiproduct fields
  like wspd_wdir.

Fixes #65.

fortran/wrf_user.f90

@@ -64,7 +64,67 @@ END SUBROUTINE DCOMPUTETK
 
 
 ! NCLFORTSTART
-SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, desiredloc, nx, ny, nz, missingval)
+SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, levels, nx, ny, nz, nlev, missingval)
+    IMPLICIT NONE
+
+    !f2py threadsafe
+    !f2py intent(in,out) :: out2d
+
+    INTEGER, INTENT(IN) :: nx, ny, nz, nlev
+    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) ::  data3d
+    REAL(KIND=8), DIMENSION(nx,ny,nlev), INTENT(OUT) :: out2d
+    REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: zdata
+    REAL(KIND=8), DIMENSION(nlev), INTENT(IN) :: levels
+    REAL(KIND=8), INTENT(IN) :: missingval
+
+! NCLEND
+
+    INTEGER :: i,j,kp,ip,im,lev
+    LOGICAL :: dointerp
+    REAL(KIND=8) :: w1,w2,desiredloc
+
+    ! does vertical coordinate increase or decrease with increasing k?
+    ! set offset appropriately
+
+    ip = 0
+    im = 1
+    IF (zdata(1,1,1) .GT. zdata(1,1,nz)) THEN
+        ip = 1
+        im = 0
+    END IF
+
+    !$OMP PARALLEL DO COLLAPSE(3) PRIVATE(i,j,lev,kp,dointerp,w1,w2,desiredloc) &
+    !$OMP FIRSTPRIVATE(ip,im) SCHEDULE(runtime)
+    DO lev = 1,nlev
+        DO i = 1,nx
+            DO j = 1,ny
+                ! Initialize to missing.  Was initially hard-coded to -999999.
+                out2d(i,j,lev) = missingval
+                dointerp = .FALSE.
+                kp = nz
+                desiredloc = levels(lev)
+
+                DO WHILE ((.NOT. dointerp) .AND. (kp >= 2))
+                    IF (((zdata(i,j,kp-im) < desiredloc) .AND. (zdata(i,j,kp-ip) > desiredloc))) THEN
+                        w2 = (desiredloc - zdata(i,j,kp-im))/(zdata(i,j,kp-ip) - zdata(i,j,kp-im))
+                        w1 = 1.D0 - w2
+                        out2d(i,j,lev) = w1*data3d(i,j,kp-im) + w2*data3d(i,j,kp-ip)
+                        dointerp = .TRUE.
+                    END IF
+                    kp = kp - 1
+                END DO
+            END DO
+        END DO
+    END DO
+    !$OMP END PARALLEL DO
+
+    RETURN
+
+END SUBROUTINE DINTERP3DZ
+
+
+! NCLFORTSTART
+SUBROUTINE DINTERP3DZ_2DLEV(data3d, out2d, zdata, levs2d, nx, ny, nz, missingval)
     IMPLICIT NONE
 
     !f2py threadsafe
@@ -74,14 +134,14 @@ SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, desiredloc, nx, ny, nz, missingval)
     REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) ::  data3d
     REAL(KIND=8), DIMENSION(nx,ny), INTENT(OUT) :: out2d
     REAL(KIND=8), DIMENSION(nx,ny,nz), INTENT(IN) :: zdata
-    REAL(KIND=8), INTENT(IN) :: desiredloc
+    REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: levs2d
     REAL(KIND=8), INTENT(IN) :: missingval
 
 ! NCLEND
 
     INTEGER :: i,j,kp,ip,im
     LOGICAL :: dointerp
-    REAL(KIND=8) :: w1,w2
+    REAL(KIND=8) :: w1,w2,desiredloc
 
     ! does vertical coordinate increase or decrease with increasing k?
     ! set offset appropriately
@@ -93,7 +153,7 @@ SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, desiredloc, nx, ny, nz, missingval)
         im = 0
     END IF
 
-    !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(i,j,kp,dointerp,w1,w2) &
+    !$OMP PARALLEL DO COLLAPSE(2) PRIVATE(i,j,kp,dointerp,w1,w2,desiredloc) &
     !$OMP FIRSTPRIVATE(ip,im) SCHEDULE(runtime)
     DO i = 1,nx
         DO j = 1,ny
@@ -101,6 +161,7 @@ SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, desiredloc, nx, ny, nz, missingval)
             out2d(i,j) = missingval
             dointerp = .FALSE.
             kp = nz
+            desiredloc = levs2d(i,j)
 
             DO WHILE ((.NOT. dointerp) .AND. (kp >= 2))
                 IF (((zdata(i,j,kp-im) < desiredloc) .AND. (zdata(i,j,kp-ip) > desiredloc))) THEN
@@ -117,7 +178,7 @@ SUBROUTINE DINTERP3DZ(data3d, out2d, zdata, desiredloc, nx, ny, nz, missingval)
 
     RETURN
 
-END SUBROUTINE DINTERP3DZ
+END SUBROUTINE DINTERP3DZ_2DLEV
 
 ! PORT DZSTAG HERE
 

src/wrf/extension.py

@@ -11,7 +11,7 @@ from wrf._wrffortran import (dcomputetk, dinterp3dz, dinterp2dxy, dinterp1d,
                           dcomputeabsvort, dlltoij, dijtoll, deqthecalc,
                           omgcalc, virtual_temp, wetbulbcalc, dcomputepw,
                           wrf_monotonic, wrf_vintrp, dcomputewspd, 
-                          dcomputewdir,
+                          dcomputewdir, dinterp3dz_2dlev,
                           fomp_set_num_threads, fomp_get_num_threads, 
                           fomp_get_max_threads, fomp_get_thread_num,
                           fomp_get_num_procs, fomp_in_parallel, 
@@ -34,7 +34,8 @@ from .decorators import (left_iteration, cast_type,
 from .util import combine_dims, npbytes_to_str, psafilepath
 from .py3compat import py3range
 from .specialdec import (uvmet_left_iter, cape_left_iter, 
-                         cloudfrac_left_iter, check_cape_args)
+                         cloudfrac_left_iter, check_cape_args,
+                         interplevel_left_iter, check_interplevel_args)
 
 class DiagnosticError(Exception):
     """Raised when an error occurs in a diagnostic routine."""
@@ -73,9 +74,9 @@ class DiagnosticError(Exception):
 # below assume that Fortran-ordered views are being used.  This allows
 # f2py to pass the array pointers directly to the Fortran routine.
 
-@check_args(0, 3, (3, 3, None, None))
-@left_iteration(3, 2, ref_var_idx=0, ignore_args=(2,3))
-@cast_type(arg_idxs=(0,1))
+@check_interplevel_args(is2dlev=False)
+@interplevel_left_iter(is2dlev=False)
+@cast_type(arg_idxs=(0,1,2))
 @extract_and_transpose()
 def _interpz3d(field3d, z, desiredloc, missingval, outview=None):
     """Wrapper for dinterp3dz.
@@ -83,9 +84,10 @@ def _interpz3d(field3d, z, desiredloc, missingval, outview=None):
     Located in wrf_user.f90.
     
     """
-    if outview is None:
-        outview = np.empty(field3d.shape[0:2], np.float64, order="F")
-        
+    if outview is None:   
+        outshape = field3d.shape[0:2] + desiredloc.shape
+        outview = np.empty(outshape, np.float64, order="F")
+    
     result = dinterp3dz(field3d, 
                         outview,
                         z, 
@@ -94,6 +96,27 @@ def _interpz3d(field3d, z, desiredloc, missingval, outview=None):
     return result
 
 
+@check_interplevel_args(is2dlev=True)
+@interplevel_left_iter(is2dlev=True)
+@cast_type(arg_idxs=(0,1,2))
+@extract_and_transpose()
+def _interpz3d_lev2d(field3d, z, lev2d, missingval, outview=None):
+    """Wrapper for dinterp3dz.
+    
+    Located in wrf_user.f90.
+    
+    """
+    if outview is None:
+        outview = np.empty(field3d.shape[0:2], np.float64, order="F")
+    
+    result = dinterp3dz_2dlev(field3d, 
+                              outview,
+                              z, 
+                              lev2d, 
+                              missingval)
+    return result
+
+
 @check_args(0, 3, (3,))
 @left_iteration(3, combine_dims([(0,-3),(1,-2)]), ref_var_idx=0, 
                   ignore_args=(1,))

src/wrf/interp.py

@@ -4,7 +4,7 @@ import numpy as np
 import numpy.ma as ma
 
 from .extension import (_interpz3d, _vertcross, _interpline, _smooth2d, 
-                        _monotonic, _vintrp)
+                        _monotonic, _vintrp, _interpz3d_lev2d)
 
 from .metadecorators import set_interp_metadata
 from .util import extract_vars, is_staggered, get_id, to_np, get_iterable
@@ -20,7 +20,7 @@ from wrf.g_pressure import get_pressure
 #  Note:  Extension decorator is good enough to handle left dims
 @set_interp_metadata("horiz")
 def interplevel(field3d, vert, desiredlev, missing=default_fill(np.float64), 
-                meta=True):
+                squeeze=True, meta=True):
     """Return the three-dimensional field interpolated to a horizontal plane 
     at the specified vertical level.
 
@@ -35,11 +35,20 @@ def interplevel(field3d, vert, desiredlev, missing=default_fill(np.float64),
             pressure or height. This array must have the same dimensionality
             as *field3d*.
         
-        desiredlev (:obj:`float`): The desired vertical level.  
+        desiredlev (:obj:`float`, 1D sequence, or :class:`numpy.ndarray`): The 
+            desired vertical level(s). This can be a single value (e.g. 500), 
+            a sequence of values (e.g. [1000, 850, 700, 500, 250]), or a 
+            multidimensional array where the right two dimensions (ny x nx) 
+            must match *field3d*, and any leftmost dimensions match 
+            field3d.shape[:-3] (e.g. planetary boundary layer).  
             Must be in the same units as the *vert* parameter.
         
         missing (:obj:`float`): The fill value to use for the output.  
             Default is :data:`wrf.default_fill(numpy.float64)`.
+            
+        squeeze (:obj:`bool`, optional): Set to False to prevent dimensions 
+            with a size of 1 from being automatically removed from the shape 
+            of the output. Default is True.
         
         meta (:obj:`bool`): Set to False to disable metadata and return 
             :class:`numpy.ndarray` instead of 
@@ -66,26 +75,34 @@ def interplevel(field3d, vert, desiredlev, missing=default_fill(np.float64),
             
             p = getvar(wrfin, "pressure")
             ht = getvar(wrfin, "z", units="dm")
+            pblh = getvar(wrfin, "PBLH")
             
             ht_500 = interplevel(ht, p, 500.0)
+            
+            
         
     
     """
-    # Some fields like uvmet have an extra left dimension for the product
-    # type, we'll handle that iteration here.
-    multi = True if field3d.ndim - vert.ndim == 1 else False
     
-    if not multi:
-        result = _interpz3d(field3d, vert, desiredlev, missing)
+    _desiredlev = np.asarray(desiredlev)
+    if _desiredlev.ndim == 0:
+        _desiredlev = np.array([desiredlev], np.float64)
+        levsare2d = False
     else:
-        outshape = field3d.shape[0:-3] + field3d.shape[-2:]
-        result = np.empty(outshape, dtype=field3d.dtype)
-            
-        for i in py3range(field3d.shape[0]):
-            result[i,:] = (
-                _interpz3d(field3d[i,:], vert, desiredlev, missing)[:])
-            
-    return ma.masked_values (result, missing)
+        levsare2d = _desiredlev.ndim >= 2
+
+    if not levsare2d:
+        result = _interpz3d(field3d, vert, _desiredlev, missing)
+    else:
+        result = _interpz3d_lev2d(field3d, vert, _desiredlev, missing)
+     
+    masked = ma.masked_values (result, missing)
+    
+    if not meta:
+        if squeeze:
+            return masked.squeeze()
+    
+    return masked
 
 
 @set_interp_metadata("cross")

src/wrf/metadecorators.py

@@ -792,16 +792,21 @@ def _set_horiz_meta(wrapped, instance, args, kwargs):
         
     """  
     argvars = from_args(wrapped, ("field3d", "vert", "desiredlev", 
-                                  "missing"), 
+                                  "missing", "squeeze"), 
                           *args, **kwargs)  
     
     field3d = argvars["field3d"]
     z = argvars["vert"]
-    desiredloc = argvars["desiredlev"]
+    desiredlev = argvars["desiredlev"]
+    _desiredlev = np.asarray(desiredlev)
     missingval = argvars["missing"]
+    squeeze = argvars["squeeze"]
     
     result = wrapped(*args, **kwargs)
     
+    levsare2d = _desiredlev.ndim >= 2
+    multiproduct = field3d.ndim - z.ndim == 1
+        
     # Defaults, in case the data isn't a DataArray
     outname = None
     outdimnames = None
@@ -812,34 +817,49 @@ def _set_horiz_meta(wrapped, instance, args, kwargs):
     vert_units = None
     if isinstance(z, DataArray):
         vert_units = z.attrs.get("units", None)
-    
-    # If we have no metadata to start with, only set the level
-    levelstr = ("{0} {1}".format(desiredloc, vert_units) 
-                if vert_units is not None 
-                else "{0}".format(desiredloc))
-    
-    name_levelstr = ("{0}_{1}".format(desiredloc, vert_units) 
-                if vert_units is not None 
-                else "{0}".format(desiredloc))
+        
     
     if isinstance(field3d, DataArray):
         outcoords = OrderedDict()
         outdimnames = list(field3d.dims)
         outcoords.update(field3d.coords)
-        outdimnames.remove(field3d.dims[-3])
+        
+        del outdimnames[-3]
+        
         try:
             del outcoords[field3d.dims[-3]]
         except KeyError:
             pass # xarray 0.9
+        
+        if not levsare2d:
+            outdimnames.insert(-2, "levels")
+            if _desiredlev.ndim == 0:
+                outcoords["levels"] = [desiredlev]
+            else:
+                outcoords["levels"] = _desiredlev
+        else:
+            if (_desiredlev.ndim == 2):
+                outcoords["levels"] = field3d.dims[-2:], _desiredlev[:]
+            else:
+                if multiproduct:
+                    d = field3d.dims[1:-3] + field3d.dims[-2:]
+                else:
+                    d = field3d.dims[0:-3] + field3d.dims[-2:]
+                outcoords["levels"] = d, _desiredlev[:]
+                
+                
+                
+                
+        
         outattrs.update(field3d.attrs)
-        outname = "{0}_{1}".format(field3d.name, name_levelstr)
+        outname = "{0}_interp".format(field3d.name)
         
     else:
-        outname = "field3d_{0}".format(name_levelstr)
-        
-    outattrs["level"] = levelstr
+        outname = "field3d_interp"
+    
     outattrs["missing_value"] = missingval
     outattrs["_FillValue"] = missingval
+    outattrs["vert_units"] = vert_units
     
     for key in ("MemoryOrder", "description"):
         try:
@@ -847,8 +867,10 @@ def _set_horiz_meta(wrapped, instance, args, kwargs):
         except KeyError:
             pass
     
-    return DataArray(result, name=outname, dims=outdimnames, 
-                     coords=outcoords, attrs=outattrs)
+    da = DataArray(result, name=outname, dims=outdimnames, 
+                   coords=outcoords, attrs=outattrs)
+    
+    return da.squeeze() if squeeze else da
 
 
 def _set_cross_meta(wrapped, instance, args, kwargs):
@@ -1646,7 +1668,7 @@ def set_interp_metadata(interp_type):
     """   
     @wrapt.decorator
     def func_wrapper(wrapped, instance, args, kwargs):
-        do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"]
+        do_meta = from_args(wrapped, ("meta"), *args, **kwargs)["meta"]
         
         if do_meta is None:
             do_meta = True

src/wrf/specialdec.py

@@ -5,6 +5,7 @@ import numpy as np
 import wrapt 
 
 from .util import iter_left_indexes, to_np
+from .py3compat import py3range
 from .config import xarray_enabled
 from .constants import default_fill
 
@@ -520,6 +521,83 @@ def cloudfrac_left_iter(alg_dtype=np.float64):
     return func_wrapper
 
 
+def interplevel_left_iter(is2dlev, alg_dtype=np.float64):
+    @wrapt.decorator
+    def func_wrapper(wrapped, instance, args, kwargs):
+        new_args = list(args)
+        new_kwargs = dict(kwargs)
+        
+        field3d = args[0]
+        z = args[1]
+        levels = args[2]
+            
+        num_left_dims = z.ndim - 3
+        orig_dtype = field3d.dtype
+        left_dims = z.shape[0:num_left_dims]
+        multiproduct = True if field3d.ndim - z.ndim == 1 else False
+        
+        # No special left side iteration, build the output from the 
+        # low, mid, high results.
+        if (num_left_dims == 0):
+            if multiproduct:
+                if not is2dlev:
+                    outshape = (field3d.shape[0:-3] + levels.shape + 
+                                field3d.shape[-2:])
+                else:
+                    outshape = (field3d.shape[0:-3] + field3d.shape[-2:])
+                
+                output = np.empty(outshape, dtype=alg_dtype)
+                for i in py3range(field3d.shape[0]):
+                    new_args[0] = field3d[i,:]
+                    new_kwargs["outview"] = output[i,:]
+                    _ = wrapped(*new_args, **new_kwargs)
+            else:
+                output = wrapped(*args, **kwargs)
+
+            return output
+        
+        if multiproduct:
+            outdims = field3d.shape[0:1] + left_dims
+        else:
+            outdims = left_dims
+            
+        extra_dims = tuple(outdims)
+        
+        if not is2dlev:
+            outdims += levels.shape
+        
+        outdims += z.shape[-2:]
+        
+        outview_array = np.empty(outdims, alg_dtype)
+
+        for left_idxs in iter_left_indexes(extra_dims):
+            
+            field_out_slice_idxs = left_idxs + (slice(None),)
+            
+            if multiproduct:
+                z_slice_idxs = left_idxs[1:] + (slice(None),)
+            else:
+                z_slice_idxs = left_idxs + (slice(None),)
+                
+            
+            new_args[0] = field3d[field_out_slice_idxs]
+            new_args[1] = z[z_slice_idxs]
+            
+            if is2dlev:
+                if levels.ndim > 2:
+                    new_args[2] = levels[z_slice_idxs]
+            
+            new_kwargs["outview"] = outview_array[field_out_slice_idxs]
+            
+            _ = wrapped(*new_args, **new_kwargs)
+            
+        output = outview_array.astype(orig_dtype)
+ 
+        return output
+    
+    return func_wrapper
+     
+
 def check_cape_args():
     """A decorator to check that the cape_3d arguments are valid.
     
@@ -574,3 +652,47 @@ def check_cape_args():
     
     return func_wrapper
 
+
+def check_interplevel_args(is2dlev):
+    """A decorator to check that the interplevel arguments are valid.
+    
+    An exception is raised when an invalid argument is found.
+            
+    Returns:
+    
+        None
+        
+    Raises:
+    
+        :class:`ValueError`: Raised when an invalid argument is detected.
+        
+    """
+    @wrapt.decorator
+    def func_wrapper(wrapped, instance, args, kwargs):
+        
+        field3d = args[0]
+        z = args[1]
+        levels = args[2]
+        
+        multiproduct = True if (field3d.ndim - z.ndim) == 1 else False
+        
+        if not multiproduct:
+            if field3d.shape != z.shape:
+                raise ValueError("arguments 0 and 1 must have the same shape")
+        else:
+            if field3d.shape[1:] != z.shape:
+                raise ValueError("argument 0 and 1 must have same rightmost "
+                                 "dimensions")
+                
+        if is2dlev:
+            if levels.ndim != 2:
+                if (levels.shape[0:-2] != z.shape[0:-3] or 
+                    levels.shape[-2:] != z.shape[-2:]):
+                    raise ValueError("argument 1 and 2 must have "
+                                     "the same leftmost and rightmost "
+                                     "dimensions")
+      
+        return wrapped(*args, **kwargs)
+    
+    return func_wrapper
+

test/utests.py

@@ -302,12 +302,46 @@ def make_interp_test(varname, wrf_in, referent, multi=False,
             ref_ht_500 = _get_refvals(referent, "z_500", repeat, multi)
             hts = getvar(in_wrfnc, "z", timeidx=timeidx)
             p = getvar(in_wrfnc, "pressure", timeidx=timeidx)
+            wspd_wdir = getvar(in_wrfnc, "wspd_wdir", timeidx=timeidx)
             
             # Make sure the numpy versions work first
             hts_500 = interplevel(to_np(hts), to_np(p), 500)
             hts_500 = interplevel(hts, p, 500)
             
-            nt.assert_allclose(to_np(hts_500), ref_ht_500)
+            print(hts_500)
+            
+            #nt.assert_allclose(to_np(hts_500), ref_ht_500)
+            
+            hts_multi= interplevel(to_np(hts), to_np(p), [1000,500,250])
+            hts_multi = interplevel(hts, p, [1000,500,250])
+            
+            print(hts_multi)
+            
+            pblh = getvar(in_wrfnc, "PBLH", timeidx=timeidx)
+            hts_pblh = interplevel(p, hts, pblh)
+            
+            print(hts_pblh)
+            
+            #nt.assert_allclose(to_np(hts_500), ref_ht_500)
+            
+            wspd_wdir_500 = interplevel(to_np(wspd_wdir), to_np(p), 500)
+            wspd_wdir_500 = interplevel(wspd_wdir, p, 500)
+            print(wspd_wdir_500)
+            
+            wspd_wdir_multi= interplevel(to_np(wspd_wdir), 
+                                         to_np(p), [1000,500,250])
+            wdpd_wdir_multi = interplevel(wspd_wdir, p, [1000,500,250])
+            
+            print(wdpd_wdir_multi)
+            
+            wspd_wdir_pblh = interplevel(to_np(wspd_wdir), to_np(hts), pblh)
+            wspd_wdir_pblh = interplevel(wspd_wdir, hts, pblh)
+            print(wspd_wdir_pblh)
+            
+            wspd_wdir_pblh = interplevel(to_np(wspd_wdir), 
+                                         to_np(hts), pblh[0,:])
+            wspd_wdir_pblh = interplevel(wspd_wdir, hts, pblh[0,:])
+            print(wspd_wdir_pblh)
             
         elif (varname == "vertcross"):
             ref_ht_cross = _get_refvals(referent, "ht_cross", repeat, multi)