From fb61c12efa50c0280f1ffe0e76ba4f792a7a9503 Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Wed, 24 Jun 2015 18:56:05 -0400 Subject: [PATCH 1/3] DOC: added comments to _pyarray_correlate in multiarraymodule.c #5954 --- numpy/core/src/multiarray/multiarraymodule.c | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 73265c3b607f..f883f3a8afb5 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1153,8 +1153,8 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, NPY_BEGIN_THREADS_DEF; - n1 = PyArray_DIMS(ap1)[0]; - n2 = PyArray_DIMS(ap2)[0]; + n1 = PyArray_DIMS(ap1)[0]; /* size of x */ + n2 = PyArray_DIMS(ap2)[0]; /* size of y */ if (n1 < n2) { ret = ap1; ap1 = ap2; @@ -1174,15 +1174,18 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, case 0: length = length - n + 1; n_left = n_right = 0; + /* mode = 'valid' */ break; case 1: n_left = (npy_intp)(n/2); n_right = n - n_left - 1; + /* mode = 'same' */ break; case 2: n_right = n - 1; n_left = n - 1; length = length + n - 1; + /* mode = 'full' */ break; default: PyErr_SetString(PyExc_ValueError, "mode must be 0, 1, or 2"); @@ -1192,6 +1195,7 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, /* * 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) { @@ -1205,18 +1209,18 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, } NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ret)); - is1 = PyArray_STRIDES(ap1)[0]; - 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; op += os; + ip1 = PyArray_DATA(ap1); /* x[0] */ + is1 = PyArray_STRIDES(ap1)[0]; /* x strides */ + ip2 = PyArray_DATA(ap2); /* y[0] */ + is2 = PyArray_STRIDES(ap2)[0]; /* y strides */ + op = PyArray_DATA(ret); /* answer data */ + os = PyArray_DESCR(ret)->elsize; /* answer strides */ } if (small_correlate(ip1, is1, n1 - n2 + 1, PyArray_TYPE(ap1), ip2, is2, n, PyArray_TYPE(ap2), From c762c56bff862f5863d9dd68a8cb17332f3414d1 Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Wed, 24 Jun 2015 19:06:04 -0400 Subject: [PATCH 2/3] ENH: added maxlag mode to numpy.correlate() and numpy.convolve() #5954 implemented maxlag int or tuple versions of mode argument to numpy.correlate and numpy.convolve. changes were implemented down through array_correlate(), array_correlate2, PyArray_Correlate(), PyArray_Correlate2() and _pyarray_correlate() in multiarraymodule.c --- numpy/core/numeric.py | 130 ++++++++++++++--- numpy/core/src/multiarray/multiarraymodule.c | 146 +++++++++++++------ 2 files changed, 211 insertions(+), 65 deletions(-) diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index f29b750f68b0..de5f12749dd2 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -8,7 +8,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_ @@ -830,7 +830,34 @@ def _mode_from_name(mode): return _mode_from_name_dict[mode.lower()[0]] return mode -def correlate(a, v, mode='valid'): +def _lags_from_mode(alen, vlen, mode): + if type(mode) is int: # maxlag + lags = (-mode+1, mode, 1) + mode = 3 + elif type(mode) is tuple: # minlag and maxlag + if len(mode) > 2: + lags = (int(mode[0]), int(mode[1]), int(mode[2])) + else: + lags = (int(mode[0]), int(mode[1]), 1) + mode = 3 + elif isinstance(mode, basestring): + mode = _mode_from_name_dict[mode.lower()[0]] + if alen < vlen: + alen, vlen = vlen, alen + inverted = 1 + else: + inverted = 0 + if mode is 0: + lags = (0, alen-vlen+1, 1) + elif mode is 1: + lags = (-int(vlen/2), alen - int(vlen/2), 1) + elif mode is 2: + lags = (-vlen+1, alen, 1) + if inverted: + lags = (-int(ceil((lags[1]-lags[0])/float(lags[2])))*lags[2]-lags[0]+lags[2], -lags[0]+lags[2], lags[2]) + return mode, lags + +def correlate(a, v, mode='valid', returns_lags=False): """ Cross-correlation of two 1-dimensional sequences. @@ -846,9 +873,15 @@ def correlate(a, v, mode='valid'): ---------- a, v : array_like Input sequences. - mode : {'valid', 'same', 'full'}, optional + mode : int, int tuple, or {'valid', 'same', 'full'}, optional Refer to the `convolve` docstring. Note that the default is `valid`, unlike `convolve`, which uses `full`. + returns_lags : 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`. @@ -857,6 +890,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 -------- @@ -876,10 +912,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], mode="full", returns_lags=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=2) 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=(-1,2,2), returns_lags=True) + (array([ 2., 3.]), array([-1, 1])) Using complex sequences: @@ -894,10 +934,13 @@ 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) + mode, lags = _lags_from_mode(len(a), len(v), mode) + if returns_lags: + return multiarray.correlate2(a, v, mode, lags[0], lags[1], lags[2]), arange(lags[0], lags[1], lags[2]) + else: + return multiarray.correlate2(a, v, mode, lags[0], lags[1], lags[2]) -def convolve(a,v,mode='full'): +def convolve(a, v, mode='full', returns_lags=False): """ Returns the discrete, linear convolution of two one-dimensional sequences. @@ -915,27 +958,53 @@ 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 + 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 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. + '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. + 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. + returns_lags : 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 -------- @@ -974,14 +1043,34 @@ 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], 'valid') - array([ 2.5]) + >>> np.convolve([1,2,3],[0,1,0.5], mode='valid', returns_lags=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=2, returns_lags=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=(-2,6,2), returns_lags=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) @@ -991,8 +1080,11 @@ def convolve(a,v,mode='full'): raise ValueError('a cannot be empty') if len(v) == 0 : raise ValueError('v cannot be empty') - mode = _mode_from_name(mode) - return multiarray.correlate(a, v[::-1], mode) + mode, lags = _lags_from_mode(len(a), len(v), mode) + if returns_lags: + return multiarray.correlate2(a, v[::-1], mode, lags[0], lags[1], lags[2]), arange(lags[0], lags[1], lags[2]) + else: + return multiarray.correlate2(a, v[::-1], mode, lags[0], lags[1], lags[2]) def outer(a, b, out=None): """ diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index f883f3a8afb5..c57e0ee9e003 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1142,17 +1142,23 @@ PyArray_CopyAndTranspose(PyObject *op) */ static PyArrayObject* _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, - int mode, int *inverted) + int mode, 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 == NPY_MAX_INTP || lagstep == 0) { + lagstep = 1; + } + n1 = PyArray_DIMS(ap1)[0]; /* size of x */ n2 = PyArray_DIMS(ap2)[0]; /* size of y */ if (n1 < n2) { @@ -1163,34 +1169,63 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, i = n1; n1 = n2; n2 = i; - *inverted = 1; - } else { - *inverted = 0; + minlag = -minlag; + maxlag = -maxlag; + lagstep = -lagstep; } - length = n1; - n = n2; + if (maxlag == NPY_MAX_INTP || minlag == NPY_MAX_INTP) { switch(mode) { + lagstep = 1; case 0: - length = length - n + 1; - n_left = n_right = 0; /* mode = 'valid' */ + minlag = 0; + maxlag = n1 - n2 + 1; break; case 1: - n_left = (npy_intp)(n/2); - n_right = n - n_left - 1; /* mode = 'same' */ + minlag = -(npy_intp)(n2/2); + maxlag = n1 + minlag; break; case 2: - n_right = n - 1; - n_left = n - 1; - length = length + n - 1; /* mode = 'full' */ + minlag = -n2 + 1; + maxlag = n1; + break; + case 3: + /* mode = 'maxlag' */ + if (maxlag == NPY_MAX_INTP) { + /* nothing good to do with situation */ + minlag = -n2 + 1; + maxlag = n1; + } + else if (minlag == NPY_MAX_INTP) { + minlag = (maxlag > n2 ? -n2+1 : -maxlag+1); + } break; default: - PyErr_SetString(PyExc_ValueError, "mode must be 0, 1, or 2"); + PyErr_SetString(PyExc_ValueError, "mode must be 0, 1, 2, or 3"); return NULL; } + } + + if (lagstep < 0) { + *inverted = 1; + i = minlag; + i1 = (npy_intp)(npy_ceil((maxlag - minlag)/(float)lagstep))*lagstep; + minlag = i1 + minlag - lagstep; + maxlag = i - lagstep; + lagstep = -lagstep; + } + else { + *inverted = 0; + } + if (maxlag <= minlag) { + length = 0; + } + else { + length = (maxlag - minlag + lagstep - 1)/lagstep; + } /* * Need to choose an output array that can hold a sum @@ -1201,6 +1236,7 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, if (ret == NULL) { return NULL; } + PyArray_FILLWBYTE(ret, 0); dot = PyArray_DESCR(ret)->f->dotfunc; if (dot == NULL) { PyErr_SetString(PyExc_ValueError, @@ -1209,36 +1245,51 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, } NPY_BEGIN_THREADS_DESCR(PyArray_DESCR(ret)); - n = n - n_left; - for (i = 0; i < n_left; i++) { - dot(ip1, is1, ip2, is2, op, n, ret); - n++; - ip2 -= is2; - op += os; ip1 = PyArray_DATA(ap1); /* x[0] */ is1 = PyArray_STRIDES(ap1)[0]; /* x strides */ ip2 = PyArray_DATA(ap2); /* y[0] */ is2 = PyArray_STRIDES(ap2)[0]; /* y strides */ op = PyArray_DATA(ret); /* answer data */ os = PyArray_DESCR(ret)->elsize; /* answer strides */ + + lag = minlag; + maxleft = (0 < maxlag ? 0 : maxlag); + for (lag = minlag; lag < maxleft; lag+=lagstep) { + /* for lags where y is left of x */ + n = n2 + lag; /* overlap is length of y - lag */ + dot(ip1, is1, ip2 - lag*is2, is2, op, n, ret); + op += os; /* iterate over answer vector */ + } + if (maxlag < n1 - n2) { + /* if maxlag doesn't take y all the way to the end of x */ + n11 = maxlag + n2; /* relevant length of x is smaller */ + } + else { + n11 = n1; } - 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); + /* starts at lag=minlag if minlag>0. + * Does lags where y entirely overlaps with x. + */ + if (lagstep == 1 && 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 { - for (i = 0; i < (n1 - n2 + 1); i++) { - dot(ip1, is1, ip2, is2, op, n, ret); - ip1 += is1; + 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; } @@ -1313,7 +1364,7 @@ _pyarray_revert(PyArrayObject *ret) * correlate(a2, a1), and conjugate the second argument for complex inputs */ NPY_NO_EXPORT PyObject * -PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) +PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode, npy_intp minlag, npy_intp maxlag, npy_intp lagstep) { PyArrayObject *ap1, *ap2, *ret = NULL; int typenum; @@ -1348,7 +1399,7 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) ap2 = cap2; } - ret = _pyarray_correlate(ap1, ap2, typenum, mode, &inverted); + ret = _pyarray_correlate(ap1, ap2, typenum, mode, &inverted, minlag, maxlag, lagstep); if (ret == NULL) { goto clean_ap2; } @@ -1381,7 +1432,7 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) * Numeric.correlate(a1,a2,mode) */ NPY_NO_EXPORT PyObject * -PyArray_Correlate(PyObject *op1, PyObject *op2, int mode) +PyArray_Correlate(PyObject *op1, PyObject *op2, int mode, npy_intp minlag, npy_intp maxlag, npy_intp lagstep) { PyArrayObject *ap1, *ap2, *ret = NULL; int typenum; @@ -1405,7 +1456,7 @@ PyArray_Correlate(PyObject *op1, PyObject *op2, int mode) goto fail; } - ret = _pyarray_correlate(ap1, ap2, typenum, mode, &unused); + ret = _pyarray_correlate(ap1, ap2, typenum, mode, &unused, minlag, maxlag, lagstep); if(ret == NULL) { goto fail; } @@ -2901,27 +2952,30 @@ array_correlate(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) { PyObject *shape, *a0; int mode = 0; - static char *kwlist[] = {"a", "v", "mode", NULL}; + npy_intp maxlag = NPY_MAX_INTP, minlag = NPY_MAX_INTP, lagstep = NPY_MAX_INTP; + 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_Correlate(a0, shape, mode); + return PyArray_Correlate(a0, shape, mode, minlag, maxlag, lagstep); } -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}; + npy_intp maxlag = NPY_MAX_INTP, minlag = NPY_MAX_INTP, lagstep = NPY_MAX_INTP; + PyObject *tmp; + 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); + return PyArray_Correlate2(a0, shape, mode, minlag, maxlag, lagstep); } static PyObject * From db91eab13e4c33d8ebc850038dc86d62b390249a Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Wed, 24 Jun 2015 19:07:22 -0400 Subject: [PATCH 3/3] ENH: unit tests for maxlag mode in numpy.correlate() and numpy.convolve() #5954 --- numpy/core/tests/test_numeric.py | 103 ++++++++++++++++++++++++++++++- 1 file changed, 101 insertions(+), 2 deletions(-) diff --git a/numpy/core/tests/test_numeric.py b/numpy/core/tests/test_numeric.py index c39900ffa7bd..1e6b41f55782 100644 --- a/numpy/core/tests/test_numeric.py +++ b/numpy/core/tests/test_numeric.py @@ -1953,21 +1953,62 @@ 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_lags=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_lags=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_lags=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_lags=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_lags=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_lags=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') @@ -1975,6 +2016,64 @@ 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, maxlag, returns_lags=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, (minlag, maxlag), returns_lags=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, (minlag, maxlag, lagstep), returns_lags=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, (minlag, maxlag, lagstep), returns_lags=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, (minlag, maxlag, lagstep), returns_lags=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, (minlag, maxlag), returns_lags=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, (maxlag, minlag), returns_lags=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, (minlag, maxlag), returns_lags=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, (minlag, maxlag), returns_lags=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')