diff --git a/src/chebyshevtransform.jl b/src/chebyshevtransform.jl index 2e58876d..2d3f17d9 100644 --- a/src/chebyshevtransform.jl +++ b/src/chebyshevtransform.jl @@ -53,22 +53,34 @@ _maybemutablecopy(x::StridedArray{T}, ::Type{T}) where {T} = x _maybemutablecopy(x, T) = Array{T}(x) @inline _plan_mul!(y::AbstractArray{T}, P::Plan{T}, x::AbstractArray) where T = mul!(y, P, _maybemutablecopy(x, T)) +function applydim!(op!, X::AbstractArray, Rpre, Rpost, ind) + for Ipost in Rpost, Ipre in Rpre + v = view(X, Ipre, ind, Ipost) + op!(v) + end + X +end +function applydim!(op!, X::AbstractArray, d::Integer, ind) + Rpre = CartesianIndices(axes(X)[1:d-1]) + Rpost = CartesianIndices(axes(X)[d+1:end]) + applydim!(op!, X, Rpre, Rpost, ind) +end for op in (:ldiv, :lmul) - op_dim_begin! = Symbol(string(op) * "_dim_begin!") - op_dim_end! = Symbol(string(op) * "_dim_end!") - op! = Symbol(string(op) * "!") + op_dim_begin! = Symbol(op, :_dim_begin!) + op_dim_end! = Symbol(op, :_dim_end!) + op! = Symbol(op, :!) @eval begin - function $op_dim_begin!(α, d::Number, y::AbstractArray{<:Any,N}) where N + function $op_dim_begin!(α, d::Number, y::AbstractArray) # scale just the d-th dimension by permuting it to the first - ỹ = PermutedDimsArray(y, _permfirst(d, N)) - $op!(α, view(ỹ, 1, ntuple(_ -> :, Val(N-1))...)) + d ∈ 1:ndims(y) || throw(ArgumentError("dimension $d must lie between 1 and $(ndims(y))")) + applydim!(v -> $op!(α, v), y, d, 1) end - function $op_dim_end!(α, d::Number, y::AbstractArray{<:Any,N}) where N + function $op_dim_end!(α, d::Number, y::AbstractArray) # scale just the d-th dimension by permuting it to the first - ỹ = PermutedDimsArray(y, _permfirst(d, N)) - $op!(α, view(ỹ, size(ỹ,1), ntuple(_ -> :, Val(N-1))...)) + d ∈ 1:ndims(y) || throw(ArgumentError("dimension $d must lie between 1 and $(ndims(y))")) + applydim!(v -> $op!(α, v), y, d, size(y, d)) end end end @@ -363,35 +375,34 @@ function plan_chebyshevutransform(x::AbstractArray{T,N}, ::Val{2}, dims...; kws. ChebyshevUTransformPlan{T,2}(FFTW.plan_r2r(x, USECONDKIND, dims...; kws...)) end - -_permfirst(d, N) = [d; 1:d-1; d+1:N] - -@inline function _chebu1_prescale!(d::Number, X::AbstractArray{T,N}) where {T,N} - X̃ = PermutedDimsArray(X, _permfirst(d, N)) - m = size(X̃,1) - X̃ .= (sinpi.(one(T)/(2m) .+ ((1:m) .- one(T))/m) ./ m) .* X̃ - X -end - -@inline function _chebu1_prescale!(d, y::AbstractArray) - for k in d - _chebu1_prescale!(k, y) +for f in [:_chebu1_prescale!, :_chebu1_postscale!, :_chebu2_prescale!, :_chebu2_postscale!, + :_ichebu1_postscale!] + _f = Symbol(:_, f) + @eval begin + @inline function $f(d::Number, X::AbstractArray) + d ∈ 1:ndims(X) || throw("dimension $d must lie between 1 and $(ndims(X))") + $_f(d, X) + X + end + @inline function $f(d, y::AbstractArray) + for k in d + $f(k, y) + end + y + end end - y end -@inline function _chebu1_postscale!(d::Number, X::AbstractArray{T,N}) where {T,N} - X̃ = PermutedDimsArray(X, _permfirst(d, N)) - m = size(X̃,1) - X̃ .= X̃ ./ (sinpi.(one(T)/(2m) .+ ((1:m) .- one(T))/m) ./ m) - X +function __chebu1_prescale!(d::Number, X::AbstractArray{T}) where {T} + m = size(X,d) + r = one(T)/(2m) .+ ((1:m) .- one(T))./m + applydim!(v -> v .*= sinpi.(r) ./ m, X, d, :) end -@inline function _chebu1_postscale!(d, y::AbstractArray) - for k in d - _chebu1_postscale!(k, y) - end - y +@inline function __chebu1_postscale!(d::Number, X::AbstractArray{T}) where {T} + m = size(X,d) + r = one(T)/(2m) .+ ((1:m) .- one(T))./m + applydim!(v -> v ./= sinpi.(r) ./ m, X, d, :) end function *(P::ChebyshevUTransformPlan{T,1,K,true,N}, x::AbstractArray{T,N}) where {T,K,N} @@ -413,35 +424,18 @@ function mul!(y::AbstractArray{T}, P::ChebyshevUTransformPlan{T,1,K,false}, x::A end -@inline function _chebu2_prescale!(d::Number, X::AbstractArray{T,N}) where {T,N} - X̃ = PermutedDimsArray(X, _permfirst(d, N)) - m = size(X̃,1) +@inline function __chebu2_prescale!(d, X::AbstractArray{T}) where {T} + m = size(X,d) c = one(T)/ (m+1) - X̃ .= sinpi.((1:m) .* c) .* X̃ - X + r = (1:m) .* c + applydim!(v -> v .*= sinpi.(r), X, d, :) end -@inline function _chebu2_prescale!(d, y::AbstractArray) - for k in d - _chebu2_prescale!(k, y) - end - y -end - - -@inline function _chebu2_postscale!(d::Number, X::AbstractArray{T,N}) where {T,N} - X̃ = PermutedDimsArray(X, _permfirst(d, N)) - m = size(X̃,1) +@inline function __chebu2_postscale!(d::Number, X::AbstractArray{T}) where {T} + m = size(X,d) c = one(T)/ (m+1) - X̃ .= X̃ ./ sinpi.((1:m) .* c) - X -end - -@inline function _chebu2_postscale!(d, y::AbstractArray) - for k in d - _chebu2_postscale!(k, y) - end - y + r = (1:m) .* c + applydim!(v -> v ./= sinpi.(r), X, d, :) end function *(P::ChebyshevUTransformPlan{T,2,K,true,N}, x::AbstractArray{T,N}) where {T,K,N} @@ -525,19 +519,10 @@ inv(P::IChebyshevUTransformPlan{T,2}) where {T} = ChebyshevUTransformPlan{T,2}(P inv(P::ChebyshevUTransformPlan{T,1}) where {T} = IChebyshevUTransformPlan{T,1}(inv(P.plan).p) inv(P::IChebyshevUTransformPlan{T,1}) where {T} = ChebyshevUTransformPlan{T,1}(inv(P.plan).p) -@inline function _ichebu1_postscale!(d::Number, X::AbstractArray{T,N}) where {T,N} - X̃ = PermutedDimsArray(X, _permfirst(d, N)) - m = size(X̃,1) - X̃ .= X̃ ./ (2 .* sinpi.(one(T)/(2m) .+ ((1:m) .- one(T))/m)) - X -end - - -@inline function _ichebu1_postscale!(d, y::AbstractArray) - for k in d - _ichebu1_postscale!(k, y) - end - y +@inline function __ichebu1_postscale!(d::Number, X::AbstractArray{T}) where {T} + m = size(X,d) + r = one(T)/(2m) .+ ((1:m) .- one(T))/m + applydim!(v -> v ./= 2 .* sinpi.(r), X, d, :) end function *(P::IChebyshevUTransformPlan{T,1,K,true}, x::AbstractArray{T}) where {T<:fftwNumber,K} diff --git a/src/toeplitzhankel.jl b/src/toeplitzhankel.jl index 78cc43fa..1c552a09 100644 --- a/src/toeplitzhankel.jl +++ b/src/toeplitzhankel.jl @@ -137,14 +137,14 @@ function *(P::ChebyshevToLegendrePlanTH, v::AbstractVector{S}) where S v end -function _cheb2leg_rescale1!(V::AbstractArray{S}) where S - m = size(V,1) - for j = CartesianIndices(tail(axes(V))) +function _cheb2leg_rescale1!(V::AbstractArray{S}, Rpre, Rpost, d) where S + m = size(V,d) + for Ipost in Rpost, Ipre in Rpre ret = zero(S) @inbounds for k = 1:2:m - ret += -V[k,j]/(k*(k-2)) + ret += -V[Ipre,k,Ipost]/(k*(k-2)) end - V[1,j] = ret + V[Ipre,1,Ipost] = ret end V end @@ -152,11 +152,13 @@ end _dropfirstdim(d::Int) = () _dropfirstdim(d::Int, m, szs...) = ((d == 1 ? 2 : 1):m, _dropfirstdim(d-1, szs...)...) -function *(P::ChebyshevToLegendrePlanTH, V::AbstractArray{<:Any,N}) where N +function *(P::ChebyshevToLegendrePlanTH, V::AbstractArray) m,n = size(V) 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))) + Rpre = CartesianIndices(axes(V)[1:d-1]) + Rpost = CartesianIndices(axes(V)[d+1:end]) + _cheb2leg_rescale1!(V, Rpre, Rpost, d) _th_applymul!(d, view(V, _dropfirstdim(d, size(V)...)...), T, L, R, tmp) end V @@ -729,4 +731,4 @@ th_cheb2jac(v, α, β, dims...) = plan_th_cheb2jac!(eltype(v), size(v), α, β, plan_th_jac2cheb!(::Type{T}, mn, α, β, dims...) where T = Jac2ChebPlanTH(plan_th_jac2jac!(T, mn, α, β, -one(α)/2, -one(α)/2, dims...)) plan_th_jac2cheb!(arr::AbstractArray{T}, α, β, dims...) where T = plan_th_jac2cheb!(T, size(arr), α, β, dims...) -th_jac2cheb(v, α, β, dims...) = plan_th_jac2cheb!(eltype(v), size(v), α, β, dims...)*copy(v) \ No newline at end of file +th_jac2cheb(v, α, β, dims...) = plan_th_jac2cheb!(eltype(v), size(v), α, β, dims...)*copy(v) diff --git a/src/toeplitzplans.jl b/src/toeplitzplans.jl index 42d24062..9be77234 100644 --- a/src/toeplitzplans.jl +++ b/src/toeplitzplans.jl @@ -61,12 +61,11 @@ function *(A::ToeplitzPlan{T,N}, X::AbstractArray{T,N}) where {T,N} # 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 .* Ỹ + applydim!(v -> v .= vc .* v, Y, d, :) end # Transform back