From 570a259f9d95712134ef46c970e2eeb8f7e8a5e2 Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Tue, 7 May 2024 23:14:29 +0200
Subject: [PATCH 1/7] improve performance of `conv`

---
 src/reformulations/conv.jl | 31 ++++++++++++++++++++++++++-----
 test/test_atoms.jl         |  5 +++++
 test/test_utilities.jl     | 24 ++++++++++++++++++------
 3 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/src/reformulations/conv.jl b/src/reformulations/conv.jl
index 591ced4b7..6254946c1 100644
--- a/src/reformulations/conv.jl
+++ b/src/reformulations/conv.jl
@@ -3,15 +3,36 @@
 # 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)
+    # It is much more efficient to construct sparse matrices
+    # this way rather than starting from `spzeros` and indexing into it.
+    Is = Int[]
+    Js = Int[]
+    Vs = eltype(h)[]
+    # build matrix by columns
+    for j in 1:n
+        append!(Is, j:(j+m-1))
+        append!(Js, repeat([j], 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]
-    end
+    length(x) > 0 ||
+        throw(ArgumentError("convolution with empty vector not supported"))
+
+    X = conv1D_matrix(x, length(y))
     return X * y
 end
 
diff --git a/test/test_atoms.jl b/test/test_atoms.jl
index 36e3c7f7c..bd27daf4e 100644
--- a/test/test_atoms.jl
+++ b/test/test_atoms.jl
@@ -1968,6 +1968,11 @@ 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..edda73350 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,20 @@ 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
+end
+
 function test_conj_issue_416()
     A = [1 1im; -1im 1]
     X = ComplexVariable(2, 2)

From efba1e141d39cb0e936832fbf9c31f7ee760d94e Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Tue, 7 May 2024 23:24:53 +0200
Subject: [PATCH 2/7] tweak

---
 src/reformulations/conv.jl | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/reformulations/conv.jl b/src/reformulations/conv.jl
index 6254946c1..a58aac412 100644
--- a/src/reformulations/conv.jl
+++ b/src/reformulations/conv.jl
@@ -16,10 +16,13 @@ function conv1D_matrix(h::AbstractVector, n::Integer)
     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, repeat([j], m))
+        append!(Js, (j for _ in 1:m))
         append!(Vs, h)
     end
     return SparseArrays.sparse(Is, Js, Vs, m + n - 1, n)

From a48dc79b91cca1f908df789d1b762d45a5e79342 Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Tue, 7 May 2024 23:31:00 +0200
Subject: [PATCH 3/7] rm comment

---
 src/reformulations/conv.jl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/reformulations/conv.jl b/src/reformulations/conv.jl
