Skip to content

Commit 14f394b

Browse files
vchuravyvchuravy
authored
Heat2D example (#20)
1 parent f0430b9 commit 14f394b

File tree

9 files changed

+554
-10
lines changed

9 files changed

+554
-10
lines changed

‎docs/Project.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ KrylovPreconditioners = "45d422c2-293f-44ce-8315-2cb988662dec"
88
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
99
Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
1010
NewtonKrylov = "0be81120-40bf-4f8b-adf0-26103efb66f1"
11+
OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
1112
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
1213
SummationByPartsOperators = "9f78cca6-572e-554e-b819-917d2f1cf240"
1314

‎docs/make.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ examples = [
2525
"Implicit -- Spring" => "spring",
2626
"Implicit -- Heat 1D" => "heat_1D",
2727
"Implicit -- Heat 1D DG" => "heat_1D_DG",
28+
"Implicit -- Heat 2D" => "heat_2D",
2829
]
2930

3031
for (_, name) in examples

‎examples/Project.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
77
KrylovPreconditioners = "45d422c2-293f-44ce-8315-2cb988662dec"
88
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
99
NewtonKrylov = "0be81120-40bf-4f8b-adf0-26103efb66f1"
10+
Observables = "510215fc-4207-5dde-b226-833fc4488ee2"
11+
OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
1012
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
1113
SummationByPartsOperators = "9f78cca6-572e-554e-b819-917d2f1cf240"
1214

‎examples/halovector.jl‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
using OffsetArrays
2+
3+
struct HaloVector{FC, D} <: AbstractVector{FC}
4+
data::D
5+
6+
function HaloVector(data::D) where {D}
7+
FC = eltype(data)
8+
return new{FC, D}(data)
9+
end
10+
end
11+
12+
function Base.similar(v::HaloVector)
13+
data = similar(v.data)
14+
return HaloVector(data)
15+
end
16+
17+
function Base.length(v::HaloVector)
18+
m, n = size(v.data)
19+
l = (m - 2) * (n - 2)
20+
return l
21+
end
22+
23+
function Base.size(v::HaloVector)
24+
l = length(v)
25+
return (l,)
26+
end
27+
28+
function Base.getindex(v::HaloVector, idx)
29+
m, n = size(v.data)
30+
row = div(idx - 1, n - 2) + 1
31+
col = mod(idx - 1, n - 2) + 1
32+
return v.data[row, col]
33+
end
34+
35+
function Base.setindex!(v::HaloVector, val, idx)
36+
m, n = size(v.data)
37+
row = div(idx - 1, n - 2) + 1
38+
col = mod(idx - 1, n - 2) + 1
39+
return v.data[row, col] = val
40+
end
41+
42+
Base.zero(v::HaloVector) = HaloVector(zero(v.data))
43+
44+
# Otherwise we get Uninitialized data in the halo...
45+
Base.copy(v::HaloVector) = HaloVector(copy(v.data))
46+
47+
48+
using Krylov
49+
import Krylov.FloatOrComplex
50+
51+
function Krylov.kdot(n::Integer, x::HaloVector{T}, y::HaloVector{T}) where {T <: FloatOrComplex}
52+
mx, nx = size(x.data)
53+
_x = x.data
54+
_y = y.data
55+
res = zero(T)
56+
for i in 1:(mx - 1)
57+
for j in 1:(nx - 1)
58+
res += _x[i, j] * _y[i, j]
59+
end
60+
end
61+
return res
62+
end
63+
64+
function Krylov.knorm(n::Integer, x::HaloVector{T}) where {T <: FloatOrComplex}
65+
mx, nx = size(x.data)
66+
_x = x.data
67+
res = zero(T)
68+
for i in 1:(mx - 1)
69+
for j in 1:(nx - 1)
70+
res += _x[i, j]^2
71+
end
72+
end
73+
return sqrt(res)
74+
end
75+
76+
function Krylov.kscal!(n::Integer, s::T, x::HaloVector{T}) where {T <: FloatOrComplex}
77+
mx, nx = size(x.data)
78+
_x = x.data
79+
for i in 1:(mx - 1)
80+
for j in 1:(nx - 1)
81+
_x[i, j] = s * _x[i, j]
82+
end
83+
end
84+
return x
85+
end
86+
87+
function Krylov.kaxpy!(n::Integer, s::T, x::HaloVector{T}, y::HaloVector{T}) where {T <: FloatOrComplex}
88+
mx, nx = size(x.data)
89+
_x = x.data
90+
_y = y.data
91+
for i in 1:(mx - 1)
92+
for j in 1:(nx - 1)
93+
_y[i, j] += s * _x[i, j]
94+
end
95+
end
96+
return y
97+
end
98+
99+
function Krylov.kaxpby!(n::Integer, s::T, x::HaloVector{T}, t::T, y::HaloVector{T}) where {T <: FloatOrComplex}
100+
mx, nx = size(x.data)
101+
_x = x.data
102+
_y = y.data
103+
for i in 1:(mx - 1)
104+
for j in 1:(nx - 1)
105+
_y[i, j] = s * _x[i, j] + t * _y[i, j]
106+
end
107+
end
108+
return y
109+
end
110+
111+
function Krylov.kcopy!(n::Integer, y::HaloVector{T}, x::HaloVector{T}) where {T <: FloatOrComplex}
112+
mx, nx = size(x.data)
113+
_x = x.data
114+
_y = y.data
115+
for i in 1:(mx - 1)
116+
for j in 1:(nx - 1)
117+
_y[i, j] = _x[i, j]
118+
end
119+
end
120+
return y
121+
end
122+
123+
function Krylov.kfill!(x::HaloVector{T}, val::T) where {T <: FloatOrComplex}
124+
mx, nx = size(x.data)
125+
_x = x.data
126+
for i in 1:(mx - 1)
127+
for j in 1:(nx - 1)
128+
_x[i, j] = val
129+
end
130+
end
131+
return x
132+
end
133+
134+
function Krylov.kref!(n::Integer, x::HaloVector{T}, y::HaloVector{T}, c::T, s::T) where {T <: FloatOrComplex}
135+
mx, nx = size(x.data)
136+
_x = x.data
137+
_y = y.data
138+
for i in 1:(mx - 1)
139+
for j in 1:(nx - 1)
140+
x_ij = _x[i, j]
141+
y_ij = _y[i, j]
142+
_x[i, j] = c * x_ij + s * y_ij
143+
_y[i, j] = conj(s) * x_ij - c * y_ij
144+
end
145+
end
146+
return x, y
147+
end

‎examples/heat_2D.jl‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
# # Implicit time-integration
2+
3+
# ## Necessary packages
4+
using NewtonKrylov
5+
using CairoMakie
6+
using OffsetArrays
7+
8+
include(joinpath(dirname(pathof(NewtonKrylov)), "..", "examples", "implicit.jl"))
9+
include(joinpath(dirname(pathof(NewtonKrylov)), "..", "examples", "halovector.jl"))
10+
11+
# ## Diffusion 2D
12+
13+
# ### Boundary
14+
15+
function bc_periodic!(u)
16+
N, M = size(u)
17+
N = N - 2
18+
M = M - 2
19+
20+
## (wrap around)
21+
u[0, :] .= u[N, :]
22+
u[N + 1, :] .= u[1, :]
23+
u[:, 0] .= u[:, N]
24+
u[:, N + 1] .= u[:, 1]
25+
return nothing
26+
end
27+
28+
function bc_zero!(u)
29+
N, M = size(u)
30+
N = N - 2
31+
M = M - 2
32+
33+
u[0, :] .= 0
34+
u[N + 1, :] .= 0
35+
u[:, 0] .= 0
36+
u[:, N + 1] .= 0
37+
return nothing
38+
end
39+
40+
41+
function diffusion!(du::HaloVector, u::HaloVector, p, t)
42+
return diffusion!(du.data, u.data, p, t)
43+
end
44+
45+
function diffusion!(du, u, (a, Δx, Δy, bc!), _)
46+
N, M = size(u)
47+
N = N - 2
48+
M = M - 2
49+
50+
## Enforce boundary conditions
51+
bc!(u)
52+
53+
for i in 1:N
54+
for j in 1:M
55+
du[i, j] = a * (
56+
(u[i + 1, j] - 2 * u[i, j] + u[i - 1, j]) / Δx^2 +
57+
(u[i, j + 1] - 2 * u[i, j] + u[i, j - 1]) / Δy^2
58+
)
59+
end
60+
end
61+
return
62+
end
63+
64+
a = 0.01
65+
66+
N = M = 40
67+
68+
Δx = 1 / (N + 1) # x-grid spacing
69+
Δy = 1 / (M + 1) # y-grid spacing
70+
71+
## TODO: This is a time-step from explicit
72+
Δt = Δx^2 * Δy^2 / (2.0 * a * (Δx^2 + Δy^2)) # Largest stable time step
73+
74+
let
75+
N = M = 12
76+
u₀ = HaloVector(OffsetArray(zeros(N + 2, M + 2), 0:(N + 1), 0:(M + 1)))
77+
jacobian(G_Euler!, diffusion!, u₀, (a, Δx, Δy, bc_zero!), Δt, 0.0)
78+
end
79+
80+
xs = 0.0:Δx:1.0
81+
ys = 0.0:Δy:1.0
82+
83+
function f(x, y)
84+
return sin(π * x) .* sin(π * y)
85+
end
86+
87+
u₀ = let
88+
_u₀ = zeros(N + 2, M + 2)
89+
_u₀ .= f.(xs, ys')
90+
HaloVector(OffsetArray(_u₀, 0:(N + 1), 0:(M + 1)))
91+
end
92+
93+
function plot_state(t, u, xs, ys)
94+
data = @lift $u.data.parent
95+
96+
xs_line = @lift @view $data[:, M ÷ 2]
97+
ys_line = @lift @view $data[N ÷ 2, :]
98+
99+
fig = Figure()
100+
101+
ax = Axis(fig[1, 1])
102+
heatmap!(ax, xs, ys, data)
103+
104+
ax1 = Axis(fig[2, 1])
105+
lines!(ax1, xs, xs_line)
106+
107+
ax2 = Axis(fig[2, 2])
108+
lines!(ax2, ys, ys_line)
109+
110+
fig[0, :] = Label(fig, @lift("t = $($t)"))
111+
112+
return fig
113+
end
114+
115+
fig = plot_state(Observable(0.0), Observable(u₀), xs, ys)
116+
117+
function create_video_implicit(filename, G!, f!, xs, ys, u, p, Δt, t_stop, frame_kwargs)
118+
ts = 0.0:Δt:t_stop
119+
_u = Observable(u)
120+
_t = Observable(first(ts))
121+
122+
fig = plot_state(_t, _u, xs, ys)
123+
return record(fig, filename; frame_kwargs...) do io
124+
function callback(__u)
125+
_u[] = u
126+
_t[] += Δt
127+
## autolimits!(ax) # update limits
128+
recordframe!(io)
129+
return yield()
130+
end
131+
solve(G!, f!, u, p, Δt, ts; callback, verbose = 1, krylov_kwargs = (; verbose = 1, reorthogonalization = true))
132+
end
133+
end
134+
135+
create_video_implicit(
136+
joinpath(@__DIR__, "implicit_euler.mp4"),
137+
G_Euler!, diffusion!, xs, ys, copy(u₀), (a, Δx, Δy, bc_zero!), Δt, 10.0, (; framerate = 30)
138+
)
139+
140+
create_video_implicit(
141+
joinpath(@__DIR__, "implicit_midpoint.mp4"),
142+
G_Midpoint!, diffusion!, xs, ys, copy(u₀), (a, Δx, Δy, bc_zero!), Δt, 10.0, (; framerate = 30)
143+
)
144+
145+
create_video_implicit(
146+
joinpath(@__DIR__, "implicit_trapezoid.mp4"),
147+
G_Trapezoid!, diffusion!, xs, ys, copy(u₀), (a, Δx, Δy, bc_zero!), Δt, 10.0, (; framerate = 30)
148+
)
149+
150+
## TODO:
151+
create_video_implicit(
152+
joinpath(@__DIR__, "implicit_euler_periodic.mp4"),
153+
G_Trapezoid!, diffusion!, xs, ys, copy(u₀), (a, Δx, Δy, bc_periodic!), Δt, 2 * Δt, (; framerate = 30)
154+
)

‎examples/implicit.jl‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
# # [Implicit schemes](@id implicit_schemes)
22
using NewtonKrylov
33

4+
using NewtonKrylov
5+
46
# ## Implicit Euler
57

68
function G_Euler!(res, uₙ, Δt, f!, du, u, p, t)
@@ -38,8 +40,8 @@ end
3840

3941
function jacobian(G!, f!, uₙ, p, Δt, t)
4042
u = copy(uₙ)
41-
du = similar(uₙ)
42-
res = similar(uₙ)
43+
du = zero(uₙ)
44+
res = zero(uₙ)
4345

4446
F!(res, u, (uₙ, Δt, du, p, t)) = G!(res, uₙ, Δt, f!, du, u, p, t)
4547

@@ -51,17 +53,23 @@ end
5153

5254
import Krylov
5355

54-
function solve(G!, f!, uₙ, p, Δt, ts; callback = _ -> nothing, verbose = 0, Workspace = Krylov.GmresWorkspace)
56+
function solve(
57+
G!, f!, uₙ, p, Δt, ts; callback = _ -> nothing,
58+
verbose = 0, Workspace = Krylov.GmresWorkspace, krylov_kwargs = (;)
59+
)
5560
u = copy(uₙ)
56-
du = similar(uₙ)
57-
res = similar(uₙ)
61+
du = zero(uₙ)
62+
res = zero(uₙ)
5863
F!(res, u, (uₙ, Δt, du, p, t)) = G!(res, uₙ, Δt, f!, du, u, p, t)
5964

6065
for t in ts
6166
if t == first(ts)
6267
continue
6368
end
64-
_, stats = newton_krylov!(F!, u, (uₙ, Δt, du, p, t), res; verbose, Workspace, tol_abs = 6.0e-6)
69+
_, stats = newton_krylov!(
70+
F!, u, (uₙ, Δt, du, p, t), res;
71+
verbose, Workspace, tol_abs = 6.0e-6, krylov_kwargs
72+
)
6573
if !stats.solved
6674
@warn "non linear solve failed marching on" t stats
6775
end

‎notebooks/Project.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
33
Krylov = "ba0b0d4f-ebba-5204-a429-3ac8c609bfb7"
44
NewtonKrylov = "0be81120-40bf-4f8b-adf0-26103efb66f1"
5+
OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
56
Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781"
67
PlutoLinks = "0ff47ea0-7a50-410d-8455-4348d5de0420"
78
PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8"

0 commit comments

Comments
 (0)