From f2b37ddb555db767dec0882587ea85f90d8f8f68 Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Wed, 8 May 2024 00:46:41 +0200
Subject: [PATCH] improve performance of `conv` (#634)

* improve performance of `conv`

* tweak

* rm comment

* format

* Apply suggestions from code review

---------

Co-authored-by: Oscar Dowson <odow@users.noreply.github.com>
---
 src/reformulations/conv.jl | 30 ++++++++++++++++++++++++++----
 test/test_atoms.jl         |  4 ++++
 test/test_utilities.jl     | 25 +++++++++++++++++++------
 3 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/src/reformulations/conv.jl b/src/reformulations/conv.jl
index 591ced4b7..9aedf5e95 100644
--- a/src/reformulations/conv.jl
+++ b/src/reformulations/conv.jl
@@ -3,15 +3,37 @@
 # Use of this source code is governed by a BSD-style license that can be found
 # in the LICENSE file or at https://opensource.org/license/bsd-2-clause
 
+"""
+    conv1D_matrix(h::AbstractVector, n::Integer) -> SparseMatrixCSC
+
+Create a sparse matrix `A` such that if `x` has length `n`,
+then we have `A * x ≈ conv1d(h, x)`.
+"""
+function conv1D_matrix(h::AbstractVector, n::Integer)
+    m = length(h)
+    Is = Int[]
+    Js = Int[]
+    Vs = eltype(h)[]
+    sizehint!(Is, n * m)
+    sizehint!(Js, n * m)
+    sizehint!(Vs, n * m)
+    # build matrix by columns
+    for j in 1:n
+        append!(Is, j:(j+m-1))
+        append!(Js, (j for _ in 1:m))
+        append!(Vs, h)
+    end
+    return SparseArrays.sparse(Is, Js, Vs, m + n - 1, n)
+end
+
 function conv(x::Value, y::AbstractExpr)
     if length(x) != size(x, 1) || size(y, 2) > 1
         error("convolution only supported between two vectors")
     end
-    m, n = length(x), size(y, 1)
-    X = spzeros(eltype(x), m + n - 1, n)
-    for i in 1:n, j in 1:m
-        X[i+j-1, i] = x[j]
+    if length(x) == 0
+        throw(ArgumentError("convolution with empty vector not supported"))
     end
+    X = conv1D_matrix(x, length(y))
     return X * y
 end
 
diff --git a/test/test_atoms.jl b/test/test_atoms.jl
index 36e3c7f7c..8d065b5f4 100644
--- a/test/test_atoms.jl
+++ b/test/test_atoms.jl
@@ -1968,6 +1968,10 @@ function test_conv()
         ErrorException("convolution only supported between two vectors"),
         conv([1, 2], Variable(2, 2)),
     )
+    @test_throws(
+        ArgumentError("convolution with empty vector not supported"),
+        conv([], Variable(2)),
+    )
     return
 end
 
diff --git a/test/test_utilities.jl b/test/test_utilities.jl
index 1407a24dc..381cb23d8 100644
--- a/test/test_utilities.jl
+++ b/test/test_utilities.jl
@@ -660,6 +660,14 @@ function test_logsumexp_stability()
     return
 end
 
+# simple 1D convolution implementation
+function _conv(h, x)
+    m = length(h)
+    n = length(x)
+    zero_pad_x(i) = 1 <= i <= n ? x[i] : 0
+    return [sum(h[j] * zero_pad_x(i - j + 1) for j in 1:m) for i in 1:m+n-1]
+end
+
 function test_conv_issue_364()
     n = 3
     m = 11
@@ -667,16 +675,21 @@ function test_conv_issue_364()
     x = rand(n)
     hvar = Variable(m)
     hvar.value = h
-    function _conv(h, x)
-        m = length(h)
-        n = length(x)
-        zero_pad_x(i) = 1 <= i <= n ? x[i] : 0
-        return [sum(h[j] * zero_pad_x(i - j + 1) for j in 1:m) for i in 1:m+n-1]
-    end
     @test evaluate(conv(hvar, x)) ≈ _conv(h, x)
     return
 end
 
+function test_conv1D_matrix()
+    for (x_len, y_len) in ((20, 5), (5, 20), (5, 5), (1, 1), (2, 3))
+        for im1 in (im, 0), im2 in (im, 0)
+            x = randn(x_len) + randn(x_len) * im1
+            y = randn(y_len) + randn(y_len) * im2
+            @test Convex.conv1D_matrix(x, length(y)) * y ≈ _conv(x, y)
+        end
+    end
+    return
+end
+
 function test_conj_issue_416()
     A = [1 1im; -1im 1]
     X = ComplexVariable(2, 2)