Merge pull request #5843 from jaimefrio/interp_bug

BUG: Fix interp issues with small arrays

1 files changed, 111 insertions, 80 deletions

diff --git a/numpy/core/src/multiarray/compiled_base.c b/numpy/core/src/multiarray/compiled_base.c
index 825ca8c7a..aaedbac61 100644
--- a/numpy/core/src/multiarray/compiled_base.c
+++ b/numpy/core/src/multiarray/compiled_base.c
@@ -492,6 +492,8 @@ fail:
     return NULL;
 }
 
+#define LIKELY_IN_CACHE_SIZE 8
+
 /** @brief find index of a sorted array such that arr[i] <= key < arr[i + 1].
  *
  * If an starting index guess is in-range, the array values around this
@@ -512,12 +514,13 @@ fail:
  * @return index
  */
 static npy_intp
-binary_search_with_guess(double key, double arr [], npy_intp len,
-                         npy_intp guess)
+binary_search_with_guess(const npy_double key, const npy_double *arr,
+                         npy_intp len, npy_intp guess)
 {
     npy_intp imin = 0;
     npy_intp imax = len;
 
+    /* Handle keys outside of the arr range first */
     if (key > arr[len - 1]) {
         return len;
     }
@@ -525,35 +528,63 @@ binary_search_with_guess(double key, double arr [], npy_intp len,
         return -1;
     }
 
-    if (guess < 0)  {
-        guess = 0;
-    }
-    else if (guess >= len - 1) {
-        guess = len - 2;
-    }
+    /*
+     * It would seem that for the following code to work, 'len' should
+     * at least be 4. But because of the way 'guess' is normalized, it
+     * will always be set to 1 if len <= 4. Given that, and that keys
+     * outside of the 'arr' bounds have already been handled, and the
+     * order in which comparisons happen below, it should become obvious
+     * that it will work with any array of at least 2 items.
+     */
+    assert (len >= 2);
 
-    /* check most likely values: guess, guess + 1, guess - 1 */
-    if ((key > arr[guess]) && (key <= arr[guess + 1])) {
-        return guess;
-    }
-    else if ((guess < len - 2) && (key > arr[guess + 1]) &&
-             (key <= arr[guess + 2])) {
-        return guess + 1;
+    if (guess > len - 3) {
+        guess = len - 3;
     }
-    else if ((guess > 1) && (key > arr[guess - 1]) &&
-             (key <= arr[guess])) {
-        return guess - 1;
+    if (guess < 1)  {
+        guess = 1;
     }
-    /* may be able to restrict bounds to range likely to be in memory */
-    if ((guess > 8) && (key >  arr[guess - 8])) {
-        imin = guess - 8;
+
+    /* check most likely values: guess - 1, guess, guess + 1 */
+    if (key <= arr[guess]) {
+        if (key <= arr[guess - 1]) {
+            imax = guess - 1;
+            /* last attempt to restrict search to items in cache */
+            if (guess > LIKELY_IN_CACHE_SIZE &&
+                        key > arr[guess - LIKELY_IN_CACHE_SIZE]) {
+                imin = guess - LIKELY_IN_CACHE_SIZE;
+            }
+        }
+        else {
+            /* key > arr[guess - 1] */
+            return guess - 1;
+        }
     }
-    if ((guess < len - 9) && (key <= arr[guess + 8])) {
-        imax = guess + 8;
+    else {
+        /* key > arr[guess] */
+        if (key <= arr[guess + 1]) {
+            return guess;
+        }
+        else {
+            /* key > arr[guess + 1] */
+            if (key <= arr[guess + 2]) {
+                return guess + 1;
+            }
+            else {
+                /* key > arr[guess + 2] */
+                imin = guess + 2;
+                /* last attempt to restrict search to items in cache */
+                if (guess < len - LIKELY_IN_CACHE_SIZE - 1 &&
+                            key <= arr[guess + LIKELY_IN_CACHE_SIZE]) {
+                    imax = guess + LIKELY_IN_CACHE_SIZE;
+                }
+            }
+        }
     }
+
     /* finally, find index by bisection */
     while (imin < imax) {
-        npy_intp imid = imin + ((imax - imin) >> 1);
+        const npy_intp imid = imin + ((imax - imin) >> 1);
         if (key >= arr[imid]) {
             imin = imid + 1;
         }
@@ -564,6 +595,8 @@ binary_search_with_guess(double key, double arr [], npy_intp len,
     return imin - 1;
 }
 
+#undef LIKELY_IN_CACHE_SIZE
+
 NPY_NO_EXPORT PyObject *
 arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
 {
@@ -571,12 +604,15 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
     PyObject *fp, *xp, *x;
     PyObject *left = NULL, *right = NULL;
     PyArrayObject *afp = NULL, *axp = NULL, *ax = NULL, *af = NULL;
-    npy_intp i, lenx, lenxp, j, jprev;
-    double lval, rval;
-    double *dy, *dx, *dz, *dres, *slopes;
+    npy_intp i, lenx, lenxp;
+    npy_double lval, rval;
+    const npy_double *dy, *dx, *dz;
+    npy_double *dres, *slopes = NULL;
 
     static char *kwlist[] = {"x", "xp", "fp", "left", "right", NULL};
 
+    NPY_BEGIN_THREADS_DEF;
+
     if (!PyArg_ParseTupleAndKeywords(args, kwdict, "OOO|OO", kwlist,
                                      &x, &xp, &fp, &left, &right)) {
         return NULL;
@@ -594,29 +630,29 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
     if (ax == NULL) {
         goto fail;
     }
-    lenxp = PyArray_DIMS(axp)[0];
+    lenxp = PyArray_SIZE(axp);
     if (lenxp == 0) {
         PyErr_SetString(PyExc_ValueError,
                 "array of sample points is empty");
         goto fail;
     }
-    if (PyArray_DIMS(afp)[0] != lenxp) {
+    if (PyArray_SIZE(afp) != lenxp) {
         PyErr_SetString(PyExc_ValueError,
                 "fp and xp are not of the same length.");
         goto fail;
     }
 
     af = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(ax),
-                                        PyArray_DIMS(ax), NPY_DOUBLE);
+                                            PyArray_DIMS(ax), NPY_DOUBLE);
     if (af == NULL) {
         goto fail;
     }
     lenx = PyArray_SIZE(ax);
 
-    dy = (double *)PyArray_DATA(afp);
-    dx = (double *)PyArray_DATA(axp);
-    dz = (double *)PyArray_DATA(ax);
-    dres = (double *)PyArray_DATA(af);
+    dy = (const npy_double *)PyArray_DATA(afp);
+    dx = (const npy_double *)PyArray_DATA(axp);
+    dz = (const npy_double *)PyArray_DATA(ax);
+    dres = (npy_double *)PyArray_DATA(af);
     /* Get left and right fill values. */
     if ((left == NULL) || (left == Py_None)) {
         lval = dy[0];
@@ -628,7 +664,7 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
         }
     }
     if ((right == NULL) || (right == Py_None)) {
-        rval = dy[lenxp-1];
+        rval = dy[lenxp - 1];
     }
     else {
         rval = PyFloat_AsDouble(right);
@@ -637,72 +673,67 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict)
         }
     }
 
-    /* only pre-calculate slopes if there are relatively few of them. */
-    j = jprev = 0;
-    if (lenxp <= lenx) {
-        slopes = (double *) PyArray_malloc((lenxp - 1)*sizeof(double));
-        if (! slopes) {
-            goto fail;
-        }
-        NPY_BEGIN_ALLOW_THREADS;
-        for (i = 0; i < lenxp - 1; i++) {
-            slopes[i] = (dy[i + 1] - dy[i])/(dx[i + 1] - dx[i]);
+    /* binary_search_with_guess needs at least a 2 item long array */
+    if (lenxp == 1) {
+        const npy_double xp_val = dx[0];
+        const npy_double fp_val = dy[0];
+
+        NPY_BEGIN_THREADS_THRESHOLDED(lenx);
+        for (i = 0; i < lenx; ++i) {
+            const npy_double x_val = dz[i];
+            dres[i] = (x_val < xp_val) ? lval :
+                                         ((x_val > xp_val) ? rval : fp_val);
         }
-        for (i = 0; i < lenx; i++) {
-            const double x = dz[i];
+        NPY_END_THREADS;
+    }
+    else {
+        npy_intp j = 0;
 
-            if (npy_isnan(x)) {
-                dres[i] = x;
-                continue;
+        /* only pre-calculate slopes if there are relatively few of them. */
+        if (lenxp <= lenx) {
+            slopes = PyArray_malloc((lenxp - 1) * sizeof(npy_double));
+            if (slopes == NULL) {
+                goto fail;
             }
+        }
 
-            j = binary_search_with_guess(x, dx, lenxp, jprev);
-            jprev = j;
-            if (j == -1) {
-                dres[i] = lval;
-            }
-            else if (j == lenxp - 1) {
-                dres[i] = dy[j];
-            }
-            else if (j == lenxp) {
-                dres[i] = rval;
-            }
-            else {
-                dres[i] = slopes[j]*(x - dx[j]) + dy[j];
+        NPY_BEGIN_THREADS;
+
+        if (slopes != NULL) {
+            for (i = 0; i < lenxp - 1; ++i) {
+                slopes[i] = (dy[i+1] - dy[i]) / (dx[i+1] - dx[i]);
             }
         }
-        NPY_END_ALLOW_THREADS;
-        PyArray_free(slopes);
-    }
-    else {
-        NPY_BEGIN_ALLOW_THREADS;
-        for (i = 0; i < lenx; i++) {
-            const double x = dz[i];
 
-            if (npy_isnan(x)) {
-                dres[i] = x;
+        for (i = 0; i < lenx; ++i) {
+            const npy_double x_val = dz[i];
+
+            if (npy_isnan(x_val)) {
+                dres[i] = x_val;
                 continue;
             }
 
-            j = binary_search_with_guess(x, dx, lenxp, jprev);
-            jprev = j;
+            j = binary_search_with_guess(x_val, dx, lenxp, j);
             if (j == -1) {
                 dres[i] = lval;
             }
-            else if (j == lenxp - 1) {
-                dres[i] = dy[j];
-            }
             else if (j == lenxp) {
                 dres[i] = rval;
             }
+            else if (j == lenxp - 1) {
+                dres[i] = dy[j];
+            }
             else {
-                const double slope = (dy[j + 1] - dy[j])/(dx[j + 1] - dx[j]);
-                dres[i] = slope*(x - dx[j]) + dy[j];
+                const npy_double slope = (slopes != NULL) ? slopes[j] :
+                                         (dy[j+1] - dy[j]) / (dx[j+1] - dx[j]);
+                dres[i] = slope*(x_val - dx[j]) + dy[j];
             }
         }
-        NPY_END_ALLOW_THREADS;
+
+        NPY_END_THREADS;
     }
 
+    PyArray_free(slopes);
     Py_DECREF(afp);
     Py_DECREF(axp);
     Py_DECREF(ax);