work on Rect transform (#133)

* work on Rect transform

* DiagTrav pad

* Make transform and evaluate Fast

* add disk sum

* Update Project.toml

* tests pass

* v0.4

Author: dlfivefifty

.github/workflows/ci.yml

@@ -10,7 +10,8 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1'
+          - '1.7'
+          - '1'          
         os:
           - ubuntu-latest
           - macOS-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "MultivariateOrthogonalPolynomials"
 uuid = "4f6956fd-4f93-5457-9149-7bfc4b2ce06d"
-version = "0.3"
+version = "0.4"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -28,20 +28,20 @@ ArrayLayouts = "0.8.6"
 BandedMatrices = "0.17"
 BlockArrays = "0.16.14"
 BlockBandedMatrices = "0.11.5"
-ClassicalOrthogonalPolynomials = "0.6"
-ContinuumArrays = "0.10.2, 0.11"
+ClassicalOrthogonalPolynomials = "0.6.9"
+ContinuumArrays = "0.11"
 DomainSets = "0.5"
-FastTransforms = "0.13, 0.14"
-FillArrays = "0.12, 0.13"
-HarmonicOrthogonalPolynomials = "0.2.7"
+FastTransforms = "0.14"
+FillArrays = "0.13"
+HarmonicOrthogonalPolynomials = "0.3"
 InfiniteArrays = "0.12"
 InfiniteLinearAlgebra = "0.6.6"
 LazyArrays = "0.22.10"
-LazyBandedMatrices = "0.8"
+LazyBandedMatrices = "0.8.4"
 QuasiArrays = "0.9"
 SpecialFunctions = "1, 2"
 StaticArrays = "1"
-julia = "1.6"
+julia = "1.7"
 
 [extras]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

src/MultivariateOrthogonalPolynomials.jl

@@ -11,19 +11,20 @@ import Base.Broadcast: Broadcasted, broadcasted, DefaultArrayStyle
 import DomainSets: boundary
 
 import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle, domain
-import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted
+import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted, plan_grid_transform, checkpoints
 
 import ArrayLayouts: MemoryLayout, sublayout, sub_materialize
 import BlockArrays: block, blockindex, BlockSlice, viewblock, blockcolsupport, AbstractBlockStyle, BlockStyle
 import BlockBandedMatrices: _BandedBlockBandedMatrix, AbstractBandedBlockBandedMatrix, _BandedMatrix, blockbandwidths, subblockbandwidths
 import LinearAlgebra: factorize
-import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout
-import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes
-import InfiniteArrays: InfiniteCardinal
+import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout, PaddedLayout
+import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes, DiagTravLayout
+import InfiniteArrays: InfiniteCardinal, OneToInf
 
-import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis
+import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis, pad, recurrencecoefficients, clenshaw
 import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan,
-                                          PartialDerivative, AngularMomentum, BlockOneTo, BlockRange1, interlace
+                                          PartialDerivative, AngularMomentum, BlockOneTo, BlockRange1, interlace,
+                                          MultivariateOPLayout
 
 export MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial,
        UnitTriangle, UnitDisk,

src/disk.jl

@@ -336,3 +336,14 @@ function \(A::Zernike{T}, wB::Weighted{V,Zernike{V}}) where {T,V}
     @assert iszero(B.a)
     ModalInterlace{TV}((Normalized.(Jacobi{TV}.(A.b, A.a:∞)) .\ HalfWeighted{:a}.(Normalized.(Jacobi{TV}.(B.b, B.a:∞)))) .* convert(TV, 2)^(-c/2), (ℵ₀,ℵ₀), (2Int(B.b), 2Int(A.a+A.b)))
 end
+
+
+###
+# sum
+###
+
+function Base._sum(P::Zernike{T}, dims) where T
+    @assert dims == 1
+    @assert P.a == P.b == 0
+    Hcat(sqrt(convert(T, π)), Zeros{T}(1,∞))
+end

src/rect.jl

@@ -5,6 +5,9 @@ end
 KronPolynomial{d,T}(a::Vararg{Any,d}) where {d,T} = KronPolynomial{d,T,typeof(a)}(a)
 KronPolynomial{d}(a::Vararg{Any,d}) where d = KronPolynomial{d,mapreduce(eltype, promote_type, a)}(a...)
 KronPolynomial(a::Vararg{Any,d}) where d = KronPolynomial{d}(a...)
+KronPolynomial{d,T}(a::AbstractVector) where {d,T} = KronPolynomial{d,T,typeof(a)}(a)
+KronPolynomial{d}(a::AbstractVector) where d = KronPolynomial{d,eltype(eltype(a))}(a)
+KronPolynomial(a::AbstractVector) = KronPolynomial{length(a)}(a)
 
 const RectPolynomial{T, PP} = KronPolynomial{2, T, PP}
 
@@ -17,8 +20,12 @@ function getindex(P::RectPolynomial{T}, xy::StaticVector{2}, Jj::BlockIndex{1}):
     a[x,J-j+1]b[y,j]
 end
 getindex(P::RectPolynomial, xy::StaticVector{2}, B::Block{1}) = [P[xy, B[j]] for j=1:Int(B)]
-getindex(P::RectPolynomial, xy::StaticVector{2}, JR::BlockOneTo) = mortar([P[xy,Block(J)] for J = 1:Int(JR[end])])
-
+function getindex(P::RectPolynomial, xy::StaticVector{2}, JR::BlockOneTo)
+    A,B = P.args
+    x,y = xy
+    N = size(JR,1)
+    DiagTrav(A[x,OneTo(N)] .* B[y,OneTo(N)]')
+end
 @simplify function *(Dx::PartialDerivative{1}, P::RectPolynomial)
     A,B = P.args
     U,M = (Derivative(axes(A,1))*A).args
@@ -36,4 +43,44 @@ function \(P::RectPolynomial, Q::RectPolynomial)
     PA,PB = P.args
     QA,QB = Q.args
     KronTrav(PA\QA, PB\QB)
+end
+
+struct ApplyPlan{T, F, Pl}
+    f::F
+    plan::Pl
+end
+
+ApplyPlan(f, P) = ApplyPlan{eltype(P), typeof(f), typeof(P)}(f, P)
+
+*(A::ApplyPlan, B::AbstractArray) = A.f(A.plan*B)
+
+function checkpoints(P::RectPolynomial)
+    x,y = checkpoints.(P.args)
+    SVector.(x, y')
+end
+
+function plan_grid_transform(P::KronPolynomial{d,<:Any,<:Fill}, B, dims=1:ndims(B)) where d
+    @assert dims == 1
+
+    T = first(P.args)
+    x, F = plan_grid_transform(T, Array{eltype(B)}(undef, Fill(blocksize(B,1),d)...))
+    @assert d == 2
+    x̃ = Vector(x)
+    SVector.(x̃, x̃'), ApplyPlan(DiagTrav, F)
+end
+
+pad(C::DiagTrav, ::BlockedUnitRange{RangeCumsum{Int,OneToInf{Int}}}) = DiagTrav(pad(C.array, ∞, ∞))
+
+QuasiArrays.mul(A::BivariateOrthogonalPolynomial, b::DiagTrav) = ApplyQuasiArray(*, A, b)
+
+function Base.unsafe_getindex(f::Mul{MultivariateOPLayout{2},<:DiagTravLayout{<:PaddedLayout}}, 𝐱::SVector)
+    P,c = f.A, f.B
+    A,B = P.args
+    x,y = 𝐱
+    clenshaw(clenshaw(paddeddata(c.array), recurrencecoefficients(A)..., x), recurrencecoefficients(B)..., y)
+end
+
+Base.@propagate_inbounds function getindex(f::Mul{MultivariateOPLayout{2},<:DiagTravLayout{<:PaddedLayout}}, x::SVector, j...)
+    @inbounds checkbounds(ApplyQuasiArray(*,f.A,f.B), x, j...)
+    Base.unsafe_getindex(f, x, j...)
 end

src/triangle.jl

@@ -570,9 +570,8 @@ getindex(P::JacobiTriangle, xy::SVector{2}, JR::BlockOneTo) =
 ###
 
 function getindex(f::ApplyQuasiVector{T,typeof(*),<:Tuple{JacobiTriangle,AbstractVector}}, xy::SVector{2})::T where T
-    P,c∞ = arguments(f)
-    c = paddeddata(c∞)
-    N = blocksize(c,1)
+    P,c = arguments(f)
+    N = Int(last(blockcolsupport(c,1)))
 
     N == 1 && return c[1]
     X = jacobimatrix(Val(1), P)
@@ -596,7 +595,7 @@ function getindex(f::ApplyQuasiVector{T,typeof(*),<:Tuple{JacobiTriangle,Abstrac
         # this also resizes γ2
         lmul!(C', γ2)
         # Need to swap sign since it should have been -C
-        γ2 .= (-).(γ2) .+ view(c,Block(n))
+        γ2 .= (-).(γ2) .+ c[Block(n)]
         γ2,γ1 = γ1,γ2
         muladd!(one(T), B', γ2, one(T), γ1)
         xy_muladd!(xy, A', γ2, one(T), γ1)
@@ -652,24 +651,23 @@ function grid(Pn::SubQuasiArray{T,2,<:JacobiTriangle,<:Tuple{Inclusion,BlockSlic
     trigrid(maximum(jr.block))
 end
 
-struct TriTransform{T}
+struct TriPlan{T}
     tri2cheb::FastTransforms.FTPlan{T,2,FastTransforms.TRIANGLE}
     grid2cheb::FastTransforms.FTPlan{T,2,FastTransforms.TRIANGLEANALYSIS}
     a::T
     b::T
     c::T
 end
 
-TriTransform{T}(F::AbstractMatrix{T}, a, b, c) where T =
-    TriTransform{T}(plan_tri2cheb(F, a, b, c), plan_tri_analysis(F), a, b, c)
+TriPlan{T}(F::AbstractMatrix{T}, a, b, c) where T =
+    TriPlan{T}(plan_tri2cheb(F, a, b, c), plan_tri_analysis(F), a, b, c)
 
-TriTransform{T}(N::Block{1}, a, b, c) where T = TriTransform{T}(Matrix{T}(undef, Int(N), Int(N)), a, b, c)
+TriPlan{T}(N::Block{1}, a, b, c) where T = TriPlan{T}(Matrix{T}(undef, Int(N), Int(N)), a, b, c)
 
-*(T::TriTransform, F::AbstractMatrix) = DiagTrav(tridenormalize!(T.tri2cheb\(T.grid2cheb*F),T.a,T.b,T.c))
+*(T::TriPlan, F::AbstractMatrix) = DiagTrav(tridenormalize!(T.tri2cheb\(T.grid2cheb*F),T.a,T.b,T.c))
 
-function factorize(V::SubQuasiArray{T,2,<:JacobiTriangle,<:Tuple{Inclusion,BlockSlice{BlockOneTo}}}) where T
-    P = parent(V)
-    _,jr = parentindices(V)
-    N = maximum(jr.block)
-    TransformFactorization(grid(V), TriTransform{T}(N, P.a, P.b, P.c))
+function plan_grid_transform(P::JacobiTriangle, arr::AbstractVector, dims=1:ndims(arr))
+    T = promote_type(eltype(P), eltype(arr))
+    N = findblock(axes(P,2), length(arr))
+    grid(P[:,Block.(OneTo(Int(N)))]), TriPlan{T}(N, P.a, P.b, P.c)
 end

test/test_disk.jl

@@ -1,6 +1,6 @@
 using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, BlockArrays, BandedMatrices, FastTransforms, LinearAlgebra, RecipesBase, Test, SpecialFunctions, LazyArrays, InfiniteArrays
 import MultivariateOrthogonalPolynomials: ModalTrav, grid, ZernikeTransform, ZernikeITransform, *, ModalInterlace
-import ClassicalOrthogonalPolynomials: HalfWeighted
+import ClassicalOrthogonalPolynomials: HalfWeighted, expand
 import ForwardDiff: hessian
 
 @testset "Disk" begin
@@ -560,4 +560,10 @@ end
             end
         end
     end
+
+    @testset "sum" begin
+        P = Zernike()
+        @test sum(expand(P, 𝐱 -> 1)) ≈ π
+        @test sum(expand(P, 𝐱 -> ((x,y) = 𝐱; exp(x*cos(y))))) ≈ 3.4898933353782744
+    end
 end

test/test_rect.jl

@@ -1,21 +1,42 @@
-using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, LinearAlgebra, Test
+using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, LinearAlgebra, BlockArrays, FillArrays, Test
+import ClassicalOrthogonalPolynomials: expand
 
 @testset "RectPolynomial" begin
     @testset "Evaluation" begin
         T = ChebyshevT()
         T² = RectPolynomial(T, T)
-        xy = SVector(0.1,0.2)
-        @test T²[xy, Block(1)[1]] == T²[xy, 1]
-        @test T²[xy, Block(1)] == T²[xy, Block.(1:1)]
-        @test T²[xy, Block(2)] == [0.1,0.2]
-        @test T²[xy, Block(3)] ≈ [cos(2*acos(0.1)), 0.1*0.2, cos(2*acos(0.2))]
+        𝐱 = SVector(0.1,0.2)
+        @test T²[𝐱, Block(1)[1]] == T²[𝐱, 1]
+        @test T²[𝐱, Block(1)] == T²[𝐱, Block.(1:1)]
+        @test T²[𝐱, Block(2)] == [0.1,0.2]
+        @test T²[𝐱, Block(3)] ≈ [cos(2*acos(0.1)), 0.1*0.2, cos(2*acos(0.2))]
 
         U = ChebyshevU()
         V = KronPolynomial(T, U)
-        @test V[xy, Block(1)[1]] == V[xy, 1]
-        @test V[xy, Block(1)] == V[xy, Block.(1:1)]
-        @test V[xy, Block(2)] == [0.1,2*0.2]
-        @test V[xy, Block(3)] ≈ [cos(2*acos(0.1)), 2*0.1*0.2, sin(3*acos(0.2))/sin(acos(0.2))]
+        @test V[𝐱, Block(1)[1]] == V[𝐱, 1]
+        @test V[𝐱, Block(1)] == V[𝐱, Block.(1:1)]
+        @test V[𝐱, Block(2)] == [0.1,2*0.2]
+        @test V[𝐱, Block(3)] ≈ [cos(2*acos(0.1)), 2*0.1*0.2, sin(3*acos(0.2))/sin(acos(0.2))]
+    end
+
+    @testset "Transform" begin
+        T = ChebyshevT()
+        T² = RectPolynomial(Fill(T, 2))
+        T²ₙ = T²[:,Block.(Base.OneTo(5))]
+        𝐱 = axes(T²ₙ,1)
+        x,y = first.(𝐱),last.(𝐱)
+        @test T²ₙ \ one.(x) == [1; zeros(14)]
+        T² \ x
+        f = expand(T², 𝐱 -> ((x,y) = 𝐱; exp(x*cos(y-0.1))))
+        @test f[SVector(0.1,0.2)] ≈ exp(0.1*cos(0.1))
+        
+        U = ChebyshevU()
+        U² = RectPolynomial(Fill(U, 2))
+
+        a,b = f.args
+        f[SVector(0.1,0.2)]
+
+        a,b = T² , (T² \ broadcast(𝐱 -> ((x,y) = 𝐱; exp(x*cos(y))), 𝐱))
     end
 
     @testset "Conversion" begin
@@ -33,8 +54,8 @@ using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticA
         T² = RectPolynomial(T, T)
         U² = RectPolynomial(U, U)
         C² = RectPolynomial(C, C)
-        xy = axes(T²,1)
-        D_x,D_y = PartialDerivative{1}(xy),PartialDerivative{2}(xy)
+        𝐱 = axes(T²,1)
+        D_x,D_y = PartialDerivative{1}(𝐱),PartialDerivative{2}(𝐱)
         D_x*T²
         D_y*T²
         U²\D_x*T²
@@ -53,8 +74,8 @@ using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticA
         W² = RectPolynomial(W, W)
         P² = RectPolynomial(P, P)
         Q² = RectPolynomial(Q, Q)
-        xy = axes(W²,1)
-        D_x,D_y = PartialDerivative{1}(xy),PartialDerivative{2}(xy)
+        𝐱 = axes(W²,1)
+        D_x,D_y = PartialDerivative{1}(𝐱),PartialDerivative{2}(𝐱)
         Δ = Q²\(D_x^2 + D_y^2)*W²
 
         K = Block.(1:200); @time L = Δ[K,K]; @time qr(L);