@@ -401,6 +401,94 @@ relative_difference = norm(J_fd - J_ad) / size(J_fd, 1)
401401
402402# This discrepancy is of the expected order of magnitude for central finite difference approximations.
403403
404+ # ## Automatic Jacobian sparsity detection and coloring
405+
406+ # When solving large sparse nonlinear ODE systems originating from spatial discretizations
407+ # with compact stencils such as the DG method with implicit time integrators,
408+ # exploiting the sparsity of the Jacobian can lead to significant speedups in the Newton-Raphson solver.
409+ # Similarly, steady-state problems can also be solved faster.
410+
411+ # [Trixi.jl](https://github.com/trixi-framework/Trixi.jl) supports efficient Jacobian computations by leveraging the
412+ # [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl)
413+ # and [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl) packages.
414+ # These tools allow to detect the sparsity pattern of the Jacobian and compute the
415+ # optional coloring vector for efficient Jacobian evaluations.
416+ # These are then handed over to the ODE solver from [OrdinaryDiffEq.jl](https://github.com/SciML/OrdinaryDiffEq.jl).
417+
418+ # Below is a minimal example in 1D, showing how to use these packages with Trixi.jl.
419+ # First, we define the the `equation` and `mesh` as for an ordinary simulation:
420+
421+ using Trixi
422+
423+ advection_velocity = 1.0
424+ equation = LinearScalarAdvectionEquation1D (advection_velocity)
425+
426+ mesh = TreeMesh ((- 1.0 ,), (1.0 ,), initial_refinement_level = 4 , n_cells_max = 10 ^ 4 )
427+
428+ # We define the basic floating point type used for the actual simulation
429+ # and construct the solver:
430+ float_type = Float64
431+ solver = DGSEM (polydeg = 3 , surface_flux = flux_godunov, RealT = float_type)
432+
433+ # Next, we set up the sparsity detection. For this we need
434+ using SparseConnectivityTracer # For Jacobian sparsity pattern
435+
436+ # We use the [global `TracerSparsityDetector()`](https://adrianhill.de/SparseConnectivityTracer.jl/stable/user/global_vs_local/) here.
437+ jac_detector = TracerSparsityDetector ()
438+
439+ # Next, we retrieve the right element type corresponding to `float_type` for the Jacobian sparsity detection.
440+ # For more details, see the API documentation of
441+ # [`jacobian_eltype`](https://adrianhill.de/SparseConnectivityTracer.jl/stable/user/api/#SparseConnectivityTracer.jacobian_eltype).
442+ jac_eltype = jacobian_eltype (float_type, jac_detector)
443+
444+ # Now we can construct the semidiscretization for sparsity detection with `jac_eltype` as the
445+ # datatype for the working arrays and helper datastructures.
446+ semi_jac_type = SemidiscretizationHyperbolic (mesh, equation,
447+ initial_condition_convergence_test, solver,
448+ uEltype = jac_eltype) # Supply sparsity detection datatype here
449+
450+ tspan = (0.0 , 1.0 ) # Re-used later in `rhs!` evaluation
451+ ode_jac_type = semidiscretize (semi_jac_type, tspan)
452+ u0_ode = ode_jac_type. u0
453+ du_ode = similar (u0_ode)
454+
455+ # Wrap the RHS for sparsity detection to match the expected signature f!(du, u) required by
456+ # [`jacobian_sparsity`](https://adrianhill.de/SparseConnectivityTracer.jl/stable/user/api/#ADTypes.jacobian_sparsity).
457+ rhs_wrapped! = (du, u) -> Trixi. rhs! (du, u, semi_jac_type, tspan[1 ])
458+ jac_prototype = jacobian_sparsity (rhs_wrapped!, du_ode, u0_ode, jac_detector)
459+
460+ # Optionally, we can also compute the coloring vector to reduce Jacobian evaluations
461+ # to `1 + maximum(coloring_vec)` for finite differencing and `maximum(coloring_vec)` for algorithmic differentiation.
462+ # For this, we need
463+ using SparseMatrixColorings
464+
465+ # We partition by columns as we are using finite differencing here.
466+ # One would also partition by columns if forward-based algorithmic differentiation were used,
467+ # and only partition by rows if reverse-mode AD were used.
468+ # See also [the documentation of the now deprecated SparseDiffTools.jl](https://github.com/JuliaDiff/SparseDiffTools.jl?tab=readme-ov-file#matrix-coloring) package,
469+ # the predecessor in spirit to SparseConnectivityTracer.jl and SparseMatrixColorings.jl, for more information.
470+ coloring_prob = ColoringProblem (; structure = :nonsymmetric , partition = :column )
471+ coloring_alg = GreedyColoringAlgorithm (; decompression = :direct )
472+ coloring_result = coloring (jac_prototype, coloring_prob, coloring_alg)
473+ coloring_vec = column_colors (coloring_result)
474+
475+ # Now, set up the actual semidiscretization for the simulation.
476+ # The datatype is automatically retrieved from the solver (in this case `float_type = Float64`).
477+ semi_float_type = SemidiscretizationHyperbolic (mesh, equation,
478+ initial_condition_convergence_test, solver)
479+ # Supply the sparse Jacobian prototype and the optional coloring vector.
480+ # Internally, an [`ODEFunction`](https://docs.sciml.ai/DiffEqDocs/stable/types/ode_types/#SciMLBase.ODEFunction)
481+ # with `jac_prototype = jac_prototype` and `colorvec = coloring_vec` is created.
482+ ode_jac_sparse = semidiscretize (semi_float_type, tspan,
483+ jac_prototype = jac_prototype,
484+ colorvec = coloring_vec)
485+
486+ # You can now solve the ODE problem efficiently with an implicit solver.
487+ # Currently we are bound to finite differencing here.
488+ using OrdinaryDiffEqSDIRK, ADTypes
489+ sol = solve (ode_jac_sparse, TRBDF2 (; autodiff = AutoFiniteDiff ()), dt = 0.1 ,
490+ save_everystep = false )
491+
404492# ## Linear systems
405493
406494# When a linear PDE is discretized using a linear scheme such as a standard DG method,
0 commit comments