ENH: add maxlag and lags parameters to correlate/convolve

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 = [

doc/release/upcoming_changes/31261.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.

doc/release/upcoming_changes/31261.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]))

doc/source/reference/routines.statistics.rst

@@ -40,6 +40,7 @@ Correlating
 
    corrcoef
    correlate
+   correlate_lags
    cov
 
 Histograms

numpy/__init__.py

@@ -203,6 +203,7 @@
         copysign,
         copyto,
         correlate,
+        correlate_lags,
         cos,
         cosh,
         count_nonzero,

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",

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",

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

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 = {

numpy/_core/include/numpy/ndarraytypes.h

@@ -269,7 +269,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 */

numpy/_core/include/numpy/numpyconfig.h

@@ -80,6 +80,7 @@
 #define NPY_2_3_API_VERSION 0x00000014
 #define NPY_2_4_API_VERSION 0x00000015
 #define NPY_2_5_API_VERSION 0x00000015
+#define NPY_2_6_API_VERSION 0x00000016
 
 
 /*
@@ -175,6 +176,8 @@
     #define NPY_FEATURE_VERSION_STRING "2.4"
 #elif NPY_FEATURE_VERSION == NPY_2_5_API_VERSION
     #define NPY_FEATURE_VERSION_STRING "2.5"
+#elif NPY_FEATURE_VERSION == NPY_2_6_API_VERSION
+    #define NPY_FEATURE_VERSION_STRING "2.6"
 #else
     #error "Missing version string define for new NumPy version."
 #endif

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.

numpy/_core/multiarray.py

@@ -37,6 +37,7 @@
     'asanyarray', 'ascontiguousarray', 'asfortranarray', 'bincount',
     'broadcast', 'busday_count', 'busday_offset', 'busdaycalendar', 'can_cast',
     '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',

numpy/_core/multiarray.pyi

@@ -130,6 +130,7 @@ __all__ = [
     "copyto",
     "correlate",
     "correlate2",
+    "correlatelags",
     "count_nonzero",
     "c_einsum",
     "datetime_as_string",
@@ -254,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]] = ...