diff --git a/examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian.jl b/examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian.jl new file mode 100644 index 00000000000..e098f1484c2 --- /dev/null +++ b/examples/tree_1d_dgsem/elixir_advection_diffusion_implicit_sparse_jacobian.jl @@ -0,0 +1,116 @@ +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) + +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) diff --git a/src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl b/src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl index 54ede387fa2..7943cbc60ea 100644 --- a/src/semidiscretization/semidiscretization_hyperbolic_parabolic.jl +++ b/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 + # 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 + # 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) diff --git a/test/test_parabolic_1d.jl b/test/test_parabolic_1d.jl index d2cc65230c4..6058880f9bf 100644 --- a/test/test_parabolic_1d.jl +++ b/test/test_parabolic_1d.jl @@ -95,6 +95,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"), diff --git a/test/test_unit.jl b/test/test_unit.jl index 54403a3e3c2..9106abc2996 100644 --- a/test/test_unit.jl +++ b/test/test_unit.jl @@ -2692,6 +2692,103 @@ 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 + + function rhs_hyperbolic_parabolic!(du_ode, u_ode, + semi::SemidiscretizationHyperbolicParabolic, t) + Trixi.@trixi_timeit timer() "rhs_hyperbolic_parabolic!" begin + # Implementation of split ODE problem in OrdinaryDiffEq + du_para = similar(du_ode) # This obviously allocates + rhs!(du_ode, u_ode, semi, t) + rhs_parabolic!(du_para, u_ode, semi, t) + + Trixi.@threaded for i in eachindex(du_ode) + # Try to enable optimizations due to `muladd` by avoiding `+=` + # https://github.com/trixi-framework/Trixi.jl/pull/2480#discussion_r2224531702 + 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