diff --git a/doc/source/reference/c-api.array.rst b/doc/source/reference/c-api.array.rst index 049c7b5379f1..fc3845f87cd3 100644 --- a/doc/source/reference/c-api.array.rst +++ b/doc/source/reference/c-api.array.rst @@ -2161,6 +2161,32 @@ Array Functions z[k] = sum_n op1[n] * conj(op2[n+k]) +.. c:function:: PyObject* PyArray_CorrelateLags(PyObject* op1, PyObject* op2, npy_intp minlag, npy_intp maxlag, npy_intp lagstep) + + Updated version of PyArray_Correlate2, which allows for custom lags in calculation. + The correlation is computed at each output point by multiplying *op1* by + a shifted version of *op2* and summing the result. + As a result of the shift, needed values outside of the defined range of + *op1* and *op2* are interpreted as zero. + + This version differs from the other PyArray_Correlate* functions in that + PyArray_CorrelateLags calculates the correlation for lags starting at minlag + and ending at (maxlag - 1), with steps of size lagstep. + The lags for which the correlation will be calculated correspond with the values in + the vector formed by numpy.arange((minlag, maxlag, lagstep)). + + .. rubric:: Notes + + The following lag values should be used to reproduce the mode parameter used by + PyArray_Correlate2 (n1 and n2 are the vector lengths): + mode | minlag | maxlag | lagstep + 0 ("valid") | 0 | n1-n2+1 | 1 + 1 ("same") | -n2/2 | n1-n2/2 | 1 + 2 ("full") | -n2+1 | n1 | 1 + + PyArray_CorrelateLags uses the same definition of correlation as PyArray_Correlate2 + in that the conjugate is taken and the argument order matters. + .. c:function:: PyObject* PyArray_Where(PyObject* condition, PyObject* x, PyObject* y) If both ``x`` and ``y`` are ``NULL``, then return diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py index 972966627b99..d6a0fbdd70e6 100644 --- a/numpy/core/code_generators/numpy_api.py +++ b/numpy/core/code_generators/numpy_api.py @@ -344,6 +344,8 @@ # End 1.9 API 'PyArray_CheckAnyScalarExact': (300, NonNull(1)), # End 1.10 API + 'PyArray_CorrelateLags': (301,), + # End 1.11 API } ufunc_types_api = { diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index b3eed97149dc..8fed76d6d6c3 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -23,7 +23,7 @@ from . import umath from .umath import (invert, sin, UFUNC_BUFSIZE_DEFAULT, ERR_IGNORE, ERR_WARN, ERR_RAISE, ERR_CALL, ERR_PRINT, ERR_LOG, - ERR_DEFAULT, PINF, NAN) + ERR_DEFAULT, PINF, NAN, ceil) from . import numerictypes from .numerictypes import longlong, intc, int_, float_, complex_, bool_ from ._internal import TooHardError @@ -377,7 +377,8 @@ def extend_all(module): def asarray(a, dtype=None, order=None): - """Convert the input to an array. + """ + Convert the input to an array. Parameters ---------- @@ -388,9 +389,8 @@ def asarray(a, dtype=None, order=None): dtype : data-type, optional By default, the data-type is inferred from the input data. order : {'C', 'F'}, optional - Whether to use row-major (C-style) or - column-major (Fortran-style) memory representation. - Defaults to 'C'. + Whether to use row-major ('C') or column-major ('F' for FORTRAN) + memory representation. Defaults to 'C'. Returns ------- @@ -448,7 +448,8 @@ class ndarray is returned. def asanyarray(a, dtype=None, order=None): - """Convert the input to an ndarray, but pass ndarray subclasses through. + """ + Convert the input to an ndarray, but pass ndarray subclasses through. Parameters ---------- @@ -459,8 +460,8 @@ def asanyarray(a, dtype=None, order=None): dtype : data-type, optional By default, the data-type is inferred from the input data. order : {'C', 'F'}, optional - Whether to use row-major (C-style) or column-major - (Fortran-style) memory representation. Defaults to 'C'. + Whether to use row-major ('C') or column-major ('F') memory + representation. Defaults to 'C'. Returns ------- @@ -814,16 +815,55 @@ def flatnonzero(a): _mode_from_name_dict = {'v': 0, 's': 1, - 'f': 2} - + 'f': 2, + 'l': 3, + } def _mode_from_name(mode): - if isinstance(mode, basestring): + if type(mode) is int: + assert(mode in _mode_from_name_dict.values()) + return mode + elif isinstance(mode, basestring): return _mode_from_name_dict[mode.lower()[0]] return mode - -def correlate(a, v, mode='valid'): +def _lags_from_lags(l): + if type(l) is int: # maxlag + lags = (-l+1, l, 1) + elif type(l) is tuple: # minlag and maxlag + if len(l) > 2: + lags = (int(l[0]), int(l[1]), int(l[2])) + else: + lags = (int(l[0]), int(l[1]), 1) + return lags + +def _lags_from_mode(alen, vlen, mode, lags=()): + inverted = 0 + if alen < vlen: + alen, vlen = vlen, alen + inverted = 1 + + if mode is 0: + mode_lags = (0, alen-vlen+1, 1) + elif mode is 1: + mode_lags = (-int(vlen/2), alen - int(vlen/2), 1) + elif mode is 2: + mode_lags = (-vlen+1, alen, 1) + elif mode is 3: + mode_lags = _lags_from_lags(lags) + + if inverted and mode != 3: + mode_lags = (-int(ceil((mode_lags[1]-mode_lags[0])/float(mode_lags[2]))) + *mode_lags[2]-mode_lags[0]+mode_lags[2], + -mode_lags[0]+mode_lags[2], mode_lags[2]) + + if lags: + lags = _lags_from_lags(lags) + assert lags == mode_lags, \ + 'Custom lags should be passed to correlate() with mode="lags".' + return mode_lags + +def correlate(a, v, mode='default', lags=(), returns_lagvector=False): """ Cross-correlation of two 1-dimensional sequences. @@ -839,9 +879,18 @@ def correlate(a, v, mode='valid'): ---------- a, v : array_like Input sequences. - mode : {'valid', 'same', 'full'}, optional + mode : {'valid', 'same', 'full', 'lags'}, optional Refer to the `convolve` docstring. Note that the default - is 'valid', unlike `convolve`, which uses 'full'. + is `valid`, unlike `convolve`, which uses `full`. + lags : int or int tuple, optional + Refer to the `convolve` docstring. + mode should be unset or set to 'lags' to use the lags argument + returns_lagvector : bool, optional + If True, the function returns a lagvector array in addition to the + cross-correlation result. The lagvector contains the indices of + the lags for which the cross-correlation was calculated. It is + the same length as the return array, and corresponds one-to-one. + False is default. old_behavior : bool `old_behavior` was removed in NumPy 1.10. If you need the old behavior, use `multiarray.correlate`. @@ -850,6 +899,9 @@ def correlate(a, v, mode='valid'): ------- out : ndarray Discrete cross-correlation of `a` and `v`. + lagvector : ndarray, optional + The indices of the lags for which the cross-correlation was calculated. + It is the same length as out, and corresponds one-to-one. See Also -------- @@ -869,10 +921,14 @@ def correlate(a, v, mode='valid'): -------- >>> np.correlate([1, 2, 3], [0, 1, 0.5]) array([ 3.5]) - >>> np.correlate([1, 2, 3], [0, 1, 0.5], "same") + >>> np.correlate([1, 2, 3], [0, 1, 0.5], mode="same") array([ 2. , 3.5, 3. ]) - >>> np.correlate([1, 2, 3], [0, 1, 0.5], "full") - array([ 0.5, 2. , 3.5, 3. , 0. ]) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], mode="full", returns_lagvector=True) + (array([ 0.5, 2. , 3.5, 3. , 0. ]), array([-2, -1, 0, 1, 2])) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], mode="lags", lags=2) + array([ 2. , 3.5, 3. ]) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], mode="lags", lags=(-1,2,2), returns_lagvector=True) + (array([ 2., 3.]), array([-1, 1])) Using complex sequences: @@ -887,11 +943,24 @@ def correlate(a, v, mode='valid'): array([ 0.0+0.j , 3.0+1.j , 1.5+1.5j, 1.0+0.j , 0.5+0.5j]) """ + if mode == 'default': + if lags: + mode = 'lags' + else: + mode = 'valid' mode = _mode_from_name(mode) - return multiarray.correlate2(a, v, mode) + if mode in (0, 1, 2) and not returns_lagvector: + return multiarray.correlate2(a, v, mode) + lags = _lags_from_mode(len(a), len(v), mode, lags) + if returns_lagvector: + return multiarray.correlate2(a, v, 3, lags[0], lags[1], lags[2]), \ + arange(lags[0], lags[1], lags[2]) + else: + return multiarray.correlate2(a, v, 3, lags[0], lags[1], lags[2]) -def convolve(a,v,mode='full'): + +def convolve(a, v, mode='full', lags=(), returns_lagvector=False): """ Returns the discrete, linear convolution of two one-dimensional sequences. @@ -909,27 +978,59 @@ def convolve(a,v,mode='full'): First one-dimensional input array. v : (M,) array_like Second one-dimensional input array. - mode : {'full', 'valid', 'same'}, optional + mode : int, int tuple, or {'full', 'valid', 'same'}, optional 'full': By default, mode is 'full'. This returns the convolution at each point of overlap, with an output shape of (N+M-1,). At the end-points of the convolution, the signals do not overlap - completely, and boundary effects may be seen. + completely, and boundary effects may be seen. This corresponds + with a lag tuple of (-M+1, N, 1) for N>M or (-N+1, M, 1) + for M>N. 'same': - Mode 'same' returns output of length ``max(M, N)``. Boundary - effects are still visible. + Mode `same` returns output of length ``max(M, N)``. Boundary + effects are still visible. This corresponds with a lag tuple of + (-M/2, N-M/2, 1) for N>M or (-M+N/2+1, N/2+1, 1) for M>N. 'valid': Mode 'valid' returns output of length ``max(M, N) - min(M, N) + 1``. The convolution product is only given for points where the signals overlap completely. Values outside - the signal boundary have no effect. + the signal boundary have no effect. This corresponds with a lag tuple + of (0, N-M+1, 1) for N>M or (-M+N, 1, 1) for M>N. + + 'lags': + Mode 'lags' uses the lags argument to define the lags for which + to perform the convolution. + lags : int or int tuple, optional + mode should be unset or set to 'lags' to use the lags argument + + int (maxlag): + This calculates the convolution for all lags starting at + (-maxlag + 1) and ending at (maxlag - 1), with steps of size 1. + See the optional returns_lagvec argument to get an array containing + lags corresponding to the convolution values in the return array. + + tuple (minlag, maxlag) or (minlag, maxlag, lagstep): + This calculates the convolution for all lags starting at + minlag and ending at (maxlag - 1), with steps of size lagstep. + The lags for which the convolution will be calculated correspond + with the values in the vector formed by numpy.arange() with the + same tuple argument. + returns_lagvector : bool, optional + If True, the function returns a lagvector array in addition to the + convolution result. The lagvector contains the indices of + the lags for which the convolution was calculated. It is + the same length as the return array, and corresponds one-to-one. + False is default. Returns ------- out : ndarray Discrete, linear convolution of `a` and `v`. + lagvector : ndarray, optional + The indices of the lags for which the convolution was calculated. + It is the same length as out, and corresponds one-to-one. See Also -------- @@ -968,25 +1069,58 @@ def convolve(a,v,mode='full'): Contains boundary effects, where zeros are taken into account: - >>> np.convolve([1,2,3],[0,1,0.5], 'same') + >>> np.convolve([1,2,3],[0,1,0.5], mode='same') array([ 1. , 2.5, 4. ]) The two arrays are of the same length, so there - is only one position where they completely overlap: + is only one position where they completely overlap, + corresponding to a lag of 0. lagvector=True causes + the function to return the lagvector corresponding + to the convolution in addition to the convolution + itself: + + >>> np.convolve([1,2,3],[0,1,0.5], mode='valid', returns_lagvector=True) + (array([ 2.5]), array([0])) - >>> np.convolve([1,2,3],[0,1,0.5], 'valid') - array([ 2.5]) + Find the convolution for lags ranging from -1 to 1 + (0 is the lag for which the left sides of the arrays + are aligned, -1 has the second vector to the left of + the first, and +1 has the second vector to the right + of the first): + + >>> np.convolve([1,2,3],[0,1,0.5], mode='lags', lags=2, returns_lagvector=True) + (array([ 1. , 2.5, 4. ]), array([-1, 0, 1])) + + Find the convolution for lags ranging from -2 to 4 + with steps of length 2 (the maxlag member of the + lag range tuple is non-inclusive, similar to np.arange()): + + >>> np.convolve([1,2,3,4,5],[0,1,0.5], mode='lags', lags=(-2,6,2), returns_lagvector=True) + (array([ 0. , 2.5, 5.5, 2.5]), array([-2, 0, 2, 4])) """ a, v = array(a, copy=False, ndmin=1), array(v, copy=False, ndmin=1) - if (len(v) > len(a)): - a, v = v, a if len(a) == 0: raise ValueError('a cannot be empty') if len(v) == 0: raise ValueError('v cannot be empty') + if (len(v) > len(a)): + a, v = v, a + if mode == 'default': + if lags: + mode = 'lags' + else: + mode = 'full' mode = _mode_from_name(mode) - return multiarray.correlate(a, v[::-1], mode) + if mode in (0, 1, 2) and not returns_lagvector: + return multiarray.correlate2(a, v[::-1], mode) + + lags = _lags_from_mode(len(a), len(v), mode, lags) + if returns_lagvector: + return multiarray.correlate2(a, v[::-1], 3, lags[0], lags[1], lags[2]),\ + arange(lags[0], lags[1], lags[2]) + else: + return multiarray.correlate2(a, v[::-1], 3, lags[0], lags[1], lags[2]) def outer(a, b, out=None): @@ -1096,7 +1230,6 @@ def alterdot(): restoredot : `restoredot` undoes the effects of `alterdot`. """ - # 2014-08-13, 1.10 warnings.warn("alterdot no longer does anything.", DeprecationWarning) @@ -1120,7 +1253,6 @@ def restoredot(): alterdot : `restoredot` undoes the effects of `alterdot`. """ - # 2014-08-13, 1.10 warnings.warn("restoredot no longer does anything.", DeprecationWarning) diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 62b5628566e9..50ed082a39fb 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1084,6 +1084,8 @@ PyArray_CopyAndTranspose(PyObject *op) return (PyObject *)ret; } + + /* * Implementation which is common between PyArray_Correlate * and PyArray_Correlate2. @@ -1091,18 +1093,25 @@ PyArray_CopyAndTranspose(PyObject *op) * inverted is set to 1 if computed correlate(ap2, ap1), 0 otherwise */ static PyArrayObject* -_pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, - int mode, int *inverted) +_pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, + int typenum, int *inverted, + npy_intp minlag, npy_intp maxlag, npy_intp lagstep) { PyArrayObject *ret; npy_intp length; - npy_intp i, n1, n2, n, n_left, n_right; + npy_intp i, i1, n1, n2, n, n11; + npy_intp lag, tmplag, maxleft, maxright; npy_intp is1, is2, os; char *ip1, *ip2, *op; PyArray_DotFunc *dot; NPY_BEGIN_THREADS_DEF; + if (lagstep == 0) { + lagstep = 1; + } + + /* size of x (n1) and y (n2) */ n1 = PyArray_DIMS(ap1)[0]; n2 = PyArray_DIMS(ap2)[0]; if (n1 < n2) { @@ -1113,40 +1122,41 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, i = n1; n1 = n2; n2 = i; + minlag = -minlag; + maxlag = -maxlag; + lagstep = -lagstep; + } + + if (lagstep < 0) { *inverted = 1; - } else { + i = minlag; + i1 = (npy_intp)(npy_ceil((maxlag - minlag)/(float)lagstep))*lagstep; + minlag = i1 + minlag - lagstep; + maxlag = i - lagstep; + lagstep = -lagstep; + } + else { *inverted = 0; } - - length = n1; - n = n2; - switch(mode) { - case 0: - length = length - n + 1; - n_left = n_right = 0; - break; - case 1: - n_left = (npy_intp)(n/2); - n_right = n - n_left - 1; - break; - case 2: - n_right = n - 1; - n_left = n - 1; - length = length + n - 1; - break; - default: - PyErr_SetString(PyExc_ValueError, "mode must be 0, 1, or 2"); - return NULL; + if (maxlag <= minlag) { + length = 0; + } + else { + length = (maxlag - minlag + lagstep - 1)/lagstep; } /* * Need to choose an output array that can hold a sum * -- use priority to determine which subtype. + * ret is the array that will be returned as the answer */ ret = new_array_for_sum(ap1, ap2, NULL, 1, &length, typenum); if (ret == NULL) { return NULL; } + if (length > 0) { + PyArray_FILLWBYTE(ret, 0); + } dot = PyArray_DESCR(ret)->f->dotfunc; if (dot == NULL) { PyErr_SetString(PyExc_ValueError, @@ -1155,36 +1165,67 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, } NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ret)); + /* p is pointer to array start and s is stride */ + /* i1 is x, i2, is y, o is answer */ + ip1 = PyArray_DATA(ap1); is1 = PyArray_STRIDES(ap1)[0]; + ip2 = PyArray_DATA(ap2); is2 = PyArray_STRIDES(ap2)[0]; op = PyArray_DATA(ret); os = PyArray_DESCR(ret)->elsize; - ip1 = PyArray_DATA(ap1); - ip2 = PyArray_BYTES(ap2) + n_left*is2; - n = n - n_left; - for (i = 0; i < n_left; i++) { - dot(ip1, is1, ip2, is2, op, n, ret); - n++; - ip2 -= is2; + + lag = minlag; + if (lag < -n2+1) { + /* if minlag is before any overlap between the vectors, + * then skip to first relevant lag + */ + op += os*((-n2+1) - lag); + lag = -n2+1; + } + maxleft = (0 < maxlag ? 0 : maxlag); + tmplag = lag; + for (lag = tmplag; lag < maxleft; lag+=lagstep) { + /* for lags where y is left of x, + * overlap is length of y - lag + */ + n = n2 + lag; + dot(ip1, is1, ip2 - lag*is2, is2, op, n, ret); + /* iterate over answer vector */ op += os; } - if (small_correlate(ip1, is1, n1 - n2 + 1, PyArray_TYPE(ap1), - ip2, is2, n, PyArray_TYPE(ap2), - op, os)) { - ip1 += is1 * (n1 - n2 + 1); - op += os * (n1 - n2 + 1); + if (maxlag < n1 - n2) { + /* if maxlag doesn't take y all the way to the end of x, + * relevant length of x is smaller + */ + n11 = maxlag + n2; } else { - for (i = 0; i < (n1 - n2 + 1); i++) { - dot(ip1, is1, ip2, is2, op, n, ret); - ip1 += is1; + n11 = n1; + } + /* starts at lag=minlag if minlag>0. + * Does lags where y entirely overlaps with x. + */ + if (lagstep == 1 && lag < maxlag && + small_correlate(ip1 + lag*is1, is1, + n11 - n2 + 1 - lag, PyArray_TYPE(ap1), + ip2, is2, n2, PyArray_TYPE(ap2), + op, os)) { + lag = n11 - n2 + 1; + op += os * (n11 - n2 + 1); + } + else if (lag < maxlag) { + tmplag = lag; + for (lag = tmplag; lag < (n11 - n2 + 1); lag+=lagstep) { + dot(ip1 + lag*is1, is1, ip2, is2, op, n2, ret); op += os; } } - for (i = 0; i < n_right; i++) { - n--; - dot(ip1, is1, ip2, is2, op, n, ret); - ip1 += is1; + maxright = (maxlag < n1 ? maxlag : n1 ); + tmplag = lag; + for (lag = tmplag; lag < maxright; lag+=lagstep) { + /* for lags where y is right of x */ + n = n1 - lag; + dot(ip1 + lag*is1, is1, ip2, is2, op, n, ret); op += os; } @@ -1246,14 +1287,70 @@ _pyarray_revert(PyArrayObject *ret) return 0; } +/* + * Generate the lags corresponding to each mode + * n1 and n2 are sizes of two vectors + * minlag, maxlag, and lagstep are edited in-place + */ +void +_lags_from_mode(int mode, npy_intp n1, npy_intp n2, + npy_intp *minlag, npy_intp *maxlag, npy_intp *lagstep) { + + int i, i1, inverted = 0; + /* place holders for minlag, maxlag, and lagstep respectively */ + npy_intp m0, m1, s; + + if (n1 < n2) { + i = n1; + n1 = n2; + n2 = i; + inverted = 1; + } + + s = 1; + switch(mode) { + case 0: + /* mode = 'valid' */ + m0 = 0; + m1 = n1 - n2 + 1; + break; + case 1: + /* mode = 'same' */ + m0 = -(npy_intp)(n2/2); + m1 = n1 + m0; + break; + case 2: + /* mode = 'full' */ + m0 = -n2 + 1; + m1 = n1; + break; + default: + PyErr_SetString(PyExc_ValueError, "mode must be 0, 1, or 2"); + } + + if (inverted) { + i = m0; + i1 = (npy_intp)(npy_ceil((m1 - m0)/(float)s))*s; + m0 = -(i1 + m0 - s); + m1 = -(i - s); + } + + /* set minlag, maxlag, and lagstep values */ + *minlag = m0; + *maxlag = m1; + *lagstep = s; +} + /*NUMPY_API - * correlate(a1,a2,mode) + * correlate(a1,a2,minlag,maxlag,lagstep) * - * This function computes the usual correlation (correlate(a1, a2) != - * correlate(a2, a1), and conjugate the second argument for complex inputs + * This function correlates a1 with a2. + * This function does not accept modes. + * See _lags_from_mode() to convert modes to lags. */ NPY_NO_EXPORT PyObject * -PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) +PyArray_CorrelateLags(PyObject *op1, PyObject *op2, + npy_intp minlag, npy_intp maxlag, npy_intp lagstep) { PyArrayObject *ap1, *ap2, *ret = NULL; int typenum; @@ -1288,14 +1385,94 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) ap2 = cap2; } - ret = _pyarray_correlate(ap1, ap2, typenum, mode, &inverted); + ret = _pyarray_correlate(ap1, ap2, typenum, &inverted, + minlag, maxlag, lagstep); + if (ret == NULL) { + goto clean_ap2; + } + + /* + * If we inverted input orders, we need to reverse the output array + * (i.e. ret = ret[::-1]) + */ + if (inverted) { + st = _pyarray_revert(ret); + if(st) { + goto clean_ret; + } + } + + Py_DECREF(ap1); + Py_DECREF(ap2); + return (PyObject *)ret; + +clean_ret: + Py_DECREF(ret); +clean_ap2: + Py_DECREF(ap2); +clean_ap1: + Py_DECREF(ap1); + return NULL; +} + +/*NUMPY_API + * correlate(a1,a2,mode) + * + * This function computes the usual correlation (correlate(a1, a2) != + * correlate(a2, a1), and conjugate the second argument for complex inputs + * + * This function requires a mode. + */ +NPY_NO_EXPORT PyObject * +PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) +{ + PyArrayObject *ap1, *ap2, *ret = NULL; + int typenum; + PyArray_Descr *typec; + npy_intp minlag, maxlag, lagstep; + npy_intp n1, n2; + int inverted, st; + + typenum = PyArray_ObjectType(op1, 0); + typenum = PyArray_ObjectType(op2, typenum); + + typec = PyArray_DescrFromType(typenum); + Py_INCREF(typec); + ap1 = (PyArrayObject *)PyArray_FromAny(op1, typec, 1, 1, + NPY_ARRAY_DEFAULT, NULL); + if (ap1 == NULL) { + Py_DECREF(typec); + return NULL; + } + ap2 = (PyArrayObject *)PyArray_FromAny(op2, typec, 1, 1, + NPY_ARRAY_DEFAULT, NULL); + if (ap2 == NULL) { + goto clean_ap1; + } + + if (PyArray_ISCOMPLEX(ap2)) { + PyArrayObject *cap2; + cap2 = (PyArrayObject *)PyArray_Conjugate(ap2, NULL); + if (cap2 == NULL) { + goto clean_ap2; + } + Py_DECREF(ap2); + ap2 = cap2; + } + + /* size of x (n1) and y (n2) */ + n1 = PyArray_DIMS(ap1)[0]; + n2 = PyArray_DIMS(ap2)[0]; + _lags_from_mode(mode, n1, n2, &minlag, &maxlag, &lagstep); + ret = _pyarray_correlate(ap1, ap2, typenum, &inverted, + minlag, maxlag, lagstep); if (ret == NULL) { goto clean_ap2; } /* - * If we inverted input orders, we need to reverse the output array (i.e. - * ret = ret[::-1]) + * If we inverted input orders, we need to reverse the output array + * (i.e. ret = ret[::-1]) */ if (inverted) { st = _pyarray_revert(ret); @@ -1327,6 +1504,8 @@ PyArray_Correlate(PyObject *op1, PyObject *op2, int mode) int typenum; int unused; PyArray_Descr *typec; + npy_intp minlag, maxlag, lagstep; + npy_intp n1, n2; typenum = PyArray_ObjectType(op1, 0); typenum = PyArray_ObjectType(op2, typenum); @@ -1345,7 +1524,12 @@ PyArray_Correlate(PyObject *op1, PyObject *op2, int mode) goto fail; } - ret = _pyarray_correlate(ap1, ap2, typenum, mode, &unused); + /* size of x (n1) and y (n2) */ + n1 = PyArray_DIMS(ap1)[0]; + n2 = PyArray_DIMS(ap2)[0]; + _lags_from_mode(mode, n1, n2, &minlag, &maxlag, &lagstep); + ret = _pyarray_correlate(ap1, ap2, typenum, &unused, + minlag, maxlag, lagstep); if(ret == NULL) { goto fail; } @@ -2844,18 +3028,33 @@ array_correlate(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) return PyArray_Correlate(a0, shape, mode); } -static PyObject* +static PyObject * array_correlate2(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) { PyObject *shape, *a0; - int mode = 0; - static char *kwlist[] = {"a", "v", "mode", NULL}; + int mode = -1; + npy_intp minlag = 0, maxlag = 0, lagstep = 0; + static char *kwlist[] = {"a", "v", "mode", "minlag", "maxlag", "lagstep", + NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|i", kwlist, - &a0, &shape, &mode)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|innn", kwlist, + &a0, &shape, &mode, &minlag, &maxlag, &lagstep)) { return NULL; } - return PyArray_Correlate2(a0, shape, mode); + if (mode == -1) { + if (minlag == 0 && maxlag == 0 && lagstep == 0) { + /* if no lag parameters passed, use default: mode = 'valid' */ + mode = 0; + } + else { + /* if lag parameters were passed, use them */ + mode = 3; + } + } + if (mode != 3) { + return PyArray_Correlate2(a0, shape, mode); + } + return PyArray_CorrelateLags(a0, shape, minlag, maxlag, lagstep); } static PyObject * diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py index dd9c83b25a4d..e7f4def24939 100644 --- a/numpy/core/tests/test_numeric.py +++ b/numpy/core/tests/test_numeric.py @@ -1041,7 +1041,6 @@ def test_sufficient_width(self): assert_equal(np.binary_repr(10, width=7), '0001010') assert_equal(np.binary_repr(-5, width=7), '1111011') - class TestBaseRepr(TestCase): def test_base3(self): assert_equal(np.base_repr(3**5, 3), '100000') @@ -2041,21 +2040,66 @@ def _setup(self, dt): self.xs = np.arange(1, 20)[::3] self.y = np.array([-1, -2, -3], dtype=dt) self.z1 = np.array([ -3., -8., -14., -20., -26., -14., -5.], dtype=dt) + self.z1s = np.array([-8., -14., -20., -26., -14.], dtype=dt) + self.z1v = np.array([-14., -20., -26.], dtype=dt) self.z1_4 = np.array([-2., -5., -8., -11., -14., -5.], dtype=dt) self.z1r = np.array([-15., -22., -22., -16., -10., -4., -1.], dtype=dt) self.z2 = np.array([-5., -14., -26., -20., -14., -8., -3.], dtype=dt) + self.z2s = np.array([-14, -26., -20., -14., -8.], dtype=dt) + self.z2v = np.array([-26., -20., -14.], dtype=dt) self.z2r = np.array([-1., -4., -10., -16., -22., -22., -15.], dtype=dt) self.zs = np.array([-3., -14., -30., -48., -66., -84., -102., -54., -19.], dtype=dt) def test_float(self): self._setup(np.float) - z = np.correlate(self.x, self.y, 'full') + z, lagvec = np.correlate(self.x, self.y, 'full', returns_lagvector=True) assert_array_almost_equal(z, self.z1) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(-self.y.size + 1, self.x.size)) + z, lagvec = np.correlate(self.x, self.y, 'same', returns_lagvector=True) + assert_array_almost_equal(z, self.z1s) + assert_array_almost_equal(z, self.z1[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + if self.x.size > self.y.size: + assert_array_equal(lagvec, + np.arange(-int(self.y.size/2), self.x.size - int(self.y.size/2))) + else: + assert_array_equal(lagvec, + np.arange(-self.y.size + int(self.x.size/2) + 1, int(self.x.size/2) + 1)) + z, lagvec = np.correlate(self.x, self.y, 'valid', returns_lagvector=True) + assert_array_almost_equal(z, self.z1v) + assert_array_almost_equal(z, self.z1[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + if self.x.size > self.y.size: + assert_array_equal(lagvec, np.arange(0, self.x.size - self.y.size + 1)) + else: + assert_array_equal(lagvec, np.arange(self.x.size - self.y.size, 1)) + z = np.correlate(self.x, self.y[:-1], 'full') assert_array_almost_equal(z, self.z1_4) - z = np.correlate(self.y, self.x, 'full') + z, lagvec = np.correlate(self.y, self.x, 'full', returns_lagvector=True) assert_array_almost_equal(z, self.z2) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(-self.x.size + 1, self.y.size)) + z, lagvec = np.correlate(self.y, self.x, 'same', returns_lagvector=True) + assert_array_almost_equal(z, self.z2s) + assert_array_almost_equal(z, self.z2[lagvec+self.x.size-1]) + assert_equal(lagvec.size, z.size) + if self.y.size > self.x.size: + assert_array_equal(lagvec, + np.arange(-int(self.x.size/2), self.y.size - int(self.x.size/2))) + else: + assert_array_equal(lagvec, + np.arange(-self.x.size + int(self.y.size/2) + 1, int(self.y.size/2) + 1)) + z, lagvec = np.correlate(self.y, self.x, 'valid', returns_lagvector=True) + assert_array_almost_equal(z, self.z2v) + assert_array_almost_equal(z, self.z2[lagvec+self.x.size-1]) + assert_equal(lagvec.size, z.size) + if self.y.size > self.x.size: + assert_array_equal(lagvec, np.arange(0, self.y.size - self.x.size + 1)) + else: + assert_array_equal(lagvec, np.arange(self.y.size - self.x.size, 1)) z = np.correlate(self.x[::-1], self.y, 'full') assert_array_almost_equal(z, self.z1r) z = np.correlate(self.y, self.x[::-1], 'full') @@ -2063,6 +2107,79 @@ def test_float(self): z = np.correlate(self.xs, self.y, 'full') assert_array_almost_equal(z, self.zs) + def test_lags(self): + self._setup(np.float) + longlen = max(self.x.size, self.y.size) + maxlag = 3 + longlen + minlag = -2 + lagstep = 2 + tmp = np.correlate(self.x, self.y, 'full') + z_full = np.zeros(tmp.size + 100) + z_full[:tmp.size] = tmp + z, lagvec = np.correlate(self.x, self.y, lags=maxlag, + returns_lagvector=True) + assert_array_equal(z[0:3], np.zeros(3)) + assert_array_equal(z[-3:], np.zeros(3)) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(-maxlag+1, maxlag)) + z, lagvec = np.correlate(self.x, self.y, mode='lags', + lags=(minlag, maxlag), returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(minlag, maxlag)) + z, lagvec = np.correlate(self.x, self.y, lags=(minlag, maxlag, lagstep), + returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(minlag, maxlag, lagstep)) + lagstep = 3 + z, lagvec = np.correlate(self.x, self.y, mode='lags', + lags=(minlag, maxlag, lagstep), + returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + minlag = 10 + maxlag = -2 + lagstep = -2 + z, lagvec = np.correlate(self.x, self.y, mode='lags', + lags=(minlag, maxlag, lagstep), + returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag), lagstep)) + maxlag = 10.2 + minlag = -2.2 + z, lagvec = np.correlate(self.x, self.y, mode='lags', + lags=(minlag, maxlag), + returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag))) + z, lagvec = np.correlate(self.x, self.y, mode='lags', + lags=(maxlag, minlag), + returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(int(maxlag), int(minlag))) + maxlag = -1 + minlag = -2 + z, lagvec = np.correlate(self.x, self.y, mode='lags', + lags=(minlag, maxlag), + returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag))) + maxlag = 2 + minlag = 4 + z, lagvec = np.correlate(self.x, self.y, mode='lags', + lags=(minlag, maxlag), + returns_lagvector=True) + assert_array_almost_equal(z, z_full[lagvec+self.y.size-1]) + assert_equal(lagvec.size, z.size) + assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag))) + # check non-integer lagsteps? + # i want to be able to switch x and y and test + def test_object(self): self._setup(Decimal) z = np.correlate(self.x, self.y, 'full') @@ -2085,7 +2202,6 @@ def test_complex(self): z = np.correlate(y, x, mode='full') assert_array_almost_equal(z, r_z) - class TestConvolve(TestCase): def test_object(self): d = [1.] * 100