Skip to content

Commit

Permalink
Support higher order tensors in ToeplitzHankel (#233)
Browse files Browse the repository at this point in the history 
* 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
  • Loading branch information
@dlfivefifty
dlfivefifty authored Nov 12, 2023
1 parent 198d220 commit dcaf8f6
Show file tree
Hide file tree
Showing 6 changed files with 215 additions and 280 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
2 changes: 1 addition & 1 deletion src/FastTransforms.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
136 changes: 43 additions & 93 deletions src/toeplitzhankel.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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))
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand Down
Loading

0 comments on commit dcaf8f6

Please sign in to comment.