Elixir: ImplicitEuler solve with sparse jacobian

Project.toml

@@ -25,6 +25,7 @@ MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 P4est = "7d669430-f675-4ae7-b43e-fab78ec5a902"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
@@ -59,12 +60,15 @@ Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
 ECOS = "e2685f51-7e38-5353-a97d-a921fd2c8199"
 Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [extensions]
 TrixiCUDAExt = "CUDA"
 TrixiConvexECOSExt = ["Convex", "ECOS"]
 TrixiMakieExt = "Makie"
 TrixiNLsolveExt = "NLsolve"
+TrixiSparseDiffToolsExt = ["SparseDiffTools", "Symbolics"]
 
 [compat]
 Accessors = "0.1.36"

examples/tree_2d_dgsem/elixir_advection_sparse_jacobian.jl

@@ -0,0 +1,78 @@
+using Trixi
+using OrdinaryDiffEq
+using SparseDiffTools, Symbolics
+
+using .TrixiSparseDiffToolsExt
+
+###############################################################################
+### semidiscretization of the linear advection equation ###
+
+advection_velocities = (1.0, 1.1)
+equations = LinearScalarAdvectionEquation2D(advection_velocities)
+
+# Create DG solver with polynomial degree = 3 and (local) Lax-Friedrichs/Rusanov flux as surface flux
+d = 3
+# `RealT = Real` requires fewer overloads than the more explicit `RealT = Num`
+# solver_real used for computing the Jacobian
+solver_real = DGSEM(polydeg = d, surface_flux = flux_lax_friedrichs, RealT = Real)
+# solver_float used for solving using the Jacobian
+solver_float = DGSEM(polydeg = d, surface_flux = flux_lax_friedrichs)
+
+coordinates_min = (-1.0, -1.0) # minimum coordinates (min(x), min(y))
+coordinates_max = (1.0, 1.0) # maximum coordinates (max(x), max(y))
+
+# Create a uniformly refined mesh with periodic boundaries
+RefinementLevel = 5
+mesh = TreeMesh(coordinates_min, coordinates_max,
+                initial_refinement_level = RefinementLevel,
+                n_cells_max = 30_000) # set maximum capacity of tree data structure
+
+# A semidiscretization collects data structures and functions for the spatial discretization
+# semi_real used for computing the Jacobian
+semi_real = SemidiscretizationHyperbolic(mesh, equations, initial_condition_convergence_test,
+                                    solver_real)
+# semi_float used for solving using the Jacobian
+semi_float = SemidiscretizationHyperbolic(mesh, equations, initial_condition_convergence_test,
+                                    solver_float)
+
+# 2^RefinementLevel = 16 elements, (d+1) local polynomial coefficients per element
+N = Int(2^RefinementLevel * (d+1)) #64
+u0_ode = zeros(N*N)
+du_ode = 80 * ones(N*N) # initialize with something
+t0 = 0.0
+tSpan = (t0, t0 + 10.0)
+
+###############################################################################
+### Compute the Jacobian with SparseDiffTools ###
+
+# Create a function with two parameters:du_ode and u0_ode
+# to fulfill the requirments of an in_place function in SparseDiffTools
+rhs = (du_ode,u0_ode)->Trixi.rhs!(du_ode, u0_ode, semi_real, t0)
+
+#From the example to detect the pattern and choose how to do the AutoDiff automatically
+sd = SymbolicsSparsityDetection()
+adtype = AutoSparseFiniteDiff()
+
+#From the example provided in SparseDiffTools. cache will reduce calculation time when Jacobian will be calculated multiple times
+sparse_cache = sparse_jacobian_cache(adtype, sd, rhs, du_ode, u0_ode)
+
+###############################################################################
+#  callback functions during the time integration
+
+ode = semidiscretize(semi_float, tSpan, sparse_cache.jac_prototype, sparse_cache.coloring.colorvec)
+
+# The AnalysisCallback allows to analyse the solution in regular intervals and prints the results
+analysis_callback = AnalysisCallback(semi_float, interval = 100)
+
+# At the beginning of the main loop, the SummaryCallback prints a summary of the simulation setup
+# and resets the timers
+summary_callback = SummaryCallback()
+
+# Create a CallbackSet to collect all callbacks such that they can be passed to the ODE solver
+callbacks = CallbackSet(summary_callback)
+
+###############################################################################
+# Run the simulation using ImplicitEuler method
+
+sol = solve(ode, ImplicitEuler(; autodiff = AutoFiniteDiff());
+            dt = 1.0, save_everystep = false, callback = callbacks);

