Hyperbolic-Parabolic elixir using SparseConnectivityTracer

examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian.jl

@@ -0,0 +1,128 @@
+using Trixi
+using SparseConnectivityTracer # For obtaining the Jacobian sparsity pattern
+using SparseMatrixColorings # For obtaining the coloring vector
+using OrdinaryDiffEqBDF, ADTypes
+
+###############################################################################
+# semidiscretization of the linear advection-diffusion equation
+
+advection_velocity = 1.5
+equations_hyperbolic = LinearScalarAdvectionEquation1D(advection_velocity)
+diffusivity() = 5.0e-2
+equations_parabolic = LaplaceDiffusion1D(diffusivity(), equations_hyperbolic)
+
+solver = DGSEM(polydeg = 3, surface_flux = flux_lax_friedrichs)
+
+coordinates_min = -1.0
+coordinates_max = 1.0
+
+mesh = TreeMesh(coordinates_min, coordinates_max,
+                initial_refinement_level = 4,
+                n_cells_max = 30_000)
+
+function initial_condition_diffusive_convergence_test(x, t,
+                                                      equation::LinearScalarAdvectionEquation1D)
+    # Store translated coordinate for easy use of exact solution
+    RealT = eltype(x)
+    x_trans = x - equation.advection_velocity * t
+
+    nu = diffusivity()
+    c = 1
+    A = 0.5f0
+    L = 2
+    f = 1.0f0 / L
+    omega = 2 * convert(RealT, pi) * f
+    scalar = c + A * sin(omega * sum(x_trans)) * exp(-2 * nu * omega^2 * t)
+    return SVector(scalar)
+end
+initial_condition = initial_condition_diffusive_convergence_test
+
+###############################################################################
+### semidiscretization for sparsity detection ###
+
+jac_detector = TracerSparsityDetector()
+# We need to construct the semidiscretization with the correct
+# sparsity-detection ready datatype, which is retrieved here
+jac_eltype = jacobian_eltype(real(solver), jac_detector)
+
+semi_jac_type = SemidiscretizationHyperbolicParabolic(mesh,
+                                                      (equations_hyperbolic,
+                                                       equations_parabolic),
+                                                      initial_condition, solver,
+                                                      uEltype = jac_eltype)
+
+tspan = (0.0, 1.5) # Re-used for wrapping `rhs_parabolic!` below
+
+# Call `semidiscretize` to create the ODE problem to have access to the
+# initial condition based on which the sparsity pattern is computed
+ode_jac_type = semidiscretize(semi_jac_type, tspan)
+u0_ode = ode_jac_type.u0
+du_ode = similar(u0_ode)
+
+###############################################################################
+### Compute the Jacobian sparsity pattern ###
+
+# Only the parabolic part of the `SplitODEProblem` is treated implicitly so we only need the parabolic Jacobian, see
+# https://docs.sciml.ai/DiffEqDocs/stable/types/split_ode_types/#SciMLBase.SplitFunction
+# Wrap the `Trixi.rhs_parabolic!` function to match the signature `f!(du, u)`, see
+# https://adrianhill.de/SparseConnectivityTracer.jl/stable/user/api/#ADTypes.jacobian_sparsity
+rhs_parabolic_wrapped! = (du_ode, u0_ode) -> Trixi.rhs_parabolic!(du_ode, u0_ode,
+                                                                  semi_jac_type, tspan[1])
+
+jac_prototype_parabolic = jacobian_sparsity(rhs_parabolic_wrapped!, du_ode, u0_ode,
+                                            jac_detector)
+
+# For most efficient solving we also want the coloring vector
+
+# We choose `nonsymmetric` `structure` because we're computing a Jacobian, which
+# is for the Upwind-alike discretization of the advection term nonsymmmetric
+# We arbitrarily choose a column-based `partitioning`. This means that we will color
+# structurally orthogonal columns the same. `row` partitioning would be equally valid here
+coloring_prob = ColoringProblem(; structure = :nonsymmetric, partition = :column)
+# The `decompression` arg specifies our evaluation scheme. The `direct` method requires solving
+# a diagonal system, whereas the `substitution` method solves a triangular system of equations
+coloring_alg = GreedyColoringAlgorithm(; decompression = :direct)
+coloring_result = coloring(jac_prototype_parabolic, coloring_prob, coloring_alg)
+coloring_vec_parabolic = column_colors(coloring_result)
+
+###############################################################################
+### sparsity-aware semidiscretization and ODE ###
+
+# Semidiscretization for actual simulation. `uEltype` is here retrieved from `solver`
+semi_float_type = SemidiscretizationHyperbolicParabolic(mesh,
+                                                        (equations_hyperbolic,
+                                                         equations_parabolic),
+                                                        initial_condition, solver)
+
+# Supply Jacobian prototype and coloring vector to the semidiscretization
+ode_jac_sparse = semidiscretize(semi_float_type, tspan,
+                                jac_prototype_parabolic = jac_prototype_parabolic,
+                                colorvec_parabolic = coloring_vec_parabolic)
+# using "dense" `ode = semidiscretize(semi_float_type, tspan)` is 4-6 times slower!
+
+###############################################################################
+### callbacks  ###
+
+summary_callback = SummaryCallback()
+
+analysis_interval = 100
+analysis_callback = AnalysisCallback(semi_float_type, interval = analysis_interval)
+
+alive_callback = AliveCallback(analysis_interval = analysis_interval)
+
+save_restart = SaveRestartCallback(interval = 100,
+                                   save_final_restart = true)
+
+# Note: No `stepsize_callback` due to (implicit) solver with adaptive timestep control
+callbacks = CallbackSet(summary_callback, analysis_callback, alive_callback, save_restart)
+
+###############################################################################
+### solve the ODE problem ###
+
+time_int_tol = 1.0e-9
+time_abs_tol = 1.0e-9
+
+sol = solve(ode_jac_sparse, SBDF2(; autodiff = AutoFiniteDiff());
+            dt = 0.01, save_everystep = false,
+            abstol = time_abs_tol, reltol = time_int_tol,
+            ode_default_options()..., callback = callbacks)

examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl

@@ -0,0 +1,27 @@
+using Trixi
+
+###############################################################################
+# create a restart file
+
+elixir_file = "elixir_advection_diffusion_implicit_sparse_jacobian.jl"
+restart_file = "restart_000000100.h5"
+
+trixi_include(@__MODULE__, joinpath(@__DIR__, elixir_file))
+
+###############################################################################
+
+restart_filename = joinpath("out", restart_file)
+tspan = (load_time(restart_filename), 2.0)
+dt_restart = load_dt(restart_filename)
+
+ode_jac_sparse = semidiscretize(semi_float_type, tspan,
+                                restart_filename,
+                                jac_prototype_parabolic = jac_prototype_parabolic,
+                                colorvec_parabolic = coloring_vec_parabolic)
+
+###############################################################################
+# run the simulation
+
+sol = solve(ode_jac_sparse,
+            SBDF2(; autodiff = AutoFiniteDiff());
+            dt = dt_restart, save_everystep = false, callback = callbacks);

src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl

@@ -264,15 +264,28 @@ function get_node_variables!(node_variables, u_ode,
 end
 
 """
-    semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan)
+    semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan;
+                   jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
+                   colorvec_parabolic::Union{AbstractVector, Nothing} = nothing)
 
 Wrap the semidiscretization `semi` as a split ODE problem in the time interval `tspan`
 that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).
 The parabolic right-hand side is the first function of the split ODE problem and
 will be used by default by the implicit part of IMEX methods from the
 SciML ecosystem.
+
+Optional keyword arguments:
+- `jac_prototype_parabolic`: Expected to come from [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl).
+  Specifies the sparsity structure of the parabolic function's Jacobian to enable e.g. efficient implicit time stepping.
+  The [SplitODEProblem](https://docs.sciml.ai/DiffEqDocs/stable/types/split_ode_types/#SciMLBase.SplitODEProblem) only expects the Jacobian
+  to be defined on the first function it takes in, which is treated implicitly. This corresponds to the parabolic right-hand side in Trixi.
+  The hyperbolic right-hand side is expected to be treated explicitly, and therefore its Jacobian is irrelevant.
+- `colorvec_parabolic`: Expected to come from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl).
+  Allows for even faster Jacobian computation. Not necessarily required when `jac_prototype_parabolic` is given.
 """
 function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan;
+                        jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
+                        colorvec_parabolic::Union{AbstractVector, Nothing} = nothing,
                         reset_threads = true)
     # Optionally reset Polyester.jl threads. See
     # https://github.com/trixi-framework/Trixi.jl/issues/1583
