Skip to content

Commit

Permalink
Update toeplitzplans.jl
Browse files Browse the repository at this point in the history
  • Loading branch information
dlfivefifty committed Nov 11, 2023
1 parent 3906bc0 commit a520c1c
Showing 1 changed file with 12 additions and 6 deletions.
18 changes: 12 additions & 6 deletions src/toeplitzplans.jl
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ ToeplitzPlan{T}(v, tmp::Array{S,N}, dft::Plan{S}, idft::Plan{S}, dims) where {T,


divdimby2(d::Int, sz1, szs...) = isone(d) ? ((sz1 + 1) ÷ 2, szs...) : (sz1, divdimby2(d-1, szs...)...)
muldimby2(d::Int, sz1, szs...) = isone(d) ? (2sz1 - 1, szs...) : (sz1, muldimby2(d-1, szs...)...)
muldimby2(d::Int, sz1, szs...) = isone(d) ? (max(0,2sz1 - 1), szs...) : (sz1, muldimby2(d-1, szs...)...)

function toeplitzplan_size(dims, szs)
ret = szs
Expand Down Expand Up @@ -89,17 +89,23 @@ function uppertoeplitz_padvec(v::AbstractVector{T}) where T
tmp
end

uppertoeplitz_vecs(v, dims::AbstractVector, szs) = [fft!(uppertoeplitz_padvec(v[1:szs[d]])) for d in dims]
safe_fft!(A) = isempty(A) ? A : fft!(A)

uppertoeplitz_vecs(v, dims::AbstractVector, szs) = [safe_fft!(uppertoeplitz_padvec(v[1:szs[d]])) for d in dims]
uppertoeplitz_vecs(v, dims::Tuple{}, szs) = ()
uppertoeplitz_vecs(v, dims::Tuple, szs) = (fft!(uppertoeplitz_padvec(v[1:szs[first(dims)]])), uppertoeplitz_vecs(v, tail(dims), szs)...)
uppertoeplitz_vecs(v, d::Int, szs) = (fft!(uppertoeplitz_padvec(v[1:szs[d]])),)
uppertoeplitz_vecs(v, dims::Tuple, szs) = (safe_fft!(uppertoeplitz_padvec(v[1:szs[first(dims)]])), uppertoeplitz_vecs(v, tail(dims), szs)...)
uppertoeplitz_vecs(v, d::Int, szs) = (safe_fft!(uppertoeplitz_padvec(v[1:szs[d]])),)


# allow FFT to work by making sure tmp is non-empty
safe_tmp(tmp::AbstractArray{<:Any,N}) where N = isempty(tmp) ? similar(tmp, ntuple(_ -> 1, Val(N))...) : tmp

function plan_uppertoeplitz!(v::AbstractVector{T}, szs::NTuple{N,Int}, dim=ntuple(identity,Val(N))) where {T,N}
S = complex(float(T))

tmp = zeros(S, to_toeplitzplan_size(dim, szs)...)
dft = plan_fft!(tmp, dim)
idft = plan_ifft!(similar(tmp), dim)
dft = plan_fft!(safe_tmp(tmp), dim)
idft = plan_ifft!(safe_tmp(similar(tmp)), dim)

return ToeplitzPlan{float(T)}(uppertoeplitz_vecs(v, dim, szs), tmp, dft, idft, dim)
end
Expand Down

0 comments on commit a520c1c

Please sign in to comment.