Using the an implicit solver based on Ariadne with Trixi.jl

using Trixi
using Theseus
using CairoMakie
using LinearAlgebra
import Ariadne: JacobianOperator

Notes: Must disable both Polyester and LoopVectorization for Enzyme to be able to differentiate Trixi.jl

LocalPreferences.jl

[Trixi]
loop_vectorization = false
backend = "static"
@assert Trixi._PREFERENCE_THREADING !== :polyester
@assert !Trixi._PREFERENCE_LOOPVECTORIZATION

Load Trixi Example

trixi_include(@__MODULE__, joinpath(examples_dir(), "tree_2d_dgsem", "elixir_advection_basic.jl"), sol = nothing);

u = copy(ode.u0)
du = zero(ode.u0)
res = zero(ode.u0)

F! = Theseus.nonlinear_problem(Theseus.ImplicitEuler(), ode.f)
J = JacobianOperator(F!, res, u, (ode.u0, 1.0, du, ode.p, 0.0, (), 1))

out = zero(u)
v = zero(u)

# precompile
mul!(u, J, v)
F!(res, u, (ode.u0, 1.0, du, ode.p, 0.0, (), 1))

@time mul!(u, J, v)
@time F!(res, u, (ode.u0, 1.0, du, ode.p, 0.0, (), 1))
[ Info: You just called `trixi_include`. Julia may now compile the code, please be patient.
  0.000227 seconds (84 allocations: 9.062 KiB)
  0.000140 seconds (7 allocations: 352 bytes)

Cost of time(mul!) ≈ 2 * time(F!)

Solve using ODE interface

sol_trbdf2 = solve(
    ode, Theseus.TRBDF2();
    dt = 1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
    ode_default_options()..., callback = callbacks,
    # verbose=1,
    krylov_algo = :gmres,
    # krylov_kwargs=(;verbose=1)
);

████████╗██████╗ ██╗██╗  ██╗██╗
╚══██╔══╝██╔══██╗██║╚██╗██╔╝██║
   ██║   ██████╔╝██║ ╚███╔╝ ██║
   ██║   ██╔══██╗██║ ██╔██╗ ██║
   ██║   ██║  ██║██║██╔╝ ██╗██║
   ╚═╝   ╚═╝  ╚═╝╚═╝╚═╝  ╚═╝╚═╝

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ SemidiscretizationHyperbolic                                                                     │
│ ════════════════════════════                                                                     │
│ #spatial dimensions: ………………………… 2                                                                │
│ mesh: ………………………………………………………………… TreeMesh{2, Trixi.SerialTree{2, Float64}} with length 341        │
│ equations: …………………………………………………… LinearScalarAdvectionEquation2D                                  │
│ initial condition: ……………………………… initial_condition_convergence_test                               │
│ boundary conditions: ………………………… Trixi.BoundaryConditionPeriodic                                  │
│ source terms: …………………………………………… nothing                                                          │
│ solver: …………………………………………………………… DG                                                               │
│ total #DOFs per field: …………………… 4096                                                             │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ TreeMesh{2, Trixi.SerialTree{2, Float64}}                                                        │
│ ═════════════════════════════════════════                                                        │
│ center: …………………………………………………………… [0.0, 0.0]                                                       │
│ length: …………………………………………………………… 2.0                                                              │
│ periodicity: ……………………………………………… (true, true)                                                     │
│ current #cells: ……………………………………… 341                                                              │
│ #leaf-cells: ……………………………………………… 256                                                              │
│ maximum #cells: ……………………………………… 30000                                                            │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ LinearScalarAdvectionEquation2D                                                                  │
│ ═══════════════════════════════                                                                  │
│ #variables: ………………………………………………… 1                                                                │
│ │ variable 1: …………………………………………… scalar                                                           │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ DG{Float64}                                                                                      │
│ ═══════════                                                                                      │
│ basis: ……………………………………………………………… LobattoLegendreBasis{Float64}(polydeg=3)                         │
│ mortar: …………………………………………………………… LobattoLegendreMortarL2{Float64}(polydeg=3)                      │
│ surface integral: ………………………………… SurfaceIntegralWeakForm                                          │
│ │ surface flux: ……………………………………… FluxLaxFriedrichs(max_abs_speed)                                 │
│ volume integral: …………………………………… VolumeIntegralWeakForm                                           │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ AnalysisCallback                                                                                 │
│ ════════════════                                                                                 │
│ interval: ……………………………………………………… 100                                                              │
│ analyzer: ……………………………………………………… LobattoLegendreAnalyzer{Float64}(polydeg=6)                      │
│ │ error 1: …………………………………………………… l2_error                                                         │
│ │ error 2: …………………………………………………… linf_error                                                       │
│ │ integral 1: …………………………………………… entropy_timederivative                                           │
│ save analysis to file: …………………… no                                                               │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ SaveSolutionCallback                                                                             │
│ ════════════════════                                                                             │
│ interval: ……………………………………………………… 100                                                              │
│ solution variables: …………………………… cons2prim                                                        │
│ save initial solution: …………………… yes                                                              │
│ save final solution: ………………………… yes                                                              │
│ output directory: ………………………………… /home/runner/work/Ariadne.jl/Ariadne.jl/docs/build/generated/out │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ StepsizeCallback                                                                                 │
│ ════════════════                                                                                 │
│ CFL Advective: ………………………………………… Returns{Float64}(1.6)                                            │
│ CFL Diffusive: ………………………………………… Returns{Float64}(0.0)                                            │
│ Interval: ……………………………………………………… 1                                                                │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ Time integration                                                                                 │
│ ════════════════                                                                                 │
│ Start time: ………………………………………………… 0.0                                                              │
│ Final time: ………………………………………………… 1.0                                                              │
│ time integrator: …………………………………… TRBDF2                                                           │
│ adaptive: ……………………………………………………… false                                                            │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘
┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ Environment information                                                                          │
│ ═══════════════════════                                                                          │
│ #threads: ……………………………………………………… 1                                                                │
│ threading backend: ……………………………… static                                                           │
│ LoopVectorization: ……………………………… disabled                                                         │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

────────────────────────────────────────────────────────────────────────────────────────────────────
 Simulation running 'LinearScalarAdvectionEquation2D' with DGSEM(polydeg=3)
────────────────────────────────────────────────────────────────────────────────────────────────────
 #timesteps:                  0                run time:       9.74106414e-01 s
 Δt:             1.00000000e+00                └── GC time:    0.00000000e+00 s (0.000%)
 sim. time:      0.00000000e+00 (0.000%)       time/DOF/rhs!:  2.73051147e-08 s
                                               PID:                   Inf s
 #DOFs per field:          4096                alloc'd memory:        779.199 MiB
 #elements:                 256

 Variable:       scalar
 L2 error:       3.94120825e-06
 Linf error:     1.23670425e-05
 ∑∂S/∂U ⋅ Uₜ :  -2.19124035e-16
────────────────────────────────────────────────────────────────────────────────────────────────────

┌ Warning: TODO forward zero-set of memorycopy used memset rather than runtime type
│ Caused by:
│ Stacktrace:
│   [1] copy
│     @ ./array.jl:350
│   [2] unaliascopy
│     @ ./abstractarray.jl:1516
│   [3] unalias
│     @ ./abstractarray.jl:1500
│   [4] broadcast_unalias
│     @ ./broadcast.jl:946
│   [5] preprocess
│     @ ./broadcast.jl:953
│   [6] preprocess_args
│     @ ./broadcast.jl:955
│   [7] preprocess
│     @ ./broadcast.jl:952
│   [8] override_bc_copyto!
│     @ ~/.julia/packages/Enzyme/S3nC6/src/compiler/interpreter.jl:818
│   [9] copyto!
│     @ ./broadcast.jl:925
│  [10] materialize!
│     @ ./broadcast.jl:883
│  [11] materialize!
│     @ ./broadcast.jl:880
│  [12] TRBDF2
│     @ ~/work/Ariadne.jl/Ariadne.jl/libs/Theseus/src/Theseus.jl:139
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/S3nC6/src/rules/llvmrules.jl:814

────────────────────────────────────────────────────────────────────────────────────────────────────
 Simulation running 'LinearScalarAdvectionEquation2D' with DGSEM(polydeg=3)
────────────────────────────────────────────────────────────────────────────────────────────────────
 #timesteps:                 18                run time:       2.87580300e+01 s
 Δt:             5.55555556e-02                └── GC time:    4.40542304e-01 s (1.532%)
 sim. time:      1.00000000e+00 (100.000%)     time/DOF/rhs!:  2.38864595e-08 s
                                               PID:            1.26007275e-05 s
 #DOFs per field:          4096                alloc'd memory:        462.266 MiB
 #elements:                 256

 Variable:       scalar
 L2 error:       1.71224434e-04
 Linf error:     2.52822142e-04
 ∑∂S/∂U ⋅ Uₜ :  -1.97551211e-09
────────────────────────────────────────────────────────────────────────────────────────────────────

────────────────────────────────────────────────────────────────────────────────────
             Trixi.jl                      Time                    Allocations
                                  ───────────────────────   ────────────────────────
        Tot / % measured:              29.0s /   4.4%           1.77GiB /   1.9%

Section                   ncalls     time    %tot     avg     alloc    %tot      avg
────────────────────────────────────────────────────────────────────────────────────
I/O                            3    942ms   74.2%   314ms   15.3MiB   45.1%  5.10MiB
  ~I/O~                        3    769ms   60.6%   256ms   11.0MiB   32.5%  3.67MiB
  save solution                2    166ms   13.1%  82.8ms   3.69MiB   10.9%  1.85MiB
  get element variables        2   6.92ms    0.5%  3.46ms    620KiB    1.8%   310KiB
  get node variables           2    863ns    0.0%   432ns     0.00B    0.0%    0.00B
  save mesh                    2    391ns    0.0%   196ns     0.00B    0.0%    0.00B
analyze solution               2    275ms   21.7%   138ms   18.6MiB   54.9%  9.31MiB
rhs!                         533   52.1ms    4.1%  97.7μs   4.78KiB    0.0%    9.19B
  volume integral            533   32.2ms    2.5%  60.4μs     0.00B    0.0%    0.00B
  surface integral           533   5.37ms    0.4%  10.1μs     0.00B    0.0%    0.00B
  interface flux             533   5.20ms    0.4%  9.75μs     0.00B    0.0%    0.00B
  prolong2interfaces         533   3.77ms    0.3%  7.07μs     0.00B    0.0%    0.00B
  Jacobian                   533   3.42ms    0.3%  6.42μs     0.00B    0.0%    0.00B
  reset ∂u/∂t                533   1.25ms    0.1%  2.35μs     0.00B    0.0%    0.00B
  ~rhs!~                     533    685μs    0.1%  1.28μs   4.78KiB    0.0%    9.19B
  prolong2mortars            533   53.9μs    0.0%   101ns     0.00B    0.0%    0.00B
  mortar flux                533   36.9μs    0.0%  69.2ns     0.00B    0.0%    0.00B
  prolong2boundaries         533   36.8μs    0.0%  69.1ns     0.00B    0.0%    0.00B
  boundary flux              533   14.5μs    0.0%  27.1ns     0.00B    0.0%    0.00B
  source terms               533   13.9μs    0.0%  26.2ns     0.00B    0.0%    0.00B
calculate dt                  19   49.1μs    0.0%  2.58μs   2.67KiB    0.0%     144B
────────────────────────────────────────────────────────────────────────────────────

Plot the solution

We have to manually convert the sol since Theseus has it's own lightweight solution type.

plot(Trixi.PlotData2DTriangulated(sol_trbdf2.u[end], sol_trbdf2.prob.p))

Solve using OrdinaryDiffEqSDIRK

import OrdinaryDiffEqSDIRK
import DifferentiationInterface: AutoFiniteDiff
sol_sdrik = solve(
    ode, OrdinaryDiffEqSDIRK.TRBDF2(autodiff = AutoFiniteDiff());
    dt = 1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
    ode_default_options()..., callback = callbacks,
    adaptive = false
);

████████╗██████╗ ██╗██╗  ██╗██╗
╚══██╔══╝██╔══██╗██║╚██╗██╔╝██║
   ██║   ██████╔╝██║ ╚███╔╝ ██║
   ██║   ██╔══██╗██║ ██╔██╗ ██║
   ██║   ██║  ██║██║██╔╝ ██╗██║
   ╚═╝   ╚═╝  ╚═╝╚═╝╚═╝  ╚═╝╚═╝

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ SemidiscretizationHyperbolic                                                                     │
│ ════════════════════════════                                                                     │
│ #spatial dimensions: ………………………… 2                                                                │
│ mesh: ………………………………………………………………… TreeMesh{2, Trixi.SerialTree{2, Float64}} with length 341        │
│ equations: …………………………………………………… LinearScalarAdvectionEquation2D                                  │
│ initial condition: ……………………………… initial_condition_convergence_test                               │
│ boundary conditions: ………………………… Trixi.BoundaryConditionPeriodic                                  │
│ source terms: …………………………………………… nothing                                                          │
│ solver: …………………………………………………………… DG                                                               │
│ total #DOFs per field: …………………… 4096                                                             │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ TreeMesh{2, Trixi.SerialTree{2, Float64}}                                                        │
│ ═════════════════════════════════════════                                                        │
│ center: …………………………………………………………… [0.0, 0.0]                                                       │
│ length: …………………………………………………………… 2.0                                                              │
│ periodicity: ……………………………………………… (true, true)                                                     │
│ current #cells: ……………………………………… 341                                                              │
│ #leaf-cells: ……………………………………………… 256                                                              │
│ maximum #cells: ……………………………………… 30000                                                            │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ LinearScalarAdvectionEquation2D                                                                  │
│ ═══════════════════════════════                                                                  │
│ #variables: ………………………………………………… 1                                                                │
│ │ variable 1: …………………………………………… scalar                                                           │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ DG{Float64}                                                                                      │
│ ═══════════                                                                                      │
│ basis: ……………………………………………………………… LobattoLegendreBasis{Float64}(polydeg=3)                         │
│ mortar: …………………………………………………………… LobattoLegendreMortarL2{Float64}(polydeg=3)                      │
│ surface integral: ………………………………… SurfaceIntegralWeakForm                                          │
│ │ surface flux: ……………………………………… FluxLaxFriedrichs(max_abs_speed)                                 │
│ volume integral: …………………………………… VolumeIntegralWeakForm                                           │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ AnalysisCallback                                                                                 │
│ ════════════════                                                                                 │
│ interval: ……………………………………………………… 100                                                              │
│ analyzer: ……………………………………………………… LobattoLegendreAnalyzer{Float64}(polydeg=6)                      │
│ │ error 1: …………………………………………………… l2_error                                                         │
│ │ error 2: …………………………………………………… linf_error                                                       │
│ │ integral 1: …………………………………………… entropy_timederivative                                           │
│ save analysis to file: …………………… no                                                               │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ SaveSolutionCallback                                                                             │
│ ════════════════════                                                                             │
│ interval: ……………………………………………………… 100                                                              │
│ solution variables: …………………………… cons2prim                                                        │
│ save initial solution: …………………… yes                                                              │
│ save final solution: ………………………… yes                                                              │
│ output directory: ………………………………… /home/runner/work/Ariadne.jl/Ariadne.jl/docs/build/generated/out │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ StepsizeCallback                                                                                 │
│ ════════════════                                                                                 │
│ CFL Advective: ………………………………………… Returns{Float64}(1.6)                                            │
│ CFL Diffusive: ………………………………………… Returns{Float64}(0.0)                                            │
│ Interval: ……………………………………………………… 1                                                                │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ Time integration                                                                                 │
│ ════════════════                                                                                 │
│ Start time: ………………………………………………… 0.0                                                              │
│ Final time: ………………………………………………… 1.0                                                              │
│ time integrator: …………………………………… TRBDF2                                                           │
│ adaptive: ……………………………………………………… false                                                            │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘
┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ Environment information                                                                          │
│ ═══════════════════════                                                                          │
│ #threads: ……………………………………………………… 1                                                                │
│ threading backend: ……………………………… static                                                           │
│ LoopVectorization: ……………………………… disabled                                                         │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

────────────────────────────────────────────────────────────────────────────────────────────────────
 Simulation running 'LinearScalarAdvectionEquation2D' with DGSEM(polydeg=3)
────────────────────────────────────────────────────────────────────────────────────────────────────
 #timesteps:                  0                run time:       8.50000000e-07 s
 Δt:             1.00000000e+00                └── GC time:    0.00000000e+00 s (0.000%)
 sim. time:      0.00000000e+00 (0.000%)       time/DOF/rhs!:  2.76345215e-08 s
                                               PID:                   Inf s
 #DOFs per field:          4096                alloc'd memory:        783.133 MiB
 #elements:                 256

 Variable:       scalar
 L2 error:       3.94120825e-06
 Linf error:     1.23670425e-05
 ∑∂S/∂U ⋅ Uₜ :  -2.19124035e-16
────────────────────────────────────────────────────────────────────────────────────────────────────


────────────────────────────────────────────────────────────────────────────────────────────────────
 Simulation running 'LinearScalarAdvectionEquation2D' with DGSEM(polydeg=3)
────────────────────────────────────────────────────────────────────────────────────────────────────
 #timesteps:                 18                run time:       2.77551628e+01 s
 Δt:             5.55555556e-02                └── GC time:    3.38724470e-02 s (0.122%)
 sim. time:      1.00000000e+00 (100.000%)     time/DOF/rhs!:  1.98079663e-08 s
                                               PID:            4.59279096e-08 s
 #DOFs per field:          4096                alloc'd memory:        902.558 MiB
 #elements:                 256

 Variable:       scalar
 L2 error:       1.71161806e-04
 Linf error:     2.52798214e-04
 ∑∂S/∂U ⋅ Uₜ :  -1.97752408e-09
────────────────────────────────────────────────────────────────────────────────────────────────────

────────────────────────────────────────────────────────────────────────────────────
             Trixi.jl                      Time                    Allocations
                                  ───────────────────────   ────────────────────────
        Tot / % measured:              27.8s /  43.1%            244MiB /   2.1%

Section                   ncalls     time    %tot     avg     alloc    %tot      avg
────────────────────────────────────────────────────────────────────────────────────
rhs!                        148k    12.0s  100.0%  81.1μs   4.78KiB    0.1%    0.03B
  volume integral           148k    7.58s   63.4%  51.4μs     0.00B    0.0%    0.00B
  interface flux            148k    1.15s    9.6%  7.76μs     0.00B    0.0%    0.00B
  surface integral          148k    1.10s    9.2%  7.46μs     0.00B    0.0%    0.00B
  Jacobian                  148k    973ms    8.1%  6.60μs     0.00B    0.0%    0.00B
  prolong2interfaces        148k    904ms    7.6%  6.13μs     0.00B    0.0%    0.00B
  ~rhs!~                    148k    118ms    1.0%   803ns   4.78KiB    0.1%    0.03B
  reset ∂u/∂t               148k    115ms    1.0%   782ns     0.00B    0.0%    0.00B
  prolong2boundaries        148k   6.72ms    0.1%  45.6ns     0.00B    0.0%    0.00B
  prolong2mortars           148k   4.73ms    0.0%  32.1ns     0.00B    0.0%    0.00B
  mortar flux               148k   3.95ms    0.0%  26.8ns     0.00B    0.0%    0.00B
  boundary flux             148k   3.32ms    0.0%  22.5ns     0.00B    0.0%    0.00B
  source terms              148k   2.84ms    0.0%  19.2ns     0.00B    0.0%    0.00B
analyze solution               2   3.62ms    0.0%  1.81ms   4.87MiB   95.9%  2.44MiB
I/O                            3   1.88ms    0.0%   627μs    203KiB    3.9%  67.8KiB
  save solution                2   1.87ms    0.0%   937μs    200KiB    3.9%   100KiB
  ~I/O~                        3   5.62μs    0.0%  1.88μs   1.70KiB    0.0%     581B
  get element variables        2   1.82μs    0.0%   910ns   1.41KiB    0.0%     720B
  get node variables           2    868ns    0.0%   434ns     0.00B    0.0%    0.00B
  save mesh                    2    140ns    0.0%  70.0ns     0.00B    0.0%    0.00B
calculate dt                  19   95.6μs    0.0%  5.03μs   2.67KiB    0.1%     144B
────────────────────────────────────────────────────────────────────────────────────

Plot the solution

plot(Trixi.PlotData2DTriangulated(sol_sdrik.u[end], sol_sdrik.prob.p))

Increase CFL numbers

trixi_include(@__MODULE__, joinpath(examples_dir(), "tree_2d_dgsem", "elixir_advection_basic.jl"), cfl = 10, sol = nothing);

sol = solve(
    ode, Theseus.ImplicitEuler();
    dt = 1.0, # solve needs some value here but it will be overwritten by the stepsize_callback
    ode_default_options()..., callback = callbacks,
    # verbose=1,
    krylov_algo = :gmres,
    # krylov_kwargs=(;verbose=1)
);

@show callbacks.discrete_callbacks[4]
┌──────────────────────────────────────────────────────────────────────────────────────────────────┐
│ StepsizeCallback                                                                                 │
│ ════════════════                                                                                 │
│ CFL Advective: ………………………………………… Returns{Int64}(10)                                               │
│ CFL Diffusive: ………………………………………… Returns{Float64}(0.0)                                            │
│ Interval: ……………………………………………………… 1                                                                │
└──────────────────────────────────────────────────────────────────────────────────────────────────┘

Plot the solution

plot(Trixi.PlotData2DTriangulated(sol.u[end], sol.prob.p))

This page was generated using Literate.jl.