Skip to content

Commit

Permalink
simplify convolve
Browse files Browse the repository at this point in the history
  • Loading branch information
AshtonSBradley committed May 20, 2024
1 parent 3d484b4 commit 3ee6ad6
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 33 deletions.
72 changes: 41 additions & 31 deletions src/analysis.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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`.
"""
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)

Expand Down
4 changes: 2 additions & 2 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 3ee6ad6

Please sign in to comment.