Support size and \ for cheb2leg (#239)

Author: dlfivefifty

Project.toml

@@ -1,6 +1,6 @@
 name = "FastTransforms"
 uuid = "057dd010-8810-581a-b7be-e3fc3b93f78c"
-version = "0.15.15"
+version = "0.15.16"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/toeplitzhankel.jl

@@ -33,6 +33,8 @@ ToeplitzHankelPlan{S,N,M}(T::TP, L::LowR, R::LowR, dims::Dims) where {S,N,M,LowR
 ToeplitzHankelPlan{S,N}(T, L, R, dims) where {S,N} = ToeplitzHankelPlan{S,N,N+1}(T, L, R, dims)
 ToeplitzHankelPlan(T::ToeplitzPlan{S,M}, L::Matrix, R::Matrix, dims=1) where {S,M} = ToeplitzHankelPlan{S,M-1,M}((T,), (L,), (R,), dims)
 
+size(TH::ToeplitzHankelPlan) = size(first(TH.T))
+
 
 _reshape_broadcast(d, R, ::Val{N}, M) where N = reshape(R,ntuple(k -> k == d ? size(R,1) : 1, Val(N))...,M)
 function _th_applymul!(d, v::AbstractArray{<:Any,N}, T, L, R, tmp) where N
@@ -52,6 +54,8 @@ function *(P::ToeplitzHankelPlan{<:Any,N}, v::AbstractArray{<:Any,N}) where N
     v
 end
 
+*(P::ToeplitzHankelPlan, v::AbstractArray) = error("plan applied to wrong-sized array")
+
 
 # partial cholesky for a Hankel matrix
 
@@ -158,6 +162,10 @@ function *(P::ChebyshevToLegendrePlanTH, V::AbstractArray{<:Any,N}) where N
     V
 end
 
+_add1tod(d::Integer, a, b...) = d == 1 ? (a+1, b...) : (a, _add1tod(d-1, b...)...)
+_add1tod(d, a, b...) = _add1tod(first(d), a, b...)
+size(P::ChebyshevToLegendrePlanTH) = Base.front(_add1tod(P.toeplitzhankel.dims, size(first(P.toeplitzhankel.T))...))
+inv(P::ChebyshevToLegendrePlanTH{T}) where T = plan_th_leg2cheb!(T, size(P), P.toeplitzhankel.dims)
 
 
 function _leg2chebTH_TLC(::Type{S}, mn, d) where S

test/toeplitzhankeltests.jl

@@ -164,4 +164,23 @@ Random.seed!(0)
             @test M ≈ Y
         end
     end
+
+    @testset "inv" begin
+        x = randn(10)
+        pl = plan_th_cheb2leg!(x)
+        @test size(pl) == (10,)
+        @test pl\(pl*x) ≈ x
+
+        X = randn(10,3)
+        for pl in (plan_th_cheb2leg!(X), plan_th_cheb2leg!(X, 1), plan_th_cheb2leg!(X, 2))
+            @test size(pl) == (10,3)
+            @test pl\(pl*copy(X)) ≈ X
+        end
+
+        X = randn(10,3,5)
+        for pl in (plan_th_cheb2leg!(X), plan_th_cheb2leg!(X, 1), plan_th_cheb2leg!(X, 2),  plan_th_cheb2leg!(X, 3))
+            @test size(pl) == (10,3,5)
+            @test pl\(pl*copy(X)) ≈ X
+        end
+    end
 end