Support expanding columns of 2D functions (#210)

* Support expanding columns of 2D functions

* use invdiagtrav

* Update Project.toml

* add QR tests

Author: dlfivefifty

Project.toml

@@ -32,19 +32,19 @@ MultivariateOrthogonalPolynomialsStatsBaseExt = "StatsBase"
 [compat]
 ArrayLayouts = "1.11"
 BandedMatrices = "1"
-BlockArrays = "1.0"
+BlockArrays = "1.7.3"
 BlockBandedMatrices = "0.13"
 ClassicalOrthogonalPolynomials = "0.15.8"
 ContinuumArrays = "0.20"
 DomainSets = "0.7"
 FastTransforms = "0.17"
 FillArrays = "1.0"
-HarmonicOrthogonalPolynomials = "0.7"
+HarmonicOrthogonalPolynomials = "0.7.1"
 InfiniteArrays = "0.15"
-InfiniteLinearAlgebra = "0.9, 0.10"
+InfiniteLinearAlgebra = "0.10"
 LazyArrays = "2.3.1"
-LazyBandedMatrices = "0.11.3"
-QuasiArrays = "0.13"
+LazyBandedMatrices = "0.11.7"
+QuasiArrays = "0.13.1"
 Random = "1"
 RecurrenceRelationships = "0.2"
 SpecialFunctions = "1, 2"

examples/christoffelsampling.jl

@@ -0,0 +1,13 @@
+using ClassicalOrthogonalPolynomials, MultivariateOrthogonalPolynomials, StatsBase, Plots
+
+
+function christoffel(A)
+    Q,R = qr(A)
+    n = size(A,2)
+    sum(expand(Q[:,k] .^2) for k=1:n)/n
+end
+
+x,y = coordinates(ChebyshevInterval() ^ 2)
+n = 3
+A = hcat([@.(cos(k*x)cos(j*y)) for k=0:n, j=0:n]...)
+K = christoffel(A)

src/MultivariateOrthogonalPolynomials.jl

@@ -1,6 +1,6 @@
 module MultivariateOrthogonalPolynomials
 using StaticArrays: iszero
-using QuasiArrays: AbstractVector
+using QuasiArrays: AbstractVector, _getindex
 using ClassicalOrthogonalPolynomials, FastTransforms, BlockBandedMatrices, BlockArrays, DomainSets,
       QuasiArrays, StaticArrays, ContinuumArrays, InfiniteArrays, InfiniteLinearAlgebra,
       LazyArrays, SpecialFunctions, LinearAlgebra, BandedMatrices, LazyBandedMatrices, ArrayLayouts,
@@ -18,10 +18,10 @@ import BlockArrays: block, blockindex, BlockSlice, viewblock, blockcolsupport, A
 import BlockBandedMatrices: _BandedBlockBandedMatrix, AbstractBandedBlockBandedMatrix, _BandedMatrix, blockbandwidths, subblockbandwidths
 import LinearAlgebra: factorize
 import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout, PaddedLayout, applylayout, LazyMatrix, ApplyMatrix
-import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes, DiagTravLayout, invdiagtrav, ApplyBandedBlockBandedLayout, krontrav
+import LazyBandedMatrices: LazyBandedBlockBandedLayout, AbstractBandedBlockBandedLayout, AbstractLazyBandedBlockBandedLayout, _krontrav_axes, DiagTravLayout, diagtrav, invdiagtrav, ApplyBandedBlockBandedLayout, krontrav
 import InfiniteArrays: InfiniteCardinal, OneToInf
 
-import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis, pad, recurrencecoefficients, clenshaw, weightedgrammatrix, Clenshaw, OPLayout
+import ClassicalOrthogonalPolynomials: jacobimatrix, Weighted, orthogonalityweight, HalfWeighted, WeightedBasis, pad, recurrencecoefficients, clenshaw, weightedgrammatrix, Clenshaw, OPLayout, normalized
 import HarmonicOrthogonalPolynomials: BivariateOrthogonalPolynomial, MultivariateOrthogonalPolynomial, Plan,
                                           AngularMomentum, angularmomentum, BlockOneTo, BlockRange1, interlace,
                                           MultivariateOPLayout, AbstractMultivariateOPLayout, MAX_PLOT_BLOCKS
@@ -38,7 +38,8 @@ export MultivariateOrthogonalPolynomial, BivariateOrthogonalPolynomial,
 
 laplacian_axis(::Inclusion{<:SVector{2}}, A; dims...) = diff(A, (2,0); dims...) + diff(A, (0, 2); dims...)
 abslaplacian_axis(::Inclusion{<:SVector{2}}, A; dims...) = -(diff(A, (2,0); dims...) + diff(A, (0, 2); dims...))
-coordinates(P) = (first.(axes(P,1)), last.(axes(P,1)))
+coordinates(P::AbstractQuasiArray) = (first.(axes(P,1)), last.(axes(P,1)))
+coordinates(P::Domain) = coordinates(Inclusion(P))
 
 function diff_layout(::AbstractBasisLayout, a, (k,j)::NTuple{2,Int}; dims...)
     (k < 0 || j < 0) && throw(ArgumentError("order must be non-negative"))

src/rect.jl

@@ -29,6 +29,8 @@ const RectPolynomial{T, PP} = KronPolynomial{2, T, PP}
 axes(P::KronPolynomial) = (Inclusion(×(map(domain, axes.(P.args, 1))...)), _krontrav_axes(axes.(P.args, 2)...))
 
 
+normalized(P::KronPolynomial) = KronPolynomial(map(normalized, P.args))
+
 function getindex(P::RectPolynomial{T}, xy::StaticVector{2}, Jj::BlockIndex{1})::T where T
     a,b = P.args
     J,j = Int(block(Jj)),blockindex(Jj)
@@ -90,6 +92,7 @@ ApplyPlan(f, P) = ApplyPlan{eltype(P), typeof(f), typeof(P)}(f, P)
 *(A::ApplyPlan, B::AbstractArray) = A.f(A.plan*B)
 
 basis_axes(d::Inclusion{<:Any,<:ProductDomain}, v) = KronPolynomial(map(d -> basis(Inclusion(d)),components(d.domain))...)
+basis_axes(d::Inclusion{<:Any,<:DomainSets.FixedIntervalProduct{N,T,D}}, v) where {N,T,D} = KronPolynomial(Fill(basis(Inclusion(D())), N))
 
 struct TensorPlan{T, Plans}
     plans::Plans
@@ -109,12 +112,20 @@ function checkpoints(P::RectPolynomial)
     SVector.(x, y')
 end
 
-function plan_transform(P::KronPolynomial{d,<:Any,<:Fill}, B::Tuple{Block{1}}, dims=1:1) where d
-    @assert dims == 1
+function plan_transform(P::KronPolynomial{d,<:Any,<:Fill}, (B,)::Tuple{Block{1}}, dims=1:1) where d
+    @assert only(dims) == 1
 
     T = first(P.args)
     @assert d == 2
-    ApplyPlan(DiagTrav, plan_transform(T, tuple(Fill(Int(B[1]),d)...)))
+    ApplyPlan(diagtrav, plan_transform(T, tuple(Fill(Int(B),d)...)))
+end
+
+function plan_transform(P::KronPolynomial{d,<:Any,<:Fill}, (B,n)::Tuple{Block{1},Int}, dims=1:1) where d
+    @assert only(dims) == 1
+
+    T = first(P.args)
+    @assert d == 2
+    ApplyPlan(A -> diagtrav(A; dims=1:d), plan_transform(T, tuple(Fill(Int(B),d)...,n), 1:d))
 end
 
 function grid(P::RectPolynomial, B::Block{1})
@@ -133,7 +144,7 @@ function plan_transform(P::KronPolynomial{d}, B::Tuple{Block{1}}, dims=1:1) wher
     @assert d == 2
     N = Int(B[1])
     Fx,Fy = plan_transform(P.args[1], (N,N), 1),plan_transform(P.args[2], (N,N), 2)
-    ApplyPlan(DiagTrav, TensorPlan(Fx,Fy))
+    ApplyPlan(diagtrav, TensorPlan(Fx,Fy))
 end
 
 applylayout(::Type{typeof(*)}, ::Lay, ::DiagTravLayout) where Lay <: AbstractBasisLayout = ExpansionLayout{Lay}()
@@ -143,6 +154,11 @@ pad(C::DiagTrav, ::BlockedOneTo{Int,RangeCumsum{Int,OneToInf{Int}}}) = DiagTrav(
 
 QuasiArrays.mul(A::BivariateOrthogonalPolynomial, b::DiagTrav) = ApplyQuasiArray(*, A, b)
 
+
+#########
+# evaluation
+########
+
 function Base.unsafe_getindex(f::Mul{KronOPLayout{2},<:DiagTravLayout{<:PaddedLayout}}, 𝐱::SVector)
     P,c = f.A, f.B
     A,B = P.args
@@ -157,9 +173,9 @@ end
 
 ## Special Legendre case
 
-function transform_ldiv(K::KronPolynomial{d,V,<:Fill{<:Legendre}}, f::Union{AbstractQuasiVector,AbstractQuasiMatrix}) where {d,V}
+function transform_ldiv(K::KronPolynomial{d,V,<:Fill{<:Legendre}}, f::AbstractQuasiVector) where {d,V}
     T = KronPolynomial{d}(Fill(ChebyshevT{V}(), size(K.args)...))
-    dat = (T \ f).array
+    dat = invdiagtrav(T \ f)
     DiagTrav(pad(FastTransforms.th_cheb2leg(paddeddata(dat)), axes(dat)...))
 end
 

test/test_rect.jl

@@ -1,7 +1,7 @@
 using MultivariateOrthogonalPolynomials, ClassicalOrthogonalPolynomials, StaticArrays, LinearAlgebra, BlockArrays, FillArrays, Base64, LazyBandedMatrices, ArrayLayouts, Random, StatsBase, Test
-using ClassicalOrthogonalPolynomials: expand, coefficients, recurrencecoefficients
+using ClassicalOrthogonalPolynomials: expand, coefficients, recurrencecoefficients, normalized
 using MultivariateOrthogonalPolynomials: weaklaplacian, ClenshawKron
-using ContinuumArrays: plotgridvalues, ExpansionLayout
+using ContinuumArrays: plotgridvalues, ExpansionLayout, basis, grid
 using Base: oneto
 
 Random.seed!(3242)
@@ -32,17 +32,47 @@ Random.seed!(3242)
         @test T²ₙ \ one.(x) == [1; zeros(14)]
         @test (T² \ x)[1:5] ≈[0;1;zeros(3)]
 
-        f = expand(T², splat((x,y) -> exp(x*cos(y-0.1))))
-        @test f[SVector(0.1,0.2)] ≈ exp(0.1*cos(0.1))
+        f = splat((x,y) -> exp(x*cos(y-0.1)))
+        𝐟 = expand(T², f)
+        @test 𝐟[SVector(0.1,0.2)] ≈ f(SVector(0.1,0.2))
 
         U² = RectPolynomial(Fill(U, 2))
-
-        @test f[SVector(0.1,0.2)] ≈ exp(0.1cos(0.1))
+        𝐟 = expand(U², f)
+        @test 𝐟[SVector(0.1,0.2)] ≈ f(SVector(0.1,0.2))
 
         TU = RectPolynomial(T,U)
-        x,F = ClassicalOrthogonalPolynomials.plan_grid_transform(TU, Block(5))
-        f = expand(TU, splat((x,y) -> exp(x*cos(y-0.1))))
-        @test f[SVector(0.1,0.2)] ≈ exp(0.1*cos(0.1))
+        𝐟 = expand(TU, f)
+        @test 𝐟[SVector(0.1,0.2)] ≈ f(SVector(0.1,0.2))
+        
+        @testset "matrix" begin
+            N = 10
+            𝐱 = grid(T², Block(N))
+            
+            @test T²[𝐱,1] == ones(N,N)
+            @test T²[𝐱,2] == first.(𝐱)
+            @test T²[𝐱,1:3] == T²[𝐱,Block.(Base.OneTo(2))] == T²[𝐱,[Block(1),Block(2)]] == [ones(N,N) ;;; first.(𝐱) ;;; last.(𝐱)]
+            @test T²[𝐱,Block(1)] == [ones(N,N) ;;;]
+            @test T²[𝐱,[1 2; 3 4]] ≈ [T²[𝐱,[1,3]] ;;;; T²[𝐱,[2,4]]]
+            
+            
+            F = plan_transform(T², Block(N))
+            @test F * f.(𝐱) ≈ transform(T², f)[Block.(1:N)] atol=1E-6
+            
+            x,y = coordinates(ChebyshevInterval()^2)
+            A = [one(x) x y]
+            F = plan_transform(T², (Block(N), 3), 1)
+            @test F * A[𝐱,:] ≈ [I(3); zeros(52,3)]
+
+            @test T² \ A ≈ [I(3); Zeros(∞,3)]
+
+            P² = RectPolynomial(Fill(Legendre(),2))
+            F = plan_transform(P², (Block(N),3), 1)
+            𝐱 = grid(P², Block(N))
+            @test F * A[𝐱,:] ≈ P²[:,Block.(Base.OneTo(N))] \ A ≈ [I(3); Zeros(52,3)]
+
+            F = plan_transform(normalized(P²), (Block(N),3), 1)
+            @test F * A[𝐱,:] ≈ normalized(P²)[:,Block.(Base.OneTo(N))] \ A ≈ [Diagonal([2, 2/sqrt(3), 2/sqrt(3)]); Zeros(52,3)]
+        end
     end
 
     @testset "Jacobi matrices" begin
@@ -254,4 +284,17 @@ Random.seed!(3242)
         @test sample(f) isa SVector
         @test sum(sample(f, 100_000))/100_000 ≈ [sum(x .* f)/sum(f),sum(y .* f)/sum(f)] rtol=1E-1
     end
+
+    @testset "qr" begin
+        x,y = coordinates(ChebyshevInterval()^2)
+        A = [one(x) cos.(x) cos.(y)]
+
+        @test A[SVector(0.1,0.2),1] ≈ 1
+        @test A[SVector(0.1,0.2),1:3] ≈ A[SVector(0.1,0.2),:] ≈ [1,cos(0.1),cos(0.2)]
+
+        Q,R = qr(A)
+        @test Q[SVector(0.1,0.2),1] ≈ 1/2
+        @test Q[SVector(0.1,0.2),2] ≈ (cos(0.1) - sin(1))/sqrt(2cos(2) + sin(2))
+        @test Q[SVector(0.1,0.2),3] ≈ (cos(0.2) - sin(1))/sqrt(2cos(2) + sin(2))
+    end
 end