Repository: numba/numba
Pull request: Support for all modes in np.correlate and np.convolve
Hello everyone,

With this PR, I hope to bring the missing features from `np.convolve` and `np.correlate` into Numba. As you can see in the edited files, the initial solution was not correct, and I had to modify a bit the `_np_correlate_core_impl` to make it work.
I also added some tests to verify that it works and it should close #7406.

Last but not least, I think that `_np_correlate_core_impl` should allow for arguments where `len(ap1) < len(ap2)`. Here, it is assumed that the arrays are ordered properly, but in my opinion that should be the job of `_np_correlate_core_impl`. Same could apply when checking for the `mode` value. Don't hesitate to give me your opinion on that?

Thanks in advance!

Problem statement:
Support for Numpy convolve mode 'valid'
## Feature request

I've been working with model simulations and there's the need for `np.convolve(a, k, mode='valid')`.

I wish to have support for this for use with `@jit`/`@njit` decorators.

I have this example implementation from one of the discussions about using Numba for such situation, but I'd rather not have to resort to it in projects which require valid mode.
```python
def numba_convolve_mode_valid_as_loop(arr, kernel):
        m = arr.size
        n = kernel.size
        out_size = m - n + 1
        out = np.empty(out_size, dtype=np.complex128)
        kernel = np.flip(kernel)  # wrong results otherwise
        for i in range(out_size):
            out[i] = np.dot(arr[i:i + n], kernel)
        return out
```

Hints from issue discussion:
Labelling as "good second issue" as most of the parts needed for the implementation are already present, but need connecting together. Se e.g. https://github.com/numba/numba/blob/master/numba/np/arraymath.py#L4230-L4251.
Thanks for pointing it out! I'd need to set a test environment for that.

A simple
```python
@njit 
def convolve_valid(a, v):
    return np.convolve(a, v, mode='valid) 
```
seems sufficient, right?

Either Numba will compile, or it will throw an error. Then the results should be easily testable with a couple more tests.

