diff --git a/doc/source/reference/c-api.array.rst b/doc/source/reference/c-api.array.rst index 90bb56b2d6f6..03d982ccb4a9 100644 --- a/doc/source/reference/c-api.array.rst +++ b/doc/source/reference/c-api.array.rst @@ -2264,6 +2264,33 @@ 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) diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py index d1406e3b288a..813caad9e7ca 100644 --- a/numpy/core/code_generators/numpy_api.py +++ b/numpy/core/code_generators/numpy_api.py @@ -346,6 +346,8 @@ # End 1.10 API 'PyArray_MapIterArrayCopyIfOverlap': (301,), # End 1.13 API + 'PyArray_CorrelateLags': (302,), + # End 1.14 API } ufunc_types_api = { diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index f8ca4942982d..5bb790602b0d 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -24,7 +24,7 @@ from . import umath from .umath import (multiply, invert, sin, UFUNC_BUFSIZE_DEFAULT, ERR_IGNORE, ERR_WARN, ERR_RAISE, ERR_CALL, ERR_PRINT, - ERR_LOG, ERR_DEFAULT, PINF, NAN) + ERR_LOG, ERR_DEFAULT, PINF, NAN, ceil) from . import numerictypes from .numerictypes import longlong, intc, int_, float_, complex_, bool_ from ._internal import TooHardError, AxisError @@ -898,16 +898,64 @@ def flatnonzero(a): _mode_from_name_dict = {'v': 0, 's': 1, - 'f': 2} - + 'f': 2, + 'l': 3, + } +keys = [key for key in _mode_from_name_dict] +for key in keys: + _mode_from_name_dict[key.upper()] = _mode_from_name_dict[key] +_mode_from_name_dict_values = _mode_from_name_dict.values() def _mode_from_name(mode): - if isinstance(mode, basestring): - return _mode_from_name_dict[mode.lower()[0]] + # guarantees that output is a value in _mode_from_name_dict + if mode in _mode_from_name_dict_values: + return mode + try: + mode = _mode_from_name_dict[mode[0]] + except KeyError: + raise ValueError("correlate/convolve mode argument must be unused or" + + " one of {'valid', 'same', 'full', 'lags'}") 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) + else: + raise ValueError("correlate/convolve lags argument must be " + + "int or int tuple.") + return lags + +def _lags_from_mode(alen, vlen, mode): + 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) + else: + raise ValueError("correlate/convolve mode argument must be unused or" + + " one of {'valid', 'same', 'full', 'lags'}") + + if inverted: + 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]) + + return mode_lags + + +def correlate(a, v, mode='default', lags=(), returns_lagvector=False): """ Cross-correlation of two 1-dimensional sequences. @@ -923,9 +971,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`. @@ -934,6 +991,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 -------- @@ -953,10 +1013,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: @@ -971,11 +1035,31 @@ 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]) """ - mode = _mode_from_name(mode) - return multiarray.correlate2(a, v, mode) + if mode == 'default': + if lags: + mode = 'lags' + else: + mode = 'valid' + mode = _mode_from_name(mode) # guaranteed a value in _mode_from_name_dict + if mode in (0, 1, 2): + if lags: + raise ValueError("correlate mode keyword argument must be 'lags'" + + " or unused if the lags keyword argument is used.") + result = multiarray.correlate2(a, v, mode) + if returns_lagvector: + alen, vlen = len(a), len(v) + lags = _lags_from_mode(alen, vlen, mode) + elif mode == 3: + lags = _lags_from_lags(lags) + result = multiarray.correlate2(a, v, 3, lags[0], lags[1], lags[2]) + + if returns_lagvector: + return result, arange(lags[0], lags[1], lags[2]) + else: + return result -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. @@ -993,27 +1077,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 -------- @@ -1052,25 +1168,64 @@ 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])) + + 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])) - >>> np.convolve([1,2,3],[0,1,0.5], 'valid') - array([ 2.5]) + 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: + alen, vlen = len(a), len(v) + if alen == 0: raise ValueError('a cannot be empty') - if len(v) == 0: + if vlen == 0: raise ValueError('v cannot be empty') - mode = _mode_from_name(mode) - return multiarray.correlate(a, v[::-1], mode) + + if mode == 'default': + if lags: + mode = 'lags' + else: + mode = 'full' + mode = _mode_from_name(mode) # guaranteed a value in _mode_from_name_dict + if mode in (0, 1, 2): + if lags: + raise ValueError("convolve mode keyword argument must be 'lags'" + + " or unused if the lags keyword argument is used.") + result = multiarray.correlate2(a, v[::-1], mode) + if returns_lagvector: + lags = _lags_from_mode(alen, vlen, mode) + elif mode == 3: + lags = _lags_from_lags(lags) + result = multiarray.correlate2(a, v[::-1], 3, lags[0], lags[1], lags[2]) + + if returns_lagvector: + return result, arange(lags[0], lags[1], lags[2]) + else: + return result def outer(a, b, out=None): diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 842e183de49c..0f23e3314731 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1130,18 +1130,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) { @@ -1152,40 +1159,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, NULL); 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, @@ -1194,36 +1202,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; } @@ -1285,14 +1324,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 computes the cross-correlation of 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; @@ -1327,14 +1422,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 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 (inverted) { st = _pyarray_revert(ret); @@ -1366,6 +1541,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); @@ -1384,7 +1561,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; } @@ -2831,25 +3013,40 @@ array_correlate(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) int mode = 0; static char *kwlist[] = {"a", "v", "mode", NULL}; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|i:correlate", kwlist, + if (!PyArg_ParseTupleAndKeywords(args, kwds, "OO|i", kwlist, &a0, &shape, &mode)) { return NULL; } 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:correlate2", 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 1e51c4a61cd5..7474d61333e2 100644 --- a/numpy/core/tests/test_numeric.py +++ b/numpy/core/tests/test_numeric.py @@ -2177,21 +2177,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') @@ -2199,6 +2244,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')