Skip to content

Commit 82efbc7

Browse files
sbairosSteveDanielDoehringranocha
authored
Hyperbolic-Parabolic elixir using SparseConnectivityTracer (#2604)
* Hyperbolic-Parabolic elixir using SCT * Add CI test * PR corrections * Adding missed comment * Navierstokes hyperbolic parabolic convergence test * Adding CI test for advection_diffusion elixir * Revert "Adding CI test for advection_diffusion elixir" This reverts commit 5c41983. * Apply suggestions from code review * Apply suggestions from code review * Update examples/tree_2d_dgsem/elixir_navierstokes_convergence_implicit_sparse_jacobian.jl * Adding CI test for advection_diffusion elixir * Revert "Adding CI test for advection_diffusion elixir" This reverts commit 5c41983. * Convert adv_diffusion to 1D and N-S to P4est * Use IMEX methods * Add demonstration of adv_velocity=0 * Fix tests * Apply suggestions from code review * PR comments * PR comments2 * Undo overwrite of Project.toml * Remove the navierstokes elixir and test * Update unit test to actually include rhs! * Update test/test_unit.jl * Apply suggestions from code review * Update test/test_parabolic_1d.jl * Update test/test_parabolic_1d.jl * Formatting changes * PR comments * Spell check / formatting * Formatting x2 * formatting x3 * formatting last time * Apply suggestions from code review * Apply suggestions from code review * Apply suggestions from code review * Fix unit test * Update test/test_unit.jl * Update test/test_unit.jl * Update test/test_unit.jl * Apply suggestions from code review * Update test/test_unit.jl * Apply suggestions from code review * Update test/test_unit.jl * Update examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian.jl * Update examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl * Apply suggestions from code review Co-authored-by: Hendrik Ranocha <[email protected]> * Apply suggestions from code review Co-authored-by: Hendrik Ranocha <[email protected]> * Update test/test_unit.jl * Update test/test_unit.jl * Update test/test_unit.jl * Apply suggestions from code review * db * update comments --------- Co-authored-by: Steve <[email protected]> Co-authored-by: Daniel Doehring <[email protected]> Co-authored-by: Hendrik Ranocha <[email protected]>
1 parent b5ecc2a commit 82efbc7

File tree

6 files changed

+368
-14
lines changed

6 files changed

+368
-14
lines changed

‎examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian.jl‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
using Trixi
2+
using SparseConnectivityTracer # For obtaining the Jacobian sparsity pattern
3+
using SparseMatrixColorings # For obtaining the coloring vector
4+
using OrdinaryDiffEqBDF, ADTypes
5+
6+
###############################################################################
7+
# semidiscretization of the linear advection-diffusion equation
8+
9+
advection_velocity = 1.5
10+
equations_hyperbolic = LinearScalarAdvectionEquation1D(advection_velocity)
11+
diffusivity() = 5.0e-2
12+
equations_parabolic = LaplaceDiffusion1D(diffusivity(), equations_hyperbolic)
13+
14+
solver = DGSEM(polydeg = 3, surface_flux = flux_lax_friedrichs)
15+
16+
coordinates_min = -1.0
17+
coordinates_max = 1.0
18+
19+
mesh = TreeMesh(coordinates_min, coordinates_max,
20+
initial_refinement_level = 4,
21+
n_cells_max = 30_000)
22+
23+
function initial_condition_diffusive_convergence_test(x, t,
24+
equation::LinearScalarAdvectionEquation1D)
25+
# Store translated coordinate for easy use of exact solution
26+
RealT = eltype(x)
27+
x_trans = x - equation.advection_velocity * t
28+
29+
nu = diffusivity()
30+
c = 1
31+
A = 0.5f0
32+
L = 2
33+
f = 1.0f0 / L
34+
omega = 2 * convert(RealT, pi) * f
35+
scalar = c + A * sin(omega * sum(x_trans)) * exp(-2 * nu * omega^2 * t)
36+
return SVector(scalar)
37+
end
38+
initial_condition = initial_condition_diffusive_convergence_test
39+
40+
###############################################################################
41+
### semidiscretization for sparsity detection ###
42+
43+
jac_detector = TracerSparsityDetector()
44+
# We need to construct the semidiscretization with the correct
45+
# sparsity-detection ready datatype, which is retrieved here
46+
jac_eltype = jacobian_eltype(real(solver), jac_detector)
47+
48+
semi_jac_type = SemidiscretizationHyperbolicParabolic(mesh,
49+
(equations_hyperbolic,
50+
equations_parabolic),
51+
initial_condition, solver,
52+
uEltype = jac_eltype)
53+
54+
tspan = (0.0, 1.5) # Re-used for wrapping `rhs_parabolic!` below
55+
56+
# Call `semidiscretize` to create the ODE problem to have access to the
57+
# initial condition based on which the sparsity pattern is computed
58+
ode_jac_type = semidiscretize(semi_jac_type, tspan)
59+
u0_ode = ode_jac_type.u0
60+
du_ode = similar(u0_ode)
61+
62+
###############################################################################
63+
### Compute the Jacobian sparsity pattern ###
64+
65+
# Only the parabolic part of the `SplitODEProblem` is treated implicitly so we only need the parabolic Jacobian, see
66+
# https://docs.sciml.ai/DiffEqDocs/stable/types/split_ode_types/#SciMLBase.SplitFunction
67+
# Wrap the `Trixi.rhs_parabolic!` function to match the signature `f!(du, u)`, see
68+
# https://adrianhill.de/SparseConnectivityTracer.jl/stable/user/api/#ADTypes.jacobian_sparsity
69+
rhs_parabolic_wrapped! = (du_ode, u0_ode) -> Trixi.rhs_parabolic!(du_ode, u0_ode,
70+
semi_jac_type, tspan[1])
71+
72+
jac_prototype_parabolic = jacobian_sparsity(rhs_parabolic_wrapped!, du_ode, u0_ode,
73+
jac_detector)
74+
75+
# For most efficient solving we also want the coloring vector
76+
77+
# We choose `nonsymmetric` `structure` because we're computing a Jacobian, which
78+
# is for the upwind-alike discretization of the advection term nonsymmmetric
79+
# We arbitrarily choose a column-based `partitioning`. This means that we will color
80+
# structurally orthogonal columns the same. `row` partitioning would be equally valid here
81+
coloring_prob = ColoringProblem(; structure = :nonsymmetric, partition = :column)
82+
# The `decompression` arg specifies our evaluation scheme. The `direct` method requires solving
83+
# a diagonal system, whereas the `substitution` method solves a triangular system of equations
84+
coloring_alg = GreedyColoringAlgorithm(; decompression = :direct)
85+
coloring_result = coloring(jac_prototype_parabolic, coloring_prob, coloring_alg)
86+
coloring_vec_parabolic = column_colors(coloring_result)
87+
88+
###############################################################################
89+
### sparsity-aware semidiscretization and ODE ###
90+
91+
# Semidiscretization for actual simulation. `uEltype` is here retrieved from `solver`
92+
semi = SemidiscretizationHyperbolicParabolic(mesh,
93+
(equations_hyperbolic,
94+
equations_parabolic),
95+
initial_condition, solver)
96+
97+
# Supply Jacobian prototype and coloring vector to the semidiscretization
98+
ode = semidiscretize(semi, tspan,
99+
jac_prototype_parabolic = jac_prototype_parabolic,
100+
colorvec_parabolic = coloring_vec_parabolic)
101+
# using "dense" `ode = semidiscretize(semi, tspan)` is 4-6 times slower!
102+
103+
###############################################################################
104+
### callbacks ###
105+
106+
summary_callback = SummaryCallback()
107+
108+
analysis_interval = 100
109+
analysis_callback = AnalysisCallback(semi, interval = analysis_interval)
110+
111+
alive_callback = AliveCallback(analysis_interval = analysis_interval)
112+
113+
save_restart = SaveRestartCallback(interval = 100,
114+
save_final_restart = true)
115+
116+
# Note: No `stepsize_callback` due to (implicit) solver with adaptive timestep control
117+
callbacks = CallbackSet(summary_callback, analysis_callback, alive_callback, save_restart)
118+
119+
###############################################################################
120+
### solve the ODE problem ###
121+
122+
sol = solve(ode, SBDF2(; autodiff = AutoFiniteDiff());
123+
ode_default_options()...,
124+
dt = 0.01,
125+
abstol = 1e-9, reltol = 1e-9,
126+
callback = callbacks)

‎examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
using Trixi
2+
3+
###############################################################################
4+
# create a restart file
5+
6+
elixir_file = "elixir_advection_diffusion_implicit_sparse_jacobian.jl"
7+
restart_file = "restart_000000100.h5"
8+
9+
trixi_include(@__MODULE__, joinpath(@__DIR__, elixir_file))
10+
11+
###############################################################################
12+
13+
restart_filename = joinpath("out", restart_file)
14+
tspan = (load_time(restart_filename), 2.0)
15+
dt_restart = load_dt(restart_filename)
16+
17+
ode = semidiscretize(semi, tspan,
18+
restart_filename,
19+
jac_prototype_parabolic = jac_prototype_parabolic,
20+
colorvec_parabolic = coloring_vec_parabolic)
21+
22+
###############################################################################
23+
# run the simulation
24+
25+
sol = solve(ode, SBDF2(; autodiff = AutoFiniteDiff());
26+
ode_default_options()...,
27+
dt = dt_restart,
28+
abstol = 1e-9, reltol = 1e-9,
29+
callback = callbacks);

‎src/semidiscretization/semidiscretization.jl‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -141,8 +141,9 @@ function semidiscretize(semi::AbstractSemidiscretization, tspan;
141141

142142
return ODEProblem{iip, specialize}(ode, u0_ode, tspan, semi)
143143
else
144-
# We could also construct an `ODEFunction` without the Jacobian here,
145-
# but we stick to the more light-weight direct in-place function `rhs!`.
144+
# We could also construct an `ODEFunction` explicitly without the Jacobian here,
145+
# but we stick to the lean direct in-place function `rhs!` and
146+
# let OrdinaryDiffEq.jl handle the rest
146147
return ODEProblem{iip, specialize}(rhs!, u0_ode, tspan, semi)
147148
end
148149
end
@@ -194,8 +195,9 @@ function semidiscretize(semi::AbstractSemidiscretization, tspan,
194195

195196
return ODEProblem{iip, specialize}(ode, u0_ode, tspan, semi)
196197
else
197-
# We could also construct an `ODEFunction` without the Jacobian here,
198-
# but we stick to the more light-weight direct in-place function `rhs!`.
198+
# We could also construct an `ODEFunction` explicitly without the Jacobian here,
199+
# but we stick to the lean direct in-place function `rhs!` and
200+
# let OrdinaryDiffEq.jl handle the rest
199201
return ODEProblem{iip, specialize}(rhs!, u0_ode, tspan, semi)
200202
end
201203
end

‎src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl‎

Lines changed: 68 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -264,15 +264,28 @@ function get_node_variables!(node_variables, u_ode,
264264
end
265265

266266
"""
267-
semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan)
267+
semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan;
268+
jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
269+
colorvec_parabolic::Union{AbstractVector, Nothing} = nothing)
268270
269271
Wrap the semidiscretization `semi` as a split ODE problem in the time interval `tspan`
270272
that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).
271273
The parabolic right-hand side is the first function of the split ODE problem and
272274
will be used by default by the implicit part of IMEX methods from the
273275
SciML ecosystem.
276+
277+
Optional keyword arguments:
278+
- `jac_prototype_parabolic`: Expected to come from [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl).
279+
Specifies the sparsity structure of the parabolic function's Jacobian to enable e.g. efficient implicit time stepping.
280+
The [`SplitODEProblem`](https://docs.sciml.ai/DiffEqDocs/stable/types/split_ode_types/#SciMLBase.SplitODEProblem) only expects the Jacobian
281+
to be defined on the first function it takes in, which is treated implicitly. This corresponds to the parabolic right-hand side in Trixi.jl.
282+
The hyperbolic right-hand side is expected to be treated explicitly, and therefore its Jacobian is irrelevant.
283+
- `colorvec_parabolic`: Expected to come from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl).
284+
Allows for even faster Jacobian computation. Not necessarily required when `jac_prototype_parabolic` is given.
274285
"""
275286
function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan;
287+
jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
288+
colorvec_parabolic::Union{AbstractVector, Nothing} = nothing,
276289
reset_threads = true)
277290
# Optionally reset Polyester.jl threads. See
278291
# https://github.com/trixi-framework/Trixi.jl/issues/1583
@@ -286,15 +299,33 @@ function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan;
286299
# mpi_isparallel() && MPI.Barrier(mpi_comm())
287300
# See https://github.com/trixi-framework/Trixi.jl/issues/328
288301
iip = true # is-inplace, i.e., we modify a vector when calling rhs_parabolic!, rhs!
289-
# Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
290-
# first function implicitly and the second one explicitly. Thus, we pass the
291-
# stiffer parabolic function first.
292-
return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
302+
303+
# Check if Jacobian prototype is provided for sparse Jacobian
304+
if jac_prototype_parabolic !== nothing
305+
# Convert `jac_prototype_parabolic` to real type, as seen here:
306+
# https://docs.sciml.ai/DiffEqDocs/stable/tutorials/advanced_ode_example/#Declaring-a-Sparse-Jacobian-with-Automatic-Sparsity-Detection
307+
parabolic_ode = SciMLBase.ODEFunction(rhs_parabolic!,
308+
jac_prototype = convert.(eltype(u0_ode),
309+
jac_prototype_parabolic),
310+
colorvec = colorvec_parabolic) # coloring vector is optional
311+
312+
# Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
313+
# first function implicitly and the second one explicitly. Thus, we pass the
314+
# (potentially) stiffer parabolic function first.
315+
return SplitODEProblem{iip}(parabolic_ode, rhs!, u0_ode, tspan, semi)
316+
else
317+
# We could also construct an `ODEFunction` explicitly without the Jacobian here,
318+
# but we stick to the lean direct in-place functions `rhs_parabolic!` and
319+
# let OrdinaryDiffEq.jl handle the rest
320+
return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
321+
end
293322
end
294323

295324
"""
296325
semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan,
297-
restart_file::AbstractString)
326+
restart_file::AbstractString;
327+
jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
328+
colorvec_parabolic::Union{AbstractVector, Nothing} = nothing)
298329
299330
Wrap the semidiscretization `semi` as a split ODE problem in the time interval `tspan`
300331
that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).
@@ -303,9 +334,20 @@ will be used by default by the implicit part of IMEX methods from the
303334
SciML ecosystem.
304335
305336
The initial condition etc. is taken from the `restart_file`.
337+
338+
Optional keyword arguments:
339+
- `jac_prototype_parabolic`: Expected to come from [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl).
340+
Specifies the sparsity structure of the parabolic function's Jacobian to enable e.g. efficient implicit time stepping.
341+
The [`SplitODEProblem`](https://docs.sciml.ai/DiffEqDocs/stable/types/split_ode_types/#SciMLBase.SplitODEProblem) only expects the Jacobian
342+
to be defined on the first function it takes in, which is treated implicitly. This corresponds to the parabolic right-hand side in Trixi.jl.
343+
The hyperbolic right-hand side is expected to be treated explicitly, and therefore its Jacobian is irrelevant.
344+
- `colorvec_parabolic`: Expected to come from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl).
345+
Allows for even faster Jacobian computation. Not necessarily required when `jac_prototype_parabolic` is given.
306346
"""
307347
function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan,
308348
restart_file::AbstractString;
349+
jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
350+
colorvec_parabolic::Union{AbstractVector, Nothing} = nothing,
309351
reset_threads = true)
310352
# Optionally reset Polyester.jl threads. See
311353
# https://github.com/trixi-framework/Trixi.jl/issues/1583
@@ -319,10 +361,26 @@ function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan,
319361
# mpi_isparallel() && MPI.Barrier(mpi_comm())
320362
# See https://github.com/trixi-framework/Trixi.jl/issues/328
321363
iip = true # is-inplace, i.e., we modify a vector when calling rhs_parabolic!, rhs!
322-
# Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
323-
# first function implicitly and the second one explicitly. Thus, we pass the
324-
# stiffer parabolic function first.
325-
return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
364+
365+
# Check if Jacobian prototype is provided for sparse Jacobian
366+
if jac_prototype_parabolic !== nothing
367+
# Convert `jac_prototype_parabolic` to real type, as seen here:
368+
# https://docs.sciml.ai/DiffEqDocs/stable/tutorials/advanced_ode_example/#Declaring-a-Sparse-Jacobian-with-Automatic-Sparsity-Detection
369+
parabolic_ode = SciMLBase.ODEFunction(rhs_parabolic!,
370+
jac_prototype = convert.(eltype(u0_ode),
371+
jac_prototype_parabolic),
372+
colorvec = colorvec_parabolic) # coloring vector is optional
373+
374+
# Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
375+
# first function implicitly and the second one explicitly. Thus, we pass the
376+
# (potentially) stiffer parabolic function first.
377+
return SplitODEProblem{iip}(parabolic_ode, rhs!, u0_ode, tspan, semi)
378+
else
379+
# We could also construct an `ODEFunction` explicitly without the Jacobian here,
380+
# but we stick to the lean direct in-place function `rhs_parabolic!` and
381+
# let OrdinaryDiffEq.jl handle the rest
382+
return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
383+
end
326384
end
327385

328386
function rhs!(du_ode, u_ode, semi::SemidiscretizationHyperbolicParabolic, t)

‎test/test_parabolic_1d.jl‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,50 @@ end
121121
@test_allocations(Trixi.rhs!, semi, sol, 1000)
122122
end
123123

124+
@trixi_testset "TreeMesh1D: elixir_advection_diffusion_implicit_sparse_jacobian.jl" begin
125+
@test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
126+
"elixir_advection_diffusion_implicit_sparse_jacobian.jl"),
127+
tspan=(0.0, 0.4),
128+
l2=[0.05240130204342638], linf=[0.07407444680136666])
129+
# Ensure that we do not have excessive memory allocations
130+
# (e.g., from type instabilities)
131+
let
132+
t = sol.t[end]
133+
u_ode = sol.u[end]
134+
du_ode = similar(u_ode)
135+
@test (@allocated Trixi.rhs!(du_ode, u_ode, semi, t)) < 1000
136+
end
137+
end
138+
139+
@trixi_testset "TreeMesh1D: elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl" begin
140+
@test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
141+
"elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl"),
142+
l2=[0.08292233849124372], linf=[0.11726345328639576])
143+
# Ensure that we do not have excessive memory allocations
144+
# (e.g., from type instabilities)
145+
let
146+
t = sol.t[end]
147+
u_ode = sol.u[end]
148+
du_ode = similar(u_ode)
149+
@test (@allocated Trixi.rhs!(du_ode, u_ode, semi, t)) < 1000
150+
end
151+
end
152+
153+
@trixi_testset "elixir_advection_implicit_sparse_jacobian_restart.jl (no colorvec)" begin
154+
@test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
155+
"elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl"),
156+
colorvec_parabolic=nothing,
157+
l2=[0.08292233849124372], linf=[0.11726345328639576])
158+
# Ensure that we do not have excessive memory allocations
159+
# (e.g., from type instabilities)
160+
let
161+
t = sol.t[end]
162+
u_ode = sol.u[end]
163+
du_ode = similar(u_ode)
164+
@test (@allocated Trixi.rhs!(du_ode, u_ode, semi, t)) < 1000
165+
end
166+
end
167+
124168
@trixi_testset "TreeMesh1D: elixir_navierstokes_convergence_periodic.jl" begin
125169
@test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
126170
"elixir_navierstokes_convergence_periodic.jl"),

0 commit comments

Comments
 (0)