feat: DBCP convolve atom

Author: haozhu10015

src/dbcp/__init__.py

@@ -1,6 +1,7 @@
+from .atoms import convolve
 from .problem import BiconvexProblem, BiconvexRelaxProblem
 
-__all__ = ["BiconvexProblem", "BiconvexRelaxProblem", "__version__"]
+__all__ = ["BiconvexProblem", "BiconvexRelaxProblem", "convolve", "__version__"]
 
 try:
     from importlib.metadata import version, PackageNotFoundError

src/dbcp/atoms.py

@@ -0,0 +1,39 @@
+import cvxpy as cp
+
+
+def convolve(x, y):
+    """Discrete convolution of two 1-D cvxpy expressions.
+
+    Suppose :math:`x` and :math:`y` are 1-D cvxpy expressions of lengths
+    :math:`m` and :math:`n`, respectively.
+    This function returns a cvxpy expression :math:`c` of length :math:`m + n - 1`, where
+
+    .. math::
+
+        c_k = \\sum_{i + j = k} x_i y_j,\\quad k = 1, \\ldots, m + n - 1.
+
+    Matches numpy.convolve for 1-D arrays.
+
+    This function extends cvxpy.convolve atom to support the convolution
+    operation between two cvxpy expressions.
+
+    Parameters
+    ----------
+    x : cp.Expression
+        A 1-D cvxpy expression.
+    y : cp.Expression
+        A 1-D cvxpy expression.
+
+    Returns
+    -------
+    cp.Expression
+        The convolution of x and y.
+    """
+    if x.ndim != 1 or y.ndim != 1:
+        raise ValueError("Both inputs must be 1-D cvxpy expressions.")
+
+    c = [0] * (x.shape[0] + y.shape[0] - 1)
+    for i, a in enumerate(y):
+        for j, b in enumerate(x):
+            c[i + j] += a * b
+    return cp.hstack(c)

tests/test_atoms.py

@@ -0,0 +1,33 @@
+import pytest
+import cvxpy as cp
+import numpy as np
+import dbcp
+
+np.random.seed(10015)
+
+
+def test_convolve():
+    # Test the convolution atom to match np.convolve
+    x0 = np.random.randn(10)
+    y0 = np.random.randn(20)
+    c0 = np.convolve(x0, y0)
+
+    x = cp.Variable(10)
+    y = cp.Variable(20)
+    c = dbcp.convolve(x, y)
+
+    x.value = x0
+    y.value = y0
+    assert np.allclose(c.value, c0)
+
+
+def test_convolve_invalid_input():
+    x = cp.Variable((3, 4))
+    y = cp.Variable(5)
+    with pytest.raises(ValueError):
+        dbcp.convolve(x, y)
+
+    x = cp.Variable(6)
+    y = cp.Variable((2, 3))
+    with pytest.raises(ValueError):
+        dbcp.convolve(x, y)