Skip to content

Commit e1db1ca

Browse files
authored
4D Chebyshev transform (#235)
1 parent 90d1dae commit e1db1ca

File tree

3 files changed

+113
-156
lines changed

3 files changed

+113
-156
lines changed

Project.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
name = "FastTransforms"
22
uuid = "057dd010-8810-581a-b7be-e3fc3b93f78c"
3-
version = "0.15.12"
3+
version = "0.15.13"
44

55
[deps]
66
AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/chebyshevtransform.jl

Lines changed: 15 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -58,76 +58,22 @@ plan_chebyshevtransform(x::AbstractArray, dims...; kws...) = plan_chebyshevtrans
5858
@inline _plan_mul!(y::AbstractArray{T}, P::Plan{T}, x::AbstractArray) where T = mul!(y, P, convert(Array{T}, x))
5959

6060

61+
for op in (:ldiv, :lmul)
62+
op_dim_begin! = Symbol(string(op) * "_dim_begin!")
63+
op_dim_end! = Symbol(string(op) * "_dim_end!")
64+
op! = Symbol(string(op) * "!")
65+
@eval begin
66+
function $op_dim_begin!(α, d::Number, y::AbstractArray{<:Any,N}) where N
67+
# scale just the d-th dimension by permuting it to the first
68+
= PermutedDimsArray(y, _permfirst(d, N))
69+
$op!(α, view(ỹ, 1, ntuple(_ -> :, Val(N-1))...))
70+
end
6171

62-
ldiv_dim_begin!(α, d::Number, y::AbstractVector) = y[1] /= α
63-
function ldiv_dim_begin!(α, d::Number, y::AbstractMatrix)
64-
if isone(d)
65-
ldiv!(α, @view(y[1,:]))
66-
else
67-
ldiv!(α, @view(y[:,1]))
68-
end
69-
end
70-
function ldiv_dim_begin!(α, d::Number, y::AbstractArray{<:Any,3})
71-
if isone(d)
72-
ldiv!(α, @view(y[1,:,:]))
73-
elseif d == 2
74-
ldiv!(α, @view(y[:,1,:]))
75-
else # d == 3
76-
ldiv!(α, @view(y[:,:,1]))
77-
end
78-
end
79-
80-
ldiv_dim_end!(α, d::Number, y::AbstractVector) = y[end] /= α
81-
function ldiv_dim_end!(α, d::Number, y::AbstractMatrix)
82-
if isone(d)
83-
ldiv!(α, @view(y[end,:]))
84-
else
85-
ldiv!(α, @view(y[:,end]))
86-
end
87-
end
88-
function ldiv_dim_end!(α, d::Number, y::AbstractArray{<:Any,3})
89-
if isone(d)
90-
ldiv!(α, @view(y[end,:,:]))
91-
elseif d == 2
92-
ldiv!(α, @view(y[:,end,:]))
93-
else # d == 3
94-
ldiv!(α, @view(y[:,:,end]))
95-
end
96-
end
97-
98-
lmul_dim_begin!(α, d::Number, y::AbstractVector) = y[1] *= α
99-
function lmul_dim_begin!(α, d::Number, y::AbstractMatrix)
100-
if isone(d)
101-
lmul!(α, @view(y[1,:]))
102-
else
103-
lmul!(α, @view(y[:,1]))
104-
end
105-
end
106-
function lmul_dim_begin!(α, d::Number, y::AbstractArray{<:Any,3})
107-
if isone(d)
108-
lmul!(α, @view(y[1,:,:]))
109-
elseif d == 2
110-
lmul!(α, @view(y[:,1,:]))
111-
else # d == 3
112-
lmul!(α, @view(y[:,:,1]))
113-
end
114-
end
115-
116-
lmul_dim_end!(α, d::Number, y::AbstractVector) = y[end] *= α
117-
function lmul_dim_end!(α, d::Number, y::AbstractMatrix)
118-
if isone(d)
119-
lmul!(α, @view(y[end,:]))
120-
else
121-
lmul!(α, @view(y[:,end]))
122-
end
123-
end
124-
function lmul_dim_end!(α, d::Number, y::AbstractArray{<:Any,3})
125-
if isone(d)
126-
lmul!(α, @view(y[end,:,:]))
127-
elseif d == 2
128-
lmul!(α, @view(y[:,end,:]))
129-
else # d == 3
130-
lmul!(α, @view(y[:,:,end]))
72+
function $op_dim_end!(α, d::Number, y::AbstractArray{<:Any,N}) where N
73+
# scale just the d-th dimension by permuting it to the first
74+
= PermutedDimsArray(y, _permfirst(d, N))
75+
$op!(α, view(ỹ, size(ỹ,1), ntuple(_ -> :, Val(N-1))...))
76+
end
13177
end
13278
end
13379

test/chebyshevtests.jl

Lines changed: 97 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -322,96 +322,107 @@ using FastTransforms, Test
322322
end
323323

324324
@testset "tensor" begin
325-
X = randn(4,5,6)
326-
= similar(X)
327-
@testset "chebyshevtransform" begin
328-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevtransform(X[:,k,j]) end
329-
@test @inferred(chebyshevtransform(X,1)) @inferred(chebyshevtransform!(copy(X),1))
330-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevtransform(X[k,:,j]) end
331-
@test chebyshevtransform(X,2) chebyshevtransform!(copy(X),2)
332-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevtransform(X[k,j,:]) end
333-
@test chebyshevtransform(X,3) chebyshevtransform!(copy(X),3)
334-
335-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevtransform(X[:,k,j],Val(2)) end
336-
@test @inferred(chebyshevtransform(X,Val(2),1)) @inferred(chebyshevtransform!(copy(X),Val(2),1))
337-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevtransform(X[k,:,j],Val(2)) end
338-
@test chebyshevtransform(X,Val(2),2) chebyshevtransform!(copy(X),Val(2),2)
339-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevtransform(X[k,j,:],Val(2)) end
340-
@test chebyshevtransform(X,Val(2),3) chebyshevtransform!(copy(X),Val(2),3)
341-
342-
@test @inferred(chebyshevtransform(X)) @inferred(chebyshevtransform!(copy(X))) chebyshevtransform(chebyshevtransform(chebyshevtransform(X,1),2),3)
343-
@test @inferred(chebyshevtransform(X,Val(2))) @inferred(chebyshevtransform!(copy(X),Val(2))) chebyshevtransform(chebyshevtransform(chebyshevtransform(X,Val(2),1),Val(2),2),Val(2),3)
344-
end
345-
346-
@testset "ichebyshevtransform" begin
347-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevtransform(X[:,k,j]) end
348-
@test @inferred(ichebyshevtransform(X,1)) @inferred(ichebyshevtransform!(copy(X),1))
349-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevtransform(X[k,:,j]) end
350-
@test ichebyshevtransform(X,2) ichebyshevtransform!(copy(X),2)
351-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevtransform(X[k,j,:]) end
352-
@test ichebyshevtransform(X,3) ichebyshevtransform!(copy(X),3)
353-
354-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevtransform(X[:,k,j],Val(2)) end
355-
@test @inferred(ichebyshevtransform(X,Val(2),1)) @inferred(ichebyshevtransform!(copy(X),Val(2),1))
356-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevtransform(X[k,:,j],Val(2)) end
357-
@test ichebyshevtransform(X,Val(2),2) ichebyshevtransform!(copy(X),Val(2),2)
358-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevtransform(X[k,j,:],Val(2)) end
359-
@test ichebyshevtransform(X,Val(2),3) ichebyshevtransform!(copy(X),Val(2),3)
360-
361-
@test @inferred(ichebyshevtransform(X)) @inferred(ichebyshevtransform!(copy(X))) ichebyshevtransform(ichebyshevtransform(ichebyshevtransform(X,1),2),3)
362-
@test @inferred(ichebyshevtransform(X,Val(2))) @inferred(ichebyshevtransform!(copy(X),Val(2))) ichebyshevtransform(ichebyshevtransform(ichebyshevtransform(X,Val(2),1),Val(2),2),Val(2),3)
363-
364-
@test ichebyshevtransform(chebyshevtransform(X)) X
365-
@test chebyshevtransform(ichebyshevtransform(X)) X
366-
end
367-
368-
@testset "chebyshevutransform" begin
369-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevutransform(X[:,k,j]) end
370-
@test @inferred(chebyshevutransform(X,1)) @inferred(chebyshevutransform!(copy(X),1))
371-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevutransform(X[k,:,j]) end
372-
@test chebyshevutransform(X,2) chebyshevutransform!(copy(X),2)
373-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevutransform(X[k,j,:]) end
374-
@test chebyshevutransform(X,3) chebyshevutransform!(copy(X),3)
375-
376-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevutransform(X[:,k,j],Val(2)) end
377-
@test @inferred(chebyshevutransform(X,Val(2),1)) @inferred(chebyshevutransform!(copy(X),Val(2),1))
378-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevutransform(X[k,:,j],Val(2)) end
379-
@test chebyshevutransform(X,Val(2),2) chebyshevutransform!(copy(X),Val(2),2)
380-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevutransform(X[k,j,:],Val(2)) end
381-
@test chebyshevutransform(X,Val(2),3) chebyshevutransform!(copy(X),Val(2),3)
382-
383-
@test @inferred(chebyshevutransform(X)) @inferred(chebyshevutransform!(copy(X))) chebyshevutransform(chebyshevutransform(chebyshevutransform(X,1),2),3)
384-
@test @inferred(chebyshevutransform(X,Val(2))) @inferred(chebyshevutransform!(copy(X),Val(2))) chebyshevutransform(chebyshevutransform(chebyshevutransform(X,Val(2),1),Val(2),2),Val(2),3)
325+
@testset "3D" begin
326+
X = randn(4,5,6)
327+
= similar(X)
328+
@testset "chebyshevtransform" begin
329+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevtransform(X[:,k,j]) end
330+
@test @inferred(chebyshevtransform(X,1)) @inferred(chebyshevtransform!(copy(X),1))
331+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevtransform(X[k,:,j]) end
332+
@test chebyshevtransform(X,2) chebyshevtransform!(copy(X),2)
333+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevtransform(X[k,j,:]) end
334+
@test chebyshevtransform(X,3) chebyshevtransform!(copy(X),3)
335+
336+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevtransform(X[:,k,j],Val(2)) end
337+
@test @inferred(chebyshevtransform(X,Val(2),1)) @inferred(chebyshevtransform!(copy(X),Val(2),1))
338+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevtransform(X[k,:,j],Val(2)) end
339+
@test chebyshevtransform(X,Val(2),2) chebyshevtransform!(copy(X),Val(2),2)
340+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevtransform(X[k,j,:],Val(2)) end
341+
@test chebyshevtransform(X,Val(2),3) chebyshevtransform!(copy(X),Val(2),3)
342+
343+
@test @inferred(chebyshevtransform(X)) @inferred(chebyshevtransform!(copy(X))) chebyshevtransform(chebyshevtransform(chebyshevtransform(X,1),2),3)
344+
@test @inferred(chebyshevtransform(X,Val(2))) @inferred(chebyshevtransform!(copy(X),Val(2))) chebyshevtransform(chebyshevtransform(chebyshevtransform(X,Val(2),1),Val(2),2),Val(2),3)
345+
end
346+
347+
@testset "ichebyshevtransform" begin
348+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevtransform(X[:,k,j]) end
349+
@test @inferred(ichebyshevtransform(X,1)) @inferred(ichebyshevtransform!(copy(X),1))
350+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevtransform(X[k,:,j]) end
351+
@test ichebyshevtransform(X,2) ichebyshevtransform!(copy(X),2)
352+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevtransform(X[k,j,:]) end
353+
@test ichebyshevtransform(X,3) ichebyshevtransform!(copy(X),3)
354+
355+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevtransform(X[:,k,j],Val(2)) end
356+
@test @inferred(ichebyshevtransform(X,Val(2),1)) @inferred(ichebyshevtransform!(copy(X),Val(2),1))
357+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevtransform(X[k,:,j],Val(2)) end
358+
@test ichebyshevtransform(X,Val(2),2) ichebyshevtransform!(copy(X),Val(2),2)
359+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevtransform(X[k,j,:],Val(2)) end
360+
@test ichebyshevtransform(X,Val(2),3) ichebyshevtransform!(copy(X),Val(2),3)
361+
362+
@test @inferred(ichebyshevtransform(X)) @inferred(ichebyshevtransform!(copy(X))) ichebyshevtransform(ichebyshevtransform(ichebyshevtransform(X,1),2),3)
363+
@test @inferred(ichebyshevtransform(X,Val(2))) @inferred(ichebyshevtransform!(copy(X),Val(2))) ichebyshevtransform(ichebyshevtransform(ichebyshevtransform(X,Val(2),1),Val(2),2),Val(2),3)
364+
365+
@test ichebyshevtransform(chebyshevtransform(X)) X
366+
@test chebyshevtransform(ichebyshevtransform(X)) X
367+
end
368+
369+
@testset "chebyshevutransform" begin
370+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevutransform(X[:,k,j]) end
371+
@test @inferred(chebyshevutransform(X,1)) @inferred(chebyshevutransform!(copy(X),1))
372+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevutransform(X[k,:,j]) end
373+
@test chebyshevutransform(X,2) chebyshevutransform!(copy(X),2)
374+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevutransform(X[k,j,:]) end
375+
@test chebyshevutransform(X,3) chebyshevutransform!(copy(X),3)
376+
377+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = chebyshevutransform(X[:,k,j],Val(2)) end
378+
@test @inferred(chebyshevutransform(X,Val(2),1)) @inferred(chebyshevutransform!(copy(X),Val(2),1))
379+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = chebyshevutransform(X[k,:,j],Val(2)) end
380+
@test chebyshevutransform(X,Val(2),2) chebyshevutransform!(copy(X),Val(2),2)
381+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = chebyshevutransform(X[k,j,:],Val(2)) end
382+
@test chebyshevutransform(X,Val(2),3) chebyshevutransform!(copy(X),Val(2),3)
383+
384+
@test @inferred(chebyshevutransform(X)) @inferred(chebyshevutransform!(copy(X))) chebyshevutransform(chebyshevutransform(chebyshevutransform(X,1),2),3)
385+
@test @inferred(chebyshevutransform(X,Val(2))) @inferred(chebyshevutransform!(copy(X),Val(2))) chebyshevutransform(chebyshevutransform(chebyshevutransform(X,Val(2),1),Val(2),2),Val(2),3)
386+
end
387+
388+
@testset "ichebyshevutransform" begin
389+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevutransform(X[:,k,j]) end
390+
@test @inferred(ichebyshevutransform(X,1)) @inferred(ichebyshevutransform!(copy(X),1))
391+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevutransform(X[k,:,j]) end
392+
@test ichebyshevutransform(X,2) ichebyshevutransform!(copy(X),2)
393+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevutransform(X[k,j,:]) end
394+
@test ichebyshevutransform(X,3) ichebyshevutransform!(copy(X),3)
395+
396+
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevutransform(X[:,k,j],Val(2)) end
397+
@test @inferred(ichebyshevutransform(X,Val(2),1)) @inferred(ichebyshevutransform!(copy(X),Val(2),1))
398+
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevutransform(X[k,:,j],Val(2)) end
399+
@test ichebyshevutransform(X,Val(2),2) ichebyshevutransform!(copy(X),Val(2),2)
400+
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevutransform(X[k,j,:],Val(2)) end
401+
@test ichebyshevutransform(X,Val(2),3) ichebyshevutransform!(copy(X),Val(2),3)
402+
403+
@test @inferred(ichebyshevutransform(X)) @inferred(ichebyshevutransform!(copy(X))) ichebyshevutransform(ichebyshevutransform(ichebyshevutransform(X,1),2),3)
404+
@test @inferred(ichebyshevutransform(X,Val(2))) @inferred(ichebyshevutransform!(copy(X),Val(2))) ichebyshevutransform(ichebyshevutransform(ichebyshevutransform(X,Val(2),1),Val(2),2),Val(2),3)
405+
406+
@test ichebyshevutransform(chebyshevutransform(X)) X
407+
@test chebyshevutransform(ichebyshevutransform(X)) X
408+
end
409+
410+
X = randn(1,1,1)
411+
@test chebyshevtransform!(copy(X), Val(1)) == ichebyshevtransform!(copy(X), Val(1)) == X
412+
@test_throws ArgumentError chebyshevtransform!(copy(X), Val(2))
413+
@test_throws ArgumentError ichebyshevtransform!(copy(X), Val(2))
385414
end
386415

387-
@testset "ichebyshevutransform" begin
388-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevutransform(X[:,k,j]) end
389-
@test @inferred(ichebyshevutransform(X,1)) @inferred(ichebyshevutransform!(copy(X),1))
390-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevutransform(X[k,:,j]) end
391-
@test ichebyshevutransform(X,2) ichebyshevutransform!(copy(X),2)
392-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevutransform(X[k,j,:]) end
393-
@test ichebyshevutransform(X,3) ichebyshevutransform!(copy(X),3)
394-
395-
for k = axes(X,2), j = axes(X,3) X̃[:,k,j] = ichebyshevutransform(X[:,k,j],Val(2)) end
396-
@test @inferred(ichebyshevutransform(X,Val(2),1)) @inferred(ichebyshevutransform!(copy(X),Val(2),1))
397-
for k = axes(X,1), j = axes(X,3) X̃[k,:,j] = ichebyshevutransform(X[k,:,j],Val(2)) end
398-
@test ichebyshevutransform(X,Val(2),2) ichebyshevutransform!(copy(X),Val(2),2)
399-
for k = axes(X,1), j = axes(X,2) X̃[k,j,:] = ichebyshevutransform(X[k,j,:],Val(2)) end
400-
@test ichebyshevutransform(X,Val(2),3) ichebyshevutransform!(copy(X),Val(2),3)
401-
402-
@test @inferred(ichebyshevutransform(X)) @inferred(ichebyshevutransform!(copy(X))) ichebyshevutransform(ichebyshevutransform(ichebyshevutransform(X,1),2),3)
403-
@test @inferred(ichebyshevutransform(X,Val(2))) @inferred(ichebyshevutransform!(copy(X),Val(2))) ichebyshevutransform(ichebyshevutransform(ichebyshevutransform(X,Val(2),1),Val(2),2),Val(2),3)
404-
405-
@test ichebyshevutransform(chebyshevutransform(X)) X
406-
@test chebyshevutransform(ichebyshevutransform(X)) X
416+
@testset "4D" begin
417+
X = randn(2,3,4,5)
418+
= similar(X)
419+
for trans in (chebyshevtransform, ichebyshevtransform, chebyshevutransform, ichebyshevutransform)
420+
for k = axes(X,2), j = axes(X,3), l = axes(X,4) X̃[:,k,j,l] = trans(X[:,k,j,l]) end
421+
@test @inferred(trans(X,1))
422+
@test @inferred(trans(X)) trans(trans(trans(trans(X,1),2),3),4)
423+
end
407424
end
408-
409-
X = randn(1,1,1)
410-
@test chebyshevtransform!(copy(X), Val(1)) == ichebyshevtransform!(copy(X), Val(1)) == X
411-
@test_throws ArgumentError chebyshevtransform!(copy(X), Val(2))
412-
@test_throws ArgumentError ichebyshevtransform!(copy(X), Val(2))
413425
end
414-
415426
@testset "Integer" begin
416427
@test chebyshevtransform([1,2,3]) == chebyshevtransform([1.,2,3])
417428
@test chebyshevtransform([1,2,3], Val(2)) == chebyshevtransform([1.,2,3], Val(2))

0 commit comments

Comments
 (0)