Merge pull request #1 from pelson/linear_extrap_and_auto_rise

Added automatic rising detection, and a linear extrapolation kernel.

Author: pelson

stratify/__init__.py

@@ -2,5 +2,6 @@
 
 from ._vinterp import (interpolate,
                        INTERPOLATE_LINEAR, INTERPOLATE_NEAREST,
-                       EXTRAPOLATE_NAN, EXTRAPOLATE_NEAREST)
+                       EXTRAPOLATE_NAN, EXTRAPOLATE_NEAREST,
+                       EXTRAPOLATE_LINEAR)
 

stratify/_vinterp.pyx

@@ -17,7 +17,7 @@ cdef extern from "math.h" nogil:
 
 __all__ = ['interpolate',
            'INTERPOLATE_LINEAR', 'INTERPOLATE_NEAREST',
-           'EXTRAPOLATE_NAN', 'EXTRAPOLATE_NEAREST']
+           'EXTRAPOLATE_NAN', 'EXTRAPOLATE_NEAREST', 'EXTRAPOLATE_LINEAR']
 
 
 # interp_kernel defines the inner part of an interpolation operation.
@@ -270,6 +270,34 @@ cdef long nearest_edge_extrap(int direction, double[:] z_src,
         fz_target[i] = fz_src[i, index]
 
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cdef long linear_extrap(int direction, double[:] z_src,
+                        double[:, :] fz_src, double level,
+                        double[:] fz_target) nogil except -1:
+    """Linear extrapolation using either the first or last 2 values."""
+    cdef unsigned int m = fz_src.shape[0]
+    cdef unsigned int n_src_pts = fz_src.shape[1]
+    cdef unsigned int p0, p1, i
+    cdef double frac
+
+    if n_src_pts < 2:
+        with gil:
+            raise ValueError('Linear extrapolation requires at least '
+                             '2 source points. Got {}.'.format(n_src_pts))
+
+    if direction < 0:
+        p0, p1 = 0, 1
+    else:
+        p0, p1 = n_src_pts - 2, n_src_pts - 1
+
+    frac = ((level - z_src[p0]) /
+            (z_src[p1] - z_src[p0]))
+
+    for i in range(m):
+       fz_target[i] = fz_src[i, p0] + frac * (fz_src[i, p1] - fz_src[i, p0])
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
 cdef long nan_data_extrap(int direction, double[:] z_src,
@@ -335,6 +363,11 @@ cdef class _NearestExtrapKernel(ExtrapKernel):
         self.kernel = nearest_edge_extrap
 
 
+cdef class _LinearExtrapKernel(ExtrapKernel):
+    def __init__(self):
+        self.kernel = linear_extrap
+
+
 cdef class _TestableDirectionExtrapKernel(ExtrapKernel):
     def __init__(self):
         self.kernel = _testable_direction_extrap
@@ -345,9 +378,10 @@ INTERPOLATE_LINEAR = _LinearInterpKernel()
 INTERPOLATE_NEAREST = _NearestInterpKernel()
 EXTRAPOLATE_NAN = _NanExtrapKernel()
 EXTRAPOLATE_NEAREST = _NearestExtrapKernel()
+EXTRAPOLATE_LINEAR = _LinearExtrapKernel()
 
 
-def interpolate(z_target, z_src, fz_src, axis=-1, rising=True,
+def interpolate(z_target, z_src, fz_src, axis=-1, rising=None,
                 interpolation=INTERPOLATE_LINEAR,
                 extrapolation=EXTRAPOLATE_NAN):
     """
@@ -382,12 +416,16 @@ def interpolate(z_target, z_src, fz_src, axis=-1, rising=True,
         the same as the shape of ``z_src``.
     axis: int (default -1)
         The axis to perform the interpolation over.
-    rising: bool (default True)
-        Whether the values of the target z coordinate are generally rising or
-        generally falling. For example, values of pressure levels will be
-        generally falling as the z coordinate increases.
+    rising: bool (default None)
+        Whether the values of the source's interpolation coordinate values
+        are generally rising or generally falling. For example, values of
+        pressure levels will be generally falling as the z coordinate
+        increases.
         This will determine whether extrapolation needs to occur for
         ``z_target`` below the first and above the last ``z_src``.
+        If rising is None, the first two interpolation coordinate values
+        will be used to determine the general direction. In most cases,
+        this is a good option.
     interpolation: :class:`.InterpKernel` instance
         The core interpolation operation to use. :attr:`.INTERPOLATE_LINEAR`
         and :attr:`_INTERPOLATE_NEAREST` are provided for convenient
@@ -407,7 +445,7 @@ cdef class _Interpolator(object):
     """
     Where the magic happens for gridwise_interp. The work of this __init__ is
     mostly for putting the input nd arrays into a 3 and 4 dimensional form for
-    convenient (read efficient) Cython form. Inline comments should help with
+    convenient (read: efficient) Cython form. Inline comments should help with
     understanding.
 
    """
@@ -420,7 +458,7 @@ cdef class _Interpolator(object):
     cpdef public _result_working_shape, result_shape, _first_value
 
     def __init__(self, z_target, z_src, fz_src, axis=-1,
-                 bint rising=True,
+                 rising=None,
                  InterpKernel interpolation=INTERPOLATE_LINEAR,
                  ExtrapKernel extrapolation=EXTRAPOLATE_NAN):
         # Cast data to numpy arrays if not already.
@@ -496,7 +534,17 @@ cdef class _Interpolator(object):
         #: The shape of the interpolated data.
         self.result_shape = tuple(result_shape)
 
-        self.rising = rising
+        if rising is None:
+            if z_src.shape[zp_axis] < 2:
+                raise ValueError('The rising keyword must be defined when '
+                                 'the size of the source array is <2 in '
+                                 'the interpolation axis.')
+            z_src_indexer = [0] * z_src.ndim
+            z_src_indexer[zp_axis] = slice(0, 2)
+            first_two = z_src[z_src_indexer]
+            rising = first_two[0] <= first_two[1] 
+
+        self.rising = bool(rising)
 
         self.interpolation = interpolation.kernel
         self.extrapolation = extrapolation.kernel
@@ -562,7 +610,7 @@ cdef class _Interpolator(object):
                     gridwise_interpolation(z_target[i, :, j], z_src[i, :, j], fz_src[:, i, :, j],
                                            self.rising,
                                            self.interpolation,
-                                          self.extrapolation,
+                                           self.extrapolation,
                                            fz_target_view[:, i, :, j])
 
         return fz_target.reshape(self.result_shape).astype(self._target_dtype)

stratify/tests/test_vinterp.py

@@ -10,7 +10,7 @@
 
 
 class TestColumnInterpolation(unittest.TestCase):
-    def interpolate(self, x_target, x_src):
+    def interpolate(self, x_target, x_src, rising=None):
         x_target = np.array(x_target)
         x_src = np.array(x_src)
         fx_src = np.empty(x_src.shape)
@@ -19,13 +19,15 @@ def interpolate(self, x_target, x_src):
         extrap_direct = vinterp._TestableDirectionExtrapKernel()
 
         r1 = stratify.interpolate(x_target, x_src, fx_src,
-                                 interpolation=index_interp,
-                                 extrapolation=extrap_direct)
+                                  rising=rising,
+                                  interpolation=index_interp,
+                                  extrapolation=extrap_direct)
 
-        r2 = stratify.interpolate(-1 * x_target, -1 * x_src, fx_src,
-                                 rising=False, interpolation=index_interp,
-                                 extrapolation=extrap_direct)
-        assert_array_equal(r1, r2)
+        if rising is not None:
+            r2 = stratify.interpolate(-1 * x_target, -1 * x_src, fx_src,
+                                     rising=not rising, interpolation=index_interp,
+                                     extrapolation=extrap_direct)
+            assert_array_equal(r1, r2)
 
         return r1
 
@@ -46,7 +48,7 @@ def test_lower_extrap_only(self):
         assert_array_equal(r, [-np.inf, -np.inf, -np.inf])
 
     def test_upper_extrap_only(self):
-        r = self.interpolate([1, 2, 3], [-4, -5])
+        r = self.interpolate([1, 2, 3], [-4, -5], rising=True)
         assert_array_equal(r, [np.inf, np.inf, np.inf])
 
     def test_extrap_on_both_sides_only(self):
@@ -65,7 +67,7 @@ def test_nan_in_target(self):
     def test_nan_in_src(self):
         msg = 'The source coordinate .* NaN'
         with self.assertRaisesRegexp(ValueError, msg):
-            self.interpolate([1], [0, np.nan])
+            self.interpolate([1], [0, np.nan], rising=True)
 
     def test_all_nan_in_src(self):
         r = self.interpolate([1, 2, 3, 4], [np.nan, np.nan, np.nan])
@@ -86,13 +88,13 @@ def test_wrong_rising_target(self):
         assert_array_equal(r, [1, np.inf])
 
     def test_wrong_rising_source(self):
-        r = self.interpolate([1, 2], [2, 1])
+        r = self.interpolate([1, 2], [2, 1], rising=True)
         assert_array_equal(r, [-np.inf, 0])
 
     def test_wrong_rising_source_and_target(self):
         # If we overshoot the first level, there is no hope,
         # so we end up extrapolating.
-        r = self.interpolate([3, 2, 1, 0], [2, 1])
+        r = self.interpolate([3, 2, 1, 0], [2, 1], rising=True)
         assert_array_equal(r, [np.inf, np.inf, np.inf, np.inf])
 
     def test_non_monotonic_coordinate_interp(self):
@@ -103,6 +105,20 @@ def test_non_monotonic_coordinate_extrap(self):
         result = self.interpolate([0, 15, 16, 17, 5, 15., 25], [10., 40, 0, 20])
         assert_array_equal(result, [-np.inf, 1, 1, 1, 2, 3, np.inf])
 
+    def test_length_one_interp(self):
+        r = self.interpolate([1], [2], rising=True)
+        assert_array_equal(r, [-np.inf])
+
+    def test_auto_rising_not_enough_values(self):
+        with self.assertRaises(ValueError):
+            r = self.interpolate([1], [2])
+
+    def test_auto_rising_equal_values(self):
+        # The code checks whether the first value is <= or equal to
+        # the second. If it didn't, we'd end up with +inf, not -inf.
+        r = self.interpolate([1], [2, 2])
+        assert_array_equal(r, [-np.inf])
+
 
 class Test_INTERPOLATE_LINEAR(unittest.TestCase):
     def interpolate(self, x_target):
@@ -129,6 +145,19 @@ def test_high_precision(self):
         assert_array_almost_equal(self.interpolate([1.123456789]),
                                   [11.23456789], decimal=6)
 
+    def test_single_point(self):
+        # Test that a single input point that falls exactly on the target
+        # level triggers a shortcut that avoids the expectation of >=2 source
+        # points.
+        interpolation = stratify.INTERPOLATE_LINEAR
+        extrapolation = vinterp._TestableDirectionExtrapKernel()
+
+        r = stratify.interpolate([2], [2], [20],
+                                 interpolation=interpolation,
+                                 extrapolation=extrapolation,
+                                 rising=True)
+        self.assertEqual(r, 20) 
+
 
 class Test_INTERPOLATE_NEAREST(unittest.TestCase):
     def interpolate(self, x_target):
@@ -140,8 +169,8 @@ def interpolate(self, x_target):
 
         # Use -2 to test negative number support.
         return stratify.interpolate(np.array(x_target) - 2, x_src - 2, fx_src,
-                                   interpolation=interpolation,
-                                   extrapolation=extrapolation)
+                                    interpolation=interpolation,
+                                    extrapolation=extrapolation)
 
     def test_on_the_mark(self):
         assert_array_equal(self.interpolate([0, 1, 2, 3, 4]),
@@ -167,8 +196,8 @@ def interpolate(self, x_target):
 
         # Use -2 to test negative number support.
         return stratify.interpolate(np.array(x_target) - 2, x_src - 2, fx_src,
-                                   interpolation=interpolation,
-                                   extrapolation=extrapolation)
+                                    interpolation=interpolation,
+                                    extrapolation=extrapolation)
 
     def test_below(self):
         assert_array_equal(self.interpolate([-1]), [np.nan])
@@ -187,8 +216,8 @@ def interpolate(self, x_target):
 
         # Use -2 to test negative number support.
         return stratify.interpolate(np.array(x_target) - 2, x_src - 2, fx_src,
-                                   interpolation=interpolation,
-                                   extrapolation=extrapolation)
+                                    interpolation=interpolation,
+                                    extrapolation=extrapolation)
 
     def test_below(self):
         assert_array_equal(self.interpolate([-1]), [0.])
@@ -197,6 +226,40 @@ def test_above(self):
         assert_array_equal(self.interpolate([5]), [40])
 
 
+class Test_EXTRAPOLATE_LINEAR(unittest.TestCase):
+    def interpolate(self, x_target):
+        interpolation = vinterp._TestableIndexInterpKernel()
+        extrapolation = stratify.EXTRAPOLATE_LINEAR
+
+        x_src = np.arange(5)
+        # To spice things up a bit, let's make x_src non-equal distance.
+        x_src[4] = 9
+        fx_src = 10 * x_src
+
+        # Use -2 to test negative number support.
+        return stratify.interpolate(np.array(x_target) - 2, x_src - 2, fx_src,
+                                    interpolation=interpolation,
+                                    extrapolation=extrapolation)
+
+    def test_below(self):
+        assert_array_equal(self.interpolate([-1]), [-10.])
+
+    def test_above(self):
+        assert_array_almost_equal(self.interpolate([15.123]), [151.23])
+
+    def test_npts(self):
+        interpolation = vinterp._TestableIndexInterpKernel()
+        extrapolation = stratify.EXTRAPOLATE_LINEAR
+
+        msg = (r'Linear extrapolation requires at least 2 '
+               r'source points. Got 1.')
+
+        with self.assertRaisesRegexp(ValueError, msg):
+            stratify.interpolate([1, 3.], [2], [20],
+                                 interpolation=interpolation,
+                                 extrapolation=extrapolation, rising=True)
+
+
 class Test__Interpolator(unittest.TestCase):
     def test_axis_m1(self):
         data = np.empty([5, 4, 23, 7, 3])