ext/TrixiSparseDiffToolsExt.jl

@@ -0,0 +1,43 @@
+# Package extension for using SparseDiffTools with Trixi for implicit solvers
+module TrixiSparseDiffToolsExt
+
+using Trixi
+
+import Base: *, zero, one # For overloading
+
+###############################################################################
+### Hacks ###
+
+# Required for setting up the Lobatto Legendre basis for abstract `Real` type
+function Trixi.eps(::Type{Real}, RealT = Float64)
+    return eps(RealT)
+end
+
+# There are several places in trixi where they do one(RealT) or zero(uEltype) where RealT or uEltype is Real
+# this just returns an Int64 1 or 0 respectively. We don't want to use ints so we override this behavior
+# Real(x::Real) = Float64(x)
+Base.one(::Type{Real}) = 1.0
+Base.zero(::Type{Real}) = 0.0
+
+# Multiplying two Matrix{Real}s gives a Matrix{Any}.
+# This causes problems when instantiating the Legendre basis.
+# Called in `calc_{forward,reverse}_{upper, lower}`
+function *(A::Matrix{Real}, B::Matrix{Real})::Matrix{Real}
+    m, n = size(A, 1), size(B, 2)
+    kA = size(A, 2)
+    kB = size(B, 1)
+    @assert kA == kB "Matrix dimensions must match for multiplication"
+
+    C = Matrix{Real}(undef, m, n)
+    for i in 1:m, j in 1:n
+        #acc::Real = zero(promote_type(typeof(A[i,1]), typeof(B[1,j])))
+        acc = zero(Real)
+        for k in 1:kA
+            acc += A[i,k] * B[k,j]
+        end
+        C[i,j] = acc
+    end
+    return C
+end
+
+end # module TrixiNLsolveExt

src/semidiscretization/semidiscretization.jl

@@ -120,6 +120,36 @@ function semidiscretize(semi::AbstractSemidiscretization, tspan;
     return ODEProblem{iip, specialize}(rhs!, u0_ode, tspan, semi)
 end
 
+"""
+    semidiscretize(semi::AbstractSemidiscretization, tspan)
+
+Wrap the semidiscretization `semi` as an ODE problem in the time interval `tspan`
+that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).
+
+The arguments `jac_prototype` and `coloring` are SparseDiffTools.jl objects which will be used 
+to reduce repeating calculations as part of the ODE solve.
+"""
+function semidiscretize(semi::AbstractSemidiscretization, tspan, jac_prototype, coloring;
+                        reset_threads = true)
+    # Optionally reset Polyester.jl threads. See
+    # https://github.com/trixi-framework/Trixi.jl/issues/1583
+    # https://github.com/JuliaSIMD/Polyester.jl/issues/30
+    if reset_threads
+        Polyester.reset_threads!()
+    end
+
+    u0_ode = compute_coefficients(first(tspan), semi)
+    # TODO: MPI, do we want to synchronize loading and print debug statements, e.g. using
+    #       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!
+    specialize = SciMLBase.FullSpecialize # specialize on rhs! and parameters (semi)
+    # See SparseDiffTools.jl docs for the typing of jac_prototype and colorvec
+    # https://github.com/JuliaDiff/SparseDiffTools.jl 
+    ode = SciMLBase.ODEFunction(rhs!, jac_prototype=jac_prototype, colorvec=coloring)
+    return ODEProblem{iip, specialize}(ode, u0_ode, tspan, semi)
+end
+
 """
     semidiscretize(semi::AbstractSemidiscretization, tspan, 
                    restart_file::AbstractString)