Support expanding columns of 2D functions

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"
 LazyArrays = "2.3.1"
 LazyBandedMatrices = "0.11.3"
-QuasiArrays = "0.13"
+QuasiArrays = "0.13.1"
 Random = "1"
 RecurrenceRelationships = "0.2"
 SpecialFunctions = "1, 2"

examples/christoffelsampling.jl

@@ -0,0 +1,4 @@
+using ClassicalOrthogonalPolynomials, MultivariateOrthogonalPolynomials, StatsBase, Plots
+
+x,y = coordinates(ChebyshevInterval() ^ 2)
+qr([one(x) x y])

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)
@@ -109,12 +111,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),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(DiagTrav, plan_transform(T, tuple(Fill(Int(B[1]),d)...)))
+    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 +143,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}()

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 MultivariateOrthogonalPolynomials: weaklaplacian, ClenshawKron
-using ContinuumArrays: plotgridvalues, ExpansionLayout
+using ContinuumArrays: plotgridvalues, ExpansionLayout, basis, grid
 using Base: oneto
 
 Random.seed!(3242)
@@ -32,17 +32,39 @@ 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)]
+        end
     end
 
     @testset "Jacobi matrices" begin
@@ -254,4 +276,20 @@ 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) x 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,0.1,0.2]
+
+        P = basis(x)
+        P\A
+
+
+        B[0.1,:]
+
+        P \ A
+    end
 end