Path: blob/main/src/semidiscretization/semidiscretization.jl
5586 views
# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).1# Since these FMAs can increase the performance of many numerical algorithms,2# we need to opt-in explicitly.3# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.4@muladd begin5#! format: noindent67"""8ndofs(semi::AbstractSemidiscretization)910Return the number of degrees of freedom associated with each scalar variable.11"""12@inline function ndofs(semi::AbstractSemidiscretization)13mesh, _, solver, cache = mesh_equations_solver_cache(semi)14return ndofs(mesh, solver, cache)15end1617"""18ndofsglobal(semi::AbstractSemidiscretization)1920Return the global number of degrees of freedom associated with each scalar variable across all MPI ranks.21This is the same as [`ndofs`](@ref) for simulations running in serial or22parallelized via threads. It will in general be different for simulations23running in parallel with MPI.24"""25@inline function ndofsglobal(semi::AbstractSemidiscretization)26mesh, _, solver, cache = mesh_equations_solver_cache(semi)27return ndofsglobal(mesh, solver, cache)28end2930"""31integrate_via_indices(func, u_ode, semi::AbstractSemidiscretization, args...; normalize=true)3233Call `func(u, i..., element, equations, solver, args...)` for all nodal indices `i..., element`34and integrate the result using a quadrature associated with the semidiscretization `semi`.3536If `normalize` is true, the result is divided by the total volume of the computational domain.37"""38function integrate_via_indices(func::Func, u_ode, semi::AbstractSemidiscretization,39args...; normalize = true) where {Func}40mesh, equations, solver, cache = mesh_equations_solver_cache(semi)4142u = wrap_array(u_ode, mesh, equations, solver, cache)43return integrate_via_indices(func, u, mesh, equations, solver, cache, args...,44normalize = normalize)45end4647"""48integrate([func=(u_node,equations)->u_node,] u_ode, semi::AbstractSemidiscretization; normalize=true)4950Call `func(u_node, equations)` for each vector of nodal variables `u_node` in `u_ode`51and integrate the result using a quadrature associated with the semidiscretization `semi`.5253If `normalize` is true, the result is divided by the total volume of the computational domain.54"""55function integrate(func::Func, u_ode, semi::AbstractSemidiscretization;56normalize = true) where {Func}57mesh, equations, solver, cache = mesh_equations_solver_cache(semi)5859u = wrap_array(u_ode, mesh, equations, solver, cache)60return integrate(func, u, mesh, equations, solver, cache, normalize = normalize)61end6263function integrate(u_ode, semi::AbstractSemidiscretization; normalize = true)64return integrate(cons2cons, u_ode, semi; normalize = normalize)65end6667# Select the right-hand side function corresponding to the semidiscretization `semi`.68@inline default_rhs(::AbstractSemidiscretization) = rhs!6970"""71calc_error_norms([func=(u_node,equations)->u_node,] u_ode, t, analyzer, semi::AbstractSemidiscretization, cache_analysis)7273Calculate discrete L2 and L∞ error norms of `func` applied to each nodal variable `u_node` in `u_ode`.74If no exact solution is available, "errors" are calculated using some reference state and can be useful75for regression tests.76"""77function calc_error_norms(u_ode, t, analyzer, semi::AbstractSemidiscretization,78cache_analysis)79return calc_error_norms(cons2cons, u_ode, t, analyzer, semi, cache_analysis)80end8182"""83semidiscretize(semi::AbstractSemidiscretization, tspan;84jac_prototype::Union{AbstractMatrix, Nothing} = nothing,85colorvec::Union{AbstractVector, Nothing} = nothing,86storage_type = nothing,87real_type = nothing)8889Wrap the semidiscretization `semi` as an ODE problem in the time interval `tspan`90that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).9192Optional keyword arguments:93- `jac_prototype`: Expected to come from [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl).94Specifies the sparsity structure of the Jacobian to enable e.g. efficient implicit time stepping.95- `colorvec`: Expected to come from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl).96Allows for even faster Jacobian computation if a sparse `jac_prototype` is given (optional).97- `storage_type` and `real_type`: Configure the underlying computational datastructures.98`storage_type` changes the fundamental array type being used, allowing the experimental use of `CuArray`99or other GPU array types. `real_type` changes the computational data type being used.100"""101function semidiscretize(semi::AbstractSemidiscretization, tspan;102jac_prototype::Union{AbstractMatrix, Nothing} = nothing,103colorvec::Union{AbstractVector, Nothing} = nothing,104reset_threads = true,105storage_type = nothing,106real_type = nothing)107# Optionally reset Polyester.jl threads. See108# https://github.com/trixi-framework/Trixi.jl/issues/1583109# https://github.com/JuliaSIMD/Polyester.jl/issues/30110if reset_threads111Polyester.reset_threads!()112end113114if !(storage_type === nothing && real_type === nothing)115if storage_type === nothing116storage_type = Array117end118if real_type === nothing119real_type = real(semi)120end121semi = trixi_adapt(storage_type, real_type, semi)122if eltype(tspan) !== real_type123tspan = convert.(real_type, tspan)124end125end126127u0_ode = compute_coefficients(first(tspan), semi) # Invoke initial condition128rhs_semi! = default_rhs(semi)129130# TODO: MPI, do we want to synchronize loading and print debug statements, e.g. using131# mpi_isparallel() && MPI.Barrier(mpi_comm())132# See https://github.com/trixi-framework/Trixi.jl/issues/328133iip = true # is-inplace, i.e., we modify a vector when calling `rhs_semi!`134specialize = SciMLBase.FullSpecialize # specialize on `rhs_semi!` and parameters (semi)135136# Check if Jacobian prototype is provided for sparse Jacobian137if jac_prototype !== nothing138# Convert `jac_prototype` to real type, as seen here:139# https://docs.sciml.ai/DiffEqDocs/stable/tutorials/advanced_ode_example/#Declaring-a-Sparse-Jacobian-with-Automatic-Sparsity-Detection140ode = SciMLBase.ODEFunction(rhs_semi!,141jac_prototype = convert.(eltype(u0_ode),142jac_prototype),143colorvec = colorvec) # coloring vector is optional144145return ODEProblem{iip, specialize}(ode, u0_ode, tspan, semi)146else147# We could also construct an `ODEFunction` explicitly without the Jacobian here,148# but we stick to the lean direct in-place function `rhs_semi!` and149# let OrdinaryDiffEq.jl handle the rest150return ODEProblem{iip, specialize}(rhs_semi!, u0_ode, tspan, semi)151end152end153154"""155semidiscretize(semi::AbstractSemidiscretization, tspan,156restart_file::AbstractString;157jac_prototype::Union{AbstractMatrix, Nothing} = nothing,158colorvec::Union{AbstractVector, Nothing} = nothing)159160Wrap the semidiscretization `semi` as an ODE problem in the time interval `tspan`161that can be passed to `solve` from the [SciML ecosystem](https://diffeq.sciml.ai/latest/).162The initial condition etc. is taken from the `restart_file`.163164Optional keyword arguments:165- `jac_prototype`: Expected to come from [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl).166Specifies the sparsity structure of the Jacobian to enable e.g. efficient implicit time stepping.167- `colorvec`: Expected to come from [SparseMatrixColorings.jl](https://github.com/gdalle/SparseMatrixColorings.jl).168Allows for even faster Jacobian computation. Not necessarily required when `jac_prototype` is given.169"""170function semidiscretize(semi::AbstractSemidiscretization, tspan,171restart_file::AbstractString;172jac_prototype::Union{AbstractMatrix, Nothing} = nothing,173colorvec::Union{AbstractVector, Nothing} = nothing,174reset_threads = true)175# Optionally reset Polyester.jl threads. See176# https://github.com/trixi-framework/Trixi.jl/issues/1583177# https://github.com/JuliaSIMD/Polyester.jl/issues/30178if reset_threads179Polyester.reset_threads!()180end181182u0_ode = load_restart_file(semi, restart_file) # Load initial condition from restart file183rhs_semi! = default_rhs(semi)184185# TODO: MPI, do we want to synchronize loading and print debug statements, e.g. using186# mpi_isparallel() && MPI.Barrier(mpi_comm())187# See https://github.com/trixi-framework/Trixi.jl/issues/328188iip = true # is-inplace, i.e., we modify a vector when calling `rhs_semi!`189specialize = SciMLBase.FullSpecialize # specialize on `rhs_semi!` and parameters (semi)190191# Check if Jacobian prototype is provided for sparse Jacobian192if jac_prototype !== nothing193# Convert `jac_prototype` to real type, as seen here:194# https://docs.sciml.ai/DiffEqDocs/stable/tutorials/advanced_ode_example/#Declaring-a-Sparse-Jacobian-with-Automatic-Sparsity-Detection195ode = SciMLBase.ODEFunction(rhs_semi!,196jac_prototype = convert.(eltype(u0_ode),197jac_prototype),198colorvec = colorvec) # coloring vector is optional199200return ODEProblem{iip, specialize}(ode, u0_ode, tspan, semi)201else202# We could also construct an `ODEFunction` explicitly without the Jacobian here,203# but we stick to the lean direct in-place function `rhs_semi!` and204# let OrdinaryDiffEq.jl handle the rest205return ODEProblem{iip, specialize}(rhs_semi!, u0_ode, tspan, semi)206end207end208209"""210compute_coefficients(func, t, semi::AbstractSemidiscretization)211212Compute the discrete coefficients of the continuous function `func` at time `t`213associated with the semidiscretization `semi`.214For example, the discrete coefficients of `func` for a discontinuous Galerkin215spectral element method ([`DGSEM`](@ref)) are the values of `func` at the216Lobatto-Legendre nodes. Similarly, a classical finite difference method will use217the values of `func` at the nodes of the grid assoociated with the semidiscretization218`semi`.219220For semidiscretizations `semi` associated with an initial condition, `func` can be omitted221to use the given initial condition at time `t`.222"""223function compute_coefficients(func, t, semi::AbstractSemidiscretization)224u_ode = allocate_coefficients(mesh_equations_solver_cache(semi)...)225# Call `compute_coefficients` defined below226compute_coefficients!(u_ode, func, t, semi)227return u_ode228end229230"""231compute_coefficients!(u_ode, func, t, semi::AbstractSemidiscretization)232233Same as [`compute_coefficients`](@ref) but stores the result in `u_ode`.234"""235function compute_coefficients!(u_ode, func, t, semi::AbstractSemidiscretization)236backend = trixi_backend(u_ode)237u = wrap_array(u_ode, semi)238# Call `compute_coefficients` defined by the solver239return compute_coefficients!(backend, u, func, t,240mesh_equations_solver_cache(semi)...)241end242243# Helper function to compute linear structure for a given semidiscretization244# `semi` and applied RHS `apply_rhs!`245function _linear_structure_from_rhs(semi::AbstractSemidiscretization, apply_rhs!)246# allocate memory247u_ode = allocate_coefficients(mesh_equations_solver_cache(semi)...)248du_ode = similar(u_ode)249250# get the right hand side from boundary conditions and optional source terms251u_ode .= zero(eltype(u_ode))252apply_rhs!(du_ode, u_ode)253b = -du_ode254255# Create a copy of `b` used internally to extract the linear part of `semi`.256# This is necessary to get everything correct when the user updates the257# returned vector `b`.258b_tmp = copy(b)259260# wrap the linear operator261A = LinearMap(length(u_ode), ismutating = true) do dest, src262apply_rhs!(dest, src)263@. dest += b_tmp264return dest265end266267return A, b268end269270"""271linear_structure(semi::AbstractSemidiscretization;272t0 = zero(real(semi)))273274Wraps the right-hand side operator of the semidiscretization `semi`275at time `t0` as an affine-linear operator given by a linear operator `A`276and a vector `b`:277```math278\\partial_t u(t) = A u(t) - b.279```280Works only for linear equations, i.e.,281equations which `have_constant_speed(equations) == True()`.282283This has the benefit of greatly reduced memory consumption compared to constructing284the full system matrix explicitly, as done for instance in285[`jacobian_fd`](@ref) and [`jacobian_ad_forward`](@ref).286287The returned linear operator `A` is a matrix-free operator which can be288supplied to iterative solvers from, e.g., [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl).289290It is also possible to use this to construct a sparse matrix without the detour of constructing291first the full Jacobian by calling292```julia293using SparseArrays294A_map, b = linear_structure(semi, t0 = t0)295A_sparse = sparse(A_map)296```297"""298function linear_structure(semi::AbstractSemidiscretization;299t0 = zero(real(semi)))300if have_constant_speed(semi.equations) == False()301throw(ArgumentError("`linear_structure` expects linear equations."))302end303304apply_rhs! = function (dest, src)305return rhs!(dest, src, semi, t0)306end307308return _linear_structure_from_rhs(semi, apply_rhs!)309end310311"""312jacobian_fd(semi::AbstractSemidiscretization;313t0 = zero(real(semi)),314u0_ode = compute_coefficients(t0, semi))315316Uses the right-hand side operator of the semidiscretization `semi`317and simple second order finite difference to compute the Jacobian `J`318of the semidiscretization `semi` at time `t0` and state `u0_ode`.319"""320function jacobian_fd(semi::AbstractSemidiscretization;321t0 = zero(real(semi)),322u0_ode = compute_coefficients(t0, semi))323# copy the initial state since it will be modified in the following324u_ode = copy(u0_ode)325du0_ode = similar(u_ode)326dup_ode = similar(u_ode)327dum_ode = similar(u_ode)328rhs_semi! = default_rhs(semi)329330# compute residual of linearization state331rhs_semi!(du0_ode, u_ode, semi, t0)332333# initialize Jacobian matrix334J = zeros(eltype(u_ode), length(u_ode), length(u_ode))335336# use second order finite difference to estimate Jacobian matrix337for idx in eachindex(u0_ode)338# determine size of fluctuation339# This is the approach used by FiniteDiff.jl to compute the340# step size, which assures that the finite difference is accurate341# for very small and very large absolute values `u0_ode[idx]`.342# See https://github.com/trixi-framework/Trixi.jl/pull/2514#issuecomment-3190534904.343absstep = sqrt(eps(typeof(u0_ode[idx])))344relstep = absstep345epsilon = max(relstep * abs(u0_ode[idx]), absstep)346347# plus fluctuation348u_ode[idx] = u0_ode[idx] + epsilon349rhs_semi!(dup_ode, u_ode, semi, t0)350351# minus fluctuation352u_ode[idx] = u0_ode[idx] - epsilon353rhs_semi!(dum_ode, u_ode, semi, t0)354355# restore linearization state356u_ode[idx] = u0_ode[idx]357358# central second order finite difference359@. J[:, idx] = (dup_ode - dum_ode) / (2 * epsilon)360end361362return J363end364365"""366jacobian_ad_forward(semi::AbstractSemidiscretization;367t0 = zero(real(semi)),368u0_ode = compute_coefficients(t0, semi))369370Uses the right-hand side operator of the semidiscretization `semi`371and forward mode automatic differentiation to compute the Jacobian `J`372of the semidiscretization `semi` at time `t0` and state `u0_ode`.373"""374function jacobian_ad_forward(semi::AbstractSemidiscretization;375t0 = zero(real(semi)),376u0_ode = compute_coefficients(t0, semi))377return jacobian_ad_forward(semi, t0, u0_ode)378end379380# The following version is for plain arrays381function jacobian_ad_forward(semi::AbstractSemidiscretization, t0, u0_ode)382du_ode = similar(u0_ode)383config = ForwardDiff.JacobianConfig(nothing, du_ode, u0_ode)384385# Use a function barrier since the generation of the `config` we use above386# is not type-stable387return _jacobian_ad_forward(semi, t0, u0_ode, du_ode, config)388end389390function _jacobian_ad_forward(semi, t0, u0_ode, du_ode, config)391new_semi = remake(semi, uEltype = eltype(config))392rhs_semi! = default_rhs(new_semi)393# Create anonymous function passed as first argument to `ForwardDiff.jacobian` to match394# `ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray,395# cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())`396J = ForwardDiff.jacobian(du_ode, u0_ode, config) do du_ode, u_ode397return rhs_semi!(du_ode, u_ode, new_semi, t0)398end399400return J401end402403# unpack u if it is wrapped in VectorOfArray (mainly for DGMulti solvers)404jacobian_ad_forward(semi::AbstractSemidiscretization, t0, u0_ode::VectorOfArray) = jacobian_ad_forward(semi,405t0,406parent(u0_ode))407408# This version is specialized to `StructArray`s used by some `DGMulti` solvers.409# We need to convert the numerical solution vectors since ForwardDiff cannot410# handle arrays of `SVector`s.411function jacobian_ad_forward(semi::AbstractSemidiscretization, t0, _u0_ode::StructArray)412u0_ode_plain = similar(_u0_ode, eltype(eltype(_u0_ode)),413(size(_u0_ode)..., nvariables(semi)))414for (v, u_v) in enumerate(StructArrays.components(_u0_ode))415u0_ode_plain[.., v] = u_v416end417du_ode_plain = similar(u0_ode_plain)418config = ForwardDiff.JacobianConfig(nothing, du_ode_plain, u0_ode_plain)419420# Use a function barrier since the generation of the `config` we use above421# is not type-stable422return _jacobian_ad_forward_structarrays(semi, t0, u0_ode_plain, du_ode_plain,423config)424end425426function _jacobian_ad_forward_structarrays(semi, t0, u0_ode_plain, du_ode_plain, config)427new_semi = remake(semi, uEltype = eltype(config))428rhs_semi! = default_rhs(new_semi)429# Create anonymous function passed as first argument to `ForwardDiff.jacobian` to match430# `ForwardDiff.jacobian(f!, y::AbstractArray, x::AbstractArray,431# cfg::JacobianConfig = JacobianConfig(f!, y, x), check=Val{true}())`432J = ForwardDiff.jacobian(du_ode_plain, u0_ode_plain,433config) do du_ode_plain, u_ode_plain434u_ode = StructArray{SVector{nvariables(semi), eltype(config)}}(ntuple(v -> view(u_ode_plain,435:,436:,437v),438nvariables(semi)))439du_ode = StructArray{SVector{nvariables(semi), eltype(config)}}(ntuple(v -> view(du_ode_plain,440:,441:,442v),443nvariables(semi)))444return rhs_semi!(du_ode, u_ode, new_semi, t0)445end446447return J448end449450# This version is specialized to arrays of `StaticArray`s used by some `DGMulti` solvers.451# We need to convert the numerical solution vectors since ForwardDiff cannot452# handle arrays of `SVector`s.453function jacobian_ad_forward(semi::AbstractSemidiscretization, t0,454_u0_ode::AbstractArray{<:SVector})455u0_ode_plain = reinterpret(eltype(eltype(_u0_ode)), _u0_ode)456du_ode_plain = similar(u0_ode_plain)457config = ForwardDiff.JacobianConfig(nothing, du_ode_plain, u0_ode_plain)458459# Use a function barrier since the generation of the `config` we use above460# is not type-stable461return _jacobian_ad_forward_staticarrays(semi, t0, u0_ode_plain, du_ode_plain,462config)463end464465function _jacobian_ad_forward_staticarrays(semi, t0, u0_ode_plain, du_ode_plain, config)466new_semi = remake(semi, uEltype = eltype(config))467rhs_semi! = default_rhs(new_semi)468J = ForwardDiff.jacobian(du_ode_plain, u0_ode_plain,469config) do du_ode_plain, u_ode_plain470u_ode = reinterpret(SVector{nvariables(semi), eltype(config)}, u_ode_plain)471du_ode = reinterpret(SVector{nvariables(semi), eltype(config)}, du_ode_plain)472return rhs_semi!(du_ode, u_ode, new_semi, t0)473end474475return J476end477478# Sometimes, it can be useful to save some (scalar) variables associated with each element,479# e.g. AMR indicators or shock indicators. Since these usually have to be re-computed480# directly before IO and do not necessarily need to be stored in memory before,481# get_element_variables!(element_variables, ..)482# is used to retrieve such up to date element variables, modifying483# `element_variables::Dict{Symbol,Any}` in place.484function get_element_variables!(element_variables, u_ode,485semi::AbstractSemidiscretization)486u = wrap_array(u_ode, semi)487return get_element_variables!(element_variables, u,488mesh_equations_solver_cache(semi)...)489end490491# Required for storing `extra_node_variables` in the `SaveSolutionCallback`.492# Not to be confused with `get_node_vars` which returns the variables of the simulated equation.493function get_node_variables!(node_variables, u_ode, semi::AbstractSemidiscretization)494return get_node_variables!(node_variables, u_ode,495mesh_equations_solver_cache(semi)...)496end497498# To implement AMR and use OrdinaryDiffEq.jl etc., we have to be a bit creative.499# Since the caches of the SciML ecosystem are immutable structs, we cannot simply500# change the underlying arrays therein. Hence, to support changing the number of501# DOFs, we need to use `resize!`. In some sense, this will force us to write more502# efficient code, since `resize!` will make use of previously allocated memory503# instead of allocating memory from scratch every time.504#505# However, multidimensional `Array`s don't support `resize!`. One option might be506# to use ElasticArrays.jl. But I don't really like that approach. Needing to use507# ElasticArray doesn't feel completely good to me, since we also want to experiment508# with other array types such as PaddedMatrices.jl, see trixi-framework/Trixi.jl#166.509# Then, we would need to wrap an Array inside something from PaddedMatrices.jl inside510# something from ElasticArrays.jl - or the other way round? Is that possible at all?511# If we go further, this looks like it could easily explode.512#513# Currently, the best option seems to be to let OrdinaryDiffEq.jl use `Vector`s,514# which can be `resize!`ed for AMR. Then, we have to wrap these `Vector`s inside515# Trixi.jl as our favorite multidimensional array type. We need to do this wrapping516# in every method exposed to OrdinaryDiffEq, i.e. in the first levels of things like517# rhs!, AMRCallback, StepsizeCallback, AnalysisCallback, SaveSolutionCallback518#519# This wrapping will also allow us to experiment more easily with additional520# kinds of wrapping, e.g. HybridArrays.jl or PaddedMatrices.jl to inform the521# compiler about the sizes of the first few dimensions in DG methods, i.e.522# nvariables(equations) and nnodes(dg).523#524# In some sense, having plain multidimensional `Array`s not support `resize!`525# isn't necessarily a bug (although it would be nice to add this possibility to526# base Julia) but can turn out to be a feature for us, because it will allow us527# more specializations.528# Since we can use multiple dispatch, these kinds of specializations can be529# tailored specifically to each combinations of mesh/solver etc.530#531# Under the hood, `wrap_array(u_ode, mesh, equations, solver, cache)` might532# (and probably will) use `unsafe_wrap`. Hence, you have to remember to533# `GC.@preserve` temporaries that are only used indirectly via `wrap_array`534# to avoid stochastic memory errors.535#536# Xref https://github.com/SciML/OrdinaryDiffEq.jl/pull/1275537function wrap_array(u_ode, semi::AbstractSemidiscretization)538return wrap_array(u_ode, mesh_equations_solver_cache(semi)...)539end540541# Like `wrap_array`, but guarantees to return a plain `Array`, which can be better542# for writing solution files etc.543function wrap_array_native(u_ode, semi::AbstractSemidiscretization)544return wrap_array_native(u_ode, mesh_equations_solver_cache(semi)...)545end546547# TODO: Taal, document interface?548# New mesh/solver combinations have to implement549# - ndofs(mesh, solver, cache)550# - ndofsgloabal(mesh, solver, cache)551# - ndims(mesh)552# - nnodes(solver)553# - real(solver)554# - create_cache(mesh, equations, solver, RealT)555# - wrap_array(u_ode, mesh, equations, solver, cache)556# - integrate(func, u, mesh, equations, solver, cache; normalize=true)557# - integrate_via_indices(func, u, mesh, equations, solver, cache, args...; normalize=true)558# - calc_error_norms(func, u, t, analyzer, mesh, equations, initial_condition, solver, cache, cache_analysis)559# - allocate_coefficients(mesh, equations, solver, cache)560# - compute_coefficients!(u, func, mesh, equations, solver, cache)561# - rhs!(du, u, t, mesh, equations, boundary_conditions, source_terms, solver, cache)562#563564end # @muladd565566567