Support zero Dirichlet conditions (#37)

* save

* Start Dirichlet

* Support Dirichlet

* tests pass

* Work on ADI

* Arrowhead Mul

* Update adi_poisson.jl

* Avoid accessing outside arrays

* Doesn't error with p = 160

* Increase coverage

Author: dlfivefifty

Project.toml

@@ -1,7 +1,7 @@
 name = "PiecewiseOrthogonalPolynomials"
 uuid = "4461d12d-4663-4550-8580-cb764c85e20f"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.1.0"
+version = "0.2.0"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
@@ -21,16 +21,16 @@ QuasiArrays = "c4ea9172-b204-11e9-377d-29865faadc5c"
 
 [compat]
 ArrayLayouts = "1.0.12"
-BandedMatrices = "0.17.30"
+BandedMatrices = "0.17.33"
 BlockArrays = "0.16.25"
-BlockBandedMatrices = "0.12"
+BlockBandedMatrices = "0.12.2"
 ClassicalOrthogonalPolynomials = "0.11"
 ContinuumArrays = "0.14"
 FillArrays = "1.0"
 InfiniteArrays = "0.12.6"
 InfiniteLinearAlgebra = "0.6.16"
 LazyArrays = "1.0"
-LazyBandedMatrices = "0.8.13"
+LazyBandedMatrices = "0.8.15"
 MatrixFactorizations = "2.0"
 QuasiArrays = "0.11"
 julia = "1.9"

examples/adi_poisson.jl

@@ -1,6 +1,7 @@
-using PiecewiseOrthogonalPolynomials, MatrixFactorizations
+using PiecewiseOrthogonalPolynomials, MatrixFactorizations, HypergeometricFunctions
 using Elliptic
 using ClassicalOrthogonalPolynomials, StaticArrays, LinearAlgebra
+using Base: oneto
 
 """
 Solve the Poisson equation with zero Dirichlet boundary conditions on the square
@@ -16,34 +17,44 @@ function mobius(z, a, b, c, d, α)
     (t₁*z + t₂)/(t₃*z + t₄)
 end
 
+ellipticK(z) = convert(eltype(α),π)/2*HypergeometricFunctions._₂F₁(one(α)/2,one(α)/2,1, z)
 
-function ADI_shifts(J, a, b, c, d, tol)
+
+function ADI_shifts(J, a, b, c, d, tol=1e-15)
     γ = (c-a)*(d-b)/((c-b)*(d-a))
     α = -1 + 2γ + 2√Complex(γ^2 - γ)
     α = Real(α)
 
-    K = Elliptic.K(1-1/α^2)
-    dn = [Elliptic.Jacobi.dn((2*j + 1)*K/(2J), 1-1/α^2) for j = 0:J-1]
+    # K = ellipticK(1-1/big(α)^2)
+    if α < 1e7
+        K = Elliptic.K(1-1/α^2)
+        dn = Elliptic.Jacobi.dn.((2*(0:J-1) + 1)*K/(2J), 1-1/α^2)
+    else
+        K = 2log(2)+log(α) + (-1+2log(2)+log(α))/α^2/4
+        m1 = 1/α^2
+        u = (1/2:J-1/2) * K/J
+        dn = @. sech(u) + m1/4 * (sinh(u)cosh(u) + u) * tanh(u) * sech(u)
+    end
 
     [mobius(-α*i, a, b, c, d, α) for i = dn], [mobius(α*i, a, b, c, d, α) for i = dn]
 end
 
-function ADI(A, B, M, F, a, b, c, d, tol)
+function ADI(A, B, C, F, a, b, c, d, tol=1e-15)
     "ADI method for solving standard sylvester AX - XB = F"
     # Modified slightly by John to allow for the mass matrix
     n = size(A)[1]
-    X = zeros((n, n))
+    X = zeros(axes(A))
 
     γ = (c-a)*(d-b)/((c-b)*(d-a))
     J = Int(ceil(log(16γ)*log(4/tol)/π^2))
     # J = 200
     p, q = ADI_shifts(J, a, b, c, d, tol)
 
     for j = 1:J
-        X = (F - (A - p[j]*M)*X)/(B - p[j]*M)
-        X = (A - q[j]*M)\(F - X*(B - q[j]*M))
+        X = ((A/p[j] - C)*X - F/p[j])/reversecholesky(Symmetric(C - B/p[j]))
+        X = reversecholesky(Symmetric(C - A/q[j]))\(X*(B/q[j] - C) - F/q[j])
     end
-    
+
     X
 end
 
@@ -56,8 +67,8 @@ function analysis_2D(f, n, p)
 
     for i = 0:n-1  # loop over cells in positive x direction
         for j = 0:n-1  # loop over cells in positive y direction
-            local f_ = z -> ((x,y)= z; f((x + i*dx - 1, y + j*dx - 1)))  # define f on reference cell
-            F[i+1:n:n*p, j+1:n:n*p] = T * f_.(SVector.(z, z')) # interpolate f into 2D tensor Legendre polynomials on reference cell
+            f_ = (x,y) -> f(x + i*dx - 1, y + j*dx - 1)  # define f on reference cell
+            F[i+1:n:n*p, j+1:n:n*p] = T * f_.(z, z') # interpolate f into 2D tensor Legendre polynomials on reference cell
         end
     end
 
@@ -67,65 +78,42 @@ end
 r = range(-1, 1, 5)
 K = length(r)-1
 
-C = ContinuousPolynomial{1}(r)
 P = ContinuousPolynomial{0}(r)
-D = Derivative(axes(C,1))
-Δ = -weaklaplacian(C)
-M = grammatrix(C)
-e1s, e2s = [], []
-p = 40 # truncation degree on each cell
-N = K+1 + K*(p+1)  # amount of basis functions in C
-
-# Truncated Laplacian + Dirichlet bcs
-
-
+Q = DirichletPolynomial(r)
+Δ = -weaklaplacian(Q)
+M = grammatrix(Q)
 
+p = 300 # truncation degree on each cell
+KR = Block.(oneto(p))
+Δₙ = Δ[KR,KR]
+Mₙ = M[KR,KR]
 
-pΔ = Matrix(Δ[Block.(1:p), Block.(1:p)]);
-pΔ[:,1] .=0; pΔ[1,:] .=0; pΔ[1,1] = 1.;
-pΔ[:,K+1] .=0; pΔ[K+1,:] .=0; pΔ[K+1,K+1] = 1.;
-
-# Truncated mass + Dirichlet bcs
-pM = Matrix(M[Block.(1:p), Block.(1:p)]);
-pM[:,1] .=0; pM[1,:] .=0; pM[1,1] = 1.;
-pM[:,K+1] .=0; pM[K+1,:] .=0; pM[K+1,K+1] = 1.;
-
-"""
-Via the standard route ADI
-"""
-# Reverse Cholesky
-rpΔ = pΔ[end:-1:1, end:-1:1]
-L = cholesky(Symmetric(rpΔ)).L
-L = L[end:-1:1, end:-1:1]
-L * L' ≈ pΔ
+U = reversecholesky(Symmetric(Δₙ)).U
 
-# Compute spectrum
-A = (L \ (L \ pM)') # = L⁻¹ pΔ L⁻ᵀ
+A = (U \ (U \ Mₙ)') # = L⁻¹ pΔ L⁻ᵀ
 e1s, e2s = eigmin(A), eigmax(A)
 
 z = SVector.(-1:0.01:1, (-1:0.01:1)')
 
 # RHS
-f(z) = ((x,y)= z; -2 .*sin.(pi*x) .* (2pi*y .*cos.(pi*y) .+ (1-pi^2*y^2) .*sin.(pi*y)))
+f = (x,y) -> -2 .*sin.(pi*x) .* (2pi*y .*cos.(pi*y) .+ (1-pi^2*y^2) .*sin.(pi*y))
 fp = analysis_2D(f, K, p)  # interpolate F into P⊗P
-Fa = P[first.(z)[:,1], Block.(1:p)] * fp  * P[first.(z)[:,1], Block.(1:p)]'
-norm(f.(z) - Fa)
+Fa = P[first.(z)[:,1], KR] * fp  * P[first.(z)[:,1], KR]'
+norm(splat(f).(z) - Fa)
 
 # weak form for RHS
-F = (C'*P)[Block.(1:p), Block.(1:p)]*fp*((C'*P)[Block.(1:p), Block.(1:p)])'  # RHS <f,v>
-F[1, :] .= 0; F[K+1, :] .= 0; F[:, 1] .= 0; F[:, K+1] .= 0  # Dirichlet bcs
+F = (Q'*P)[KR, KR]*fp*((Q'*P)[KR, KR])'  # RHS <f,v>
 
-tol = 1e-15 # ADI tolerance
-A, B, a, b, c, d = pM, -pM, e1s, e2s, -e2s, -e1s    
-X = ADI(A, B, pΔ, F, a, b, c, d, tol)
+A, B, a, b, c, d = Mₙ, -Mₙ, e1s, e2s, -e2s, -e1s
+@time X = ADI(A, B, Δₙ, F, a, b, c, d)
 
 # X = UΔ
-U = (L' \ (L \ X'))'
+Y = (U' \ (U \ X'))'
 
 u_exact = z -> ((x,y)= z; sin.(π*x)*sin.(π*y)*y^2)
-Ua = C[first.(z)[:,1], Block.(1:p)] * U  * C[first.(z)[:,1], Block.(1:p)]'
+Ua = Q[first.(z)[:,1], Block.(1:p)] * Y  * Q[first.(z)[:,1], Block.(1:p)]'
 
-norm(u_exact.(z) - Ua) # ℓ^∞ error.
+@test u_exact.(z) ≈ Ua # ℓ^∞ error.
 
 """
 Via (5.3) and (5.6) of Kars' thesis.
@@ -153,7 +141,7 @@ F = (C'*P)[Block.(1:p), Block.(1:p)]*fp*((C'*P)[Block.(1:p), Block.(1:p)])'  # R
 F[1, :] .= 0; F[K+1, :] .= 0; F[:, 1] .= 0; F[:, K+1] .= 0  # Dirichlet bcs
 
 tol = 1e-15 # ADI tolerance
