/
Poisson.jl
169 lines (147 loc) · 5.3 KB
/
Poisson.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
"""
Poisson{N,M}
Composite type for conservative variable coefficient Poisson equations:
∮ds β ∂x/∂n = σ
The resulting linear system is
Ax = [L+D+L']x = z
where A is symmetric, block-tridiagonal and extremely sparse. Moreover,
`D[I]=-∑ᵢ(L[I,i]+L'[I,i])`. This means matrix storage, multiplication,
ect can be easily implemented and optimized without external libraries.
To help iteratively solve the system above, the Poisson structure holds
helper arrays for `inv(D)`, the error `ϵ`, and residual `r=z-Ax`. An iterative
solution method then estimates the error `ϵ=̃A⁻¹r` and increments `x+=ϵ`, `r-=Aϵ`.
"""
abstract type AbstractPoisson{T,S,V} end
struct Poisson{T,S<:AbstractArray{T},V<:AbstractArray{T}} <: AbstractPoisson{T,S,V}
L :: V # Lower diagonal coefficients
D :: S # Diagonal coefficients
iD :: S # 1/Diagonal
x :: S # approximate solution
ϵ :: S # increment/error
r :: S # residual
z :: S # source
n :: Vector{Int16} # pressure solver iterations
perdir :: NTuple # direction of periodic boundary condition
function Poisson(x::AbstractArray{T},L::AbstractArray{T},z::AbstractArray{T};perdir=(0,)) where T
@assert axes(x) == axes(z) && axes(x) == Base.front(axes(L)) && last(axes(L)) == eachindex(axes(x))
r = similar(x); fill!(r,0)
ϵ,D,iD = copy(r),copy(r),copy(r)
set_diag!(D,iD,L)
new{T,typeof(x),typeof(L)}(L,D,iD,x,ϵ,r,z,[],perdir)
end
end
function set_diag!(D,iD,L)
@inside D[I] = diag(I,L)
@inside iD[I] = abs2(D[I])<2eps(eltype(D)) ? 0. : inv(D[I])
end
update!(p::Poisson) = set_diag!(p.D,p.iD,p.L)
@fastmath @inline function diag(I::CartesianIndex{d},L) where {d}
s = zero(eltype(L))
for i in 1:d
s -= @inbounds(L[I,i]+L[I+δ(i,I),i])
end
return s
end
"""
mult!(p::Poisson,x)
Efficient function for Poisson matrix-vector multiplication.
Fills `p.z = p.A x` with 0 in the ghost cells.
"""
function mult!(p::Poisson,x)
@assert axes(p.z)==axes(x)
fill!(p.z,0)
@inside p.z[I] = mult(I,p.L,p.D,x)
return p.z
end
@fastmath @inline function mult(I::CartesianIndex{d},L,D,x) where {d}
s = @inbounds(x[I]*D[I])
for i in 1:d
s += @inbounds(x[I-δ(i,I)]*L[I,i]+x[I+δ(i,I)]*L[I+δ(i,I),i])
end
return s
end
"""
residual!(p::Poisson)
Computes the resiual `r = z-Ax` and corrects it such that
`r = 0` if `iD==0` which ensures local satisfiability
and
`sum(r) = 0` which ensures global satisfiability.
The global correction is done by adjusting all points uniformly,
minimizing the local effect. Other approaches are possible.
Note: These corrections mean `x` is not strictly solving `Ax=z`, but
without the corrections, no solution exists.
"""
function residual!(p::Poisson)
@inside p.r[I] = ifelse(p.iD[I]==0,0,p.z[I]-mult(I,p.L,p.D,p.x))
s = sum(p.r)/length(p.r[inside(p.r)])
abs(s) <= 2eps(eltype(s)) && return
@inside p.r[I] = p.r[I]-s
end
increment!(p::Poisson) = @loop (p.r[I] = p.r[I]-mult(I,p.L,p.D,p.ϵ);
p.x[I] = p.x[I]+p.ϵ[I]) over I ∈ inside(p.x)
"""
Jacobi!(p::Poisson; it=1)
Jacobi smoother run `it` times.
Note: This runs for general backends, but is _very_ slow to converge.
"""
@fastmath Jacobi!(p;it=1) = for _ ∈ 1:it
@inside p.ϵ[I] = p.r[I]*p.iD[I]
BC!(p.ϵ;perdir=p.perdir)
increment!(p)
end
using LinearAlgebra: ⋅
"""
pcg!(p::Poisson; it=6)
Conjugate-Gradient smoother with Jacobi preditioning. Runs at most `it` iterations,
but will exit early if the Gram-Schmidt update parameter `|α| < 1%` or `|r D⁻¹ r| < 1e-8`.
Note: This runs for general backends and is the default smoother.
"""
function pcg!(p::Poisson{T};it=6) where T
x,r,ϵ,z = p.x,p.r,p.ϵ,p.z
@inside z[I] = ϵ[I] = r[I]*p.iD[I]
insideI = inside(x) # [insideI]
rho = T(r⋅z)
abs(rho)<10eps(T) && return
for i in 1:it
BC!(ϵ;perdir=p.perdir)
@inside z[I] = mult(I,p.L,p.D,ϵ)
alpha = rho/T(z[insideI]⋅ϵ[insideI])
@loop (x[I] += alpha*ϵ[I];
r[I] -= alpha*z[I]) over I ∈ inside(x)
(i==it || abs(alpha)<1e-2) && return
@inside z[I] = r[I]*p.iD[I]
rho2 = T(r⋅z)
abs(rho2)<10eps(T) && return
beta = rho2/rho
@inside ϵ[I] = beta*ϵ[I]+z[I]
rho = rho2
end
end
smooth!(p) = pcg!(p)
L₂(p::Poisson) = p.r ⋅ p.r # special method since outside(p.r)≡0
L∞(p::Poisson) = maximum(abs.(p.r))
"""
solver!(A::Poisson;log,tol,itmx)
Approximate iterative solver for the Poisson matrix equation `Ax=b`.
- `A`: Poisson matrix with working arrays.
- `A.x`: Solution vector. Can start with an initial guess.
- `A.z`: Right-Hand-Side vector. Will be overwritten!
- `A.n[end]`: stores the number of iterations performed.
- `log`: If `true`, this function returns a vector holding the `L₂`-norm of the residual at each iteration.
- `tol`: Convergence tolerance on the `L₂`-norm residual.
- `itmx`: Maximum number of iterations.
"""
function solver!(p::Poisson;log=false,tol=1e-4,itmx=1e3)
BC!(p.x;perdir=p.perdir)
residual!(p); r₂ = L₂(p)
log && (res = [r₂])
nᵖ=0
while r₂>tol && nᵖ<itmx
smooth!(p); r₂ = L₂(p)
log && push!(res,r₂)
nᵖ+=1
end
BC!(p.x;perdir=p.perdir)
push!(p.n,nᵖ)
log && return res
end