Use transform for plotting (#127)

* Use transform for plotting

* Separate weighted and regular Zernike plot

* Update test_disk.jl

* test weighted zernike plot

Author: dlfivefifty

examples/diskhelmholtz.jl

@@ -1,5 +1,5 @@
 using MultivariateOrthogonalPolynomials, FastTransforms, BlockBandedMatrices, Plots, LinearAlgebra, StaticArrays
-plotly()
+pyplot()
 
 
 ####
@@ -25,7 +25,7 @@ L = Δ + k^2 * S # discretisation of Helmholtz
 f = @.(cos(x * exp(y)))
 
 u = W * (L \ (Z \ f))
-surface(u)
+contourf(u)
 
 
 # One can also fix the discretisation size
@@ -36,16 +36,14 @@ Wₙ = W[:,Block.(1:N)]
 Lₙ = L[Block.(1:N),Block.(1:N)]
 
 u = Wₙ * (Lₙ \ (Zₙ \ f))
-surface(u)
+contourf(u)
 
 
 # We can also do eigenvalues of the Laplacian
-
+N = 20
 Δₙ = Δ[Block.(1:N),Block.(1:N)]
 Sₙ = S[Block.(1:N),Block.(1:N)]
 
-BandedMatrix(Δₙ)
-
 λ,Q = eigen(Symmetric(Matrix(Δₙ)), Symmetric(Matrix(Sₙ)))
 
-surface(Wₙ * Q[:,end])
+contourf(Wₙ * Q[:,end-10])

src/MultivariateOrthogonalPolynomials.jl

@@ -10,10 +10,10 @@ import Base: axes, in, ==, *, ^, \, copy, OneTo, getindex, size, oneto, all, res
 import DomainSets: boundary
 
 import QuasiArrays: LazyQuasiMatrix, LazyQuasiArrayStyle
-import ContinuumArrays: @simplify, Weight, grid, plotgrid, TransformFactorization, ExpansionLayout
+import ContinuumArrays: @simplify, Weight, weight, grid, plotgrid, TransformFactorization, ExpansionLayout, plotvalues, unweighted
 
 import ArrayLayouts: MemoryLayout, sublayout, sub_materialize
-import BlockArrays: block, blockindex, BlockSlice, viewblock
+import BlockArrays: block, blockindex, BlockSlice, viewblock, blockcolsupport
 import BlockBandedMatrices: _BandedBlockBandedMatrix, AbstractBandedBlockBandedMatrix, _BandedMatrix, blockbandwidths, subblockbandwidths
 import LinearAlgebra: factorize
 import LazyArrays: arguments, paddeddata, LazyArrayStyle, LazyLayout

src/disk.jl

@@ -286,6 +286,24 @@ function plotgrid(S::SubQuasiArray{<:Any,2,<:Zernike})
     plotgrid(Z[kr,Block.(OneTo(Int(findblock(axes(Z,2),maximum(jr)))))])
 end
 
+
+function plotvalues(u::ApplyQuasiVector{T,typeof(*),<:Tuple{Zernike, AbstractVector}}, x) where T
+    Z,c = u.args
+    CS = blockcolsupport(c)
+    N = Int(last(CS)) ÷ 2 + 1 # polynomial degree
+    F = ZernikeITransform{T}(2N, Z.a, Z.b)
+    C = F * c[Block.(OneTo(2N))] # transform to grid
+    [permutedims(u[x[1,:]]); # evaluate on edge of disk
+     C C[:,1];
+     fill(u[x[end,1]], 1, size(x,2))] # evaluate at origin and repeat
+end
+
+function plotvalues(u::ApplyQuasiVector{T,typeof(*),<:Tuple{Weighted{<:Any,<:Zernike}, AbstractVector}}, x) where T
+    U = plotvalues(unweighted(u), x)
+    w = weight(u.args[1])
+    w[x] .* U
+end
+
 struct ZernikeTransform{T} <: Plan{T}
     N::Int
     disk2cxf::FastTransforms.FTPlan{T,2,FastTransforms.DISK}

test/test_disk.jl

@@ -311,7 +311,13 @@ import ForwardDiff: hessian
         u = Z * [1; 2; zeros(∞)];
         rep = RecipesBase.apply_recipe(Dict{Symbol, Any}(), u)
         g = MultivariateOrthogonalPolynomials.plotgrid(Z[:,1:3])
-        @test rep[1].args == (first.(g),last.(g),u[g])
+        @test all(rep[1].args .≈ (first.(g),last.(g),u[g]))
+
+        W = Weighted(Zernike(1))
+        u = W * [1; 2; zeros(∞)];
+        rep = RecipesBase.apply_recipe(Dict{Symbol, Any}(), u)
+        g = MultivariateOrthogonalPolynomials.plotgrid(W[:,1:3])
+        @test all(rep[1].args .≈ (first.(g),last.(g),u[g]))
     end
 end