2nd kind matrix transforms

Author: dlfivefifty

src/chebyshevtransform.jl

@@ -89,19 +89,41 @@ function mul!(y::AbstractArray{T}, P::ChebyshevTransformPlan{T,1,K,false}, x::Ab
     _cheb1_rescale!(P.plan.region, y)
 end
 
-function *(P::ChebyshevTransformPlan{T,2,K,true}, x::AbstractVector{T}) where {T,K}
+
+_cheb2_rescale!(_, y::AbstractVector) = (y[1] /= 2; y[end] /= 2; ldiv!(length(y)-1, y))
+
+function _cheb2_rescale!(d::Number, y::AbstractMatrix{T}) where T
+    if isone(d)
+        ldiv!(2, @view(y[1,:]))
+        ldiv!(2, @view(y[end,:]))
+    else
+        ldiv!(2, @view(y[:,1]))
+        ldiv!(2, @view(y[:,end]))
+    end
+    ldiv!(size(y,d)-1, y)
+end
+
+# TODO: higher dimensional arrays
+function _cheb2_rescale!(d::UnitRange, y::AbstractMatrix{T}) where T
+    @assert d == 1:2
+    ldiv!(2, @view(y[1,:]))
+    ldiv!(2, @view(y[end,:]))
+    ldiv!(2, @view(y[:,1]))
+    ldiv!(2, @view(y[:,end]))
+    ldiv!(prod(size(y) .- 1), y)
+end
+
+function *(P::ChebyshevTransformPlan{T,2,K,true}, x::AbstractArray{T}) where {T,K}
     n = length(x)
     y = P.plan*x # will be  === x if in-place
-    y[1] /= 2; y[end] /= 2
-    lmul!(inv(convert(T,n-1)),y)
+    _cheb2_rescale!(P.plan.region, y)
 end
 
-function mul!(y::AbstractVector{T}, P::ChebyshevTransformPlan{T,2,K,false}, x::AbstractVector) where {T,K}
+function mul!(y::AbstractArray{T}, P::ChebyshevTransformPlan{T,2,K,false}, x::AbstractArray) where {T,K}
     n = length(x)
     length(y) == n || throw(DimensionMismatch("output must match dimension"))
     _plan_mul!(y, P.plan, x)
-    y[1] /= 2; y[end] /= 2
-    lmul!(inv(convert(T,n-1)),y)
+    _cheb2_rescale!(P.plan.region, y)
 end
 
 *(P::ChebyshevTransformPlan{T,kind,K,false}, x::AbstractArray{T}) where {T,kind,K} =

test/chebyshevtests.jl

@@ -176,8 +176,11 @@ using FastTransforms, Test
         X = randn(4,5)
         @test @inferred(chebyshevtransform(X,1)) ≈ @inferred(chebyshevtransform!(copy(X),1)) ≈ hcat(chebyshevtransform.([X[:,k] for k=axes(X,2)])...)
         @test chebyshevtransform(X,2) ≈ chebyshevtransform!(copy(X),2) ≈ hcat(chebyshevtransform.([X[k,:] for k=axes(X,1)])...)'
+        @test @inferred(chebyshevtransform(X,Val(2),1)) ≈ @inferred(chebyshevtransform!(copy(X),Val(2),1)) ≈ hcat(chebyshevtransform.([X[:,k] for k=axes(X,2)],Val(2))...)
+        @test chebyshevtransform(X,Val(2),2) ≈ chebyshevtransform!(copy(X),Val(2),2) ≈ hcat(chebyshevtransform.([X[k,:] for k=axes(X,1)],Val(2))...)'
 
-        @test @inferred(chebyshevtransform(X)) == @inferred(chebyshevtransform!(copy(X))) == chebyshevtransform(chebyshevtransform(X,1),2)
+        @test @inferred(chebyshevtransform(X)) ≈ @inferred(chebyshevtransform!(copy(X))) ≈ chebyshevtransform(chebyshevtransform(X,1),2)
+        @test @inferred(chebyshevtransform(X,Val(2))) ≈ @inferred(chebyshevtransform!(copy(X),Val(2))) ≈ chebyshevtransform(chebyshevtransform(X,Val(2),1),Val(2),2)
 
 
         X = randn(1,1)