Skip to content

Commit adf63f4

Browse files
committed
Remove uses of PermutedDimsArray
1 parent 9f135a8 commit adf63f4

File tree

1 file changed

+55
-69
lines changed

1 file changed

+55
-69
lines changed

src/chebyshevtransform.jl

Lines changed: 55 additions & 69 deletions
Original file line numberDiff line numberDiff line change
@@ -54,22 +54,31 @@ end
5454
@inline _plan_mul!(y::AbstractArray{T}, P::Plan{T}, x::StridedArray{T}) where T = mul!(y, P, x)
5555
@inline _plan_mul!(y::AbstractArray{T}, P::Plan{T}, x::AbstractArray) where T = mul!(y, P, convert(Array{T}, x))
5656

57+
function applydim!(op!, X::AbstractArray, Rpre, Rpost, ind)
58+
for Ipost in Rpost, Ipre in Rpre
59+
v = view(X, Ipre, ind, Ipost)
60+
op!(v)
61+
end
62+
X
63+
end
5764

5865
for op in (:ldiv, :lmul)
59-
op_dim_begin! = Symbol(string(op) * "_dim_begin!")
60-
op_dim_end! = Symbol(string(op) * "_dim_end!")
61-
op! = Symbol(string(op) * "!")
66+
op_dim_begin! = Symbol(op, :_dim_begin!)
67+
op_dim_end! = Symbol(op, :_dim_end!)
68+
op! = Symbol(op, :!)
6269
@eval begin
63-
function $op_dim_begin!(α, d::Number, y::AbstractArray{<:Any,N}) where N
70+
function $op_dim_begin!(α, d::Number, y::AbstractArray)
6471
# scale just the d-th dimension by permuting it to the first
65-
= PermutedDimsArray(y, _permfirst(d, N))
66-
$op!(α, view(ỹ, 1, ntuple(_ -> :, Val(N-1))...))
72+
Rpre = CartesianIndices(axes(y)[1:d-1])
73+
Rpost = CartesianIndices(axes(y)[d+1:end])
74+
applydim!(v -> $op!(α, v), y, Rpre, Rpost, 1)
6775
end
6876

69-
function $op_dim_end!(α, d::Number, y::AbstractArray{<:Any,N}) where N
77+
function $op_dim_end!(α, d::Number, y::AbstractArray)
7078
# scale just the d-th dimension by permuting it to the first
71-
= PermutedDimsArray(y, _permfirst(d, N))
72-
$op!(α, view(ỹ, size(ỹ,1), ntuple(_ -> :, Val(N-1))...))
79+
Rpre = CartesianIndices(axes(y)[1:d-1])
80+
Rpost = CartesianIndices(axes(y)[d+1:end])
81+
applydim!(v -> $op!(α, v), y, Rpre, Rpost, size(y, d))
7382
end
7483
end
7584
end
@@ -366,32 +375,35 @@ end
366375

367376
_permfirst(d, N) = [d; 1:d-1; d+1:N]
368377

369-
@inline function _chebu1_prescale!(d::Number, X::AbstractArray{T,N}) where {T,N}
370-
= PermutedDimsArray(X, _permfirst(d, N))
371-
m = size(X̃,1)
372-
X̃ .= (sinpi.(one(T)/(2m) .+ ((1:m) .- one(T))/m) ./ m) .*
373-
X
374-
end
375-
376-
@inline function _chebu1_prescale!(d, y::AbstractArray)
377-
for k in d
378-
_chebu1_prescale!(k, y)
378+
for f in [:_chebu1_prescale!, :_chebu1_postscale!, :_chebu2_prescale!, :_chebu2_postscale!,
379+
:_ichebu1_postscale!]
380+
_f = Symbol(:_, f)
381+
@eval begin
382+
@inline function $f(d::Number, X::AbstractArray)
383+
Rpre = CartesianIndices(axes(X)[1:d-1])
384+
Rpost = CartesianIndices(axes(X)[d+1:end])
385+
$_f(d, X, Rpre, Rpost)
386+
X
387+
end
388+
@inline function $f(d, y::AbstractArray)
389+
for k in d
390+
$f(k, y)
391+
end
392+
y
393+
end
379394
end
380-
y
381395
end
382396

383-
@inline function _chebu1_postscale!(d::Number, X::AbstractArray{T,N}) where {T,N}
384-
= PermutedDimsArray(X, _permfirst(d, N))
385-
m = size(X̃,1)
386-
X̃ .=./ (sinpi.(one(T)/(2m) .+ ((1:m) .- one(T))/m) ./ m)
387-
X
397+
function __chebu1_prescale!(d::Number, X::AbstractArray{T}, Rpre, Rpost) where {T}
398+
m = size(X,d)
399+
r = one(T)/(2m) .+ ((1:m) .- one(T))./m
400+
applydim!(v -> v .*= sinpi.(r) ./ m, X, Rpre, Rpost, :)
388401
end
389402