@@ -286,15 +299,32 @@ function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan;
     #       mpi_isparallel() && MPI.Barrier(mpi_comm())
     #       See https://github.com/trixi-framework/Trixi.jl/issues/328
     iip = true # is-inplace, i.e., we modify a vector when calling rhs_parabolic!, rhs!
-    # Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
-    # first function implicitly and the second one explicitly. Thus, we pass the
-    # stiffer parabolic function first.
-    return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
+
+    # Check if Jacobian prototype is provided for sparse Jacobian
+    if jac_prototype_parabolic !== nothing
+        # Convert `jac_prototype_parabolic` to real type, as seen here:
+        # https://docs.sciml.ai/DiffEqDocs/stable/tutorials/advanced_ode_example/#Declaring-a-Sparse-Jacobian-with-Automatic-Sparsity-Detection
+        parabolic_ode = SciMLBase.ODEFunction(rhs_parabolic!,
+                                              jac_prototype = convert.(eltype(u0_ode),
+                                                                       jac_prototype_parabolic),
+                                              colorvec = colorvec_parabolic) # coloring vector is optional
+
+        # Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
+        # first function implicitly and the second one explicitly. Thus, we pass the
+        # (potentially) stiffer parabolic function first.
+        return SplitODEProblem{iip}(parabolic_ode, rhs!, u0_ode, tspan, semi)
+    else
+        # We could also construct an `ODEFunction` without the Jacobian here,
+        # but we stick to the more light-weight direct in-place function `rhs_parabolic!`.
+        return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
+    end
 end
 
 """
     semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan,
-                   restart_file::AbstractString)
+                   restart_file::AbstractString;
+                   jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
+                   colorvec_parabolic::Union{AbstractVector, Nothing} = nothing)
 
 Wrap the semidiscretization `semi` as a split ODE problem in the time interval `tspan`
 that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).
@@ -303,9 +333,20 @@ will be used by default by the implicit part of IMEX methods from the
 SciML ecosystem.
 
 The initial condition etc. is taken from the `restart_file`.
+
+Optional keyword arguments:
+- `jac_prototype_parabolic`: Expected to come from [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl).
+  Specifies the sparsity structure of the parabolic function's Jacobian to enable e.g. efficient implicit time stepping.
+  The [SplitODEProblem](https://docs.sciml.ai/DiffEqDocs/stable/types/split_ode_types/#SciMLBase.SplitODEProblem) only expects the Jacobian
+  to be defined on the first function it takes in, which is treated implicitly. This corresponds to the parabolic right-hand side in Trixi.
+  The hyperbolic right-hand side is expected to be treated explicitly, and therefore its Jacobian is irrelevant.
+- `colorvec_parabolic`: Expected to come from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl).
+  Allows for even faster Jacobian computation. Not necessarily required when `jac_prototype_parabolic` is given.
 """
 function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan,
                         restart_file::AbstractString;
+                        jac_prototype_parabolic::Union{AbstractMatrix, Nothing} = nothing,
+                        colorvec_parabolic::Union{AbstractVector, Nothing} = nothing,
                         reset_threads = true)
     # Optionally reset Polyester.jl threads. See
     # https://github.com/trixi-framework/Trixi.jl/issues/1583