-A, B, a, b, c, d = pΔ, -pΔ, e1s, e2s, -e2s, -e1s    
+A, B, a, b, c, d = pΔ, -pΔ, e1s, e2s, -e2s, -e1s
 X = ADI(A, B, pM, F, a, b, c, d, tol)
 
 # X = UM

src/PiecewiseOrthogonalPolynomials.jl

@@ -1,7 +1,7 @@
 module PiecewiseOrthogonalPolynomials
 using ClassicalOrthogonalPolynomials, LinearAlgebra, BlockArrays, BlockBandedMatrices, BandedMatrices, ContinuumArrays, QuasiArrays, LazyArrays, LazyBandedMatrices, FillArrays, MatrixFactorizations, ArrayLayouts
 
-import ArrayLayouts: sublayout, sub_materialize, symmetriclayout, transposelayout, SymmetricLayout, HermitianLayout, TriangularLayout, layout_getindex, materialize!, MatLdivVec, AbstractStridedLayout, triangulardata
+import ArrayLayouts: sublayout, sub_materialize, symmetriclayout, transposelayout, SymmetricLayout, HermitianLayout, TriangularLayout, layout_getindex, materialize!, MatLdivVec, AbstractStridedLayout, triangulardata, MatMulMatAdd, MatMulVecAdd, _fill_lmul!
 import BandedMatrices: _BandedMatrix
 import BlockArrays: BlockSlice, block, blockindex, blockvec
 import BlockBandedMatrices: _BandedBlockBandedMatrix, AbstractBandedBlockBandedMatrix, subblockbandwidths, blockbandwidths, AbstractBandedBlockBandedLayout, layout_replace_in_print_matrix
@@ -13,11 +13,13 @@ import Base: axes, getindex, +, -, *, /, ==, \, OneTo, oneto, replace_in_print_m
 import LinearAlgebra: BlasInt
 import MatrixFactorizations: reversecholcopy
 
-export PiecewisePolynomial, ContinuousPolynomial, Derivative, Block, weaklaplacian, grammatrix
+export PiecewisePolynomial, ContinuousPolynomial, DirichletPolynomial, Derivative, Block, weaklaplacian, grammatrix
 
+include("arrowhead.jl")
 include("piecewisepolynomial.jl")
 include("continuouspolynomial.jl")
-include("arrowhead.jl")
+include("dirichlet.jl")
+
 
 
 end # module