diff options
| -rw-r--r-- | numpy/core/src/multiarray/compiled_base.c | 191 |
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); |