diff --git a/doc/release/1.12.0-notes.rst b/doc/release/1.12.0-notes.rst index 98dd3b31697e..07871c10ceef 100644 --- a/doc/release/1.12.0-notes.rst +++ b/doc/release/1.12.0-notes.rst @@ -132,6 +132,11 @@ New nanfunctions ``nancumsum`` and ``nancumprod`` added Nanfunctions ``nancumsum`` and ``nancumprod`` have been added to compute ``cumsum`` and ``cumprod`` by ignoring nans. +``np.interp`` can now interpolate complex values +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +``np.lib.interp(x, xp, fp)`` now allows the interpolated array ``fp`` +to be complex and will interpolate at ``complex128`` precision. + Improvements ============ diff --git a/numpy/core/src/multiarray/compiled_base.c b/numpy/core/src/multiarray/compiled_base.c index 136d0859ea0b..213d74f44031 100644 --- a/numpy/core/src/multiarray/compiled_base.c +++ b/numpy/core/src/multiarray/compiled_base.c @@ -664,6 +664,182 @@ arr_interp(PyObject *NPY_UNUSED(self), PyObject *args, PyObject *kwdict) return NULL; } +/* As for arr_interp but for complex fp values */ +NPY_NO_EXPORT PyObject * +arr_interp_complex(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; + + const npy_double *dx, *dz; + const npy_cdouble *dy; + npy_cdouble lval, rval; + npy_cdouble *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; + } + + afp = (PyArrayObject *)PyArray_ContiguousFromAny(fp, NPY_CDOUBLE, 1, 1); + + if (afp == NULL) { + return NULL; + } + + axp = (PyArrayObject *)PyArray_ContiguousFromAny(xp, NPY_DOUBLE, 1, 1); + if (axp == NULL) { + goto fail; + } + ax = (PyArrayObject *)PyArray_ContiguousFromAny(x, NPY_DOUBLE, 1, 0); + if (ax == NULL) { + goto fail; + } + lenxp = PyArray_SIZE(axp); + if (lenxp == 0) { + PyErr_SetString(PyExc_ValueError, + "array of sample points is empty"); + goto fail; + } + if (PyArray_SIZE(afp) != lenxp) { + PyErr_SetString(PyExc_ValueError, + "fp and xp are not of the same length."); + goto fail; + } + + lenx = PyArray_SIZE(ax); + dx = (const npy_double *)PyArray_DATA(axp); + dz = (const npy_double *)PyArray_DATA(ax); + + af = (PyArrayObject *)PyArray_SimpleNew(PyArray_NDIM(ax), + PyArray_DIMS(ax), NPY_CDOUBLE); + if (af == NULL) { + goto fail; + } + + dy = (const npy_cdouble *)PyArray_DATA(afp); + dres = (npy_cdouble *)PyArray_DATA(af); + /* Get left and right fill values. */ + if ((left == NULL) || (left == Py_None)) { + lval = dy[0]; + } + else { + lval.real = PyComplex_RealAsDouble(left); + if ((lval.real == -1) && PyErr_Occurred()) { + goto fail; + } + lval.imag = PyComplex_ImagAsDouble(left); + if ((lval.imag == -1) && PyErr_Occurred()) { + goto fail; + } + } + + if ((right == NULL) || (right == Py_None)) { + rval = dy[lenxp - 1]; + } + else { + rval.real = PyComplex_RealAsDouble(right); + if ((rval.real == -1) && PyErr_Occurred()) { + goto fail; + } + rval.imag = PyComplex_ImagAsDouble(right); + if ((rval.imag == -1) && PyErr_Occurred()) { + goto fail; + } + } + + /* binary_search_with_guess needs at least a 3 item long array */ + if (lenxp == 1) { + const npy_double xp_val = dx[0]; + const npy_cdouble 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); + } + NPY_END_THREADS; + } + else { + npy_intp j = 0; + + /* only pre-calculate slopes if there are relatively few of them. */ + if (lenxp <= lenx) { + slopes = PyArray_malloc((lenxp - 1) * sizeof(npy_cdouble)); + if (slopes == NULL) { + goto fail; + } + } + + NPY_BEGIN_THREADS; + + if (slopes != NULL) { + for (i = 0; i < lenxp - 1; ++i) { + const double inv_dx = 1.0 / (dx[i+1] - dx[i]); + slopes[i].real = (dy[i+1].real - dy[i].real) * inv_dx; + slopes[i].imag = (dy[i+1].imag - dy[i].imag) * inv_dx; + } + } + + for (i = 0; i < lenx; ++i) { + const npy_double x_val = dz[i]; + + if (npy_isnan(x_val)) { + dres[i].real = x_val; + dres[i].imag = 0.0; + continue; + } + + j = binary_search_with_guess(x_val, dx, lenxp, j); + if (j == -1) { + dres[i] = lval; + } + else if (j == lenxp) { + dres[i] = rval; + } + else if (j == lenxp - 1) { + dres[i] = dy[j]; + } + else { + if (slopes!=NULL) { + dres[i].real = slopes[j].real*(x_val - dx[j]) + dy[j].real; + dres[i].imag = slopes[j].imag*(x_val - dx[j]) + dy[j].imag; + } + else { + const npy_double inv_dx = 1.0 / (dx[j+1] - dx[j]); + dres[i].real = (dy[j+1].real - dy[j].real)*(x_val - dx[j])* + inv_dx + dy[j].real; + dres[i].imag = (dy[j+1].imag - dy[j].imag)*(x_val - dx[j])* + inv_dx + dy[j].imag; + } + } + } + + NPY_END_THREADS; + } + PyArray_free(slopes); + + Py_DECREF(afp); + Py_DECREF(axp); + Py_DECREF(ax); + return (PyObject *)af; + +fail: + Py_XDECREF(afp); + Py_XDECREF(axp); + Py_XDECREF(ax); + Py_XDECREF(af); + return NULL; +} + /* * Converts a Python sequence into 'count' PyArrayObjects * diff --git a/numpy/core/src/multiarray/compiled_base.h b/numpy/core/src/multiarray/compiled_base.h index 19e3778adfc3..51508531c5a3 100644 --- a/numpy/core/src/multiarray/compiled_base.h +++ b/numpy/core/src/multiarray/compiled_base.h @@ -11,6 +11,8 @@ arr_digitize(PyObject *, PyObject *, PyObject *kwds); NPY_NO_EXPORT PyObject * arr_interp(PyObject *, PyObject *, PyObject *); NPY_NO_EXPORT PyObject * +arr_interp_complex(PyObject *, PyObject *, PyObject *); +NPY_NO_EXPORT PyObject * arr_ravel_multi_index(PyObject *, PyObject *, PyObject *); NPY_NO_EXPORT PyObject * arr_unravel_index(PyObject *, PyObject *, PyObject *); diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index e2731068b5ef..62b5628566e9 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4233,6 +4233,8 @@ static struct PyMethodDef array_module_methods[] = { METH_VARARGS | METH_KEYWORDS, NULL}, {"interp", (PyCFunction)arr_interp, METH_VARARGS | METH_KEYWORDS, NULL}, + {"interp_complex", (PyCFunction)arr_interp_complex, + METH_VARARGS | METH_KEYWORDS, NULL}, {"ravel_multi_index", (PyCFunction)arr_ravel_multi_index, METH_VARARGS | METH_KEYWORDS, NULL}, {"unravel_index", (PyCFunction)arr_unravel_index, diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py index 8f15fc547b74..3533a59fcafd 100644 --- a/numpy/lib/function_base.py +++ b/numpy/lib/function_base.py @@ -23,8 +23,10 @@ from numpy.core.numerictypes import typecodes, number from numpy.lib.twodim_base import diag from .utils import deprecate -from numpy.core.multiarray import _insert, add_docstring -from numpy.core.multiarray import digitize, bincount, interp as compiled_interp +from numpy.core.multiarray import ( + _insert, add_docstring, digitize, bincount, + interp as compiled_interp, interp_complex as compiled_interp_complex + ) from numpy.core.umath import _add_newdoc_ufunc as add_newdoc_ufunc from numpy.compat import long from numpy.compat.py3k import basestring @@ -1663,13 +1665,13 @@ def interp(x, xp, fp, left=None, right=None, period=None): `period` is not specified. Otherwise, `xp` is internally sorted after normalizing the periodic boundaries with ``xp = xp % period``. - fp : 1-D sequence of floats + fp : 1-D sequence of float or complex The y-coordinates of the data points, same length as `xp`. - left : float, optional + left : optional float or complex corresponding to fp Value to return for `x < xp[0]`, default is `fp[0]`. - right : float, optional + right : optional float or complex corresponding to fp Value to return for `x > xp[-1]`, default is `fp[-1]`. period : None or float, optional @@ -1681,7 +1683,7 @@ def interp(x, xp, fp, left=None, right=None, period=None): Returns ------- - y : float or ndarray + y : float or complex (corresponding to fp) or ndarray The interpolated values, same shape as `x`. Raises @@ -1732,14 +1734,31 @@ def interp(x, xp, fp, left=None, right=None, period=None): >>> np.interp(x, xp, fp, period=360) array([7.5, 5., 8.75, 6.25, 3., 3.25, 3.5, 3.75]) + Complex interpolation + >>> x = [1.5, 4.0] + >>> xp = [2,3,5] + >>> fp = [1.0j, 0, 2+3j] + >>> np.interp(x, xp, fp) + array([ 0.+1.j , 1.+1.5j]) + """ + + fp = np.asarray(fp) + + if np.iscomplexobj(fp): + interp_func = compiled_interp_complex + input_dtype = np.complex128 + else: + interp_func = compiled_interp + input_dtype = np.float64 + if period is None: if isinstance(x, (float, int, number)): - return compiled_interp([x], xp, fp, left, right).item() + return interp_func([x], xp, fp, left, right).item() elif isinstance(x, np.ndarray) and x.ndim == 0: - return compiled_interp([x], xp, fp, left, right).item() + return interp_func([x], xp, fp, left, right).item() else: - return compiled_interp(x, xp, fp, left, right) + return interp_func(x, xp, fp, left, right) else: if period == 0: raise ValueError("period must be a non-zero value") @@ -1752,7 +1771,8 @@ def interp(x, xp, fp, left=None, right=None, period=None): x = [x] x = np.asarray(x, dtype=np.float64) xp = np.asarray(xp, dtype=np.float64) - fp = np.asarray(fp, dtype=np.float64) + fp = np.asarray(fp, dtype=input_dtype) + if xp.ndim != 1 or fp.ndim != 1: raise ValueError("Data points must be 1-D sequences") if xp.shape[0] != fp.shape[0]: @@ -1765,12 +1785,12 @@ def interp(x, xp, fp, left=None, right=None, period=None): fp = fp[asort_xp] xp = np.concatenate((xp[-1:]-period, xp, xp[0:1]+period)) fp = np.concatenate((fp[-1:], fp, fp[0:1])) + if return_array: - return compiled_interp(x, xp, fp, left, right) + return interp_func(x, xp, fp, left, right) else: - return compiled_interp(x, xp, fp, left, right).item() - - + return interp_func(x, xp, fp, left, right).item() + def angle(z, deg=0): """ Return the angle of the complex argument. diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py index 0442792947fa..0f71393ad148 100644 --- a/numpy/lib/tests/test_function_base.py +++ b/numpy/lib/tests/test_function_base.py @@ -2235,6 +2235,28 @@ def test_scalar_interpolation_point(self): assert_almost_equal(np.interp(x0, x, y), x0) x0 = np.nan assert_almost_equal(np.interp(x0, x, y), x0) + + def test_complex_interp(self): + # test complex interpolation + x = np.linspace(0, 1, 5) + y = np.linspace(0, 1, 5) + (1 + np.linspace(0, 1, 5))*1.0j + x0 = 0.3 + y0 = x0 + (1+x0)*1.0j + assert_almost_equal(np.interp(x0, x, y), y0) + # test complex left and right + x0 = -1 + left = 2 + 3.0j + assert_almost_equal(np.interp(x0, x, y, left=left), left) + x0 = 2.0 + right = 2 + 3.0j + assert_almost_equal(np.interp(x0, x, y, right=right), right) + # test complex periodic + x = [-180, -170, -185, 185, -10, -5, 0, 365] + xp = [190, -190, 350, -350] + fp = [5+1.0j, 10+2j, 3+3j, 4+4j] + y = [7.5+1.5j, 5.+1.0j, 8.75+1.75j, 6.25+1.25j, 3.+3j, 3.25+3.25j, + 3.5+3.5j, 3.75+3.75j] + assert_almost_equal(np.interp(x, xp, fp, period=360), y) def test_zero_dimensional_interpolation_point(self): x = np.linspace(0, 1, 5)