Hyperbolic-Parabolic elixir using SparseConnectivityTracer

examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian.jl

@@ -0,0 +1,113 @@
+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 ###
+
+# Wrap the `Trixi.rhs!` 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
+
+coloring_prob = ColoringProblem(; structure = :nonsymmetric, partition = :column)
+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)
+
+# Note: No `stepsize_callback` due to (implicit) solver with adaptive timestep control
+callbacks = CallbackSet(summary_callback, analysis_callback, alive_callback)
+
+###############################################################################
+### 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)

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

@@ -113,6 +113,21 @@ 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_navierstokes_convergence_periodic.jl" begin
     @test_trixi_include(joinpath(EXAMPLES_DIR,
                                  "elixir_navierstokes_convergence_periodic.jl"),

test/test_unit.jl

@@ -2744,6 +2744,101 @@ end
     @test isapprox(jac_finite_diff, Matrix(jac_sparse_finite_diff); rtol = 5e-8)
     @test isapprox(sparse(jac_finite_diff), jac_sparse_finite_diff; rtol = 5e-8)
 end
+
+@testset "Parabolic-Hyperbolic Problem Sparsity Pattern" begin
+
+    # Poor-mans rebuild of `SplitODEProblem` from SciML
+    function rhs_hyperbolic_parabolic!(du_ode, u_ode,
+                                    semi::SemidiscretizationHyperbolicParabolic, t)
+        Trixi.@trixi_timeit timer() "rhs_hyperbolic_parabolic!" begin
+            du_para = similar(du_ode) # This obviously allocates, but fine for this test
+            rhs!(du_ode, u_ode, semi, t)
+            rhs_parabolic!(du_para, u_ode, semi, t)
+
+            Trixi.@threaded for i in eachindex(du_ode)
+                du_ode[i] = du_ode[i] + du_para[i]
+            end
+        end
+        return nothing
+    end
+
+    ###############################################################################
+    ### equations, solver, mesh ###
+
+    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)
+
+    ###############################################################################
+    ### 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)
+
+    # Semidiscretization for sparsity pattern detection
+    semi_jac_type = SemidiscretizationHyperbolicParabolic(mesh, (equations_hyperbolic, equations_parabolic),
+                                                 initial_condition_convergence_test,
+                                                 solver,
+                                                 uEltype = jac_eltype) # Need to supply Jacobian element type
+
+    tspan = (0.0, 1.5) # Re-used for wrapping `rhs` 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
+    # Although we do sparsity detection on the entire RHS (since semi_jac_type depends on both equations and 
+    # equations_parabolic), this is equivalent to doing sparsity detection on the diffusion problem alone
+    # (see next section for a demonstration of this)
+
+    # 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)
+    ###############################################################################
+    ### Compare sparsity pattern based on hyperbolic-parabolic ###
+    ### problem with sparsity pattern of parabolic-only problem ###
+
+    # We construct a semidiscretization just as we did previously, but this time with advection_velocity=0
+    equations_hyperbolic_no_advection = LinearScalarAdvectionEquation1D(0)
+    equations_parabolic_no_advection = LaplaceDiffusion1D(diffusivity(), equations_hyperbolic_no_advection)
+
+    semi_jac_type_no_advection = SemidiscretizationHyperbolicParabolic(mesh,
+                                                (equations_hyperbolic_no_advection, 
+                                                equations_parabolic_no_advection),
+                                                initial_condition_convergence_test, solver,
+                                                uEltype = jac_eltype)
+    ode_jac_type_no_advection = semidiscretize(semi_jac_type_no_advection, tspan)
+
+    # Do sparsity detection on our semidiscretization with advection turned off
+    rhs_hyper_para! = (du_ode, u0_ode) -> rhs_hyperbolic_parabolic!(du_ode, u0_ode, semi_jac_type_no_advection, tspan[1])
+
+    jac_prototype_parabolic_no_advection = jacobian_sparsity(rhs_parabolic_wrapped!, du_ode, u0_ode, jac_detector)
+
+    # Given that the stencil for parabolic solvers are always larger than those of hyperbolic solvers,
+    # the sparsity detection will never depend on the hyperbolic part of the problem
+    @test jac_prototype_parabolic == jac_prototype_parabolic_no_advection
+end
 end
 
 end #module