From 368d4a22fcf32621eae1aca9de70500ee313a8d3 Mon Sep 17 00:00:00 2001
From: ZongyuLi-umich <zonyul@umich.edu>
Date: Thu, 14 Oct 2021 19:57:40 -0400
Subject: [PATCH] ker3

---
 src/fft_convolve.jl     | 16 ++++++++--------
 test/adjoint-fftconv.jl |  6 +++---
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/src/fft_convolve.jl b/src/fft_convolve.jl
index 8d57b8a9..5ff47d4a 100644
--- a/src/fft_convolve.jl
+++ b/src/fft_convolve.jl
@@ -144,13 +144,13 @@ end
 
 
 """
-    fft_conv!(output, image3, ker, plans)
-In-place version of convolving a 3D `image3` with `ker`
+    fft_conv!(output, image3, ker3, plans)
+In-place version of convolving a 3D `image3` with a 3D kernel `ker3`
 """
 function fft_conv!(
     output::AbstractArray{<:RealU,3},
     image3::AbstractArray{<:RealU,3},
-    ker::AbstractMatrix{<:RealU},
+    ker3::AbstractArray{<:RealU,3},
     plans::Vector{<:PlanPSF},
     )
 
@@ -159,7 +159,7 @@ function fft_conv!(
     fun = y -> fft_conv!(
             (@view output[:, y, :]),
             (@view image3[:, y, :]),
-            ker,
+            (@view ker3[:, :, y]),
             plans[Threads.threadid()],
             )
 
@@ -171,13 +171,13 @@ end
 
 
 """
-    fft_conv_adj!(output, image3, ker, plans)
-In-place version of adjoint of convolving a 3D `image3` with `ker`
+    fft_conv_adj!(output, image3, ker3, plans)
+In-place version of adjoint of convolving a 3D `image3` with a 3D kernel `ker3`
 """
 function fft_conv_adj!(
     output::AbstractArray{<:RealU,3},
     image3::AbstractArray{<:RealU,3},
-    ker::AbstractMatrix{<:RealU},
+    ker3::AbstractArray{<:RealU,3},
     plans::Vector{<:PlanPSF},
 )
 
@@ -186,7 +186,7 @@ function fft_conv_adj!(
     fun = y -> fft_conv_adj!(
             (@view output[:, y, :]),
             (@view image3[:, y, :]),
-            ker,
+            (@view ker3[:, :, y]),
             plans[Threads.threadid()],
             )
 
diff --git a/test/adjoint-fftconv.jl b/test/adjoint-fftconv.jl
index e9f7bb83..0855cbd0 100644
--- a/test/adjoint-fftconv.jl
+++ b/test/adjoint-fftconv.jl
@@ -23,11 +23,11 @@ end
     T = Float32
     plan = plan_psf(nx, nz, nx_psf; T)
     image3 = rand(T, nx, nx, nz)
-    ker = ones(T, nx_psf, nx_psf) / (nx_psf)^2
+    ker3 = ones(T, nx_psf, nx_psf, nx) / (nx_psf)^2
     result = similar(image3)
-    fft_conv!(result, image3, ker, plan)
+    fft_conv!(result, image3, ker3, plan)
     @test maximum(result) ≤ 1
-    fft_conv_adj!(result, image3, ker, plan)
+    fft_conv_adj!(result, image3, ker3, plan)
     @test maximum(result) ≤ 1.5 # boundary is the sum of replicate padding
 end