feat: implement convolve function

Author: rwnobrega

site/toc.yaml

@@ -103,6 +103,7 @@ Utilities:
   - int_to_bits
   - sampling_rate_compress
   - sampling_rate_expand
+  - convolve
   - fourier_transform
   - boxplus
   - gaussian_q

src/komm/__init__.py

@@ -125,7 +125,7 @@
 from ._util.bit_operations import bits_to_int, int_to_bits
 from ._util.information_theory import binary_entropy, binary_entropy_inv, entropy, relative_entropy
 from ._util.sequences import autocorrelation, cyclic_autocorrelation
-from ._util.signal_processing import fourier_transform, sampling_rate_compress, sampling_rate_expand
+from ._util.signal_processing import convolve, fourier_transform, sampling_rate_compress, sampling_rate_expand
 from ._util.special_functions import boxplus, gaussian_q, gaussian_q_inv, marcum_q
 
 # isort: on
@@ -241,6 +241,7 @@
     "binary_entropy_inv",
     "entropy",
     "relative_entropy",
+    "convolve",
     "fourier_transform",
     "sampling_rate_compress",
     "sampling_rate_expand",

src/komm/_util/signal_processing.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 import numpy.typing as npt
@@ -118,6 +118,60 @@ def sampling_rate_expand(
     return output
 
 
+def convolve(
+    input1: npt.ArrayLike,
+    input2: npt.ArrayLike,
+    axis: int = -1,
+    mode: Literal["full", "valid", "same"] = "full",
+):
+    r"""
+    Convolves two signals. This is like [`numpy.convolve`](https://numpy.org/doc/stable/reference/generated/numpy.convolve.html), but one of the inputs may be multidimensional (the other input must still be $1$-dimensional). The convolution is applied independently to each array along a specified axis.
+
+    Note:
+        This is a simple wrapper around `numpy.convolve`.
+
+    Parameters:
+        input1: The first input array.
+
+        input2: The second input array.
+
+        axis: The axis along which the convolution is applied in the multidimensional input. Default is the last axis.
+
+        mode: Convolution mode (same as `numpy.convolve`).
+
+    Examples:
+        >>> komm.convolve([1, 1], [1, 1])
+        array([1, 2, 1])
+
+        >>> komm.convolve([1, 1], [[1, 1, 0], [0, 1, 1]])
+        array([[1, 2, 1, 0],
+               [0, 1, 2, 1]])
+
+        >>> komm.convolve([[1, 1, 0], [0, 1, 1]], [1, 1])
+        array([[1, 2, 1, 0],
+               [0, 1, 2, 1]])
+
+        >>> komm.convolve([1, 1], [[1, 1, 0], [0, 1, 1]], axis=0)
+        array([[1, 1, 0],
+               [1, 2, 1],
+               [0, 1, 1]])
+
+        >>> komm.convolve([[1, 1]], [[1, 1]])
+        Traceback (most recent call last):
+        ...
+        ValueError: at least one of the inputs must be 1-dimensional
+    """
+    input1 = np.asarray(input1)
+    input2 = np.asarray(input2)
+    if input1.ndim == 1:
+        x, arr = input1, input2
+    elif input2.ndim == 1:
+        x, arr = input2, input1
+    else:
+        raise ValueError("at least one of the inputs must be 1-dimensional")
+    return np.apply_along_axis(lambda y: np.convolve(x, y, mode), axis, arr)
+
+
 def fourier_transform(
     waveform: npt.ArrayLike,
     time_step: float,