Merge pull request #242 from jishnub/scalepermdim

Remove uses of PermutedDimsArray in pre and post-scaling
JuliaApproximation · Jun 8, 2024 · 9dead75 · 9dead75
2 parents 17f3e59 + da5f214
commit 9dead75
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 82 deletions.
diff --git 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
-            ỹ = PermutedDimsArray(y, _permfirst(d, N))
-            $op!(α, view(ỹ, 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
-            ỹ = PermutedDimsArray(y, _permfirst(d, N))
-            $op!(α, view(ỹ, size(ỹ,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
@@ -137,26 +137,28 @@ 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
 
 _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)
+th_jac2cheb(v, α, β, dims...) = plan_th_jac2cheb!(eltype(v), size(v), α, β, dims...)*copy(v)
diff --git 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)
-        Ỹ = PermutedDimsArray(Y, _permfirst(d, N))
-        Ỹ .= vc .* Ỹ
+        applydim!(v -> v .= vc .* v, Y, d, :)
     end
 
     # Transform back