390-
@inline function _chebu1_postscale!(d, y::AbstractArray)
391-
for k in d
392-
_chebu1_postscale!(k, y)
393-
end
394-
y
403+
@inline function __chebu1_postscale!(d::Number, X::AbstractArray{T}, Rpre, Rpost) where {T}
404+
m = size(X,d)
405+
r = one(T)/(2m) .+ ((1:m) .- one(T))./m
406+
applydim!(v -> v ./= sinpi.(r) ./ m, X, Rpre, Rpost, :)
395407
end
396408

397409
function *(P::ChebyshevUTransformPlan{T,1,K,true,N}, x::AbstractArray{T,N}) where {T,K,N}
@@ -413,35 +425,18 @@ function mul!(y::AbstractArray{T}, P::ChebyshevUTransformPlan{T,1,K,false}, x::A
413425
end
414426

415427

416-
@inline function _chebu2_prescale!(d::Number, X::AbstractArray{T,N}) where {T,N}
417-
= PermutedDimsArray(X, _permfirst(d, N))
418-
m = size(X̃,1)
428+
@inline function __chebu2_prescale!(d, X::AbstractArray{T}, Rpre, Rpost) where {T}
429+
m = size(X,d)
419430
c = one(T)/ (m+1)
420-
X̃ .= sinpi.((1:m) .* c) .*
421-
X
422-
end
423-
424-
@inline function _chebu2_prescale!(d, y::AbstractArray)
425-
for k in d
426-
_chebu2_prescale!(k, y)
427-
end
428-
y
431+
r = (1:m) .* c
432+
applydim!(v -> v .*= sinpi.(r), X, Rpre, Rpost, :)
429433
end
430434

431-
432-
@inline function _chebu2_postscale!(d::Number, X::AbstractArray{T,N}) where {T,N}
433-
= PermutedDimsArray(X, _permfirst(d, N))
434-
m = size(X̃,1)
435+
@inline function __chebu2_postscale!(d::Number, X::AbstractArray{T}, Rpre, Rpost) where {T}
436+
m = size(X,d)
435437
c = one(T)/ (m+1)
436-
X̃ .=./ sinpi.((1:m) .* c)
437-
X
438-
end
439-
440-
@inline function _chebu2_postscale!(d, y::AbstractArray)
441-
for k in d
442-
_chebu2_postscale!(k, y)
443-
end
444-
y
438+
r = (1:m) .* c
439+
applydim!(v -> v ./= sinpi.(r), X, Rpre, Rpost, :)
445440
end
446441

447442
function *(P::ChebyshevUTransformPlan{T,2,K,true,N}, x::AbstractArray{T,N}) where {T,K,N}
@@ -525,19 +520,10 @@ inv(P::IChebyshevUTransformPlan{T,2}) where {T} = ChebyshevUTransformPlan{T,2}(P
525520
inv(P::ChebyshevUTransformPlan{T,1}) where {T} = IChebyshevUTransformPlan{T,1}(inv(P.plan).p)
526521
inv(P::IChebyshevUTransformPlan{T,1}) where {T} = ChebyshevUTransformPlan{T,1}(inv(P.plan).p)
527522

528-
@inline function _ichebu1_postscale!(d::Number, X::AbstractArray{T,N}) where {T,N}
529-
= PermutedDimsArray(X, _permfirst(d, N))
530-
m = size(X̃,1)
531-
X̃ .=./ (2 .* sinpi.(one(T)/(2m) .+ ((1:m) .- one(T))/m))
532-
X
533-
end
534-
535-
536-
@inline function _ichebu1_postscale!(d, y::AbstractArray)
537-
for k in d
538-
_ichebu1_postscale!(k, y)
539-
end
540-
y
523+
@inline function __ichebu1_postscale!(d::Number, X::AbstractArray{T}, Rpre, Rpost) where {T}
524+
m = size(X,d)
525+
r = one(T)/(2m) .+ ((1:m) .- one(T))/m
526+
applydim!(v -> v ./= 2 .* sinpi.(r), X, Rpre, Rpost, :)
541527
end
542528

543529
function *(P::IChebyshevUTransformPlan{T,1,K,true}, x::AbstractArray{T}) where {T<:fftwNumber,K}
@@ -742,4 +728,4 @@ for pln in (:plan_chebyshevtransform!, :plan_chebyshevtransform,
742728
$pln(x::AbstractArray, dims...; kws...) = $pln(x, Val(1), dims...; kws...)
743729
$pln(::Type{T}, szs, dims...; kwds...) where T = $pln(Array{T}(undef, szs...), dims...; kwds...)
744730
end
745-
end
731+
end

0 commit comments

Comments
 (0)