That said, I'd need to compile and use my personal version of Numba for those, and I'm definitely pretty lost in that regard. I currently installed it via `pip`.
As a new contributor, some instructions for getting set up for development are [here](https://numba.readthedocs.io/en/stable/developer/contributing.html#getting-set-up)  and are a great place to start.
Thank you!!
I'll take a look and update if I can't make it.
Hello @danilo-bc, are you still working on this ? Otherwise, I would gladly give it a try!
> Hello @danilo-bc, are you still working on this ? Otherwise, I would gladly give it a try!

Please feel free to take it over. I've had a few unexpected problems in the meantime.
> > Hello @danilo-bc, are you still working on this ? Otherwise, I would gladly give it a try!
> 
> Please feel free to take it over. I've had a few unexpected problems in the meantime.

Thanks for the quick reply! I will try it myself then.

Commit message: chore(convolve): testing exception

Patch to review:
diff --git a/docs/source/reference/numpysupported.rst b/docs/source/reference/numpysupported.rst
index 9e62a41d091..ac4edc5f8cc 100644
--- a/docs/source/reference/numpysupported.rst
+++ b/docs/source/reference/numpysupported.rst
@@ -453,10 +453,10 @@ The following top-level functions are supported:
 * :func:`numpy.broadcast_shapes`
 * :func:`numpy.column_stack`
 * :func:`numpy.concatenate` (only supports tuple arguments)
-* :func:`numpy.convolve` (only the 2 first arguments)
+* :func:`numpy.convolve` (all arguments)
 * :func:`numpy.copy` (only the first argument)
 * :func:`numpy.corrcoef` (only the 3 first arguments, requires SciPy)
-* :func:`numpy.correlate` (only the 2 first arguments)
+* :func:`numpy.correlate` (all arguments)
 * :func:`numpy.count_nonzero` (axis only supports scalar values)
 * :func:`numpy.cov` (only the 5 first arguments)
 * :func:`numpy.cross` (only the 2 first arguments; at least one of the input
diff --git a/numba/np/arraymath.py b/numba/np/arraymath.py
index 537c226ce83..3e17e676787 100644
--- a/numba/np/arraymath.py
+++ b/numba/np/arraymath.py
@@ -4088,16 +4088,6 @@ def _np_correlate_core(ap1, ap2, mode, direction):
     pass
 
 
-class _corr_conv_Mode(IntEnum):
-    """
-    Enumerated modes for correlate/convolve as per:
-    https://github.com/numpy/numpy/blob/ac6b1a902b99e340cf7eeeeb7392c91e38db9dd8/numpy/core/numeric.py#L862-L870    # noqa: E501
-    """
-    VALID = 0
-    SAME = 1
-    FULL = 2
-
-
 @overload(_np_correlate_core)
 def _np_correlate_core_impl(ap1, ap2, mode, direction):
     a_dt = as_dtype(ap1.dtype)
@@ -4105,37 +4095,44 @@ def _np_correlate_core_impl(ap1, ap2, mode, direction):
     dt = np.promote_types(a_dt, b_dt)
     innerprod = _get_inner_prod(ap1.dtype, ap2.dtype)
 
-    Mode = _corr_conv_Mode
-
     def impl(ap1, ap2, mode, direction):
         # Implementation loosely based on `_pyarray_correlate` from
         # https://github.com/numpy/numpy/blob/3bce2be74f228684ca2895ad02b63953f37e2a9d/numpy/core/src/multiarray/multiarraymodule.c#L1191    # noqa: E501
-        # For "Mode":
-        # Convolve uses 'full' by default, this is denoted by the number 2
-        # Correlate uses 'valid' by default, this is denoted by the number 0
+        # For "mode":
+        # Convolve uses 'full' by default.
+        # Correlate uses 'valid' by default.
         # For "direction", +1 to write the return values out in order 0->N
         # -1 to write them out N->0.
 
-        if not (mode == Mode.VALID or mode == Mode.FULL):
-            raise ValueError("Invalid mode")
-
         n1 = len(ap1)
         n2 = len(ap2)
+
+        if n1 < n2:
+            # This should never occur when called by np.convolve because
+            # _np_correlate.impl swaps arguments based on length.
+            # The same applies for np.correlate.
+            raise ValueError("'len(ap1)' must greater than 'len(ap2)'")
+
         length = n1
         n = n2
-        if mode == Mode.VALID: # mode == valid == 0, correlate default
+        if mode == "valid":
             length = length - n + 1
             n_left = 0
             n_right = 0
-        elif mode == Mode.FULL: # mode == full == 2, convolve default
+        elif mode == "full":
             n_right = n - 1
             n_left = n - 1
             length = length + n - 1
+        elif mode == "same":
+            n_left = n // 2
+            n_right = n - n_left - 1
         else:
-            raise ValueError("Invalid mode")
+            raise ValueError(
+                "Invalid 'mode', "
+                "valid are 'full', 'same', 'valid'"
+            )
 
         ret = np.zeros(length, dt)
-        n = n - n_left
 
         if direction == 1:
             idx = 0
@@ -4147,23 +4144,26 @@ def impl(ap1, ap2, mode, direction):
             raise ValueError("Invalid direction")
 
         for i in range(n_left):
-            ret[idx] = innerprod(ap1[:idx + 1], ap2[-(idx + 1):])
+            k = i + n - n_left
+            ret[idx] = innerprod(ap1[:k], ap2[-k:])
             idx = idx + inc
 
         for i in range(n1 - n2 + 1):
             ret[idx] = innerprod(ap1[i : i + n2], ap2)
             idx = idx + inc
 
-        for i in range(n_right, 0, -1):
-            ret[idx] = innerprod(ap1[-i:], ap2[:i])
+        for i in range(n_right):
+            k = n - i - 1
+            ret[idx] = innerprod(ap1[-k:], ap2[:k])
             idx = idx + inc
+
         return ret
 
     return impl
 
 
 @overload(np.correlate)
-def _np_correlate(a, v):
+def _np_correlate(a, v, mode="valid"):
     _assert_1d(a, 'np.correlate')
     _assert_1d(v, 'np.correlate')
 
@@ -4175,8 +4175,6 @@ def op_conj(x):
     def op_nop(x):
         return x
 
-    Mode = _corr_conv_Mode
-
     if a.dtype in types.complex_domain:
         if v.dtype in types.complex_domain:
             a_op = op_nop
@@ -4192,29 +4190,31 @@ def op_nop(x):
             a_op = op_conj
             b_op = op_nop
 
-    def impl(a, v):
+    def impl(a, v, mode="valid"):
         la = len(a)
         lv = len(v)
+
         if la == 0:
             raise ValueError("'a' cannot be empty")
         if lv == 0:
             raise ValueError("'v' cannot be empty")
+
         if la < lv:
-            return _np_correlate_core(b_op(v), a_op(a), Mode.VALID, -1)
+            return _np_correlate_core(b_op(v), a_op(a), mode, -1)
         else:
-            return _np_correlate_core(a_op(a), b_op(v), Mode.VALID, 1)
+            return _np_correlate_core(a_op(a), b_op(v), mode, 1)
 
     return impl
 
 
 @overload(np.convolve)
-def np_convolve(a, v):
+def np_convolve(a, v, mode="full"):
     _assert_1d(a, 'np.convolve')
     _assert_1d(v, 'np.convolve')
 
     Mode = _corr_conv_Mode
 
-    def impl(a, v):
+    def impl(a, v, mode="full"):
         la = len(a)
         lv = len(v)
 
@@ -4224,9 +4224,9 @@ def impl(a, v):
             raise ValueError("'v' cannot be empty")
 
         if la < lv:
-            return _np_correlate_core(v, a[::-1], Mode.FULL, 1)
+            return _np_correlate_core(v, a[::-1], corr_mode, 1)
         else:
-            return _np_correlate_core(a, v[::-1], Mode.FULL, 1)
+            return _np_correlate_core(a, v[::-1], corr_mode, 1)
 
     return impl
 
diff --git a/numba/tests/test_np_functions.py b/numba/tests/test_np_functions.py
index 093769c438c..1f7e30905b5 100644
--- a/numba/tests/test_np_functions.py
+++ b/numba/tests/test_np_functions.py
@@ -187,12 +187,12 @@ def split(a, indices, axis=0):
     return np.split(a, indices, axis=axis)
 
 
-def correlate(a, v):
-    return np.correlate(a, v)
+def correlate(a, v, mode="valid"):
+    return np.correlate(a, v, mode=mode)
 
 
-def convolve(a, v):
-    return np.convolve(a, v)
+def convolve(a, v, mode="full"):
+    return np.convolve(a, v, mode=mode)
 
 
 def tri_n(N):
@@ -1339,8 +1339,11 @@ def _test_correlate_convolve(self, pyfunc):
         lengths = (1, 2, 3, 7)
         dts = [np.int8, np.int32, np.int64, np.float32, np.float64,
                np.complex64, np.complex128]
+        modes = ["full", "valid", "same"]
 
-        for dt1, dt2, n, m in itertools.product(dts, dts, lengths, lengths):
+        for dt1, dt2, n, m, mode in itertools.product(
+            dts, dts, lengths, lengths, modes
+        ):
             a = np.arange(n, dtype=dt1)
             v = np.arange(m, dtype=dt2)
 
@@ -1349,8 +1352,9 @@ def _test_correlate_convolve(self, pyfunc):
             if np.issubdtype(dt2, np.complexfloating):
                 v = (v + 1j * v).astype(dt2)
 
-            expected = pyfunc(a, v)
-            got = cfunc(a, v)
+            expected = pyfunc(a, v, mode=mode)
+            got = cfunc(a, v, mode=mode)
+
             self.assertPreciseEqual(expected, got)
 
         _a = np.arange(12).reshape(4, 3)
@@ -1380,6 +1384,11 @@ def _test_correlate_convolve_exceptions(self, fn):
             else:
                 self.assertIn("'v' cannot be empty", str(raises.exception))
 
+        with self.assertRaises(ValueError) as raises:
+            cfunc(_b, _b, mode="invalid mode")
+
+            self.assertIn("Invalid 'mode'", str(raises.exception))
+
     def test_correlate_exceptions(self):
         # correlate supported 0 dimension arrays until 1.18
         self._test_correlate_convolve_exceptions(correlate)


Code review comments:
diff_hunk: @@ -4192,29 +4190,31 @@ def op_nop(x):
             a_op = op_conj
             b_op = op_nop
 
-    def impl(a, v):
+    def impl(a, v, mode="valid"):
         la = len(a)
         lv = len(v)
+
         if la == 0:
             raise ValueError("'a' cannot be empty")
         if lv == 0:
             raise ValueError("'v' cannot be empty")
+
         if la < lv:
-            return _np_correlate_core(b_op(v), a_op(a), Mode.VALID, -1)
+            return _np_correlate_core(b_op(v), a_op(a), mode, -1)
         else:
-            return _np_correlate_core(a_op(a), b_op(v), Mode.VALID, 1)
+            return _np_correlate_core(a_op(a), b_op(v), mode, 1)
 
     return impl
 
 
 @overload(np.convolve)
-def np_convolve(a, v):
+def np_convolve(a, v, mode="full"):
     _assert_1d(a, 'np.convolve')
     _assert_1d(v, 'np.convolve')
 
     Mode = _corr_conv_Mode
path: numba/np/arraymath.py
text: @user1:
```suggestion
```

Seems like accidentally leftover code.

@author:
Good catch! Thanks