index a58aac412..a351d0f21 100644
--- a/src/reformulations/conv.jl
+++ b/src/reformulations/conv.jl
@@ -11,8 +11,6 @@ then we have `A * x ≈ conv1d(h, x)`.
 """
 function conv1D_matrix(h::AbstractVector, n::Integer)
     m = length(h)
-    # It is much more efficient to construct sparse matrices
-    # this way rather than starting from `spzeros` and indexing into it.
     Is = Int[]
     Js = Int[]
     Vs = eltype(h)[]

From 0d51fac9e2a890a4b0f865fefd17d877d91ac28b Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Tue, 7 May 2024 23:33:46 +0200
Subject: [PATCH 4/7] format

---
 src/reformulations/conv.jl | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/reformulations/conv.jl b/src/reformulations/conv.jl
index a351d0f21..12200dc66 100644
--- a/src/reformulations/conv.jl
+++ b/src/reformulations/conv.jl
@@ -14,9 +14,9 @@ function conv1D_matrix(h::AbstractVector, n::Integer)
     Is = Int[]
     Js = Int[]
     Vs = eltype(h)[]
-    sizehint!(Is, n*m)
-    sizehint!(Js, n*m)
-    sizehint!(Vs, n*m)
+    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))

From c22eca4023114c01da81bee34d56ab46cc97a63b Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Wed, 8 May 2024 00:38:41 +0200
Subject: [PATCH 5/7] implement 2D convolution

---
 src/reformulations/conv.jl | 95 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)

diff --git a/src/reformulations/conv.jl b/src/reformulations/conv.jl
index 12200dc66..4395e0b06 100644
--- a/src/reformulations/conv.jl
+++ b/src/reformulations/conv.jl
@@ -26,6 +26,17 @@ function conv1D_matrix(h::AbstractVector, n::Integer)
     return SparseArrays.sparse(Is, Js, Vs, m + n - 1, n)
 end
 
+"""
+    conv(x, y)
+
+Produces the discrete convolution of vectors `x` of length `m` and `y` of length `n`. That is, if `z = conv(x,y)`, then `z` has length `m+n-1`, and `z` has entries
+
+```julia
+z[i] = sum(x[j] * get(y, i - j + 1, 0) for j in 1:m)
+```
+
+Note that `conv` is symmetric in `x` and `y`: `conv(x,y) == conv(y, x)` (in exact arithmetic).
+"""
 function conv(x::Value, y::AbstractExpr)
     if length(x) != size(x, 1) || size(y, 2) > 1
         error("convolution only supported between two vectors")
@@ -38,3 +49,87 @@ function conv(x::Value, y::AbstractExpr)
 end
 
 conv(x::AbstractExpr, y::Value) = conv(y, x)
+
+# direct non-variable implementation for reference
+function conv(x::AbstractVector, y::AbstractVector)
+    T = promote_type(eltype(x), eltype(y))
+    m = length(x)
+    n = length(y)
+    z = zeros(T, m + n - 1)
+    for i in eachindex(z)
+        z[i] = sum(x[j] * get(y, i - j + 1, 0) for j in 1:m)
+    end
+    return z
+end
+
+#####
+##### 2D discrete convolution
+#####
+
+# We reformulate the problem into a 1D convolution following the approach from:
+# https://en.wikipedia.org/wiki/Multidimensional_discrete_convolution#Multidimensional_convolution_with_one-dimensional_convolution_methods
+
+# In particular, we keep separate the matrix representation of the 1D convolution, since this could be re-used among multiple 2D convolutions.
+
+function conv2D_kernel(X::AbstractMatrix, sz::Tuple{Int,Int})
+    M, N = size(X)
+    K, L = sz
+    Z_rows = M + K - 1
+    Z_cols = N + L - 1
+    Xpad = zeros(eltype(X), Z_rows, Z_cols)
+    Xpad[1:size(X, 1), 1:size(X, 2)] .= X
+    Xv = vec(Xpad)[1:(N-1)*(M+K-1)+M]
+    return conv1D_matrix(Xv, (L - 1) * (M + K - 1) + K)
+end
+
+function conv2D(
+    X_size::Tuple,
+    X_kernel::AbstractMatrix{T},
+    Y::Convex.AbstractExpr,
+) where {T}
+    M, N = X_size
+    K, L = size(Y)
+    Z_rows = M + K - 1
+    Z_cols = N + L - 1
+    bottom = spzeros(T, length(size(Y, 1)+1:Z_rows), Z_cols)
+    right_side = spzeros(T, size(Y, 1), length(size(Y, 2)+1:Z_cols))
+    Y_padded = [
+        Y right_side
+        bottom
+    ]
+    Y_vector = vec(Y_padded)[1:(L-1)*(M+K-1)+K]
+    Z_vector = X_kernel * Y_vector
+    Z = reshape(Z_vector, Z_rows, Z_cols)
+    return Z
+end
+
+"""
+    conv2D(X, Y)
+
+Performs a 2D discrete convolution of `X` and `Y`. Assuming `size(X) = (M,N)` and `size(Y) = (K, L)`, then `Z = conv2D(X,Y)` has `size(Z) = (M + K - 1, N + L - 1)`, with entries
+
+```julia
+Z[i, j] = sum(X[m, n] * get(Y, (i-m+1, j-n+1), zero(T)) for m in 1:M, n in 1:N)
+```
+
+"""
+function conv2D(X::AbstractMatrix, Y::Convex.AbstractExpr)
+    return conv2D(size(X), conv2D_kernel(X, size(Y)), Y)
+end
+
+conv2D(Y::Convex.AbstractExpr, X::AbstractMatrix) = conv2D(X, Y)
+
+# direct non-variable implementation for reference
+function conv2D(X::AbstractMatrix, Y::AbstractMatrix)
+    T = promote_type(eltype(X), eltype(Y))
+    M, N = size(X)
+    K, L = size(Y)
+    Z = zeros(T, M + K - 1, N + L - 1)
+    for i in 1:M+K-1, j in 1:N+L-1
+        Z[i, j] = sum(
+            X[m, n] * get(Y, (i - m + 1, j - n + 1), zero(T)) for m in 1:M,
+            n in 1:N
+        )
+    end
+    return Z
+end

From b6f7eec5a1b4336cd7995f82df595cd113b7d287 Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Fri, 10 May 2024 00:50:55 +0200
Subject: [PATCH 6/7] wip

---
 src/Convex.jl      |  1 +
 test/test_atoms.jl | 12 ++++++++++++
 2 files changed, 13 insertions(+)

diff --git a/src/Convex.jl b/src/Convex.jl
index a4aff3901..9c2681377 100644
--- a/src/Convex.jl
+++ b/src/Convex.jl
@@ -13,6 +13,7 @@ import OrderedCollections
 import SparseArrays
 
 export conv,
+    conv2D,
     dotsort,
     entropy,
     entropy_elementwise,
diff --git a/test/test_atoms.jl b/test/test_atoms.jl
index 8d065b5f4..9728826fd 100644
--- a/test/test_atoms.jl
+++ b/test/test_atoms.jl
@@ -1975,6 +1975,18 @@ function test_conv()
     return
 end
 
+function test_conv2d()
+    target = """
+    variables: x1, x2
+    minobjective: [1.0 * x1, 2.0 * x1 + 1.0 * x2, 2.0 * x2]
+    """
+    _test_reformulation(target) do context
+        return conv2D(Variable(2, 2), [1 2; 3 4])
+    end
+
+
+end
+
 ### reformulations/dot
 
 function test_dot()

From 223ceb5b808aac5c5bbd346de4f45339551ac57d Mon Sep 17 00:00:00 2001
From: Eric Hanson <5846501+ericphanson@users.noreply.github.com>
Date: Sat, 11 May 2024 14:57:04 +0200
Subject: [PATCH 7/7] wip

---
 test/test_atoms.jl | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/test/test_atoms.jl b/test/test_atoms.jl
index 4e8504e2f..d369f085a 100644
--- a/test/test_atoms.jl
+++ b/test/test_atoms.jl
@@ -1980,17 +1980,17 @@ function test_conv()
     return
 end
 
-function test_conv2d()
-    target = """
-    variables: x1, x2
-    minobjective: [1.0 * x1, 2.0 * x1 + 1.0 * x2, 2.0 * x2]
-    """
-    _test_reformulation(target) do context
-        return conv2D(Variable(2, 2), [1 2; 3 4])
-    end
+# function test_conv2d()
+#     target = """
+#     variables: x1, x2
+#     minobjective: [1.0 * x1, 2.0 * x1 + 1.0 * x2, 2.0 * x2, 3.0 * x1 + 0.0 * x2 + 1.0 * v[3]]
+#     """
+#     _test_reformulation(target) do context
+#         return conv2D(Variable(2, 2), [1 3; 2 4])
+#     end
 
 
-end
+# end
 
 ### reformulations/dot