From d95b01ba0154c8cbcfa52afd62b78673fd8dbfbd Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Tue, 14 Apr 2026 20:23:47 +0000 Subject: [PATCH 01/13] committing progress --- numpy/_core/code_generators/cversions.txt | 3 + numpy/_core/code_generators/numpy_api.py | 1 + numpy/_core/include/numpy/ndarraytypes.h | 3 +- numpy/_core/meson.build | 3 +- numpy/_core/multiarray.py | 2 +- numpy/_core/multiarray.pyi | 1 + numpy/_core/numeric.py | 227 ++++++++++-- numpy/_core/src/multiarray/conversion_utils.c | 12 +- numpy/_core/src/multiarray/multiarraymodule.c | 322 +++++++++++++++--- numpy/_core/tests/test_numeric.py | 42 ++- 10 files changed, 531 insertions(+), 85 deletions(-) diff --git a/numpy/_core/code_generators/cversions.txt b/numpy/_core/code_generators/cversions.txt index b058875d0455..25a44713ae96 100644 --- a/numpy/_core/code_generators/cversions.txt +++ b/numpy/_core/code_generators/cversions.txt @@ -85,3 +85,6 @@ # General loop registration for ufuncs, sort, and argsort # Version 21 (NumPy 2.5.0) No change 0x00000015 = fbd24fc5b2ba4f7cd3606ec6128de7a5 +# Version 22 (NumPy 2.6.0) +# Added PyArray_CorrelateLags +0x00000016 = c5f7b81de5f99f0df1320b4a5e64ea12 \ No newline at end of file diff --git a/numpy/_core/code_generators/numpy_api.py b/numpy/_core/code_generators/numpy_api.py index c2b471c71757..90027982dd39 100644 --- a/numpy/_core/code_generators/numpy_api.py +++ b/numpy/_core/code_generators/numpy_api.py @@ -409,6 +409,7 @@ def get_annotations(): # End 2.0 API # NpyIterGetTransferFlags (slot 223) added. # End 2.3 API + 'PyArray_CorrelateLags': (369, MinVersion("2.6")), } ufunc_types_api = { diff --git a/numpy/_core/include/numpy/ndarraytypes.h b/numpy/_core/include/numpy/ndarraytypes.h index f740788f3720..d671ec2ad2aa 100644 --- a/numpy/_core/include/numpy/ndarraytypes.h +++ b/numpy/_core/include/numpy/ndarraytypes.h @@ -270,7 +270,8 @@ typedef enum { typedef enum { NPY_VALID=0, NPY_SAME=1, - NPY_FULL=2 + NPY_FULL=2, + NPY_LAGS=3 } NPY_CORRELATEMODE; /* The special not-a-time (NaT) value */ diff --git a/numpy/_core/meson.build b/numpy/_core/meson.build index aa4da9c11146..082eaffcd60e 100644 --- a/numpy/_core/meson.build +++ b/numpy/_core/meson.build @@ -52,7 +52,8 @@ C_ABI_VERSION = '0x02000000' # 0x00000014 - 2.3.x # 0x00000015 - 2.4.x # 0x00000015 - 2.5.x -C_API_VERSION = '0x00000015' +# 0x00000016 - 2.6.x +C_API_VERSION = '0x00000016' # Check whether we have a mismatch between the set C API VERSION and the # actual C API VERSION. Will raise a MismatchCAPIError if so. diff --git a/numpy/_core/multiarray.py b/numpy/_core/multiarray.py index 54d240c89e3e..06d3f1cf7d6e 100644 --- a/numpy/_core/multiarray.py +++ b/numpy/_core/multiarray.py @@ -36,7 +36,7 @@ '_monotonicity', 'add_docstring', 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'bincount', 'broadcast', 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast', - 'compare_chararrays', 'concatenate', 'copyto', 'correlate', 'correlate2', + 'compare_chararrays', 'concatenate', 'copyto', 'correlate', 'correlate2', 'correlatelags', 'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data', 'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype', 'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat', diff --git a/numpy/_core/multiarray.pyi b/numpy/_core/multiarray.pyi index e516c3cdab72..7b8e072d9810 100644 --- a/numpy/_core/multiarray.pyi +++ b/numpy/_core/multiarray.pyi @@ -128,6 +128,7 @@ __all__ = [ "copyto", "correlate", "correlate2", + "correlatelags" "count_nonzero", "c_einsum", "datetime_as_string", diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index d4e1685501d7..d4ed05eebb8c 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -718,12 +718,72 @@ def flatnonzero(a): return np.nonzero(np.ravel(a))[0] -def _correlate_dispatcher(a, v, mode=None): +def _mode_from_name(mode): + if type(mode) not in (str, int): + raise TypeError("correlate/convolve mode argument must be unused or" + + " one of {'valid', 'same', 'full', 'lags'}") + + mode_map = { + 'valid': 0, + 'same': 1, + 'full': 2, + 'lags': 3 + } + + if mode in mode_map.values(): + return mode + + if mode in mode_map: + return mode_map[mode] + + raise ValueError("correlate/convolve mode argument must be unused or" + + " one of {'valid', 'same', 'full', 'lags'}") + + +def _lags_from_lags(lag): + if type(lag) is int: # maxlag + lags = (-lag+1, lag, 1) + elif type(lag) is tuple: # minlag and maxlag + if len(lag) > 2: + lags = (int(lag[0]), int(lag[1]), int(lag[2])) + else: + lags = (int(lag[0]), int(lag[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 == 0: + mode_lags = (0, alen-vlen+1, 1) + elif mode == 1: + mode_lags = (-int(vlen/2), alen - int(vlen/2), 1) + elif mode == 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(math.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_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): return (a, v) @array_function_dispatch(_correlate_dispatcher) -def correlate(a, v, mode='valid'): +def correlate(a, v, mode='default', lags=None, returns_lagvector=False): r""" Cross-correlation of two 1-dimensional sequences. @@ -739,14 +799,26 @@ 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 'default', 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. Returns ------- 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 -------- @@ -774,13 +846,16 @@ def correlate(a, v, mode='valid'): Examples -------- - >>> import numpy as np >>> np.correlate([1, 2, 3], [0, 1, 0.5]) array([3.5]) - >>> np.correlate([1, 2, 3], [0, 1, 0.5], "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="same") + array([ 2. , 3.5, 3. ]) + >>> 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: @@ -795,15 +870,36 @@ 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]) """ - return multiarray.correlate2(a, v, mode) + if mode == 'default': + if lags is not None: + 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 is not None: + 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_throuple = _lags_from_mode(alen, vlen, mode) + elif mode == 3: + lags_throuple = _lags_from_lags(lags) + result = multiarray.correlatelags(a, v, lags_throuple[0], lags_throuple[1], lags_throuple[2]) + + if returns_lagvector: + return result, arange(lags_throuple[0], lags_throuple[1], lags_throuple[2]) + else: + return result -def _convolve_dispatcher(a, v, mode=None): +def _convolve_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): return (a, v) @array_function_dispatch(_convolve_dispatcher) -def convolve(a, v, mode='full'): +def convolve(a, v, mode='full', lags=None, returns_lagvector=False): """ Returns the discrete, linear convolution of two one-dimensional sequences. @@ -821,27 +917,60 @@ 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 -------- @@ -882,24 +1011,66 @@ 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: - - >>> np.convolve([1,2,3],[0,1,0.5], 'valid') - array([2.5]) - + 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])) + + 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=None, ndmin=1), array(v, copy=None, ndmin=1) - 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') - if len(v) > len(a): + + if vlen > alen: a, v = v, a - return multiarray.correlate(a, v[::-1], mode) + + if mode == 'default': + if lags is not None: + 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 is not None: + 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.correlatelags(a, v[::-1], lags[0], lags[1], lags[2]) + + if returns_lagvector: + return result, arange(lags[0], lags[1], lags[2]) + else: + return result def _outer_dispatcher(a, b, out=None): diff --git a/numpy/_core/src/multiarray/conversion_utils.c b/numpy/_core/src/multiarray/conversion_utils.c index 164aa2e4c8b4..da59d7666436 100644 --- a/numpy/_core/src/multiarray/conversion_utils.c +++ b/numpy/_core/src/multiarray/conversion_utils.c @@ -862,6 +862,10 @@ static int correlatemode_parser(char const *str, Py_ssize_t length, void *data) *val = NPY_FULL; is_exact = (length == 4 && strcmp(str, "full") == 0); } + else if (str[0] == 'L' || str[0] == 'l') { + *val = NPY_LAGS; + is_exact = (length == 4 && strcmp(str, "lags") == 0); + } else { return -1; } @@ -871,7 +875,7 @@ static int correlatemode_parser(char const *str, Py_ssize_t length, void *data) */ if (!is_exact) { PyErr_SetString(PyExc_ValueError, - "Use one of 'valid', 'same', or 'full' for convolve/correlate mode"); + "Use one of 'valid', 'same', 'full', or 'lags' for convolve/correlate mode"); return -1; } @@ -887,7 +891,7 @@ PyArray_CorrelatemodeConverter(PyObject *object, NPY_CORRELATEMODE *val) if (PyUnicode_Check(object)) { return string_converter_helper( object, (void *)val, correlatemode_parser, "mode", - "must be one of 'valid', 'same', or 'full'"); + "must be one of 'valid', 'same', 'full', or 'lags'"); } else { @@ -898,14 +902,14 @@ PyArray_CorrelatemodeConverter(PyObject *object, NPY_CORRELATEMODE *val) "convolve/correlate mode not understood"); return NPY_FAIL; } - if (number <= (int) NPY_FULL + if (number <= (int) NPY_LAGS && number >= (int) NPY_VALID) { *val = (NPY_CORRELATEMODE) number; return NPY_SUCCEED; } else { PyErr_Format(PyExc_ValueError, - "integer convolve/correlate mode must be 0, 1, or 2"); + "integer convolve/correlate mode must be 0, 1, 2, or 3"); return NPY_FAIL; } } diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c index 8bede253a22f..bb94393bcb04 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.c +++ b/numpy/_core/src/multiarray/multiarraymodule.c @@ -1132,26 +1132,34 @@ PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyArrayObject* out) return NULL; } - /* * Implementation which is common between PyArray_Correlate - * and PyArray_Correlate2. + * PyArray_Correlate2, and PyArray_CorrelateLags. * * 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) { + printf("_pyarray_correlate: minlag: %ld, maxlag: %ld, lagstep: %ld\n", minlag, maxlag, 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 == 0) { @@ -1169,41 +1177,41 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, ret = NULL; i = n1; n1 = n2; - n2 = i; + 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; } + /* Zero-initialize so out-of-bounds lags are correctly zero */ + memset(PyArray_DATA(ret), 0, PyArray_NBYTES(ret)); dot = PyDataType_GetArrFuncs(PyArray_DESCR(ret))->dotfunc; if (dot == NULL) { PyErr_SetString(PyExc_ValueError, @@ -1213,40 +1221,70 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, int typenum, int needs_pyapi = PyDataType_FLAGCHK(PyArray_DESCR(ret), NPY_NEEDS_PYAPI); 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_ITEMSIZE(ret); - 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); + + 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); if (needs_pyapi && PyErr_Occurred()) { goto done; } - n++; - ip2 -= is2; + /* 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) && (!needs_pyapi || !PyErr_Occurred()); - 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)) { + op += os * (n11 - n2 + 1 - lag); + lag = n11 - n2 + 1; + } + else if (lag < maxlag) { + tmplag = lag; + for (lag = tmplag; lag < (n11 - n2 + 1) && (!needs_pyapi || !PyErr_Occurred()); lag+=lagstep) { + dot(ip1 + lag*is1, is1, ip2, is2, op, n2, ret); op += os; } } - for (i = 0; i < n_right && (!needs_pyapi || !PyErr_Occurred()); i++) { - n--; - dot(ip1, is1, ip2, is2, op, n, ret); - ip1 += is1; + maxright = (maxlag < n1 ? maxlag : n1 ); + tmplag = lag; + for (lag = tmplag; lag < maxright && (!needs_pyapi || !PyErr_Occurred()); 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; } @@ -1310,6 +1348,61 @@ _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"); + return; + } + + 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) * @@ -1322,6 +1415,8 @@ 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; int st; @@ -1358,7 +1453,13 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) ap2 = cap2; } - ret = _pyarray_correlate(ap1, ap2, typenum, mode, &inverted); + /* 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); + printf("PyArray_Correlate2: mode: %d, n1: %ld, n2: %ld, minlag: %ld, maxlag: %ld, lagstep: %ld\n", mode, n1, n2, minlag, maxlag, lagstep); + ret = _pyarray_correlate(ap1, ap2, typenum, &inverted, + minlag, maxlag, lagstep); if (ret == NULL) { goto clean_ap2; } @@ -1387,6 +1488,86 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) return NULL; } +/*NUMPY_API + * correlate(a1,a2,minlag,maxlag,lagstep) + * + * 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_CorrelateLags(PyObject *op1, PyObject *op2, + npy_intp minlag, npy_intp maxlag, npy_intp lagstep) +{ + PyArrayObject *ap1, *ap2, *ret = NULL; + int typenum; + PyArray_Descr *typec; + int inverted; + int st; + + typenum = PyArray_ObjectType(op1, NPY_NOTYPE); + if (typenum == NPY_NOTYPE) { + return NULL; + } + typenum = PyArray_ObjectType(op2, typenum); + if (typenum == NPY_NOTYPE) { + return NULL; + } + + 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; + } + + 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 * Numeric.correlate(a1,a2,mode) */ @@ -1397,6 +1578,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, NPY_NOTYPE); if (typenum == NPY_NOTYPE) { @@ -1421,7 +1604,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; } @@ -3031,6 +3219,13 @@ array_correlate(PyObject *NPY_UNUSED(dummy), NULL, NULL, NULL) < 0) { return NULL; } + if (mode == 3) { + PyErr_SetString(PyExc_ValueError, + "correlate() accepts only modes 0, 1, and 2 " + "(valid, same and full). " + "Use correlatelags() for mode 3 (lags)."); + return NULL; + } return PyArray_Correlate(a0, shape, mode); } @@ -3049,9 +3244,35 @@ array_correlate2(PyObject *NPY_UNUSED(dummy), NULL, NULL, NULL) < 0) { return NULL; } + printf("array_correlate2: mode: %d\n", mode); return PyArray_Correlate2(a0, shape, mode); } +static PyObject* +array_correlatelags(PyObject *NPY_UNUSED(dummy), + PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) +{ + PyObject *shape, *a0; + npy_intp minlag = 0, maxlag = 0, lagstep = 0; + NPY_PREPARE_ARGPARSER; + + if (npy_parse_arguments("correlate_lags", args, len_args, kwnames, + "a", NULL, &a0, + "v", NULL, &shape, + "minlag", &PyArray_IntpFromPyIntConverter, &minlag, + "maxlag", &PyArray_IntpFromPyIntConverter, &maxlag, + "lagstep", &PyArray_IntpFromPyIntConverter, &lagstep, + NULL, NULL, NULL) < 0) { + return NULL; + } + printf("array_correlate_lags: minlag: %ld, maxlag: %ld, lagstep: %ld\n", minlag, maxlag, lagstep); + if (minlag == 0 && maxlag == 0 && lagstep == 0) { + /* if no lag parameters passed, use default: mode = 'valid' */ + return PyArray_Correlate2(a0, shape, 0); + } + return PyArray_CorrelateLags(a0, shape, minlag, maxlag, lagstep); +} + static PyObject * array_arange(PyObject *NPY_UNUSED(ignored), PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames) @@ -4591,6 +4812,9 @@ static struct PyMethodDef array_module_methods[] = { {"correlate2", (PyCFunction)array_correlate2, METH_FASTCALL | METH_KEYWORDS, NULL}, + {"correlatelags", + (PyCFunction)array_correlatelags, + METH_FASTCALL | METH_KEYWORDS, NULL}, {"frombuffer", (PyCFunction)array_frombuffer, METH_VARARGS | METH_KEYWORDS, NULL}, diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py index 9e71b7c6b1b8..5e2c052970e0 100644 --- a/numpy/_core/tests/test_numeric.py +++ b/numpy/_core/tests/test_numeric.py @@ -3603,8 +3603,10 @@ def _setup(self, dt): def test_float(self): self._setup(float) - z = np.correlate(self.x, self.y, 'full') + z, lags = np.correlate(self.x, self.y, 'full', returns_lagvector=True) assert_array_almost_equal(z, self.z1) + assert_array_equal(lags, np.arange(-2, 5)) + assert len(z) == len(lags) 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') @@ -3615,6 +3617,44 @@ def test_float(self): assert_array_almost_equal(z, self.z2r) z = np.correlate(self.xs, self.y, 'full') assert_array_almost_equal(z, self.zs) + # Test 'valid' mode - middle part where there's full overlap + z, lags = np.correlate(self.x, self.y, 'valid', returns_lagvector=True) + assert_array_almost_equal(z, self.z1[2:5]) # [-14., -20., -26.] + assert_array_equal(lags, np.arange(0, 3)) + assert len(z) == len(lags) + # Test 'same' mode - same length as first input + z, lags = np.correlate(self.x, self.y, 'same', returns_lagvector=True) + assert_array_almost_equal(z, self.z1[1:6]) # [-8., -14., -20., -26., -14.] + assert_array_equal(lags, np.arange(-1, 4)) + assert len(z) == len(lags) + + def test_int_lags(self): + self._setup(float) + # Test 'lags' mode + lagsin = 2 + print(f"int lags: {lagsin=}") + z, lags = np.correlate(self.x, self.y, 'lags', lags=lagsin, returns_lagvector=True) + assert_array_almost_equal(z, self.z1[1:4]) # [-8., 14., -20.] + assert_array_equal(lags, np.arange(-lagsin+1, lagsin)) + assert len(z) == len(lags) + + def test_two_tuple_lags(self): + self._setup(float) + lagsin = (-1, 3) + print(f"two tuple lags: {lagsin=}") + z, lags = np.correlate(self.x, self.y, 'lags', lags=lagsin, returns_lagvector=True) + assert_array_almost_equal(z, self.z1[1:5]) # [-8., -14., -20., -26.] + assert_array_equal(lags, np.arange(*lagsin)) + assert len(z) == len(lags) + + def test_three_tuple_lags(self): + self._setup(float) + lagsin = (-2, 5, 2) + print(f"three tuple lags: {lagsin=}") + z, lags = np.correlate(self.x, self.y, 'lags', lags=lagsin, returns_lagvector=True) + assert_array_almost_equal(z, self.z1[0::2]) # [-3, -14., -26., -5.] + assert_array_equal(lags, np.arange(*lagsin)) + assert len(z) == len(lags) def test_object(self): self._setup(Decimal) From d2ce49f80a0796aa30ac51d76df15a88c2518bc6 Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Wed, 15 Apr 2026 16:42:23 +0000 Subject: [PATCH 02/13] fixed bug and memory unsafety, made inversion entirely internal to _pyarray_correlate --- numpy/_core/numeric.py | 168 ++++++++++-------- numpy/_core/src/multiarray/multiarraymodule.c | 159 ++++++++--------- 2 files changed, 166 insertions(+), 161 deletions(-) diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index d4ed05eebb8c..8dd65dbc68a4 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -718,64 +718,75 @@ def flatnonzero(a): return np.nonzero(np.ravel(a))[0] +_CorrModeDefault = object() + +_CORR_MODE_MAP = { + 'valid': 0, + 'same': 1, + 'full': 2, + 'lags': 3, +} + + def _mode_from_name(mode): - if type(mode) not in (str, int): - raise TypeError("correlate/convolve mode argument must be unused or" + - " one of {'valid', 'same', 'full', 'lags'}") - - mode_map = { - 'valid': 0, - 'same': 1, - 'full': 2, - 'lags': 3 - } - - if mode in mode_map.values(): - return mode - - if mode in mode_map: - return mode_map[mode] - - raise ValueError("correlate/convolve mode argument must be unused or" + - " one of {'valid', 'same', 'full', 'lags'}") + if isinstance(mode, int): + if mode in _CORR_MODE_MAP.values(): + return mode + elif isinstance(mode, str): + if mode in _CORR_MODE_MAP: + return _CORR_MODE_MAP[mode] + else: + raise TypeError("correlate/convolve mode must be a string or int, " + "one of 'valid', 'same', 'full', 'lags'") + raise ValueError("correlate/convolve mode must be " + "one of 'valid', 'same', 'full', 'lags'") def _lags_from_lags(lag): - if type(lag) is int: # maxlag - lags = (-lag+1, lag, 1) - elif type(lag) is tuple: # minlag and maxlag - if len(lag) > 2: - lags = (int(lag[0]), int(lag[1]), int(lag[2])) + if isinstance(lag, (int, nt.integer)): + return (-lag + 1, int(lag), 1) + elif isinstance(lag, (tuple, list)): + if len(lag) == 3: + minlag, maxlag, lagstep = int(lag[0]), int(lag[1]), int(lag[2]) + elif len(lag) == 2: + minlag, maxlag, lagstep = int(lag[0]), int(lag[1]), 1 else: - lags = (int(lag[0]), int(lag[1]), 1) + raise ValueError("lags tuple must have 2 or 3 elements") + if lagstep == 0: + raise ValueError("lagstep must not be zero") + return (minlag, maxlag, lagstep) else: - raise ValueError("correlate/convolve lags argument must be " + - "int or int tuple.") - return lags + raise TypeError("lags must be an int or a tuple of ints") def _lags_from_mode(alen, vlen, mode): - inverted = 0 + """ + Compute the lag range for a given mode. + + Produces lags aligned with the caller's original (alen, vlen) order + (i.e. matches what ``np.correlate(a, v, mode)`` returns element-by-element). + """ + inverted = False if alen < vlen: alen, vlen = vlen, alen - inverted = 1 + inverted = True if mode == 0: - mode_lags = (0, alen-vlen+1, 1) + m0, m1 = 0, alen - vlen + 1 elif mode == 1: - mode_lags = (-int(vlen/2), alen - int(vlen/2), 1) + half = int(vlen / 2) + m0, m1 = -half, alen - half elif mode == 2: - mode_lags = (-vlen+1, alen, 1) + m0, m1 = -vlen + 1, alen else: - raise ValueError("correlate/convolve mode argument must be unused or" + - " one of {'valid', 'same', 'full', 'lags'}") + raise ValueError("correlate/convolve mode must be " + "one of 'valid', 'same', 'full'") if inverted: - mode_lags = (-int(math.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]) + count = m1 - m0 # lagstep is always 1 for mode-based lags + m0, m1 = -(m0 + count - 1), -m0 + 1 - return mode_lags + return (m0, m1, 1) def _correlate_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): @@ -783,7 +794,8 @@ def _correlate_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): @array_function_dispatch(_correlate_dispatcher) -def correlate(a, v, mode='default', lags=None, returns_lagvector=False): +def correlate(a, v, mode=_CorrModeDefault, lags=None, + returns_lagvector=False): r""" Cross-correlation of two 1-dimensional sequences. @@ -801,7 +813,9 @@ def correlate(a, v, mode='default', lags=None, returns_lagvector=False): Input sequences. mode : {'valid', 'same', 'full', 'lags'}, optional Refer to the `convolve` docstring. Note that the default - is 'default', unlike `convolve`, which uses 'full'. + is 'valid', unlike `convolve`, which uses 'full'. + If ``lags`` is provided and ``mode`` is not specified, + ``mode`` defaults to 'lags'. lags : int or int tuple, optional Refer to the `convolve` docstring. mode should be unset or set to 'lags' to use the lags argument @@ -870,28 +884,28 @@ def correlate(a, v, mode='default', lags=None, returns_lagvector=False): 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 is not None: - mode = 'lags' - else: - mode = 'valid' - mode = _mode_from_name(mode) # guaranteed a value in _mode_from_name_dict + # Resolve mode: if unset, infer from whether lags was provided + if mode is _CorrModeDefault: + mode = 'lags' if lags is not None else 'valid' + mode = _mode_from_name(mode) + if mode in (0, 1, 2): if lags is not None: - raise ValueError("correlate mode keyword argument must be 'lags'" + - " or unused if the lags keyword argument is used.") + raise ValueError("lags cannot be used with mode " + "'valid', 'same', or 'full'") result = multiarray.correlate2(a, v, mode) if returns_lagvector: - alen, vlen = len(a), len(v) - lags_throuple = _lags_from_mode(alen, vlen, mode) + lags_tuple = _lags_from_mode(len(a), len(v), mode) elif mode == 3: - lags_throuple = _lags_from_lags(lags) - result = multiarray.correlatelags(a, v, lags_throuple[0], lags_throuple[1], lags_throuple[2]) + if lags is None: + raise ValueError("lags argument is required for mode='lags'") + lags_tuple = _lags_from_lags(lags) + result = multiarray.correlatelags( + a, v, lags_tuple[0], lags_tuple[1], lags_tuple[2]) if returns_lagvector: - return result, arange(lags_throuple[0], lags_throuple[1], lags_throuple[2]) - else: - return result + return result, arange(lags_tuple[0], lags_tuple[1], lags_tuple[2]) + return result def _convolve_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): @@ -899,7 +913,8 @@ def _convolve_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): @array_function_dispatch(_convolve_dispatcher) -def convolve(a, v, mode='full', lags=None, returns_lagvector=False): +def convolve(a, v, mode=_CorrModeDefault, lags=None, + returns_lagvector=False): """ Returns the discrete, linear convolution of two one-dimensional sequences. @@ -917,7 +932,7 @@ def convolve(a, v, mode='full', lags=None, returns_lagvector=False): First one-dimensional input array. v : (M,) array_like Second one-dimensional input array. - mode : int, int tuple, or {'full', 'valid', 'same'}, optional + mode : {'full', 'valid', 'same', 'lags'}, 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 @@ -930,7 +945,7 @@ def convolve(a, v, mode='full', lags=None, returns_lagvector=False): 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 @@ -941,6 +956,8 @@ def convolve(a, v, mode='full', lags=None, returns_lagvector=False): 'lags': Mode 'lags' uses the lags argument to define the lags for which to perform the convolution. + If ``lags`` is provided and ``mode`` is not specified, + ``mode`` defaults to 'lags'. lags : int or int tuple, optional Mode should be unset or set to 'lags' to use the lags argument. @@ -1038,7 +1055,7 @@ def convolve(a, v, mode='full', lags=None, returns_lagvector=False): 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=None, ndmin=1), array(v, copy=None, ndmin=1) alen, vlen = len(a), len(v) @@ -1049,28 +1066,29 @@ def convolve(a, v, mode='full', lags=None, returns_lagvector=False): if vlen > alen: a, v = v, a - - if mode == 'default': - if lags is not None: - mode = 'lags' - else: - mode = 'full' - mode = _mode_from_name(mode) # guaranteed a value in _mode_from_name_dict + + # Resolve mode: if unset, infer from whether lags was provided + if mode is _CorrModeDefault: + mode = 'lags' if lags is not None else 'full' + mode = _mode_from_name(mode) + if mode in (0, 1, 2): if lags is not None: - raise ValueError("convolve mode keyword argument must be 'lags' " + - "or unused if the lags keyword argument is used.") + raise ValueError("lags cannot be used with mode " + "'valid', 'same', or 'full'") result = multiarray.correlate2(a, v[::-1], mode) if returns_lagvector: - lags = _lags_from_mode(alen, vlen, mode) + lags_tuple = _lags_from_mode(alen, vlen, mode) elif mode == 3: - lags = _lags_from_lags(lags) - result = multiarray.correlatelags(a, v[::-1], lags[0], lags[1], lags[2]) + if lags is None: + raise ValueError("lags argument is required for mode='lags'") + lags_tuple = _lags_from_lags(lags) + result = multiarray.correlatelags( + a, v[::-1], lags_tuple[0], lags_tuple[1], lags_tuple[2]) if returns_lagvector: - return result, arange(lags[0], lags[1], lags[2]) - else: - return result + return result, arange(lags_tuple[0], lags_tuple[1], lags_tuple[2]) + return result def _outer_dispatcher(a, b, out=None): diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c index bb94393bcb04..7de431ddc0f2 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.c +++ b/numpy/_core/src/multiarray/multiarraymodule.c @@ -1132,26 +1132,33 @@ PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyArrayObject* out) return NULL; } +/* Forward declaration -- definition is below _pyarray_correlate */ +static int _pyarray_revert(PyArrayObject *ret); + /* - * Implementation which is common between PyArray_Correlate - * PyArray_Correlate2, and PyArray_CorrelateLags. + * Core correlation computation. + * + * Handles all input forms internally: + * - Swaps arrays if n1 < n2 (and negates lags) + * - Normalizes negative lagstep to positive with reversed range + * - Reverses output if internal swaps changed the lag orientation * - * inverted is set to 1 if computed correlate(ap2, ap1), 0 otherwise + * Callers may pass any valid combination of (ap1, ap2, minlag, maxlag, lagstep) + * and receive a result aligned with their original array/lag order. */ static PyArrayObject* _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, - int typenum, int *inverted, + int typenum, npy_intp minlag, npy_intp maxlag, npy_intp lagstep) { - printf("_pyarray_correlate: minlag: %ld, maxlag: %ld, lagstep: %ld\n", minlag, maxlag, lagstep); - - PyArrayObject *ret; + PyArrayObject *ret, *swap; npy_intp length; npy_intp i, i1, n1, n2, n, n11; - npy_intp lag, tmplag, maxleft, maxright; + npy_intp lag, tmplag, maxleft, maxright, phase2_end; npy_intp is1, is2, os; char *ip1, *ip2, *op; PyArray_DotFunc *dot; + int inverted; NPY_BEGIN_THREADS_DEF; @@ -1170,30 +1177,34 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, PyErr_SetString(PyExc_ValueError, "second array argument cannot be empty"); return NULL; } + + /* Ensure n1 >= n2 by swapping. Swapping negates the lag orientation. */ if (n1 < n2) { - ret = ap1; + swap = ap1; ap1 = ap2; - ap2 = ret; - ret = NULL; + ap2 = swap; i = n1; n1 = n2; - n2 = i; + n2 = i; minlag = -minlag; maxlag = -maxlag; lagstep = -lagstep; } + /* Normalize negative lagstep: reverse the range to use positive step. + * The output will need to be reversed before returning. */ if (lagstep < 0) { - *inverted = 1; + inverted = 1; i = minlag; i1 = (npy_intp)(npy_ceil((maxlag - minlag)/(float)lagstep))*lagstep; - minlag = i1 + minlag - lagstep; + minlag = i1 + minlag - lagstep; maxlag = i - lagstep; lagstep = -lagstep; } else { - *inverted = 0; + inverted = 0; } + if (maxlag <= minlag) { length = 0; } @@ -1210,8 +1221,11 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, if (ret == NULL) { return NULL; } - /* Zero-initialize so out-of-bounds lags are correctly zero */ - memset(PyArray_DATA(ret), 0, PyArray_NBYTES(ret)); + /* Zero-initialize so out-of-bounds lags are correctly zero. + * Skip for object arrays where zeroing would create NULL pointers. */ + if (!PyDataType_FLAGCHK(PyArray_DESCR(ret), NPY_NEEDS_PYAPI)) { + memset(PyArray_DATA(ret), 0, PyArray_NBYTES(ret)); + } dot = PyDataType_GetArrFuncs(PyArray_DESCR(ret))->dotfunc; if (dot == NULL) { PyErr_SetString(PyExc_ValueError, @@ -1252,29 +1266,21 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, /* iterate over answer vector */ op += os; } - 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 { - n11 = n1; - } - /* starts at lag=minlag if minlag>0. - * Does lags where y entirely overlaps with x. - */ - if (lagstep == 1 && lag < maxlag && + n11 = n1; + /* Phase 2: lags where y entirely overlaps with x, i.e. lag in [0, n1-n2]. + * Upper bound is the smaller of maxlag and the full-overlap boundary. */ + phase2_end = maxlag < (n11 - n2 + 1) ? maxlag : (n11 - n2 + 1); + if (lagstep == 1 && lag < phase2_end && small_correlate(ip1 + lag*is1, is1, - n11 - n2 + 1 - lag, PyArray_TYPE(ap1), + phase2_end - lag, PyArray_TYPE(ap1), ip2, is2, n2, PyArray_TYPE(ap2), op, os)) { - op += os * (n11 - n2 + 1 - lag); - lag = n11 - n2 + 1; + op += os * (phase2_end - lag); + lag = phase2_end; } - else if (lag < maxlag) { + else if (lag < phase2_end) { tmplag = lag; - for (lag = tmplag; lag < (n11 - n2 + 1) && (!needs_pyapi || !PyErr_Occurred()); lag+=lagstep) { + for (lag = tmplag; lag < phase2_end && (!needs_pyapi || !PyErr_Occurred()); lag+=lagstep) { dot(ip1 + lag*is1, is1, ip2, is2, op, n2, ret); op += os; } @@ -1294,6 +1300,14 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, goto clean_ret; } + /* If we inverted the range, reverse the output so it aligns with the + * caller's lag orientation. */ + if (inverted) { + if (_pyarray_revert(ret) != 0) { + goto clean_ret; + } + } + return ret; clean_ret: @@ -1349,22 +1363,22 @@ _pyarray_revert(PyArrayObject *ret) } /* - * Generate the lags corresponding to each mode - * n1 and n2 are sizes of two vectors - * minlag, maxlag, and lagstep are edited in-place + * Generate the lags corresponding to each mode. + * Produces lags aligned with the caller's original (n1, n2) order -- + * callers can pass unswapped dimensions and use the result directly. + * minlag, maxlag, and lagstep are edited in-place. */ -void +static 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; + npy_intp tmp, i, i1; + int inverted = 0; if (n1 < n2) { - i = n1; + tmp = n1; n1 = n2; - n2 = i; + n2 = tmp; inverted = 1; } @@ -1390,14 +1404,15 @@ _lags_from_mode(int mode, npy_intp n1, npy_intp n2, return; } + /* If we swapped to get n1 >= n2, transform the lags back to the + * caller's original orientation. */ if (inverted) { i = m0; i1 = (npy_intp)(npy_ceil((m1 - m0)/(float)s))*s; - m0 = -(i1 + m0 - s); + m0 = -(i1 + m0 - s); m1 = -(i - s); } - /* set minlag, maxlag, and lagstep values */ *minlag = m0; *maxlag = m1; *lagstep = s; @@ -1417,8 +1432,6 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) PyArray_Descr *typec; npy_intp minlag, maxlag, lagstep; npy_intp n1, n2; - int inverted; - int st; typenum = PyArray_ObjectType(op1, NPY_NOTYPE); if (typenum == NPY_NOTYPE) { @@ -1453,34 +1466,22 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) 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); - printf("PyArray_Correlate2: mode: %d, n1: %ld, n2: %ld, minlag: %ld, maxlag: %ld, lagstep: %ld\n", mode, n1, n2, minlag, maxlag, lagstep); - ret = _pyarray_correlate(ap1, ap2, typenum, &inverted, + if (PyErr_Occurred()) { + goto clean_ap2; + } + ret = _pyarray_correlate(ap1, ap2, typenum, 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: @@ -1491,7 +1492,9 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) /*NUMPY_API * correlate(a1,a2,minlag,maxlag,lagstep) * - * This function computes the cross-correlation of a1 with a2. + * This function computes the cross-correlation of a1 with a2 at the + * specified lag range. The arrays and lag parameters may be provided in + * any valid form; normalization is handled internally. * This function does not accept modes. * See _lags_from_mode() to convert modes to lags. */ @@ -1502,8 +1505,6 @@ PyArray_CorrelateLags(PyObject *op1, PyObject *op2, PyArrayObject *ap1, *ap2, *ret = NULL; int typenum; PyArray_Descr *typec; - int inverted; - int st; typenum = PyArray_ObjectType(op1, NPY_NOTYPE); if (typenum == NPY_NOTYPE) { @@ -1538,29 +1539,16 @@ PyArray_CorrelateLags(PyObject *op1, PyObject *op2, ap2 = cap2; } - ret = _pyarray_correlate(ap1, ap2, typenum, &inverted, + ret = _pyarray_correlate(ap1, ap2, typenum, 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: @@ -1576,7 +1564,6 @@ PyArray_Correlate(PyObject *op1, PyObject *op2, int mode) { PyArrayObject *ap1, *ap2, *ret = NULL; int typenum; - int unused; PyArray_Descr *typec; npy_intp minlag, maxlag, lagstep; npy_intp n1, n2; @@ -1604,11 +1591,13 @@ PyArray_Correlate(PyObject *op1, PyObject *op2, int mode) goto fail; } - /* 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, + if (PyErr_Occurred()) { + goto fail; + } + ret = _pyarray_correlate(ap1, ap2, typenum, minlag, maxlag, lagstep); if (ret == NULL) { goto fail; @@ -3244,7 +3233,6 @@ array_correlate2(PyObject *NPY_UNUSED(dummy), NULL, NULL, NULL) < 0) { return NULL; } - printf("array_correlate2: mode: %d\n", mode); return PyArray_Correlate2(a0, shape, mode); } @@ -3265,7 +3253,6 @@ array_correlatelags(PyObject *NPY_UNUSED(dummy), NULL, NULL, NULL) < 0) { return NULL; } - printf("array_correlate_lags: minlag: %ld, maxlag: %ld, lagstep: %ld\n", minlag, maxlag, lagstep); if (minlag == 0 && maxlag == 0 && lagstep == 0) { /* if no lag parameters passed, use default: mode = 'valid' */ return PyArray_Correlate2(a0, shape, 0); From ce94ff5c4efa107ed5813facbc191a434d31cf7f Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Wed, 15 Apr 2026 17:01:57 +0000 Subject: [PATCH 03/13] add tests and adjust signature to include lags and maxlag arguments --- numpy/_core/numeric.py | 179 ++++++++++++------- numpy/_core/tests/test_numeric.py | 283 ++++++++++++++++++++++++++++-- 2 files changed, 380 insertions(+), 82 deletions(-) diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index 8dd65dbc68a4..dcc5e134a174 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -742,21 +742,52 @@ def _mode_from_name(mode): "one of 'valid', 'same', 'full', 'lags'") +def _lags_from_maxlag(maxlag): + """Convert a maxlag int to a (minlag, maxlag_exclusive, step) tuple. + + ``maxlag=M`` requests the symmetric inclusive range ``[-M, M]``, + matching MATLAB's ``xcorr(x, y, M)`` convention (2*M+1 lags total). + """ + if not isinstance(maxlag, (int, nt.integer)): + raise TypeError("maxlag must be an integer") + m = int(maxlag) + if m < 0: + raise ValueError("maxlag must be non-negative") + return (-m, m + 1, 1) + + def _lags_from_lags(lag): - if isinstance(lag, (int, nt.integer)): - return (-lag + 1, int(lag), 1) - elif isinstance(lag, (tuple, list)): - if len(lag) == 3: - minlag, maxlag, lagstep = int(lag[0]), int(lag[1]), int(lag[2]) - elif len(lag) == 2: - minlag, maxlag, lagstep = int(lag[0]), int(lag[1]), 1 - else: - raise ValueError("lags tuple must have 2 or 3 elements") - if lagstep == 0: + """Convert a lags specification to (minlag, maxlag_exclusive, lagstep). + + Accepts a ``range``, ``slice`` (with explicit start/stop), or a 1-D + array_like containing an arithmetic progression of integers. + """ + if isinstance(lag, range): + return (lag.start, lag.stop, lag.step) + if isinstance(lag, slice): + if lag.start is None or lag.stop is None: + raise ValueError("lags slice must have explicit start and stop") + step = 1 if lag.step is None else int(lag.step) + if step == 0: raise ValueError("lagstep must not be zero") - return (minlag, maxlag, lagstep) - else: - raise TypeError("lags must be an int or a tuple of ints") + return (int(lag.start), int(lag.stop), step) + + arr = asanyarray(lag) + if arr.ndim != 1 or arr.size == 0: + raise ValueError("lags must be a 1-D non-empty sequence") + if not nt.issubdtype(arr.dtype, nt.integer): + raise TypeError("lags values must be integers") + if arr.size == 1: + return (int(arr[0]), int(arr[0]) + 1, 1) + step = int(arr[1] - arr[0]) + if step == 0: + raise ValueError("lagstep must not be zero") + expected = arr[0] + step * arange(arr.size) + if not (arr == expected).all(): + raise ValueError( + "lags array must be an arithmetic progression " + "(use a range or slice for non-arithmetic patterns)") + return (int(arr[0]), int(arr[-1]) + step, step) def _lags_from_mode(alen, vlen, mode): @@ -789,12 +820,13 @@ def _lags_from_mode(alen, vlen, mode): return (m0, m1, 1) -def _correlate_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): +def _correlate_dispatcher(a, v, mode=None, *, maxlag=None, lags=None, + returns_lagvector=None): return (a, v) @array_function_dispatch(_correlate_dispatcher) -def correlate(a, v, mode=_CorrModeDefault, lags=None, +def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, returns_lagvector=False): r""" Cross-correlation of two 1-dimensional sequences. @@ -814,11 +846,17 @@ def correlate(a, v, mode=_CorrModeDefault, lags=None, mode : {'valid', 'same', 'full', 'lags'}, optional Refer to the `convolve` docstring. Note that the default is 'valid', unlike `convolve`, which uses 'full'. - If ``lags`` is provided and ``mode`` is not specified, + If ``maxlag`` or ``lags`` is provided and ``mode`` is not specified, ``mode`` defaults to 'lags'. - lags : int or int tuple, optional - Refer to the `convolve` docstring. - mode should be unset or set to 'lags' to use the lags argument + maxlag : int, optional + Compute the cross-correlation at lags ``-maxlag, -maxlag+1, ..., maxlag`` + (a symmetric inclusive window of ``2*maxlag+1`` lags around 0). + Mutually exclusive with ``lags``. + lags : range, slice, or 1-D array_like of int, optional + Explicit lag specification. Accepts a Python ``range`` object, a + ``slice`` with explicit start and stop, or a 1-D array_like containing + an arithmetic progression of integer lag indices. + Mutually exclusive with ``maxlag``. 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 @@ -866,9 +904,10 @@ def correlate(a, v, mode=_CorrModeDefault, lags=None, array([ 2. , 3.5, 3. ]) >>> 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) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], maxlag=1) array([ 2. , 3.5, 3. ]) - >>> np.correlate([1, 2, 3], [0, 1, 0.5], mode="lags", lags=(-1,2,2), returns_lagvector=True) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], lags=range(-1, 2, 2), + ... returns_lagvector=True) (array([ 2., 3.]), array([-1, 1])) Using complex sequences: @@ -884,22 +923,31 @@ def correlate(a, v, mode=_CorrModeDefault, lags=None, array([ 0.0+0.j , 3.0+1.j , 1.5+1.5j, 1.0+0.j , 0.5+0.5j]) """ - # Resolve mode: if unset, infer from whether lags was provided + if maxlag is not None and lags is not None: + raise TypeError("cannot specify both maxlag and lags") + lags_given = maxlag is not None or lags is not None + + # Resolve mode: if unset, infer from whether maxlag/lags was provided if mode is _CorrModeDefault: - mode = 'lags' if lags is not None else 'valid' + mode = 'lags' if lags_given else 'valid' mode = _mode_from_name(mode) if mode in (0, 1, 2): - if lags is not None: - raise ValueError("lags cannot be used with mode " - "'valid', 'same', or 'full'") + if lags_given: + raise ValueError( + "maxlag/lags cannot be used with mode " + "'valid', 'same', or 'full'") result = multiarray.correlate2(a, v, mode) if returns_lagvector: lags_tuple = _lags_from_mode(len(a), len(v), mode) elif mode == 3: - if lags is None: - raise ValueError("lags argument is required for mode='lags'") - lags_tuple = _lags_from_lags(lags) + if not lags_given: + raise ValueError( + "maxlag or lags is required for mode='lags'") + if maxlag is not None: + lags_tuple = _lags_from_maxlag(maxlag) + else: + lags_tuple = _lags_from_lags(lags) result = multiarray.correlatelags( a, v, lags_tuple[0], lags_tuple[1], lags_tuple[2]) @@ -908,12 +956,13 @@ def correlate(a, v, mode=_CorrModeDefault, lags=None, return result -def _convolve_dispatcher(a, v, mode=None, lags=None, returns_lagvector=None): +def _convolve_dispatcher(a, v, mode=None, *, maxlag=None, lags=None, + returns_lagvector=None): return (a, v) @array_function_dispatch(_convolve_dispatcher) -def convolve(a, v, mode=_CorrModeDefault, lags=None, +def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, returns_lagvector=False): """ Returns the discrete, linear convolution of two one-dimensional sequences. @@ -954,26 +1003,20 @@ def convolve(a, v, mode=_CorrModeDefault, lags=None, 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. - If ``lags`` is provided and ``mode`` is not specified, - ``mode`` defaults to 'lags'. - - 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. + Mode 'lags' uses ``maxlag`` or ``lags`` to define the lags for + which to perform the convolution. + If ``maxlag`` or ``lags`` is provided and ``mode`` is not + specified, ``mode`` defaults to 'lags'. + + maxlag : int, optional + Compute the convolution at lags ``-maxlag, -maxlag+1, ..., maxlag`` + (a symmetric inclusive window of ``2*maxlag+1`` lags around 0). + Mutually exclusive with ``lags``. + lags : range, slice, or 1-D array_like of int, optional + Explicit lag specification. Accepts a Python ``range`` object, a + ``slice`` with explicit start and stop, or a 1-D array_like containing + an arithmetic progression of integer lag indices. + Mutually exclusive with ``maxlag``. returns_lagvector : bool, optional If True, the function returns a lagvector array in addition to the convolution result. The lagvector contains the indices of @@ -1047,13 +1090,14 @@ def convolve(a, v, mode=_CorrModeDefault, lags=None, 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) + >>> np.convolve([1,2,3],[0,1,0.5], maxlag=1, 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) + with steps of length 2: + + >>> np.convolve([1,2,3,4,5], [0,1,0.5], lags=range(-2, 6, 2), + ... returns_lagvector=True) (array([ 0. , 2.5, 5.5, 2.5]), array([-2, 0, 2, 4])) """ @@ -1067,22 +1111,31 @@ def convolve(a, v, mode=_CorrModeDefault, lags=None, if vlen > alen: a, v = v, a - # Resolve mode: if unset, infer from whether lags was provided + if maxlag is not None and lags is not None: + raise TypeError("cannot specify both maxlag and lags") + lags_given = maxlag is not None or lags is not None + + # Resolve mode: if unset, infer from whether maxlag/lags was provided if mode is _CorrModeDefault: - mode = 'lags' if lags is not None else 'full' + mode = 'lags' if lags_given else 'full' mode = _mode_from_name(mode) if mode in (0, 1, 2): - if lags is not None: - raise ValueError("lags cannot be used with mode " - "'valid', 'same', or 'full'") + if lags_given: + raise ValueError( + "maxlag/lags cannot be used with mode " + "'valid', 'same', or 'full'") result = multiarray.correlate2(a, v[::-1], mode) if returns_lagvector: lags_tuple = _lags_from_mode(alen, vlen, mode) elif mode == 3: - if lags is None: - raise ValueError("lags argument is required for mode='lags'") - lags_tuple = _lags_from_lags(lags) + if not lags_given: + raise ValueError( + "maxlag or lags is required for mode='lags'") + if maxlag is not None: + lags_tuple = _lags_from_maxlag(maxlag) + else: + lags_tuple = _lags_from_lags(lags) result = multiarray.correlatelags( a, v[::-1], lags_tuple[0], lags_tuple[1], lags_tuple[2]) diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py index 5e2c052970e0..c7a4b8e4875b 100644 --- a/numpy/_core/tests/test_numeric.py +++ b/numpy/_core/tests/test_numeric.py @@ -3628,34 +3628,202 @@ def test_float(self): assert_array_equal(lags, np.arange(-1, 4)) assert len(z) == len(lags) - def test_int_lags(self): + def test_maxlag(self): self._setup(float) - # Test 'lags' mode - lagsin = 2 - print(f"int lags: {lagsin=}") - z, lags = np.correlate(self.x, self.y, 'lags', lags=lagsin, returns_lagvector=True) - assert_array_almost_equal(z, self.z1[1:4]) # [-8., 14., -20.] - assert_array_equal(lags, np.arange(-lagsin+1, lagsin)) + # maxlag=1 gives the symmetric inclusive window [-1, 0, 1] + z, lags = np.correlate(self.x, self.y, maxlag=1, + returns_lagvector=True) + assert_array_almost_equal(z, self.z1[1:4]) # [-8., -14., -20.] + assert_array_equal(lags, np.arange(-1, 2)) assert len(z) == len(lags) - - def test_two_tuple_lags(self): + + def test_lags_range(self): self._setup(float) - lagsin = (-1, 3) - print(f"two tuple lags: {lagsin=}") - z, lags = np.correlate(self.x, self.y, 'lags', lags=lagsin, returns_lagvector=True) + z, lags = np.correlate(self.x, self.y, lags=range(-1, 3), + returns_lagvector=True) assert_array_almost_equal(z, self.z1[1:5]) # [-8., -14., -20., -26.] - assert_array_equal(lags, np.arange(*lagsin)) + assert_array_equal(lags, np.arange(-1, 3)) assert len(z) == len(lags) - def test_three_tuple_lags(self): + def test_lags_range_with_step(self): self._setup(float) - lagsin = (-2, 5, 2) - print(f"three tuple lags: {lagsin=}") - z, lags = np.correlate(self.x, self.y, 'lags', lags=lagsin, returns_lagvector=True) - assert_array_almost_equal(z, self.z1[0::2]) # [-3, -14., -26., -5.] - assert_array_equal(lags, np.arange(*lagsin)) + z, lags = np.correlate(self.x, self.y, lags=range(-2, 5, 2), + returns_lagvector=True) + assert_array_almost_equal(z, self.z1[0::2]) # [-3., -14., -26., -5.] + assert_array_equal(lags, np.arange(-2, 5, 2)) assert len(z) == len(lags) + def test_lags_slice(self): + self._setup(float) + z, lags = np.correlate(self.x, self.y, lags=slice(-1, 3), + returns_lagvector=True) + assert_array_almost_equal(z, self.z1[1:5]) + assert_array_equal(lags, np.arange(-1, 3)) + + def test_lags_array(self): + self._setup(float) + # 1-D array of lags (must be an arithmetic progression) + z, lags = np.correlate(self.x, self.y, + lags=np.array([-2, 0, 2, 4]), + returns_lagvector=True) + assert_array_almost_equal(z, self.z1[0::2]) + assert_array_equal(lags, np.array([-2, 0, 2, 4])) + + def test_maxlag_and_lags_mutually_exclusive(self): + self._setup(float) + with assert_raises(TypeError): + np.correlate(self.x, self.y, maxlag=1, lags=range(-1, 2)) + + def test_lags_with_non_lags_mode(self): + self._setup(float) + with assert_raises(ValueError): + np.correlate(self.x, self.y, mode='full', maxlag=1) + with assert_raises(ValueError): + np.correlate(self.x, self.y, mode='valid', lags=range(0, 2)) + + def test_lags_mode_without_lags(self): + self._setup(float) + with assert_raises(ValueError): + np.correlate(self.x, self.y, mode='lags') + + def test_lags_non_arithmetic_array_rejected(self): + self._setup(float) + with assert_raises(ValueError): + np.correlate(self.x, self.y, lags=np.array([0, 1, 3])) + + # --- Lag range geometry tests --- + # n1=5 (self.x), n2=3 (self.y), z1 covers lags [-2..4] + + def test_lags_left_partial_only(self): + self._setup(float) + # All requested lags < 0 (left partial overlap region only) + z = np.correlate(self.x, self.y, lags=range(-2, 0)) + assert_array_almost_equal(z, self.z1[0:2]) # [-3, -8] + + def test_lags_right_partial_only(self): + self._setup(float) + # All requested lags >= n1-n2+1=3 (right partial overlap region only) + z = np.correlate(self.x, self.y, lags=range(3, 5)) + assert_array_almost_equal(z, self.z1[5:7]) # [-14, -5] + + def test_lags_full_overlap_only(self): + self._setup(float) + # Only full-overlap lags [0, n1-n2] = [0, 2] + z = np.correlate(self.x, self.y, lags=range(0, 3)) + assert_array_almost_equal(z, self.z1[2:5]) + + def test_lags_beyond_overlap(self): + self._setup(float) + # Lags entirely outside any overlap -> all zeros + z = np.correlate(self.x, self.y, lags=range(-10, -5)) + assert_array_almost_equal(z, np.zeros(5)) + z = np.correlate(self.x, self.y, lags=range(10, 15)) + assert_array_almost_equal(z, np.zeros(5)) + + def test_lags_partial_beyond_overlap(self): + self._setup(float) + # Lags partially outside overlap -> zeros at extremes + z = np.correlate(self.x, self.y, lags=range(-5, 8)) + expected = np.concatenate([np.zeros(3), self.z1, np.zeros(3)]) + assert_array_almost_equal(z, expected) + + def test_lags_single_lag(self): + self._setup(float) + z, lags = np.correlate(self.x, self.y, lags=range(0, 1), + returns_lagvector=True) + assert_array_almost_equal(z, [self.z1[2]]) + assert_array_equal(lags, np.array([0])) + + def test_lags_negative_step(self): + self._setup(float) + # Negative step: reversed lag order + z, lags = np.correlate(self.x, self.y, lags=slice(2, -3, -1), + returns_lagvector=True) + # slice(2,-3,-1) -> [2, 1, 0, -1, -2] + # corresponds to z1 indices [4, 3, 2, 1, 0] + assert_array_almost_equal(z, self.z1[4::-1]) + assert_array_equal(lags, np.arange(2, -3, -1)) + + def test_lags_empty(self): + self._setup(float) + z, lags = np.correlate(self.x, self.y, lags=range(0, 0), + returns_lagvector=True) + assert len(z) == 0 + assert len(lags) == 0 + + # --- Array geometry tests --- + + def test_lags_v_longer_than_a(self): + self._setup(float) + # When len(a) < len(v): correlate(y, x) is the time-reversal of + # correlate(x, y) for real inputs. + z = np.correlate(self.y, self.x, lags=range(-4, 3)) + assert_array_almost_equal(z, self.z1[::-1]) + + def test_lags_equal_length(self): + a = np.array([1, 2, 3], dtype=float) + v = np.array([4, 5, 6], dtype=float) + z_mode = np.correlate(a, v, 'full') + z_lags = np.correlate(a, v, lags=range(-2, 3)) + assert_array_almost_equal(z_mode, z_lags) + + def test_lags_autocorrelation(self): + # a == v: result should be symmetric for real arrays + a = np.array([1.0, 2.0, 3.0]) + z = np.correlate(a, a, lags=range(-2, 3)) + assert_array_almost_equal(z, z[::-1]) + + # --- Equivalence between mode and explicit lag range --- + + def test_lags_matches_full_mode(self): + self._setup(float) + z_mode = np.correlate(self.x, self.y, 'full') + z_lags = np.correlate(self.x, self.y, lags=range(-2, 5)) + assert_array_almost_equal(z_mode, z_lags) + + def test_lags_matches_same_mode(self): + self._setup(float) + z_mode = np.correlate(self.x, self.y, 'same') + z_lags = np.correlate(self.x, self.y, lags=range(-1, 4)) + assert_array_almost_equal(z_mode, z_lags) + + def test_lags_matches_valid_mode(self): + self._setup(float) + z_mode = np.correlate(self.x, self.y, 'valid') + z_lags = np.correlate(self.x, self.y, lags=range(0, 3)) + assert_array_almost_equal(z_mode, z_lags) + + # --- Type/dtype tests --- + + def test_complex_with_lags(self): + x = np.array([1, 2, 3, 4 + 1j], dtype=complex) + y = np.array([-1, -2j, 3 + 1j], dtype=complex) + # Lags-mode equals subset of full-mode for complex arrays + z_full = np.correlate(y, x, mode='full') + z_lags = np.correlate(y, x, lags=range(-3, 3)) + assert_array_almost_equal(z_full, z_lags) + + def test_object_with_lags(self): + self._setup(Decimal) + z_full = np.correlate(self.x, self.y, mode='full') + z_lags = np.correlate(self.x, self.y, lags=range(-2, 5)) + assert_array_almost_equal(z_full, z_lags) + + # --- Default mode resolution --- + + def test_default_mode_resolution(self): + self._setup(float) + # No mode, no lags -> 'valid' + assert_array_equal(np.correlate(self.x, self.y), + np.correlate(self.x, self.y, mode='valid')) + # No mode, with maxlag -> 'lags' + assert_array_equal(np.correlate(self.x, self.y, maxlag=1), + np.correlate(self.x, self.y, mode='lags', maxlag=1)) + # No mode, with lags -> 'lags' + assert_array_equal( + np.correlate(self.x, self.y, lags=range(-1, 2)), + np.correlate(self.x, self.y, mode='lags', lags=range(-1, 2))) + def test_object(self): self._setup(Decimal) z = np.correlate(self.x, self.y, 'full') @@ -3737,6 +3905,83 @@ def test_convolve_empty_input_error_message(self): with pytest.raises(ValueError, match="v cannot be empty"): np.convolve(np.array([1, 2]), np.array([])) + # --- Convolve lags tests (mirror correlate API) --- + + def test_convolve_maxlag(self): + a = np.array([1, 2, 3], dtype=float) + v = np.array([0, 1, 0.5], dtype=float) + z, lags = np.convolve(a, v, maxlag=1, returns_lagvector=True) + assert_array_almost_equal(z, [1.0, 2.5, 4.0]) + assert_array_equal(lags, np.arange(-1, 2)) + + def test_convolve_lags_range(self): + a = np.array([1, 2, 3], dtype=float) + v = np.array([0, 1, 0.5], dtype=float) + z, lags = np.convolve(a, v, lags=range(-1, 2), + returns_lagvector=True) + assert_array_almost_equal(z, [1.0, 2.5, 4.0]) + assert_array_equal(lags, np.arange(-1, 2)) + + def test_convolve_lags_range_with_step(self): + a = np.array([1, 2, 3, 4, 5], dtype=float) + v = np.array([0, 1, 0.5], dtype=float) + z, lags = np.convolve(a, v, lags=range(-2, 6, 2), + returns_lagvector=True) + assert_array_almost_equal(z, [0.0, 2.5, 5.5, 2.5]) + assert_array_equal(lags, np.arange(-2, 6, 2)) + + def test_convolve_returns_lagvector(self): + a = [1, 2, 3] + v = [0, 1, 0.5] + z, lags = np.convolve(a, v, mode='full', returns_lagvector=True) + assert len(z) == len(lags) + assert_array_equal(lags, np.arange(-2, 3)) + + def test_convolve_lags_matches_full_mode(self): + a = np.array([1, 2, 3, 4, 5], dtype=float) + v = np.array([6, 7, 8], dtype=float) + z_mode = np.convolve(a, v, 'full') + z_lags = np.convolve(a, v, lags=range(-2, 5)) + assert_array_almost_equal(z_mode, z_lags) + + def test_convolve_vs_correlate_reversed(self): + # convolve(a, v, lags=L) == correlate(a, v[::-1], lags=L) + a = np.array([1, 2, 3, 4, 5], dtype=float) + v = np.array([6, 7, 8], dtype=float) + z_conv = np.convolve(a, v, lags=range(-3, 6)) + z_corr = np.correlate(a, v[::-1], lags=range(-3, 6)) + assert_array_almost_equal(z_conv, z_corr) + + def test_convolve_default_mode_resolution(self): + a = [1, 2, 3] + v = [0, 1, 0.5] + # No mode, no lags -> 'full' + assert_array_equal(np.convolve(a, v), np.convolve(a, v, mode='full')) + # No mode, with maxlag -> 'lags' + assert_array_equal(np.convolve(a, v, maxlag=1), + np.convolve(a, v, mode='lags', maxlag=1)) + # No mode, with lags -> 'lags' + assert_array_equal(np.convolve(a, v, lags=range(-1, 2)), + np.convolve(a, v, mode='lags', + lags=range(-1, 2))) + + def test_convolve_maxlag_lags_mutually_exclusive(self): + with assert_raises(TypeError): + np.convolve([1, 2, 3], [4, 5, 6], + maxlag=1, lags=range(-1, 2)) + + def test_convolve_lags_with_non_lags_mode(self): + with assert_raises(ValueError): + np.convolve([1, 2, 3], [4, 5, 6], mode='full', maxlag=1) + with assert_raises(ValueError): + np.convolve([1, 2, 3], [4, 5, 6], mode='full', + lags=range(0, 2)) + + def test_convolve_lags_mode_without_lags(self): + with assert_raises(ValueError): + np.convolve([1, 2, 3], [4, 5, 6], mode='lags') + + class TestArgwhere: @pytest.mark.parametrize('nd', [0, 1, 2]) From 5d06e1ce9b1ef039fdac683fdf15b632f973b3aa Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Wed, 15 Apr 2026 17:35:56 +0000 Subject: [PATCH 04/13] improve benchmarks and follow development guidelines --- benchmarks/benchmarks/bench_core.py | 37 +++++++++++++++++++ .../upcoming_changes/maxlag2025.c_api.rst | 8 ++++ .../maxlag2025.new_feature.rst | 24 ++++++++++++ numpy/_core/numeric.py | 9 +++-- numpy/_core/src/multiarray/multiarraymodule.c | 10 +++-- numpy/_core/tests/test_numeric.py | 19 +++++----- 6 files changed, 91 insertions(+), 16 deletions(-) create mode 100644 doc/release/upcoming_changes/maxlag2025.c_api.rst create mode 100644 doc/release/upcoming_changes/maxlag2025.new_feature.rst diff --git a/benchmarks/benchmarks/bench_core.py b/benchmarks/benchmarks/bench_core.py index ea7aae007fdc..7bc2f6527ab4 100644 --- a/benchmarks/benchmarks/bench_core.py +++ b/benchmarks/benchmarks/bench_core.py @@ -150,6 +150,43 @@ def time_convolve(self, size1, size2, mode): np.convolve(self.x1, self.x2, mode=mode) +class CorrConvLags(Benchmark): + # Benchmarks for the maxlag / lags API on np.correlate / np.convolve. + # Demonstrates the speedup of computing only a window of lags vs. + # computing 'full' and slicing. + params = [[1000, int(1e5)], # size1 (longer signal) + [10, 100, 1000]] # size2 (shorter signal / kernel) + param_names = ['size1', 'size2'] + + def setup(self, size1, size2): + self.x1 = np.linspace(0, 1, num=size1) + self.x2 = np.cos(np.linspace(0, 2 * np.pi, num=size2)) + # A typical "small lag window" use case: just lags around 0. + self.maxlag_small = 5 + # And a moderate window. + self.maxlag_med = max(10, size2 // 2) + + def time_correlate_maxlag_small(self, size1, size2): + np.correlate(self.x1, self.x2, maxlag=self.maxlag_small) + + def time_correlate_maxlag_med(self, size1, size2): + np.correlate(self.x1, self.x2, maxlag=self.maxlag_med) + + def time_correlate_lags_range(self, size1, size2): + np.correlate(self.x1, self.x2, + lags=range(-self.maxlag_small, self.maxlag_small + 1)) + + def time_full_then_slice(self, size1, size2): + # The pre-existing way users had to do this: compute everything, + # throw away most of it. + full = np.correlate(self.x1, self.x2, mode='full') + center = len(full) // 2 + _ = full[center - self.maxlag_small:center + self.maxlag_small + 1] + + def time_convolve_maxlag_small(self, size1, size2): + np.convolve(self.x1, self.x2, maxlag=self.maxlag_small) + + class CountNonzero(Benchmark): param_names = ['numaxes', 'size', 'dtype'] params = [ diff --git a/doc/release/upcoming_changes/maxlag2025.c_api.rst b/doc/release/upcoming_changes/maxlag2025.c_api.rst new file mode 100644 index 000000000000..ad63af0b5b0a --- /dev/null +++ b/doc/release/upcoming_changes/maxlag2025.c_api.rst @@ -0,0 +1,8 @@ +New ``PyArray_CorrelateLags`` C API function +-------------------------------------------- +A new public C API function `PyArray_CorrelateLags` has been added for +computing one-dimensional cross-correlation at a specified lag range +``[minlag, maxlag)`` with a given ``lagstep``. The function accepts arrays +in any order and any valid lag-range orientation; internal swapping and +output reversal are handled transparently so the result is aligned with the +caller's input order and lag direction. diff --git a/doc/release/upcoming_changes/maxlag2025.new_feature.rst b/doc/release/upcoming_changes/maxlag2025.new_feature.rst new file mode 100644 index 000000000000..adb23da091e2 --- /dev/null +++ b/doc/release/upcoming_changes/maxlag2025.new_feature.rst @@ -0,0 +1,24 @@ +``maxlag`` and ``lags`` parameters for `numpy.correlate` and `numpy.convolve` +----------------------------------------------------------------------------- +`numpy.correlate` and `numpy.convolve` now accept ``maxlag`` and ``lags`` +keyword arguments to compute the result at a restricted set of lag values +instead of the full output, plus a ``returns_lagvector`` flag to return the +corresponding lag indices alongside the result. + +* ``maxlag=M`` (int) computes a symmetric inclusive window + ``[-M, M]`` (``2*M+1`` lags), matching MATLAB's ``xcorr(x, y, M)`` + convention. +* ``lags=`` accepts a Python ``range``, a ``slice`` with explicit start/stop, + or a 1-D integer array_like containing an arithmetic progression of lag + indices. +* ``returns_lagvector=True`` causes both functions to return a + ``(result, lagvector)`` tuple where ``lagvector`` enumerates the lags + corresponding to ``result``. + +These parameters are mutually exclusive and require ``mode='lags'`` (which is +the default when either is supplied). + +Example:: + + >>> np.correlate([1, 2, 3], [0, 1, 0.5], maxlag=1, returns_lagvector=True) + (array([2. , 3.5, 3. ]), array([-1, 0, 1])) diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index dcc5e134a174..8d9438f2af3d 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -849,13 +849,14 @@ def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, If ``maxlag`` or ``lags`` is provided and ``mode`` is not specified, ``mode`` defaults to 'lags'. maxlag : int, optional - Compute the cross-correlation at lags ``-maxlag, -maxlag+1, ..., maxlag`` - (a symmetric inclusive window of ``2*maxlag+1`` lags around 0). + Compute the cross-correlation at lags + ``-maxlag, -maxlag+1, ..., maxlag`` (a symmetric inclusive window + of ``2*maxlag+1`` lags around 0). Mutually exclusive with ``lags``. lags : range, slice, or 1-D array_like of int, optional Explicit lag specification. Accepts a Python ``range`` object, a - ``slice`` with explicit start and stop, or a 1-D array_like containing - an arithmetic progression of integer lag indices. + ``slice`` with explicit start and stop, or a 1-D array_like + containing an arithmetic progression of integer lag indices. Mutually exclusive with ``maxlag``. returns_lagvector : bool, optional If True, the function returns a lagvector array in addition to the diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c index 7de431ddc0f2..e3032a3913bd 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.c +++ b/numpy/_core/src/multiarray/multiarraymodule.c @@ -1280,14 +1280,18 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, } else if (lag < phase2_end) { tmplag = lag; - for (lag = tmplag; lag < phase2_end && (!needs_pyapi || !PyErr_Occurred()); lag+=lagstep) { + for (lag = tmplag; + lag < phase2_end && (!needs_pyapi || !PyErr_Occurred()); + lag += lagstep) { dot(ip1 + lag*is1, is1, ip2, is2, op, n2, ret); op += os; } } - maxright = (maxlag < n1 ? maxlag : n1 ); + maxright = (maxlag < n1) ? maxlag : n1; tmplag = lag; - for (lag = tmplag; lag < maxright && (!needs_pyapi || !PyErr_Occurred()); lag+=lagstep) { + for (lag = tmplag; + lag < maxright && (!needs_pyapi || !PyErr_Occurred()); + lag += lagstep) { /* for lags where y is right of x */ n = n1 - lag; dot(ip1 + lag*is1, is1, ip2, is2, op, n, ret); diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py index c7a4b8e4875b..b90acabb3bd2 100644 --- a/numpy/_core/tests/test_numeric.py +++ b/numpy/_core/tests/test_numeric.py @@ -3671,24 +3671,24 @@ def test_lags_array(self): def test_maxlag_and_lags_mutually_exclusive(self): self._setup(float) - with assert_raises(TypeError): + with pytest.raises(TypeError, match="cannot specify both"): np.correlate(self.x, self.y, maxlag=1, lags=range(-1, 2)) def test_lags_with_non_lags_mode(self): self._setup(float) - with assert_raises(ValueError): + with pytest.raises(ValueError, match="cannot be used with mode"): np.correlate(self.x, self.y, mode='full', maxlag=1) - with assert_raises(ValueError): + with pytest.raises(ValueError, match="cannot be used with mode"): np.correlate(self.x, self.y, mode='valid', lags=range(0, 2)) def test_lags_mode_without_lags(self): self._setup(float) - with assert_raises(ValueError): + with pytest.raises(ValueError, match="is required for mode='lags'"): np.correlate(self.x, self.y, mode='lags') def test_lags_non_arithmetic_array_rejected(self): self._setup(float) - with assert_raises(ValueError): + with pytest.raises(ValueError, match="arithmetic progression"): np.correlate(self.x, self.y, lags=np.array([0, 1, 3])) # --- Lag range geometry tests --- @@ -3966,19 +3966,20 @@ def test_convolve_default_mode_resolution(self): lags=range(-1, 2))) def test_convolve_maxlag_lags_mutually_exclusive(self): - with assert_raises(TypeError): + with pytest.raises(TypeError, match="cannot specify both"): np.convolve([1, 2, 3], [4, 5, 6], maxlag=1, lags=range(-1, 2)) def test_convolve_lags_with_non_lags_mode(self): - with assert_raises(ValueError): + with pytest.raises(ValueError, match="cannot be used with mode"): np.convolve([1, 2, 3], [4, 5, 6], mode='full', maxlag=1) - with assert_raises(ValueError): + with pytest.raises(ValueError, match="cannot be used with mode"): np.convolve([1, 2, 3], [4, 5, 6], mode='full', lags=range(0, 2)) def test_convolve_lags_mode_without_lags(self): - with assert_raises(ValueError): + with pytest.raises(ValueError, + match="is required for mode='lags'"): np.convolve([1, 2, 3], [4, 5, 6], mode='lags') From f4860c532a70e51ed066811a0b00ee629ec2887d Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Wed, 15 Apr 2026 19:01:29 +0000 Subject: [PATCH 05/13] comments --- numpy/_core/src/multiarray/multiarraymodule.c | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c index da3e30c1c34a..a29ffb7ce35e 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.c +++ b/numpy/_core/src/multiarray/multiarraymodule.c @@ -1136,7 +1136,8 @@ PyArray_MatrixProduct2(PyObject *op1, PyObject *op2, PyArrayObject* out) static int _pyarray_revert(PyArrayObject *ret); /* - * Core correlation computation. + * Core correlation computation - common between PyArray_Correlate + * and PyArray_Correlate2 and PyArray_CorrelateLags. * * Handles all input forms internally: * - Swaps arrays if n1 < n2 (and negates lags) @@ -1147,8 +1148,7 @@ static int _pyarray_revert(PyArrayObject *ret); * and receive a result aligned with their original array/lag order. */ static PyArrayObject* -_pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, - PyArray_Descr *typec, +_pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, PyArray_Descr *typec, npy_intp minlag, npy_intp maxlag, npy_intp lagstep) { PyArrayObject *ret, *swap; @@ -1252,12 +1252,12 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, op += os*((-n2+1) - lag); lag = -n2+1; } + + /* Phase 1: lags where y is left of x, i.e. lag in [-(n2-1), 0). + * Overlap length is n2 + lag. */ 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); if (needs_pyapi && PyErr_Occurred()) { @@ -1266,9 +1266,10 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, /* iterate over answer vector */ op += os; } - n11 = n1; + /* Phase 2: lags where y entirely overlaps with x, i.e. lag in [0, n1-n2]. * Upper bound is the smaller of maxlag and the full-overlap boundary. */ + n11 = n1; phase2_end = maxlag < (n11 - n2 + 1) ? maxlag : (n11 - n2 + 1); if (lagstep == 1 && lag < phase2_end && small_correlate(ip1 + lag*is1, is1, @@ -1287,12 +1288,14 @@ _pyarray_correlate(PyArrayObject *ap1, PyArrayObject *ap2, op += os; } } + + /* Phase 3: lags where y is right of x, i.e. lag in [n1-n2+1, n1). + * Overlap length is n1 - lag. */ maxright = (maxlag < n1) ? maxlag : n1; tmplag = lag; for (lag = tmplag; lag < maxright && (!needs_pyapi || !PyErr_Occurred()); 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; From e1c3ae25d06011026669b633b8ae6e47479558ce Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Thu, 16 Apr 2026 21:17:49 +0000 Subject: [PATCH 06/13] missing comma --- numpy/_core/multiarray.py | 3 ++- numpy/_core/multiarray.pyi | 3 ++- numpy/_core/tests/test_numeric.py | 14 +++++++------- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/numpy/_core/multiarray.py b/numpy/_core/multiarray.py index cc11802ccbfa..c876a9f45c61 100644 --- a/numpy/_core/multiarray.py +++ b/numpy/_core/multiarray.py @@ -36,7 +36,8 @@ '_monotonicity', 'add_docstring', 'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray', 'asfortranarray', 'bincount', 'broadcast', 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast', - 'compare_chararrays', 'concatenate', 'copyto', 'correlate', 'correlate2', 'correlatelags', + 'compare_chararrays', 'concatenate', 'copyto', 'correlate', 'correlate2', + 'correlatelags', 'count_nonzero', 'c_einsum', 'datetime_as_string', 'datetime_data', 'dot', 'dragon4_positional', 'dragon4_scientific', 'dtype', 'empty', 'empty_like', 'error', 'flagsobj', 'flatiter', 'format_longfloat', diff --git a/numpy/_core/multiarray.pyi b/numpy/_core/multiarray.pyi index 14f2aeb395c2..c75e6adb16ab 100644 --- a/numpy/_core/multiarray.pyi +++ b/numpy/_core/multiarray.pyi @@ -130,7 +130,7 @@ __all__ = [ "copyto", "correlate", "correlate2", - "correlatelags" + "correlatelags", "count_nonzero", "c_einsum", "datetime_as_string", @@ -255,6 +255,7 @@ _place: Final[Callable[..., object]] = ... _reconstruct: Final[Callable[..., object]] = ... _vec_string: Final[Callable[..., object]] = ... correlate2: Final[Callable[..., object]] = ... +correlatelags: Final[Callable[..., object]] = ... dragon4_positional: Final[Callable[..., object]] = ... dragon4_scientific: Final[Callable[..., object]] = ... interp_complex: Final[Callable[..., object]] = ... diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py index e57bee9182d3..da9a10cc2317 100644 --- a/numpy/_core/tests/test_numeric.py +++ b/numpy/_core/tests/test_numeric.py @@ -3635,7 +3635,7 @@ def test_float(self): assert_array_almost_equal(z, self.z1[1:6]) # [-8., -14., -20., -26., -14.] assert_array_equal(lags, np.arange(-1, 4)) assert len(z) == len(lags) - + def test_maxlag(self): self._setup(float) # maxlag=1 gives the symmetric inclusive window [-1, 0, 1] @@ -3687,7 +3687,7 @@ def test_lags_with_non_lags_mode(self): with pytest.raises(ValueError, match="cannot be used with mode"): np.correlate(self.x, self.y, mode='full', maxlag=1) with pytest.raises(ValueError, match="cannot be used with mode"): - np.correlate(self.x, self.y, mode='valid', lags=range(0, 2)) + np.correlate(self.x, self.y, mode='valid', lags=range(2)) def test_lags_mode_without_lags(self): self._setup(float) @@ -3717,7 +3717,7 @@ def test_lags_right_partial_only(self): def test_lags_full_overlap_only(self): self._setup(float) # Only full-overlap lags [0, n1-n2] = [0, 2] - z = np.correlate(self.x, self.y, lags=range(0, 3)) + z = np.correlate(self.x, self.y, lags=range(3)) assert_array_almost_equal(z, self.z1[2:5]) def test_lags_beyond_overlap(self): @@ -3737,7 +3737,7 @@ def test_lags_partial_beyond_overlap(self): def test_lags_single_lag(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, lags=range(0, 1), + z, lags = np.correlate(self.x, self.y, lags=range(1), returns_lagvector=True) assert_array_almost_equal(z, [self.z1[2]]) assert_array_equal(lags, np.array([0])) @@ -3754,7 +3754,7 @@ def test_lags_negative_step(self): def test_lags_empty(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, lags=range(0, 0), + z, lags = np.correlate(self.x, self.y, lags=range(0), returns_lagvector=True) assert len(z) == 0 assert len(lags) == 0 @@ -3798,7 +3798,7 @@ def test_lags_matches_same_mode(self): def test_lags_matches_valid_mode(self): self._setup(float) z_mode = np.correlate(self.x, self.y, 'valid') - z_lags = np.correlate(self.x, self.y, lags=range(0, 3)) + z_lags = np.correlate(self.x, self.y, lags=range(3)) assert_array_almost_equal(z_mode, z_lags) # --- Type/dtype tests --- @@ -3983,7 +3983,7 @@ def test_convolve_lags_with_non_lags_mode(self): np.convolve([1, 2, 3], [4, 5, 6], mode='full', maxlag=1) with pytest.raises(ValueError, match="cannot be used with mode"): np.convolve([1, 2, 3], [4, 5, 6], mode='full', - lags=range(0, 2)) + lags=range(2)) def test_convolve_lags_mode_without_lags(self): with pytest.raises(ValueError, From f3ca620e6625675fa98daa7f8ea218efee7a1dbe Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Fri, 17 Apr 2026 01:03:53 +0000 Subject: [PATCH 07/13] Rename upcoming_changes files to match PR number #31261. Update pythoncapi-compat to match upstream --- .../upcoming_changes/{maxlag2025.c_api.rst => 31261.c_api.rst} | 0 .../{maxlag2025.new_feature.rst => 31261.new_feature.rst} | 0 numpy/_core/src/common/pythoncapi-compat | 2 +- 3 files changed, 1 insertion(+), 1 deletion(-) rename doc/release/upcoming_changes/{maxlag2025.c_api.rst => 31261.c_api.rst} (100%) rename doc/release/upcoming_changes/{maxlag2025.new_feature.rst => 31261.new_feature.rst} (100%) diff --git a/doc/release/upcoming_changes/maxlag2025.c_api.rst b/doc/release/upcoming_changes/31261.c_api.rst similarity index 100% rename from doc/release/upcoming_changes/maxlag2025.c_api.rst rename to doc/release/upcoming_changes/31261.c_api.rst diff --git a/doc/release/upcoming_changes/maxlag2025.new_feature.rst b/doc/release/upcoming_changes/31261.new_feature.rst similarity index 100% rename from doc/release/upcoming_changes/maxlag2025.new_feature.rst rename to doc/release/upcoming_changes/31261.new_feature.rst diff --git a/numpy/_core/src/common/pythoncapi-compat b/numpy/_core/src/common/pythoncapi-compat index 90c06a4cae55..8636bccf29ad 160000 --- a/numpy/_core/src/common/pythoncapi-compat +++ b/numpy/_core/src/common/pythoncapi-compat @@ -1 +1 @@ -Subproject commit 90c06a4cae557bdbfa4f231a781d2b5c1a8f6d1c +Subproject commit 8636bccf29adfa23463f810b3c2830f7cff1e933 From 2a67bbe88a3743ebe28dce4a6099a76dea6d3c1f Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Fri, 17 Apr 2026 01:05:21 +0000 Subject: [PATCH 08/13] update submodules --- numpy/_core/src/npysort/x86-simd-sort | 2 +- vendored-meson/meson | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/numpy/_core/src/npysort/x86-simd-sort b/numpy/_core/src/npysort/x86-simd-sort index 6a7a01da4b0d..5adb33411f3c 160000 --- a/numpy/_core/src/npysort/x86-simd-sort +++ b/numpy/_core/src/npysort/x86-simd-sort @@ -1 +1 @@ -Subproject commit 6a7a01da4b0dfde108aa626a2364c954e2c50fe1 +Subproject commit 5adb33411f3cea8bdbafa9d91bd75bc4bf19c7dd diff --git a/vendored-meson/meson b/vendored-meson/meson index e72c717199fa..5d5a3d478da1 160000 --- a/vendored-meson/meson +++ b/vendored-meson/meson @@ -1 +1 @@ -Subproject commit e72c717199fa18d34020c7c97f9de3f388c5e055 +Subproject commit 5d5a3d478da115c812be77afa651db2492d52171 From 6447640bf64b1253c3555cdf3272fd41b4b13167 Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Fri, 17 Apr 2026 02:01:34 +0000 Subject: [PATCH 09/13] add flag for conjugation so that convolve doesn't use conjugation --- numpy/_core/numeric.py | 5 +- numpy/_core/src/multiarray/multiarraymodule.c | 49 ++++++++++++++----- numpy/_core/tests/test_numeric.py | 16 ++++++ 3 files changed, 55 insertions(+), 15 deletions(-) diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index 556857f836dd..b2bf3f753e97 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -1121,7 +1121,7 @@ def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, raise ValueError( "maxlag/lags cannot be used with mode " "'valid', 'same', or 'full'") - result = multiarray.correlate2(a, v[::-1], mode) + result = multiarray.correlate(a, v[::-1], mode) if returns_lagvector: lags_tuple = _lags_from_mode(alen, vlen, mode) elif mode == 3: @@ -1133,7 +1133,8 @@ def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, else: lags_tuple = _lags_from_lags(lags) result = multiarray.correlatelags( - a, v[::-1], lags_tuple[0], lags_tuple[1], lags_tuple[2]) + a, v[::-1], lags_tuple[0], lags_tuple[1], lags_tuple[2], + conjugate=False) if returns_lagvector: return result, arange(lags_tuple[0], lags_tuple[1], lags_tuple[2]) diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c index a29ffb7ce35e..1799fdaabc7b 100644 --- a/numpy/_core/src/multiarray/multiarraymodule.c +++ b/numpy/_core/src/multiarray/multiarraymodule.c @@ -1428,8 +1428,13 @@ _lags_from_mode(int mode, npy_intp n1, npy_intp n2, /*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 + * Computes the cross-correlation of a1 with a2 using a fixed output mode + * ('valid', 'same', or 'full'). For complex inputs, the second array is + * conjugated before the computation (i.e. correlation semantics). + * + * This is the mode-based counterpart of PyArray_CorrelateLags with + * conjugate != 0. For convolution (no conjugation), use + * PyArray_Correlate instead. */ NPY_NO_EXPORT PyObject * PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) @@ -1503,17 +1508,28 @@ PyArray_Correlate2(PyObject *op1, PyObject *op2, int mode) } /*NUMPY_API - * correlate(a1,a2,minlag,maxlag,lagstep) + * correlate(a1,a2,minlag,maxlag,lagstep,conjugate) + * + * Computes the sliding inner product of a1 with a2 at the specified lag + * range [minlag, maxlag) with step lagstep. + * This function accepts vectors in either order. This function + * does not accept modes; use _lags_from_mode() to convert a mode integer to + * (minlag, maxlag, lagstep). + * + * The conjugate flag selects the operation semantics: + * + * conjugate != 0 -- correlation semantics (same as PyArray_Correlate2): + * out[k] = sum_n a1[n+k] * conj(a2[n]) * - * This function computes the cross-correlation of a1 with a2 at the - * specified lag range. The arrays and lag parameters may be provided in - * any valid form; normalization is handled internally. - * This function does not accept modes. - * See _lags_from_mode() to convert modes to lags. + * conjugate == 0 -- no conjugation (same as PyArray_Correlate): + * out[k] = sum_n a1[n+k] * a2[n] + * For convolution the caller must pass a2 already + * reversed (a2[::-1]) before calling this function. */ NPY_NO_EXPORT PyObject * PyArray_CorrelateLags(PyObject *op1, PyObject *op2, - npy_intp minlag, npy_intp maxlag, npy_intp lagstep) + npy_intp minlag, npy_intp maxlag, npy_intp lagstep, + int conjugate) { PyArrayObject *ap1, *ap2, *ret = NULL; PyArray_Descr *typec = NULL; @@ -1546,7 +1562,7 @@ PyArray_CorrelateLags(PyObject *op1, PyObject *op2, goto clean_ap1; } - if (PyArray_ISCOMPLEX(ap2)) { + if (conjugate && PyArray_ISCOMPLEX(ap2)) { PyArrayObject *cap2; cap2 = (PyArrayObject *)PyArray_Conjugate(ap2, NULL); if (cap2 == NULL) { @@ -1577,6 +1593,11 @@ PyArray_CorrelateLags(PyObject *op1, PyObject *op2, /*NUMPY_API * Numeric.correlate(a1,a2,mode) + * + * Legacy correlation function that does NOT conjugate the second array for + * complex inputs. Equivalent to PyArray_CorrelateLags with conjugate == 0 + * (convolution semantics) after the appropriate lag range is computed via + * _lags_from_mode(). Prefer PyArray_Correlate2 for standard correlation. */ NPY_NO_EXPORT PyObject * PyArray_Correlate(PyObject *op1, PyObject *op2, int mode) @@ -3265,6 +3286,7 @@ array_correlatelags(PyObject *NPY_UNUSED(dummy), { PyObject *shape, *a0; npy_intp minlag = 0, maxlag = 0, lagstep = 0; + npy_bool conjugate = 1; NPY_PREPARE_ARGPARSER; if (npy_parse_arguments("correlate_lags", args, len_args, kwnames, @@ -3272,14 +3294,15 @@ array_correlatelags(PyObject *NPY_UNUSED(dummy), {"v", NULL, &shape}, {"minlag", &PyArray_IntpFromPyIntConverter, &minlag}, {"maxlag", &PyArray_IntpFromPyIntConverter, &maxlag}, - {"lagstep", &PyArray_IntpFromPyIntConverter, &lagstep}) < 0) { + {"lagstep", &PyArray_IntpFromPyIntConverter, &lagstep}, + {"|conjugate", &PyArray_BoolConverter, &conjugate}) < 0) { return NULL; } if (minlag == 0 && maxlag == 0 && lagstep == 0) { /* if no lag parameters passed, use default: mode = 'valid' */ return PyArray_Correlate2(a0, shape, 0); - } - return PyArray_CorrelateLags(a0, shape, minlag, maxlag, lagstep); + } + return PyArray_CorrelateLags(a0, shape, minlag, maxlag, lagstep, conjugate); } static PyObject * diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py index da9a10cc2317..9d5f79689d7f 100644 --- a/numpy/_core/tests/test_numeric.py +++ b/numpy/_core/tests/test_numeric.py @@ -3990,6 +3990,22 @@ def test_convolve_lags_mode_without_lags(self): match="is required for mode='lags'"): np.convolve([1, 2, 3], [4, 5, 6], mode='lags') + def test_complex_no_conjugation(self): + # Regression test: convolve must NOT conjugate either input. + # correlate(a, v) conjugates v; convolve(a, v) must not. + # For real inputs conjugation is a no-op, so only complex exposes + # the bug. Expected values computed as scalar polynomial product: + # (1+1j)*x^1 + 2*x^0 times 3*x^1 + (4+1j)*x^0 + a = np.array([1 + 1j, 2]) + v = np.array([3, 4 + 1j]) + expected = np.array([3 + 3j, 9 + 5j, 8 + 2j]) + assert_array_equal(np.convolve(a, v), expected) + # Same check via the lags path (maxlag=1 covers lags [-1, 0, 1], + # which is the full output for two length-2 inputs) + result, lags = np.convolve(a, v, maxlag=1, returns_lagvector=True) + assert_array_equal(result, expected) + assert_array_equal(lags, np.arange(-1, 2)) + class TestArgwhere: From 2ef7d331a72d90c2340468a150c2026ab3ccb6c3 Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Sun, 19 Apr 2026 01:14:16 +0000 Subject: [PATCH 10/13] fix stubs --- numpy/_core/numeric.pyi | 74 ++++++++++++++++++++++++++++++++++------- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/numpy/_core/numeric.pyi b/numpy/_core/numeric.pyi index 1dd34d6a4fcd..9c448e015107 100644 --- a/numpy/_core/numeric.pyi +++ b/numpy/_core/numeric.pyi @@ -1048,45 +1048,95 @@ def flatnonzero(a: ArrayLike) -> _Array1D[np.intp]: ... # NOTE: we ignore UP047 because inlining `_AnyScalarT` would result in a lot of code duplication +type _LagsArg = int | tuple[int, int] | tuple[int, int, int] | range | slice | ArrayLike | None + # keep in sync with `convolve` and `ma.core.correlate` @overload def correlate( # noqa: UP047 - a: _ArrayLike1D[_AnyNumericScalarT], v: _ArrayLike1D[_AnyNumericScalarT], mode: _CorrelateMode = "valid" + a: _ArrayLike1D[_AnyNumericScalarT], v: _ArrayLike1D[_AnyNumericScalarT], + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., ) -> _Array1D[_AnyNumericScalarT]: ... @overload -def correlate(a: _ArrayLike1DBool_co, v: _ArrayLike1DBool_co, mode: _CorrelateMode = "valid") -> _Array1D[np.bool]: ... +def correlate( + a: _ArrayLike1DBool_co, v: _ArrayLike1DBool_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., +) -> _Array1D[np.bool]: ... @overload -def correlate(a: _ArrayLike1DInt_co, v: _ArrayLike1DInt_co, mode: _CorrelateMode = "valid") -> _Array1D[np.int_ | Any]: ... +def correlate( + a: _ArrayLike1DInt_co, v: _ArrayLike1DInt_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., +) -> _Array1D[np.int_ | Any]: ... @overload -def correlate(a: _ArrayLike1DFloat_co, v: _ArrayLike1DFloat_co, mode: _CorrelateMode = "valid") -> _Array1D[np.float64 | Any]: ... +def correlate( + a: _ArrayLike1DFloat_co, v: _ArrayLike1DFloat_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., +) -> _Array1D[np.float64 | Any]: ... @overload def correlate( - a: _ArrayLike1DNumber_co, v: _ArrayLike1DNumber_co, mode: _CorrelateMode = "valid" + a: _ArrayLike1DNumber_co, v: _ArrayLike1DNumber_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., ) -> _Array1D[np.complex128 | Any]: ... @overload def correlate( - a: _ArrayLike1DTD64_co, v: _ArrayLike1DTD64_co, mode: _CorrelateMode = "valid" + a: _ArrayLike1DTD64_co, v: _ArrayLike1DTD64_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., ) -> _Array1D[np.timedelta64 | Any]: ... +@overload +def correlate( + a: ArrayLike, v: ArrayLike, + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[True], +) -> tuple[np.ndarray[Any, Any], _Array1D[np.intp]]: ... # keep in sync with `correlate` @overload def convolve( # noqa: UP047 - a: _ArrayLike1D[_AnyNumericScalarT], v: _ArrayLike1D[_AnyNumericScalarT], mode: _CorrelateMode = "valid" + a: _ArrayLike1D[_AnyNumericScalarT], v: _ArrayLike1D[_AnyNumericScalarT], + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., ) -> _Array1D[_AnyNumericScalarT]: ... @overload -def convolve(a: _ArrayLike1DBool_co, v: _ArrayLike1DBool_co, mode: _CorrelateMode = "valid") -> _Array1D[np.bool]: ... +def convolve( + a: _ArrayLike1DBool_co, v: _ArrayLike1DBool_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., +) -> _Array1D[np.bool]: ... @overload -def convolve(a: _ArrayLike1DInt_co, v: _ArrayLike1DInt_co, mode: _CorrelateMode = "valid") -> _Array1D[np.int_ | Any]: ... +def convolve( + a: _ArrayLike1DInt_co, v: _ArrayLike1DInt_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., +) -> _Array1D[np.int_ | Any]: ... @overload -def convolve(a: _ArrayLike1DFloat_co, v: _ArrayLike1DFloat_co, mode: _CorrelateMode = "valid") -> _Array1D[np.float64 | Any]: ... +def convolve( + a: _ArrayLike1DFloat_co, v: _ArrayLike1DFloat_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., +) -> _Array1D[np.float64 | Any]: ... @overload def convolve( - a: _ArrayLike1DNumber_co, v: _ArrayLike1DNumber_co, mode: _CorrelateMode = "valid" + a: _ArrayLike1DNumber_co, v: _ArrayLike1DNumber_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., ) -> _Array1D[np.complex128 | Any]: ... @overload def convolve( - a: _ArrayLike1DTD64_co, v: _ArrayLike1DTD64_co, mode: _CorrelateMode = "valid" + a: _ArrayLike1DTD64_co, v: _ArrayLike1DTD64_co, + mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[False] = ..., ) -> _Array1D[np.timedelta64 | Any]: ... +@overload +def convolve( + a: ArrayLike, v: ArrayLike, + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., + returns_lagvector: L[True], +) -> tuple[np.ndarray[Any, Any], _Array1D[np.intp]]: ... # keep roughly in sync with `convolve` and `correlate`, but for 2-D output and an additional `out` overload, # and also keep in sync with `ma.core.outer` (minus `out`) From 5d9723fe559daf80b5cb1cdb3e07326bbc4b93f5 Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Mon, 20 Apr 2026 05:51:26 +0000 Subject: [PATCH 11/13] remove returns_lagvector and replace with correlate_lags() --- numpy/__init__.py | 1 + numpy/__init__.pyi | 3 +- numpy/_core/__init__.pyi | 2 + numpy/_core/numeric.py | 191 +++++++++++++++++++----------- numpy/_core/numeric.pyi | 67 +++++------ numpy/_core/tests/test_numeric.py | 63 +++++----- numpy/matlib.pyi | 1 + 7 files changed, 195 insertions(+), 133 deletions(-) diff --git a/numpy/__init__.py b/numpy/__init__.py index ce0452ffe8d0..ee1514b14b5e 100644 --- a/numpy/__init__.py +++ b/numpy/__init__.py @@ -203,6 +203,7 @@ copysign, copyto, correlate, + correlate_lags, cos, cosh, count_nonzero, diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index 4875e64dccc1..678fbfbe6712 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -353,6 +353,7 @@ from numpy._core.numeric import ( argwhere, flatnonzero, correlate, + correlate_lags, convolve, outer, tensordot, @@ -610,7 +611,7 @@ __all__ = [ "count_nonzero", "empty", "broadcast", "dtype", "fromstring", "fromfile", "frombuffer", "from_dlpack", "where", "argwhere", "copyto", "concatenate", "lexsort", "astype", "can_cast", "promote_types", "min_scalar_type", "result_type", - "isfortran", "empty_like", "zeros_like", "ones_like", "correlate", "convolve", + "isfortran", "empty_like", "zeros_like", "ones_like", "correlate", "correlate_lags", "convolve", "inner", "dot", "outer", "vdot", "roll", "rollaxis", "moveaxis", "cross", "tensordot", "little_endian", "fromiter", "array_equal", "array_equiv", "indices", "fromfunction", "isclose", "isscalar", "binary_repr", "base_repr", "ones", diff --git a/numpy/_core/__init__.pyi b/numpy/_core/__init__.pyi index ce5427bbfcd9..5ad001ff44c9 100644 --- a/numpy/_core/__init__.pyi +++ b/numpy/_core/__init__.pyi @@ -95,6 +95,7 @@ from .numeric import ( convolve, copyto, correlate, + correlate_lags, count_nonzero, cross, dot, @@ -424,6 +425,7 @@ __all__ = [ "copysign", "copyto", "correlate", + "correlate_lags", "cos", "cosh", "count_nonzero", diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index b2bf3f753e97..ec16367e363a 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -78,7 +78,7 @@ 'argwhere', 'copyto', 'concatenate', 'lexsort', 'astype', 'can_cast', 'promote_types', 'min_scalar_type', 'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like', - 'correlate', 'convolve', 'inner', 'dot', 'outer', 'vdot', 'roll', + 'correlate', 'correlate_lags', 'convolve', 'inner', 'dot', 'outer', 'vdot', 'roll', 'rollaxis', 'moveaxis', 'cross', 'tensordot', 'little_endian', 'fromiter', 'array_equal', 'array_equiv', 'indices', 'fromfunction', 'isclose', 'isscalar', 'binary_repr', 'base_repr', 'ones', @@ -815,14 +815,95 @@ def _lags_from_mode(alen, vlen, mode): return (m0, m1, 1) -def _correlate_dispatcher(a, v, mode=None, *, maxlag=None, lags=None, - returns_lagvector=None): +def _correlate_lags_dispatcher(a_len, v_len, mode=None, *, maxlag=None, + lags=None): + return () + + +@array_function_dispatch(_correlate_lags_dispatcher) +def correlate_lags(a_len, v_len, mode=_CorrModeDefault, *, maxlag=None, + lags=None): + """ + Return the lag indices for a `~numpy.correlate` call with the same + parameters. + + Parameters + ---------- + a_len : int + Length of the first input sequence (``len(a)`` in ``correlate(a, v, ...)``). + v_len : int + Length of the second input sequence (``len(v)`` in ``correlate(a, v, ...)``). + mode : {'valid', 'same', 'full'}, optional + The same mode that would be passed to `~numpy.correlate`. + If ``maxlag`` or ``lags`` is provided and ``mode`` is not specified, + ``mode`` defaults to 'lags'. + maxlag : int, optional + The same ``maxlag`` that would be passed to `~numpy.correlate`. + Mutually exclusive with ``lags``. + lags : range, slice, or 1-D array_like of int, optional + The same ``lags`` that would be passed to `~numpy.correlate`. + Mutually exclusive with ``maxlag``. + + Returns + ------- + lags : ndarray + Array of lag indices corresponding element-by-element to the output + of ``np.correlate(a, v, mode, maxlag=maxlag, lags=lags)``. + + See Also + -------- + correlate : Cross-correlation of two 1-D sequences. + + Examples + -------- + >>> import numpy as np + >>> np.correlate_lags(5, 3, mode='full') + array([-2, -1, 0, 1, 2, 3, 4]) + >>> np.correlate_lags(5, 3, mode='valid') + array([0, 1, 2]) + >>> np.correlate_lags(5, 3, maxlag=2) + array([-2, -1, 0, 1, 2]) + + Pair the lag vector with a correlation result: + + >>> a, v = [1, 2, 3, 4, 5], [0, 1, 0.5] + >>> r = np.correlate(a, v, mode='full') + >>> lag_vec = np.correlate_lags(len(a), len(v), mode='full') + >>> r[lag_vec == 0] + array([2.]) + """ + if maxlag is not None and lags is not None: + raise TypeError("cannot specify both maxlag and lags") + lags_given = maxlag is not None or lags is not None + + if mode is _CorrModeDefault: + mode = 'lags' if lags_given else 'valid' + mode = _mode_from_name(mode) + + if mode in (0, 1, 2): + if lags_given: + raise ValueError( + "maxlag/lags cannot be used with mode " + "'valid', 'same', or 'full'") + lags_tuple = _lags_from_mode(a_len, v_len, mode) + elif mode == 3: + if not lags_given: + raise ValueError( + "maxlag or lags is required for mode='lags'") + if maxlag is not None: + lags_tuple = _lags_from_maxlag(maxlag) + else: + lags_tuple = _lags_from_lags(lags) + + return arange(lags_tuple[0], lags_tuple[1], lags_tuple[2]) + + +def _correlate_dispatcher(a, v, mode=None, *, maxlag=None, lags=None): return (a, v) @array_function_dispatch(_correlate_dispatcher) -def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, - returns_lagvector=False): +def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None): r""" Cross-correlation of two 1-dimensional sequences. @@ -853,23 +934,15 @@ def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, ``slice`` with explicit start and stop, or a 1-D array_like containing an arithmetic progression of integer lag indices. Mutually exclusive with ``maxlag``. - 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. Returns ------- 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 -------- + correlate_lags : Return the lag indices corresponding to a `correlate` call. convolve : Discrete, linear convolution of two one-dimensional sequences. scipy.signal.correlate : uses FFT which has superior performance on large arrays. @@ -898,13 +971,18 @@ def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, array([3.5]) >>> 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_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], maxlag=1) array([ 2. , 3.5, 3. ]) - >>> np.correlate([1, 2, 3], [0, 1, 0.5], lags=range(-1, 2, 2), - ... returns_lagvector=True) - (array([ 2., 3.]), array([-1, 1])) + >>> np.correlate([1, 2, 3], [0, 1, 0.5], lags=range(-1, 2, 2)) + array([ 2., 3.]) + + Pair the result with its lag indices using `correlate_lags`: + + >>> a, v = [1, 2, 3], [0, 1, 0.5] + >>> r = np.correlate(a, v, mode="full") + >>> lag_vec = np.correlate_lags(len(a), len(v), mode="full") + >>> lag_vec + array([-2, -1, 0, 1, 2]) Using complex sequences: @@ -933,9 +1011,7 @@ def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, raise ValueError( "maxlag/lags cannot be used with mode " "'valid', 'same', or 'full'") - result = multiarray.correlate2(a, v, mode) - if returns_lagvector: - lags_tuple = _lags_from_mode(len(a), len(v), mode) + return multiarray.correlate2(a, v, mode) elif mode == 3: if not lags_given: raise ValueError( @@ -944,22 +1020,16 @@ def correlate(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, lags_tuple = _lags_from_maxlag(maxlag) else: lags_tuple = _lags_from_lags(lags) - result = multiarray.correlatelags( + return multiarray.correlatelags( a, v, lags_tuple[0], lags_tuple[1], lags_tuple[2]) - if returns_lagvector: - return result, arange(lags_tuple[0], lags_tuple[1], lags_tuple[2]) - return result - -def _convolve_dispatcher(a, v, mode=None, *, maxlag=None, lags=None, - returns_lagvector=None): +def _convolve_dispatcher(a, v, mode=None, *, maxlag=None, lags=None): return (a, v) @array_function_dispatch(_convolve_dispatcher) -def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, - returns_lagvector=False): +def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None): """ Returns the discrete, linear convolution of two one-dimensional sequences. @@ -1013,23 +1083,14 @@ def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, ``slice`` with explicit start and stop, or a 1-D array_like containing an arithmetic progression of integer lag indices. Mutually exclusive with ``maxlag``. - 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 -------- + correlate_lags : Return the lag indices corresponding to a `correlate` call. scipy.signal.fftconvolve : Convolve two arrays using the Fast Fourier Transform. scipy.linalg.toeplitz : Used to construct the convolution operator. @@ -1070,31 +1131,31 @@ def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, >>> 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, - 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: + The two arrays are of the same length, so there is only one position + where they completely overlap, corresponding to a lag of 0: - >>> 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], mode='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): + (lag 0 aligns the left sides of the arrays; negative lags shift the + second array left, positive lags shift it right): + + >>> np.convolve([1,2,3],[0,1,0.5], maxlag=1) + array([ 1. , 2.5, 4. ]) - >>> np.convolve([1,2,3],[0,1,0.5], maxlag=1, 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 2: - Find the convolution for lags ranging from -2 to 4 - with steps of length 2: + >>> np.convolve([1,2,3,4,5], [0,1,0.5], lags=range(-2, 6, 2)) + array([ 0. , 2.5, 5.5, 2.5]) - >>> np.convolve([1,2,3,4,5], [0,1,0.5], lags=range(-2, 6, 2), - ... returns_lagvector=True) - (array([ 0. , 2.5, 5.5, 2.5]), array([-2, 0, 2, 4])) + Pair the result with lag indices using `correlate_lags`: + + >>> a, v = [1, 2, 3], [0, 1, 0.5] + >>> r = np.convolve(a, v, mode='valid') + >>> lag_vec = np.correlate_lags(len(a), len(v), mode='valid') + >>> lag_vec + array([0]) """ a, v = array(a, copy=None, ndmin=1), array(v, copy=None, ndmin=1) @@ -1121,9 +1182,7 @@ def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, raise ValueError( "maxlag/lags cannot be used with mode " "'valid', 'same', or 'full'") - result = multiarray.correlate(a, v[::-1], mode) - if returns_lagvector: - lags_tuple = _lags_from_mode(alen, vlen, mode) + return multiarray.correlate(a, v[::-1], mode) elif mode == 3: if not lags_given: raise ValueError( @@ -1132,14 +1191,10 @@ def convolve(a, v, mode=_CorrModeDefault, *, maxlag=None, lags=None, lags_tuple = _lags_from_maxlag(maxlag) else: lags_tuple = _lags_from_lags(lags) - result = multiarray.correlatelags( + return multiarray.correlatelags( a, v[::-1], lags_tuple[0], lags_tuple[1], lags_tuple[2], conjugate=False) - if returns_lagvector: - return result, arange(lags_tuple[0], lags_tuple[1], lags_tuple[2]) - return result - def _outer_dispatcher(a, b, out=None): return (a, b, out) diff --git a/numpy/_core/numeric.pyi b/numpy/_core/numeric.pyi index 9c448e015107..ef7fd2671881 100644 --- a/numpy/_core/numeric.pyi +++ b/numpy/_core/numeric.pyi @@ -403,6 +403,7 @@ __all__ = [ "copysign", "copyto", "correlate", + "correlate_lags", "cos", "cosh", "count_nonzero", @@ -1054,89 +1055,83 @@ type _LagsArg = int | tuple[int, int] | tuple[int, int, int] | range | slice | A @overload def correlate( # noqa: UP047 a: _ArrayLike1D[_AnyNumericScalarT], v: _ArrayLike1D[_AnyNumericScalarT], - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[_AnyNumericScalarT]: ... @overload def correlate( a: _ArrayLike1DBool_co, v: _ArrayLike1DBool_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.bool]: ... @overload def correlate( a: _ArrayLike1DInt_co, v: _ArrayLike1DInt_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.int_ | Any]: ... @overload def correlate( a: _ArrayLike1DFloat_co, v: _ArrayLike1DFloat_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.float64 | Any]: ... @overload def correlate( a: _ArrayLike1DNumber_co, v: _ArrayLike1DNumber_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.complex128 | Any]: ... @overload def correlate( a: _ArrayLike1DTD64_co, v: _ArrayLike1DTD64_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.timedelta64 | Any]: ... -@overload -def correlate( - a: ArrayLike, v: ArrayLike, - mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[True], -) -> tuple[np.ndarray[Any, Any], _Array1D[np.intp]]: ... + +def correlate_lags( + a_len: SupportsIndex, v_len: SupportsIndex, + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., +) -> _Array1D[np.intp]: ... # keep in sync with `correlate` @overload def convolve( # noqa: UP047 a: _ArrayLike1D[_AnyNumericScalarT], v: _ArrayLike1D[_AnyNumericScalarT], - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[_AnyNumericScalarT]: ... @overload def convolve( a: _ArrayLike1DBool_co, v: _ArrayLike1DBool_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.bool]: ... @overload def convolve( a: _ArrayLike1DInt_co, v: _ArrayLike1DInt_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.int_ | Any]: ... @overload def convolve( a: _ArrayLike1DFloat_co, v: _ArrayLike1DFloat_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.float64 | Any]: ... @overload def convolve( a: _ArrayLike1DNumber_co, v: _ArrayLike1DNumber_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.complex128 | Any]: ... @overload def convolve( a: _ArrayLike1DTD64_co, v: _ArrayLike1DTD64_co, - mode: _CorrelateMode = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[False] = ..., + mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., + lags: _LagsArg = ..., ) -> _Array1D[np.timedelta64 | Any]: ... -@overload -def convolve( - a: ArrayLike, v: ArrayLike, - mode: _CorrelateMode | L["lags"] = ..., *, maxlag: int | None = ..., lags: _LagsArg = ..., - returns_lagvector: L[True], -) -> tuple[np.ndarray[Any, Any], _Array1D[np.intp]]: ... # keep roughly in sync with `convolve` and `correlate`, but for 2-D output and an additional `out` overload, # and also keep in sync with `ma.core.outer` (minus `out`) diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py index 9d5f79689d7f..cfbe57d37c4d 100644 --- a/numpy/_core/tests/test_numeric.py +++ b/numpy/_core/tests/test_numeric.py @@ -3611,7 +3611,8 @@ def _setup(self, dt): def test_float(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, 'full', returns_lagvector=True) + z = np.correlate(self.x, self.y, 'full') + lags = np.correlate_lags(len(self.x), len(self.y), 'full') assert_array_almost_equal(z, self.z1) assert_array_equal(lags, np.arange(-2, 5)) assert len(z) == len(lags) @@ -3626,12 +3627,14 @@ def test_float(self): z = np.correlate(self.xs, self.y, 'full') assert_array_almost_equal(z, self.zs) # Test 'valid' mode - middle part where there's full overlap - z, lags = np.correlate(self.x, self.y, 'valid', returns_lagvector=True) + z = np.correlate(self.x, self.y, 'valid') + lags = np.correlate_lags(len(self.x), len(self.y), 'valid') assert_array_almost_equal(z, self.z1[2:5]) # [-14., -20., -26.] assert_array_equal(lags, np.arange(0, 3)) assert len(z) == len(lags) # Test 'same' mode - same length as first input - z, lags = np.correlate(self.x, self.y, 'same', returns_lagvector=True) + z = np.correlate(self.x, self.y, 'same') + lags = np.correlate_lags(len(self.x), len(self.y), 'same') assert_array_almost_equal(z, self.z1[1:6]) # [-8., -14., -20., -26., -14.] assert_array_equal(lags, np.arange(-1, 4)) assert len(z) == len(lags) @@ -3639,41 +3642,41 @@ def test_float(self): def test_maxlag(self): self._setup(float) # maxlag=1 gives the symmetric inclusive window [-1, 0, 1] - z, lags = np.correlate(self.x, self.y, maxlag=1, - returns_lagvector=True) + z = np.correlate(self.x, self.y, maxlag=1) + lags = np.correlate_lags(len(self.x), len(self.y), maxlag=1) assert_array_almost_equal(z, self.z1[1:4]) # [-8., -14., -20.] assert_array_equal(lags, np.arange(-1, 2)) assert len(z) == len(lags) def test_lags_range(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, lags=range(-1, 3), - returns_lagvector=True) + z = np.correlate(self.x, self.y, lags=range(-1, 3)) + lags = np.correlate_lags(len(self.x), len(self.y), lags=range(-1, 3)) assert_array_almost_equal(z, self.z1[1:5]) # [-8., -14., -20., -26.] assert_array_equal(lags, np.arange(-1, 3)) assert len(z) == len(lags) def test_lags_range_with_step(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, lags=range(-2, 5, 2), - returns_lagvector=True) + z = np.correlate(self.x, self.y, lags=range(-2, 5, 2)) + lags = np.correlate_lags(len(self.x), len(self.y), lags=range(-2, 5, 2)) assert_array_almost_equal(z, self.z1[0::2]) # [-3., -14., -26., -5.] assert_array_equal(lags, np.arange(-2, 5, 2)) assert len(z) == len(lags) def test_lags_slice(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, lags=slice(-1, 3), - returns_lagvector=True) + z = np.correlate(self.x, self.y, lags=slice(-1, 3)) + lags = np.correlate_lags(len(self.x), len(self.y), lags=slice(-1, 3)) assert_array_almost_equal(z, self.z1[1:5]) assert_array_equal(lags, np.arange(-1, 3)) def test_lags_array(self): self._setup(float) # 1-D array of lags (must be an arithmetic progression) - z, lags = np.correlate(self.x, self.y, - lags=np.array([-2, 0, 2, 4]), - returns_lagvector=True) + z = np.correlate(self.x, self.y, lags=np.array([-2, 0, 2, 4])) + lags = np.correlate_lags( + len(self.x), len(self.y), lags=np.array([-2, 0, 2, 4])) assert_array_almost_equal(z, self.z1[0::2]) assert_array_equal(lags, np.array([-2, 0, 2, 4])) @@ -3737,16 +3740,17 @@ def test_lags_partial_beyond_overlap(self): def test_lags_single_lag(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, lags=range(1), - returns_lagvector=True) + z = np.correlate(self.x, self.y, lags=range(1)) + lags = np.correlate_lags(len(self.x), len(self.y), lags=range(1)) assert_array_almost_equal(z, [self.z1[2]]) assert_array_equal(lags, np.array([0])) def test_lags_negative_step(self): self._setup(float) # Negative step: reversed lag order - z, lags = np.correlate(self.x, self.y, lags=slice(2, -3, -1), - returns_lagvector=True) + z = np.correlate(self.x, self.y, lags=slice(2, -3, -1)) + lags = np.correlate_lags( + len(self.x), len(self.y), lags=slice(2, -3, -1)) # slice(2,-3,-1) -> [2, 1, 0, -1, -2] # corresponds to z1 indices [4, 3, 2, 1, 0] assert_array_almost_equal(z, self.z1[4::-1]) @@ -3754,8 +3758,8 @@ def test_lags_negative_step(self): def test_lags_empty(self): self._setup(float) - z, lags = np.correlate(self.x, self.y, lags=range(0), - returns_lagvector=True) + z = np.correlate(self.x, self.y, lags=range(0)) + lags = np.correlate_lags(len(self.x), len(self.y), lags=range(0)) assert len(z) == 0 assert len(lags) == 0 @@ -3918,30 +3922,32 @@ def test_convolve_empty_input_error_message(self): def test_convolve_maxlag(self): a = np.array([1, 2, 3], dtype=float) v = np.array([0, 1, 0.5], dtype=float) - z, lags = np.convolve(a, v, maxlag=1, returns_lagvector=True) + z = np.convolve(a, v, maxlag=1) + lags = np.correlate_lags(len(a), len(v), maxlag=1) assert_array_almost_equal(z, [1.0, 2.5, 4.0]) assert_array_equal(lags, np.arange(-1, 2)) def test_convolve_lags_range(self): a = np.array([1, 2, 3], dtype=float) v = np.array([0, 1, 0.5], dtype=float) - z, lags = np.convolve(a, v, lags=range(-1, 2), - returns_lagvector=True) + z = np.convolve(a, v, lags=range(-1, 2)) + lags = np.correlate_lags(len(a), len(v), lags=range(-1, 2)) assert_array_almost_equal(z, [1.0, 2.5, 4.0]) assert_array_equal(lags, np.arange(-1, 2)) def test_convolve_lags_range_with_step(self): a = np.array([1, 2, 3, 4, 5], dtype=float) v = np.array([0, 1, 0.5], dtype=float) - z, lags = np.convolve(a, v, lags=range(-2, 6, 2), - returns_lagvector=True) + z = np.convolve(a, v, lags=range(-2, 6, 2)) + lags = np.correlate_lags(len(a), len(v), lags=range(-2, 6, 2)) assert_array_almost_equal(z, [0.0, 2.5, 5.5, 2.5]) assert_array_equal(lags, np.arange(-2, 6, 2)) - def test_convolve_returns_lagvector(self): + def test_correlate_lags_full_mode(self): a = [1, 2, 3] v = [0, 1, 0.5] - z, lags = np.convolve(a, v, mode='full', returns_lagvector=True) + z = np.convolve(a, v, mode='full') + lags = np.correlate_lags(len(a), len(v), mode='full') assert len(z) == len(lags) assert_array_equal(lags, np.arange(-2, 3)) @@ -4002,7 +4008,8 @@ def test_complex_no_conjugation(self): assert_array_equal(np.convolve(a, v), expected) # Same check via the lags path (maxlag=1 covers lags [-1, 0, 1], # which is the full output for two length-2 inputs) - result, lags = np.convolve(a, v, maxlag=1, returns_lagvector=True) + result = np.convolve(a, v, maxlag=1) + lags = np.correlate_lags(len(a), len(v), maxlag=1) assert_array_equal(result, expected) assert_array_equal(lags, np.arange(-1, 2)) diff --git a/numpy/matlib.pyi b/numpy/matlib.pyi index 06a6806e373a..e081612d7878 100644 --- a/numpy/matlib.pyi +++ b/numpy/matlib.pyi @@ -114,6 +114,7 @@ from numpy import ( # noqa: F401 core, corrcoef, correlate, + correlate_lags, cos, cosh, count_nonzero, From 1626e2bd73522b7050a183ffaee7f7735e95f0ee Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Mon, 20 Apr 2026 14:05:27 +0000 Subject: [PATCH 12/13] fix docs --- doc/source/reference/routines.statistics.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/source/reference/routines.statistics.rst b/doc/source/reference/routines.statistics.rst index cd93e60253fb..3c3eddda8271 100644 --- a/doc/source/reference/routines.statistics.rst +++ b/doc/source/reference/routines.statistics.rst @@ -40,6 +40,7 @@ Correlating corrcoef correlate + correlate_lags cov Histograms From e6b600d060d00b3cc6023d284de352c60e650372 Mon Sep 17 00:00:00 2001 From: Honi Sanders Date: Mon, 20 Apr 2026 19:12:58 +0000 Subject: [PATCH 13/13] fix docstring --- numpy/_core/numeric.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py index ec16367e363a..028fe835758e 100644 --- a/numpy/_core/numeric.py +++ b/numpy/_core/numeric.py @@ -870,7 +870,7 @@ def correlate_lags(a_len, v_len, mode=_CorrModeDefault, *, maxlag=None, >>> r = np.correlate(a, v, mode='full') >>> lag_vec = np.correlate_lags(len(a), len(v), mode='full') >>> r[lag_vec == 0] - array([2.]) + array([3.5]) """ if maxlag is not None and lags is not None: raise TypeError("cannot specify both maxlag and lags")