From dcaf8f6338a35f0f0eb512dd5d5f323178dbc179 Mon Sep 17 00:00:00 2001 From: Sheehan Olver Date: Sun, 12 Nov 2023 21:45:18 +0000 Subject: [PATCH] Support higher order tensors in ToeplitzHankel (#233) * Support higher order tensors in ToeplitzHankel * Update toeplitzhankel.jl * Generalise ToeplitzPlan * generalis plan_upper for matrices * consolidate plan_uppertoeplitz * 4-tensor tests * Update toeplitzplans.jl * generalise _th_applymul! * Tesnor leg2cheb * tests pass * cheb2leg tensor support * fixes * v0.15.12 * Update toeplitzplans.jl --- Project.toml | 2 +- src/FastTransforms.jl | 2 +- src/toeplitzhankel.jl | 136 +++++++-------------- src/toeplitzplans.jl | 231 ++++++++++-------------------------- test/toeplitzhankeltests.jl | 33 ++++++ test/toeplitzplanstests.jl | 91 +++++++++++--- 6 files changed, 215 insertions(+), 280 deletions(-) diff --git a/Project.toml b/Project.toml index 3de40b67..7b29f896 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FastTransforms" uuid = "057dd010-8810-581a-b7be-e3fc3b93f78c" -version = "0.15.11" +version = "0.15.12" [deps] AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" diff --git a/src/FastTransforms.jl b/src/FastTransforms.jl index f9d99e74..e9bf6fac 100644 --- a/src/FastTransforms.jl +++ b/src/FastTransforms.jl @@ -8,7 +8,7 @@ using FastGaussQuadrature, FillArrays, LinearAlgebra, @reexport using GenericFFT import Base: convert, unsafe_convert, eltype, ndims, adjoint, transpose, show, - *, \, inv, length, size, view, getindex + *, \, inv, length, size, view, getindex, tail, OneTo import Base.GMP: Limb diff --git a/src/toeplitzhankel.jl b/src/toeplitzhankel.jl index 4e0ab3db..fef934f6 100644 --- a/src/toeplitzhankel.jl +++ b/src/toeplitzhankel.jl @@ -16,75 +16,42 @@ so that `L[:,k] = DL*C[:,k]` and `R[:,k] = DR*C[:,k]`. This allows a Cholesky decomposition in 𝒪(K²N) operations and 𝒪(KN) storage, K = log N log ɛ⁻¹. The tuple storage allows plans applied to each dimension. """ -struct ToeplitzHankelPlan{S, N, M, N1, TP<:ToeplitzPlan{S,N1}} <: Plan{S} - T::NTuple{M,TP} - L::NTuple{M,Matrix{S}} - R::NTuple{M,Matrix{S}} - tmp::Array{S,N1} - dims::NTuple{M,Int} - function ToeplitzHankelPlan{S,N,M,N1,TP}(T::NTuple{M,TP}, L, R, dims) where {S,TP,N,N1,M} +struct ToeplitzHankelPlan{S, N, N1, LowR, TP, Dims} <: Plan{S} + T::TP # A length M Vector or Tuple of ToeplitzPlan + L::LowR # A length M Vector or Tuple of Matrices storing low rank factors of L + R::LowR # A length M Vector or Tuple of Matrices storing low rank factors of R + tmp::Array{S,N1} # A larger dimensional array to transform each scaled array all-at-once + dims::Dims # A length M Vector or Tuple of Int storing the dimensions acted on + function ToeplitzHankelPlan{S,N,N1,LowR,TP,Dims}(T::TP, L::LowR, R::LowR, dims) where {S,N,N1,LowR,TP,Dims} tmp = Array{S}(undef, max.(size.(T)...)...) - new{S,N,M,N1,TP}(T, L, R, tmp, dims) + new{S,N,N1,LowR,TP,Dims}(T, L, R, tmp, dims) end - ToeplitzHankelPlan{S,N,M,N1,TP}(T::NTuple{M,TP}, L, R, dims::Int) where {S,TP,N,N1,M} = - ToeplitzHankelPlan{S,N,M,N1,TP}(T, L, R, (dims,)) end -ToeplitzHankelPlan(T::ToeplitzPlan{S,2}, L::Matrix, R::Matrix, dims=1) where S = - ToeplitzHankelPlan{S, 1, 1, 2, typeof(T)}((T,), (L,), (R,), dims) -ToeplitzHankelPlan(T::ToeplitzPlan{S,3}, L::Matrix, R::Matrix, dims) where S = - ToeplitzHankelPlan{S, 2, 1,3, typeof(T)}((T,), (L,), (R,), dims) +ToeplitzHankelPlan{S,N,M}(T::TP, L::LowR, R::LowR, dims::Dims) where {S,N,M,LowR,TP,Dims} = ToeplitzHankelPlan{S,N,M,LowR,TP,Dims}(T, L, R, dims) +ToeplitzHankelPlan{S,N}(T, L, R, dims) where {S,N} = ToeplitzHankelPlan{S,N,N+1}(T, L, R, dims) +ToeplitzHankelPlan(T::ToeplitzPlan{S,M}, L::Matrix, R::Matrix, dims=1) where {S,M} = ToeplitzHankelPlan{S,M-1,M}((T,), (L,), (R,), dims) -ToeplitzHankelPlan(T::NTuple{2,TP}, L::Tuple, R::Tuple, dims) where {S,TP<:ToeplitzPlan{S,3}} = - ToeplitzHankelPlan{S, 2,2,3, TP}(T, L, R, dims) - -function *(P::ToeplitzHankelPlan{<:Any,1}, v::AbstractVector) - (R,),(L,),(T,),tmp = P.R,P.L,P.T,P.tmp - tmp .= R .* v - T * tmp - tmp .= L .* tmp - sum!(v, tmp) -end - -function _th_applymul1!(v, T, L, R, tmp) - N = size(R,2) - m,n = size(v) - tmp[1:m,1:n,1:N] .= reshape(R,size(R,1),1,N) .* v - T * view(tmp,1:m,1:n,1:N) - view(tmp,1:m,1:n,1:N) .*= reshape(L,size(L,1),1,N) - sum!(v, view(tmp,1:m,1:n,1:N)) -end - -function _th_applymul2!(v, T, L, R, tmp) - N = size(R,2) - m,n = size(v) - tmp[1:m,1:n,1:N] .= reshape(R,1,size(R,1),N) .* v - T * view(tmp,1:m,1:n,1:N) - view(tmp,1:m,1:n,1:N) .*= reshape(L,1,size(L,1),N) - sum!(v, view(tmp,1:m,1:n,1:N)) +_reshape_broadcast(d, R, ::Val{N}, M) where N = reshape(R,ntuple(k -> k == d ? size(R,1) : 1, Val(N))...,M) +function _th_applymul!(d, v::AbstractArray{<:Any,N}, T, L, R, tmp) where N + M = size(R,2) + ax = (axes(v)..., OneTo(M)) + tmp[ax...] .= _reshape_broadcast(d, R, Val(N), M) .* v + T * view(tmp, ax...) + view(tmp,ax...) .*= _reshape_broadcast(d, L, Val(N), M) + sum!(v, view(tmp,ax...)) end -function *(P::ToeplitzHankelPlan{<:Any,2,1}, v::AbstractMatrix) - (R,),(L,),(T,),tmp = P.R,P.L,P.T,P.tmp - if P.dims == (1,) - _th_applymul1!(v, T, L, R, tmp) - else - _th_applymul2!(v, T, L, R, tmp) +function *(P::ToeplitzHankelPlan{<:Any,N}, v::AbstractArray{<:Any,N}) where N + for (R,L,T,d) in zip(P.R,P.L,P.T,P.dims) + _th_applymul!(d, v, T, L, R, P.tmp) end v end -function *(P::ToeplitzHankelPlan{<:Any,2,2}, v::AbstractMatrix) - (R1,R2),(L1,L2),(T1,T2),tmp = P.R,P.L,P.T,P.tmp - - _th_applymul1!(v, T1, L1, R1, tmp) - _th_applymul2!(v, T2, L2, R2, tmp) - - v -end # partial cholesky for a Hankel matrix @@ -166,9 +133,9 @@ function *(P::ChebyshevToLegendrePlanTH, v::AbstractVector{S}) where S v end -function _cheb2leg_rescale1!(V::AbstractMatrix{S}) where S - m,n = size(V) - for j = 1:n +function _cheb2leg_rescale1!(V::AbstractArray{S}) where S + m = size(V,1) + for j = CartesianIndices(tail(axes(V))) ret = zero(S) @inbounds for k = 1:2:m ret += -V[k,j]/(k*(k-2)) @@ -178,24 +145,15 @@ function _cheb2leg_rescale1!(V::AbstractMatrix{S}) where S V end +_dropfirstdim(d::Int) = () +_dropfirstdim(d::Int, m, szs...) = ((d == 1 ? 2 : 1):m, _dropfirstdim(d-1, szs...)...) -function *(P::ChebyshevToLegendrePlanTH, V::AbstractMatrix) +function *(P::ChebyshevToLegendrePlanTH, V::AbstractArray{<:Any,N}) where N m,n = size(V) - dims = P.toeplitzhankel.dims - if dims == (1,) - _cheb2leg_rescale1!(V) - P.toeplitzhankel*view(V,2:m,:) - elseif dims == (2,) - _cheb2leg_rescale1!(transpose(V)) - P.toeplitzhankel*view(V,:,2:n) - else - @assert dims == (1,2) - (R1,R2),(L1,L2),(T1,T2),tmp = P.toeplitzhankel.R,P.toeplitzhankel.L,P.toeplitzhankel.T,P.toeplitzhankel.tmp - - _cheb2leg_rescale1!(V) - _th_applymul1!(view(V,2:m,:), T1, L1, R1, tmp) - _cheb2leg_rescale1!(transpose(V)) - _th_applymul2!(view(V,:,2:n), T2, L2, R2, tmp) + tmp = P.toeplitzhankel.tmp + for (d,R,L,T) in zip(P.toeplitzhankel.dims,P.toeplitzhankel.R,P.toeplitzhankel.L,P.toeplitzhankel.T) + _cheb2leg_rescale1!(PermutedDimsArray(V, _permfirst(d, N))) + _th_applymul!(d, view(V, _dropfirstdim(d, size(V)...)...), T, L, R, tmp) end V end @@ -226,18 +184,14 @@ function _leg2chebuTH_TLC(::Type{S}, mn, d) where {S} (T, (1:n) .* C, C) end - for f in (:leg2cheb, :leg2chebu) plan = Symbol("plan_th_", f, "!") TLC = Symbol("_", f, "TH_TLC") @eval begin - $plan(::Type{S}, mn::Tuple, dims::Int) where {S} = ToeplitzHankelPlan($TLC(S, mn, dims)..., dims) - - function $plan(::Type{S}, mn::NTuple{2,Int}, dims::NTuple{2,Int}) where {S} - @assert dims == (1,2) - T1,L1,C1 = $TLC(S, mn, 1) - T2,L2,C2 = $TLC(S, mn, 2) - ToeplitzHankelPlan((T1,T2), (L1,L2), (C1,C2), dims) + $plan(::Type{S}, mn::NTuple{N,Int}, dims::Int) where {S,N} = ToeplitzHankelPlan($TLC(S, mn, dims)..., dims) + function $plan(::Type{S}, mn::NTuple{N,Int}, dims) where {S,N} + TLCs = $TLC.(S, Ref(mn), dims) + ToeplitzHankelPlan{S,N}(map(first, TLCs), map(TLC -> TLC[2], TLCs), map(last, TLCs), dims) end end end @@ -265,13 +219,11 @@ function _cheb2legTH_TLC(::Type{S}, mn, d) where S T, DL .* C, DR .* C end -plan_th_cheb2leg!(::Type{S}, mn::Tuple, dims::Int) where {S} = ChebyshevToLegendrePlanTH(ToeplitzHankelPlan(_cheb2legTH_TLC(S, mn, dims)..., dims)) +plan_th_cheb2leg!(::Type{S}, mn::NTuple{N,Int}, dims::Int) where {S,N} = ChebyshevToLegendrePlanTH(ToeplitzHankelPlan(_cheb2legTH_TLC(S, mn, dims)..., dims)) -function plan_th_cheb2leg!(::Type{S}, mn::NTuple{2,Int}, dims::NTuple{2,Int}) where {S} - @assert dims == (1,2) - T1,L1,C1 = _cheb2legTH_TLC(S, mn, 1) - T2,L2,C2 = _cheb2legTH_TLC(S, mn, 2) - ChebyshevToLegendrePlanTH(ToeplitzHankelPlan((T1,T2), (L1,L2), (C1,C2), dims)) +function plan_th_cheb2leg!(::Type{S}, mn::NTuple{N,Int}, dims) where {S,N} + TLCs = _cheb2legTH_TLC.(S, Ref(mn), dims) + ChebyshevToLegendrePlanTH(ToeplitzHankelPlan{S,N}(map(first, TLCs), map(TLC -> TLC[2], TLCs), map(last, TLCs), dims)) end @@ -337,7 +289,7 @@ _good_plan_th_ultra2ultra!(::Type{S}, mn, λ₁, λ₂, dims::Int) where S = Toe function _good_plan_th_ultra2ultra!(::Type{S}, mn::NTuple{2,Int}, λ₁, λ₂, dims::NTuple{2,Int}) where S T1,L1,C1 = _ultra2ultraTH_TLC(S, mn, λ₁, λ₂, 1) T2,L2,C2 = _ultra2ultraTH_TLC(S, mn, λ₁, λ₂, 2) - ToeplitzHankelPlan((T1,T2), (L1,L2), (C1,C2), dims) + ToeplitzHankelPlan{S,2}((T1,T2), (L1,L2), (C1,C2), dims) end @@ -515,7 +467,7 @@ _good_plan_th_jac2jac!(::Type{S}, mn, α, β, γ, δ, dims::Int) where S = Toepl function _good_plan_th_jac2jac!(::Type{S}, mn::NTuple{2,Int}, α, β, γ, δ, dims::NTuple{2,Int}) where S T1,L1,C1 = _jac2jacTH_TLC(S, mn, α, β, γ, δ, 1) T2,L2,C2 = _jac2jacTH_TLC(S, mn, α, β, γ, δ, 2) - ToeplitzHankelPlan((T1,T2), (L1,L2), (C1,C2), dims) + ToeplitzHankelPlan{S,2}((T1,T2), (L1,L2), (C1,C2), dims) end @@ -685,10 +637,8 @@ end for f in (:th_leg2cheb, :th_cheb2leg, :th_leg2chebu) plan = Symbol("plan_", f, "!") @eval begin - $plan(::Type{S}, mn::NTuple{N,Int}, dims::UnitRange) where {N,S} = $plan(S, mn, tuple(dims...)) - $plan(::Type{S}, mn::Tuple{Int}, dims::Tuple{Int}=(1,)) where {S} = $plan(S, mn, dims...) - $plan(::Type{S}, (m,n)::NTuple{2,Int}) where {S} = $plan(S, (m,n), (1,2)) $plan(arr::AbstractArray{T}, dims...) where T = $plan(T, size(arr), dims...) + $plan(::Type{S}, mn::NTuple{N,Int}) where {S,N} = $plan(S, mn, ntuple(identity,Val(N))) $f(v, dims...) = $plan(eltype(v), size(v), dims...)*copy(v) end end diff --git a/src/toeplitzplans.jl b/src/toeplitzplans.jl index ab08d1f7..42d24062 100644 --- a/src/toeplitzplans.jl +++ b/src/toeplitzplans.jl @@ -1,39 +1,46 @@ using FFTW import FFTW: plan_r2r! -struct ToeplitzPlan{T, N, M, S, VECS<:Tuple{Vararg{Vector{S}}}, P<:Plan{S}, Pi<:Plan{S}} <: Plan{T} - vectors::VECS + +""" + ToeplitzPlan + +applies Toeplitz matrices fast along each dimension. +""" + +struct ToeplitzPlan{T, N, Dims, S, VECS, P<:Plan{S}, Pi<:Plan{S}} <: Plan{T} + vectors::VECS # Vector or Tuple of storage tmp::Array{S,N} dft::P idft::Pi - dims::NTuple{M,Int} + dims::Dims end -ToeplitzPlan{T}(v::AbstractVector, tmp, dft, idft, dims) where T = ToeplitzPlan{T}((v,), tmp, dft, idft, dims) -ToeplitzPlan{T}(v::Tuple{Vararg{Vector{S}}}, tmp::Array{S,N}, dft::Plan{S}, idft::Plan{S}, dims::NTuple{M,Int}) where {T,S,N,M} = ToeplitzPlan{T,N,M,S,typeof(v),typeof(dft), typeof(idft)}(v, tmp, dft, idft, dims) -ToeplitzPlan{T}(v::Tuple{Vararg{Vector{S}}}, tmp::Array{S,N}, dft::Plan{S}, idft::Plan{S}, dims::Int) where {T,S,N} = ToeplitzPlan{T}(v, tmp, dft, idft, (dims,)) +ToeplitzPlan{T}(v, tmp::Array{S,N}, dft::Plan{S}, idft::Plan{S}, dims) where {T,S,N} = ToeplitzPlan{T,N,typeof(dims),S,typeof(v),typeof(dft), typeof(idft)}(v, tmp, dft, idft, dims) + -size(A::ToeplitzPlan{<:Any,1}) = ((length(A.tmp)+1) ÷ 2,) -function size(A::ToeplitzPlan{<:Any,2,1}) - if A.dims == (1,) - ((size(A.tmp,1)+1) ÷ 2, size(A.tmp,2)) - else # A.dims == (2,) - (size(A.tmp,1), (size(A.tmp,2)+1) ÷ 2) +divdimby2(d::Int, sz1, szs...) = isone(d) ? ((sz1 + 1) ÷ 2, szs...) : (sz1, divdimby2(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 + for d in dims + ret = divdimby2(d, ret...) end + ret end -function size(A::ToeplitzPlan{<:Any,3,1}) - if A.dims == (1,) - ((size(A.tmp,1)+1) ÷ 2, size(A.tmp,2), size(A.tmp,3)) - elseif A.dims == (2,) - (size(A.tmp,1), (size(A.tmp,2)+1) ÷ 2, size(A.tmp,3)) - else - (size(A.tmp,1), size(A.tmp,2), (size(A.tmp,3)+1) ÷ 2) +function to_toeplitzplan_size(dims, szs) + ret = szs + for d in dims + ret = muldimby2(d, ret...) end + ret end -size(A::ToeplitzPlan{<:Any,2,2}) = ((size(A.tmp,1)+1) ÷ 2, (size(A.tmp,2)+1) ÷ 2) +size(A::ToeplitzPlan) = toeplitzplan_size(A.dims, size(A.tmp)) + # based on ToeplitzMatrices.jl """ @@ -44,101 +51,31 @@ Return real-valued part of `x` if `T` is a type of a real number, and `x` otherw maybereal(::Type, x) = x maybereal(::Type{<:Real}, x) = real(x) -function *(A::ToeplitzPlan{T,1}, x::AbstractVector{T}) where T - vc,tmp,dft,idft = A.vectors[1],A.tmp, A.dft,A.idft - S = eltype(tmp) - N = length(tmp) - n = length(x) - if 2n-1 ≠ N - throw(DimensionMismatch("Toeplitz plan does not match size of input")) - end - copyto!(view(tmp, 1:n), x) - fill!(view(tmp, n+1:N), zero(S)) - dft * tmp - tmp .*= vc - idft * tmp - @inbounds for k = 1:n - x[k] = maybereal(T, tmp[k]) - end - x -end +function *(A::ToeplitzPlan{T,N}, X::AbstractArray{T,N}) where {T,N} + vcs,Y,dft,idft,dims = A.vectors,A.tmp, A.dft,A.idft,A.dims -function *(A::ToeplitzPlan{T,2,1, S}, x::AbstractMatrix{T}) where {T,S} - vc,tmp,dft,idft = A.vectors[1],A.tmp, A.dft, A.idft - M,N = size(tmp) - m,n = size(x) + isempty(X) && return X - if isempty(x) - return x - end + fill!(Y, zero(eltype(Y))) + copyto!(view(Y, axes(X)...), X) - if A.dims == (1,) - copyto!(view(tmp, 1:m, :), x) - fill!(view(tmp, m+1:M, :), zero(S)) - if !isempty(tmp) - dft * tmp - end - tmp .= vc .* tmp - else - @assert A.dims == (2,) - copyto!(view(tmp, :, 1:n), x) - fill!(view(tmp, :, n+1:N), zero(S)) - dft * tmp - tmp .= tmp .* transpose(vc) + # Fourier transform each dimension + dft * Y + + # Multiply by a diagonal matrix along each dimension by permuting + # to first dimension + for (vc,d) in zip(vcs,dims) + Ỹ = PermutedDimsArray(Y, _permfirst(d, N)) + Ỹ .= vc .* Ỹ end - idft * tmp - x .= maybereal.(T, view(tmp,1:m,1:n)) -end + # Transform back + idft * Y -function *(A::ToeplitzPlan{T,2,2, S}, X::AbstractMatrix{T}) where {T,S} - vcs,tmp,dft,idft = A.vectors,A.tmp, A.dft,A.idft - vc1,vc2 = vcs - M,N = size(tmp) - m,n = size(X) - - @assert A.dims == (1,2) - copyto!(view(tmp, 1:m, 1:n), X) - fill!(view(tmp, m+1:M, :), zero(S)) - fill!(view(tmp, 1:m, n+1:N), zero(S)) - dft * tmp - tmp .= vc1 .* tmp .* transpose(vc2) - idft * tmp - @inbounds for k = 1:m, j = 1:n - X[k,j] = maybereal(T, tmp[k,j]) - end + X .= maybereal.(T, view(Y, axes(X)...)) X end -function *(A::ToeplitzPlan{T,3,1, S}, x::AbstractArray{T,3}) where {T,S} - vc,tmp,dft,idft = A.vectors[1],A.tmp, A.dft,A.idft - M,N,L = size(tmp) - m,n,l = size(x) - - if A.dims == (1,) - copyto!(view(tmp, 1:m, :, :), x) - fill!(view(tmp, m+1:M, :, :), zero(S)) - dft * tmp - tmp .= vc .* tmp - elseif A.dims == (2,) - copyto!(view(tmp, :, 1:n, :), x) - fill!(view(tmp, :, n+1:N, :), zero(S)) - dft * tmp - tmp .= tmp .* transpose(vc) - else - copyto!(view(tmp, :, :, 1:l), x) - fill!(view(tmp, :, :, l+1:L), zero(S)) - dft * tmp - tmp .= tmp .* reshape(vc, 1, 1, L) - end - idft * tmp - @inbounds for k = 1:m, j = 1:n, ℓ = 1:l - x[k,j,ℓ] = maybereal(T, tmp[k,j,ℓ]) - end - x -end - - function uppertoeplitz_padvec(v::AbstractVector{T}) where T n = length(v) @@ -151,73 +88,25 @@ function uppertoeplitz_padvec(v::AbstractVector{T}) where T tmp end -function plan_uppertoeplitz!(v::AbstractVector{T}) where T - tmp = uppertoeplitz_padvec(v) - dft = plan_fft!(tmp) - idft = plan_ifft!(similar(tmp)) - return ToeplitzPlan{float(T)}(dft * tmp, similar(tmp), dft, idft, (1,)) -end +safe_fft!(A) = isempty(A) ? A : fft!(A) -# TODO: support different transforms -# function plan_uppertoeplitz!(v1::AbstractVector{T}, v2::AbstractVector{T}) where T -# S = float(T) -# m,n = length(v1), length(v2) -# tmp = zeros(S, 2m-1, 2n-1) -# pv1 = uppertoeplitz_padvec(v1) -# pv2 = uppertoeplitz_padvec(v2) -# dft = plan_r2r!(tmp, FFTW.R2HC) -# return ToeplitzPlan((r2r!(pv1, FFTW.R2HC), r2r!(pv2, FFTW.R2HC)), tmp, dft, 1:2) -# end - -function plan_uppertoeplitz!(v::AbstractVector{T}, szs::NTuple{2,Int}, dim::Int) where T - S = complex(float(T)) - m,n = szs - if isone(dim) - tmp = zeros(S, max(0,2m-1), n) - pv = uppertoeplitz_padvec(v[1:m]) - else # dim == 2 - tmp = zeros(S, m, max(0,2n-1)) - pv = uppertoeplitz_padvec(v[1:n]) - end - if isempty(tmp) - # dummy plans just to create type - dft = plan_fft!(similar(tmp, 1, 1), dim) - idft = plan_ifft!(similar(tmp, 1, 1), dim) - ToeplitzPlan{float(T)}(pv, tmp, dft, idft, dim) - else - dft = plan_fft!(tmp, dim) - idft = plan_ifft!(similar(tmp), dim) - return ToeplitzPlan{float(T)}(fft!(pv), tmp, dft, idft, dim) - end -end +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) = (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]])),) -function plan_uppertoeplitz!(v::AbstractVector{T}, szs::NTuple{3,Int}, dim::Int) where T - S = complex(float(T)) - m,n,l = szs - if isone(dim) - tmp = zeros(S, 2m-1, n, l) - pv = uppertoeplitz_padvec(v[1:m]) - elseif dim == 2 - tmp = zeros(S, m, 2n-1, l) - pv = uppertoeplitz_padvec(v[1:n]) - else - @assert dim == 3 - tmp = zeros(S, m, n, 2l-1) - pv = uppertoeplitz_padvec(v[1:l]) - end - dft = plan_fft!(tmp, dim) - idft = plan_ifft!(similar(tmp), dim) - return ToeplitzPlan{float(T)}(fft!(pv), tmp, dft, idft, dim) -end -function plan_uppertoeplitz!(v::AbstractVector{T}, szs::NTuple{2,Int}, dim=(1,2)) where T - @assert dim == (1,2) +# 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)) - m,n = szs - tmp = zeros(S, 2m-1, 2n-1) - pv1 = uppertoeplitz_padvec(v[1:m]) - pv2 = uppertoeplitz_padvec(v[1:n]) - dft = plan_fft!(tmp, dim) - idft = plan_ifft!(similar(tmp), dim) - return ToeplitzPlan{float(T)}((fft!(pv1), fft!(pv2)), tmp, dft, idft, dim) + + tmp = zeros(S, to_toeplitzplan_size(dim, szs)...) + 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 + +plan_uppertoeplitz!(v::AbstractVector{T}) where T = plan_uppertoeplitz!(v, size(v)) diff --git a/test/toeplitzhankeltests.jl b/test/toeplitzhankeltests.jl index 19274747..b90ce94b 100644 --- a/test/toeplitzhankeltests.jl +++ b/test/toeplitzhankeltests.jl @@ -130,4 +130,37 @@ Random.seed!(0) @test norm(v - th_cheb2leg(th_leg2cheb(v)), Inf) ≤ 1E-13 @test norm(v - th_cheb2leg(th_leg2cheb(v)))/norm(v) ≤ 1E-14 end + + @testset "tensor" begin + X = randn(5,4,3) + for trans in (th_leg2cheb, th_cheb2leg) + Y = trans(X, 1) + for ℓ = 1:size(X,3) + @test Y[:,:,ℓ] ≈ trans(X[:,:,ℓ],1) + end + Y = trans(X, 2) + for ℓ = 1:size(X,3) + @test Y[:,:,ℓ] ≈ trans(X[:,:,ℓ],2) + end + Y = trans(X, 3) + for j = 1:size(X,2) + @test Y[:,j,:] ≈ trans(X[:,j,:],2) + end + + Y = trans(X, (1,3)) + for j = 1:size(X,2) + @test Y[:,j,:] ≈ trans(X[:,j,:]) + end + + Y = trans(X, 1:3) + M = copy(X) + for j = 1:size(X,3) + M[:,:,j] = trans(M[:,:,j]) + end + for k = 1:size(X,1), j=1:size(X,2) + M[k,j,:] = trans(M[k,j,:]) + end + @test M ≈ Y + end + end end \ No newline at end of file diff --git a/test/toeplitzplanstests.jl b/test/toeplitzplanstests.jl index e56d8c3e..6ea6a095 100644 --- a/test/toeplitzplanstests.jl +++ b/test/toeplitzplanstests.jl @@ -34,23 +34,86 @@ import FastTransforms: plan_uppertoeplitz! @testset "Tensor" begin T = [1 2 3; 0 1 2; 0 0 1] - X = randn(3,3,3) - P = plan_uppertoeplitz!([1,2,3], size(X), 1) - PX = P * copy(X) - for ℓ = 1:size(X,3) - @test PX[:,:,ℓ] ≈ T*X[:,:,ℓ] - end + @testset "3D" begin + X = randn(3,3,3) + P = plan_uppertoeplitz!([1,2,3], size(X), 1) + PX = P * copy(X) + for ℓ = 1:size(X,3) + @test PX[:,:,ℓ] ≈ T*X[:,:,ℓ] + end - P = plan_uppertoeplitz!([1,2,3], size(X), 2) - PX = P * copy(X) - for ℓ = 1:size(X,3) - @test PX[:,:,ℓ] ≈ X[:,:,ℓ]*T' + P = plan_uppertoeplitz!([1,2,3], size(X), 2) + PX = P * copy(X) + for ℓ = 1:size(X,3) + @test PX[:,:,ℓ] ≈ X[:,:,ℓ]*T' + end + + P = plan_uppertoeplitz!([1,2,3], size(X), 3) + PX = P * copy(X) + for j = 1:size(X,2) + @test PX[:,j,:] ≈ X[:,j,:]*T' + end + + P = plan_uppertoeplitz!([1,2,3], size(X), (1,3)) + PX = P * copy(X) + for j = 1:size(X,2) + @test PX[:,j,:] ≈ T*X[:,j,:]*T' + end + + P = plan_uppertoeplitz!([1,2,3], size(X), 1:3) + PX = P * copy(X) + M = copy(X) + for j = 1:size(X,3) + M[:,:,j] = T*M[:,:,j]*T' + end + for k = 1:size(X,1) + M[k,:,:] = M[k,:,:]*T' + end + @test M ≈ PX end - P = plan_uppertoeplitz!([1,2,3], size(X), 3) - PX = P * copy(X) - for j = 1:size(X,2) - @test PX[:,j,:] ≈ X[:,j,:]*T' + @testset "4D" begin + X = randn(3,3,3,3) + P = plan_uppertoeplitz!([1,2,3], size(X), 1) + PX = P * copy(X) + for ℓ = 1:size(X,3), m = 1:size(X,4) + @test PX[:,:,ℓ,m] ≈ T*X[:,:,ℓ,m] + end + + P = plan_uppertoeplitz!([1,2,3], size(X), 2) + PX = P * copy(X) + for ℓ = 1:size(X,3), m = 1:size(X,4) + @test PX[:,:,ℓ,m] ≈ X[:,:,ℓ,m]*T' + end + + P = plan_uppertoeplitz!([1,2,3], size(X), 3) + PX = P * copy(X) + for j = 1:size(X,2), m = 1:size(X,4) + @test PX[:,j,:,m] ≈ X[:,j,:,m]*T' + end + + P = plan_uppertoeplitz!([1,2,3], size(X), 4) + PX = P * copy(X) + for k = 1:size(X,1), j = 1:size(X,2) + @test PX[k,j,:,:] ≈ X[k,j,:,:]*T' + end + + P = plan_uppertoeplitz!([1,2,3], size(X), (1,3)) + PX = P * copy(X) + for j = 1:size(X,2), m=1:size(X,4) + @test PX[:,j,:,m] ≈ T*X[:,j,:,m]*T' + end + + P = plan_uppertoeplitz!([1,2,3], size(X), 1:4) + PX = P * copy(X) + M = copy(X) + for ℓ = 1:size(X,3), m = 1:size(X,4) + M[:,:,ℓ,m] = T*M[:,:,ℓ,m]*T' + end + for k = 1:size(X,1), j = 1:size(X,2) + M[k,j,:,:] = T*M[k,j,:,:]*T' + end + @test M ≈ PX end end