@@ -319,10 +360,25 @@ function semidiscretize(semi::SemidiscretizationHyperbolicParabolic, tspan,
     #       mpi_isparallel() && MPI.Barrier(mpi_comm())
     #       See https://github.com/trixi-framework/Trixi.jl/issues/328
     iip = true # is-inplace, i.e., we modify a vector when calling rhs_parabolic!, rhs!
-    # Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
-    # first function implicitly and the second one explicitly. Thus, we pass the
-    # stiffer parabolic function first.
-    return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
+
+    # Check if Jacobian prototype is provided for sparse Jacobian
+    if jac_prototype_parabolic !== nothing
+        # Convert `jac_prototype_parabolic` to real type, as seen here:
+        # https://docs.sciml.ai/DiffEqDocs/stable/tutorials/advanced_ode_example/#Declaring-a-Sparse-Jacobian-with-Automatic-Sparsity-Detection
+        parabolic_ode = SciMLBase.ODEFunction(rhs_parabolic!,
+                                              jac_prototype = convert.(eltype(u0_ode),
+                                                                       jac_prototype_parabolic),
+                                              colorvec = colorvec_parabolic) # coloring vector is optional
+
+        # Note that the IMEX time integration methods of OrdinaryDiffEq.jl treat the
+        # first function implicitly and the second one explicitly. Thus, we pass the
+        # (potentially) stiffer parabolic function first.
+        return SplitODEProblem{iip}(parabolic_ode, rhs!, u0_ode, tspan, semi)
+    else
+        # We could also construct an `ODEFunction` without the Jacobian here,
+        # but we stick to the more light-weight direct in-place function `rhs_parabolic!`.
+        return SplitODEProblem{iip}(rhs_parabolic!, rhs!, u0_ode, tspan, semi)
+    end
 end
 
 function rhs!(du_ode, u_ode, semi::SemidiscretizationHyperbolicParabolic, t)

test/test_parabolic_1d.jl

@@ -121,6 +121,50 @@ end
     @test_allocations(Trixi.rhs!, semi, sol, 1000)
 end
 
+@trixi_testset "TreeMesh1D: elixir_advection_diffusion_implicit_sparse_jacobian.jl" begin
+    @test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
+                                 "elixir_advection_diffusion_implicit_sparse_jacobian.jl"),
+                        tspan=(0.0, 0.4),
+                        l2=[0.05240130204342638], linf=[0.07407444680136666])
+    # Ensure that we do not have excessive memory allocations
+    # (e.g., from type instabilities)
+    let
+        t = sol.t[end]
+        u_ode = sol.u[end]
+        du_ode = similar(u_ode)
+        @test (@allocated Trixi.rhs!(du_ode, u_ode, semi_float_type, t)) < 1000
+    end
+end
+
+@trixi_testset "TreeMesh1D: elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl" begin
+    @test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
+                                 "elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl"),
+                        l2=[0.08292233849124372], linf=[0.11726345328639576])
+    # Ensure that we do not have excessive memory allocations
+    # (e.g., from type instabilities)
+    let
+        t = sol.t[end]
+        u_ode = sol.u[end]
+        du_ode = similar(u_ode)
+        @test (@allocated Trixi.rhs!(du_ode, u_ode, semi_float_type, t)) < 1000
+    end
+end
+
+@trixi_testset "elixir_advection_implicit_sparse_jacobian_restart.jl (no colorvec)" begin
+    @test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
+                                 "elixir_advection_diffusion_implicit_sparse_jacobian_restart.jl"),
+                        colorvec_parabolic=nothing,
+                        l2=[0.08292233849124372], linf=[0.11726345328639576])
+    # Ensure that we do not have excessive memory allocations
+    # (e.g., from type instabilities)
+    let
+        t = sol.t[end]
+        u_ode = sol.u[end]
+        du_ode = similar(u_ode)
+        @test (@allocated Trixi.rhs!(du_ode, u_ode, semi_float_type, t)) < 1000
+    end
+end
+
 @trixi_testset "TreeMesh1D: elixir_navierstokes_convergence_periodic.jl" begin
     @test_trixi_include(joinpath(EXAMPLES_DIR, "tree_1d_dgsem",
                                  "elixir_navierstokes_convergence_periodic.jl"),