From 3ee6ad69eb495528189a451bdd7964bc0ba18094 Mon Sep 17 00:00:00 2001
From: Ashton Bradley <ashton.bradley@gmail.com>
Date: Mon, 20 May 2024 15:09:55 +1200
Subject: [PATCH] simplify convolve

---
 src/analysis.jl  | 72 +++++++++++++++++++++++++++---------------------
 test/runtests.jl |  4 +--
 2 files changed, 43 insertions(+), 33 deletions(-)

diff --git a/src/analysis.jl b/src/analysis.jl
index 677db00..f12959b 100644
--- a/src/analysis.jl
+++ b/src/analysis.jl
@@ -101,8 +101,7 @@ end
 """
 	Wi,Wc = helmholtz(wx,...,kx,...)
 
-Computes a 2 or 3 dimensional Helmholtz decomposition of the vector field with components `wx`, `wy`, or `wx`, `wy`, `wz`. 
-`psi` is passed to provide requisite arrays in `k`-space.
+Computes a 2 or 3 dimensional Helmholtz decomposition of the vector field with components `wx`, `wy`, or `wx`, `wy`, `wz`. `psi` is passed to provide requisite arrays in `k`-space.
 Returned fields `Wi`, `Wc` are tuples of cartesian components of incompressible and compressible respectively.
 """
 function helmholtz(wx, wy, kx, ky)
@@ -172,35 +171,45 @@ function energydecomp(psi::Psi{3})
     return et, ei, ec
 end
 
-"""
-	zeropad(A)
+# """
+# 	zeropad(A)
 
-Zero-pad the array `A` to twice the size with the same element type as `A`.
-"""
-function zeropad(A)
-    S = size(A)
-    if any(isodd.(S))
-        error("Array dims not divisible by 2")
-    end
-    nO = 2 .* S
-    nI = S .÷ 2
+# Zero-pad the array `A` to twice the size with the same element type as `A`.
+# """
+# function zeropad(A)
+#     S = size(A)
+#     if any(isodd.(S))
+#         error("Array dims not divisible by 2")
+#     end
+#     nO = 2 .* S
+#     nI = S .÷ 2
+
+#     outer = []
+#     inner = []
+
+#     for no in nO
+#         push!(outer,(1:no))
+#     end
 
-    outer = []
-    inner = []
+#     for ni in nI
+#         push!(inner,(ni+1:ni+2*ni))
+#     end
 
-    for no in nO
-        push!(outer,(1:no))
-    end
+#     return PaddedView(zero(eltype(A)),A,Tuple(outer),Tuple(inner)) |> collect
+# end
 
-    for ni in nI
-        push!(inner,(ni+1:ni+2*ni))
-    end
+"""
+    zeropad(A)
 
-    return PaddedView(zero(eltype(A)),A,Tuple(outer),Tuple(inner)) |> collect
+Zero-pad the array `A` to twice the size with the same element type as `A`, for fft convolution, correlations.
+"""
+function zeropad(A)
+    outer = 2 .* size(A)
+    return PaddedView(zero(eltype(A)),A,Tuple(outer)) |> collect
 end
 
 """
-	log10range(a,b,n)
+log10range(a,b,n)
 
 Create a vector that is linearly spaced in log space, containing `n` values bracketed by `a` and `b`.
 """
@@ -226,11 +235,12 @@ function convolve(ψ1,ψ2,X,K)
 	ϕ1 = zeropad(conj.(ψ1))
     ϕ2 = zeropad(ψ2)
 
-	χ1 = fft(ϕ1)*prod(DX) |> fftshift   
-	χ2 = fft(ϕ2)*prod(DX) |> fftshift
-	return ifft(χ1.*χ2)*prod(DK)*(2*pi)^(n/2) |> fftshift
+	χ1 = fft(ϕ1)*prod(DX)  
+	χ2 = fft(ϕ2)*prod(DX) 
+	return ifft(χ1.*χ2)*prod(DK)*(2*pi)^(n/2)  
 end
 
+
 @doc raw"""
 	auto_correlate(ψ,X,K)
 
@@ -250,8 +260,8 @@ function auto_correlate(ψ,X,K)
     n = length(X)
     DX,DK = fft_differentials(X,K)
     ϕ = zeropad(ψ)
-	χ = fft(ϕ)*prod(DX) |> fftshift
-	return ifft(abs2.(χ))*prod(DK)*(2*pi)^(n/2) |> fftshift
+	χ = fft(ϕ)*prod(DX) 
+	return ifft(abs2.(χ))*prod(DK)*(2*pi)^(n/2) 
 end
 
 auto_correlate(psi::Psi{D}) where D = auto_correlate(psi.ψ,psi.X,psi.K)
@@ -276,9 +286,9 @@ function cross_correlate(ψ1,ψ2,X,K)
     DX,DK = fft_differentials(X,K)
     ϕ1 = zeropad(ψ1)
     ϕ2 = zeropad(ψ2)
-	χ1 = fft(ϕ1)*prod(DX) |> fftshift
-    χ2 = fft(ϕ2)*prod(DX) |> fftshift
-	return ifft(conj(χ1).*χ2)*prod(DK)*(2*pi)^(n/2) |> fftshift
+	χ1 = fft(ϕ1)*prod(DX)  
+    χ2 = fft(ϕ2)*prod(DX)  
+	return ifft(conj(χ1).*χ2)*prod(DK)*(2*pi)^(n/2)  
 end
 cross_correlate(psi1::Psi{D},psi2::Psi{D}) where D = cross_correlate(psi1.ψ,psi2.ψ,psi1.X,psi1.K)
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 84788b7..346da39 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -36,7 +36,7 @@ using Test
     
     Natoms = sum(abs2.(ψ))*dx^2
     AC = auto_correlate(ψ,X,K)
-    Nr,Nim = AC[n+1,n+1] .|> (real,imag) 
+    Nr,Nim = AC[1,1] .|> (real,imag) 
 
     @test Natoms ≈ Nr
     @test Nim ≈ 0.0 
@@ -84,7 +84,7 @@ end
     
     Natoms = sum(abs2.(ψ))*dx^3
     AC = auto_correlate(ψ,X,K)
-    Nr,Nim = AC[n+1,n+1,n+1] .|> (real,imag) 
+    Nr,Nim = AC[1,1,1] .|> (real,imag) 
 
     @test Natoms ≈ Nr
     @test Nim ≈ 0.0