analysis_covariant.jl

@muladd begin
#! format: noindent

# When the equations are of type AbstractCovariantEquations, the functions which we would 
# like to integrate depend on the solution as well as the auxiliary variables
function Trixi.integrate(func::Func, u,
                         mesh::Union{TreeMesh{2}, StructuredMesh{2},
                                     StructuredMeshView{2},
                                     UnstructuredMesh2D, P4estMesh{2}, T8codeMesh{2}},
                         equations::AbstractCovariantEquations{2},
                         dg::Union{DGSEM, FDSBP},
                         cache; normalize = true) where {Func}
    (; aux_node_vars) = cache.auxiliary_variables

    Trixi.integrate_via_indices(u, mesh, equations, dg, cache;
                                normalize = normalize) do u, i, j, element, equations,
                                                          dg
        u_local = Trixi.get_node_vars(u, equations, dg, i, j, element)
        aux_local = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
        return func(u_local, aux_local, equations)
    end
end

function Trixi.integrate(func::Func, u,
                         mesh::DGMultiMesh,
                         equations::AbstractCovariantEquations,
                         dg::DGMulti, cache; normalize = true) where {Func}
    rd = dg.basis
    md = mesh.md
    @unpack u_values, aux_quad_values = cache

    # interpolate u to quadrature points
    Trixi.apply_to_each_field(Trixi.mul_by!(rd.Vq), u_values, u)

    integral = zero(func(u_values[1], equations))
    total_volume = zero(sum(rd.wq))
    for element in Trixi.eachelement(dg, cache)
        weights = area_element.(aux_quad_values[:, element], equations) .* rd.wq
        integral += sum(weights .* func.(u_values[:, element], equations))
        total_volume += sum(weights)
    end
    if normalize == true
        integral /= total_volume
    end
    return integral
end

# For the covariant form, we want to integrate using the exact area element 
# J = √G = (det(AᵀA))^(1/2), which is stored in cache.auxiliary_variables, not the approximate 
# area element used in the Cartesian formulation, which stored in cache.elements
function Trixi.integrate_via_indices(func::Func, u,
                                     mesh::Union{StructuredMesh{2},
                                                 StructuredMeshView{2},
                                                 UnstructuredMesh2D, P4estMesh{2},
                                                 T8codeMesh{2}},
                                     equations::AbstractCovariantEquations{2},
                                     dg::DGSEM, cache, args...;
                                     normalize = true) where {Func}
    (; weights) = dg.basis
    (; aux_node_vars) = cache.auxiliary_variables

    # Initialize integral with zeros of the right shape
    integral = zero(func(u, 1, 1, 1, equations, dg, args...))
    total_volume = zero(real(mesh))

    # Use quadrature to numerically integrate over entire domain
    for element in eachelement(dg, cache)
        for j in eachnode(dg), i in eachnode(dg)
            aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
            J = area_element(aux_node, equations)
            integral += weights[i] * weights[j] * J *
                        func(u, i, j, element, equations, dg, args...)
            total_volume += weights[i] * weights[j] * J
        end
    end

    # Normalize with total volume
    if normalize
        integral = integral / total_volume
    end

    return integral
end

# Entropy time derivative for cons2entropy function which depends on auxiliary variables
function Trixi.analyze(::typeof(Trixi.entropy_timederivative), du, u, t,
                       mesh::P4estMesh{2}, equations::AbstractCovariantEquations{2},
                       dg::Union{DGSEM, FDSBP}, cache)
    (; aux_node_vars) = cache.auxiliary_variables

    # Calculate ∫(∂S/∂u ⋅ ∂u/∂t)dΩ
    Trixi.integrate_via_indices(u, mesh, equations, dg, cache,
                                du) do u, i, j, element, equations, dg, du_node
        # Get auxiliary variables, solution variables, and time derivative at given node
        aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
        u_node = Trixi.get_node_vars(u, equations, dg, i, j, element)
        du_node = Trixi.get_node_vars(du, equations, dg, i, j, element)

        # compute ∂S/∂u ⋅ ∂u/∂t, where the entropy variables ∂S/∂u depend on the solution 
        # and auxiliary variables
        dot(cons2entropy(u_node, aux_node, equations), du_node)
    end
end

# Entropy time derivative for cons2entropy function which depends on auxiliary variables
function Trixi.analyze(::typeof(Trixi.entropy_timederivative), du, u, t,
                       mesh::DGMultiMesh, equations::AbstractCovariantEquations,
                       dg::DGMulti, cache)
    rd = dg.basis
    md = mesh.md
    @unpack u_values, aux_quad_values = cache

    # interpolate u, du to quadrature points
    du_values = similar(u_values)
    Trixi.apply_to_each_field(Trixi.mul_by!(rd.Vq), du_values, du)
    Trixi.apply_to_each_field(Trixi.mul_by!(rd.Vq), u_values, u)

    # compute ∫v(u) * du/dt = ∫dS/dt. We can directly compute v(u) instead of computing the entropy
    # projection here, since the RHS will be projected to polynomials of degree N and testing with
    # the L2 projection of v(u) would be equivalent to testing with v(u) due to the moment-preserving
    # property of the L2 projection.
    dS_dt = zero(eltype(first(du)))
    for i in Base.OneTo(length(md.wJq))
        ref_index = mod(i - 1, rd.Nq) + 1
        node_weight = rd.wq[ref_index] * area_element(aux_quad_values[i], equations)
        dS_dt += dot(cons2entropy(u_values[i], aux_quad_values[i], equations),
                     du_values[i]) * node_weight
    end
    return dS_dt
end

# L2 and Linf error calculation for the covariant form
function Trixi.calc_error_norms(func, u, t, analyzer, mesh::P4estMesh{2},
                                equations::AbstractCovariantEquations{2},
                                initial_condition, dg::DGSEM, cache, cache_analysis)
    (; weights) = dg.basis
    (; node_coordinates) = cache.elements
    (; aux_node_vars) = cache.auxiliary_variables

    # Set up data structures
    l2_error = zero(func(Trixi.get_node_vars(u, equations, dg, 1, 1, 1), equations))
    linf_error = copy(l2_error)
    total_volume = zero(real(mesh))

    # Iterate over all elements for error calculations
    for element in eachelement(dg, cache)

        # Calculate errors at each volume quadrature node
        for j in eachnode(dg), i in eachnode(dg)
            x_node = Trixi.get_node_coords(node_coordinates, equations, dg, i, j,
                                           element)

            # Convert exact solution into contravariant components using geometric
            # information stored in aux vars
            aux_node = get_node_aux_vars(aux_node_vars, equations, dg, i, j, element)
            u_exact = initial_condition(x_node, t, aux_node, equations)

            # Compute the difference as usual
            u_numerical = Trixi.get_node_vars(u, equations, dg, i, j, element)
            diff = func(u_exact, equations) - func(u_numerical, equations)

            # For the L2 error, integrate with respect to area element stored in aux vars 
            J = area_element(aux_node, equations)
            l2_error += diff .^ 2 * (weights[i] * weights[j] * J)

            # Compute Linf error as usual
            linf_error = @. max(linf_error, abs(diff))

            # Increment total volume according to the volume element stored in aux vars
            total_volume += weights[i] * weights[j] * J
        end
    end

    # For L2 error, divide by total volume
    l2_error = @. sqrt(l2_error / total_volume)

    return l2_error, linf_error
end

# L2 and Linf error calculation for the covariant form
function Trixi.calc_error_norms(func, u, t, analyzer,
                                mesh::DGMultiMesh{NDIMS_AMBIENT},
                                equations::AbstractCovariantEquations,
                                initial_condition,
                                dg::DGMulti{NDIMS}, cache,
                                cache_analysis) where {NDIMS, NDIMS_AMBIENT}
    rd = dg.basis
    md = mesh.md
    @unpack u_values, aux_quad_values = cache

    # interpolate u to quadrature points
    Trixi.apply_to_each_field(Trixi.mul_by!(rd.Vq), u_values, u)

    component_l2_errors = zero(eltype(u_values))
    component_linf_errors = zero(eltype(u_values))
    total_volume = zero(eltype(u_values[1]))
    for i in Trixi.each_quad_node_global(mesh, dg, cache)
        u_exact = initial_condition(SVector(getindex.(md.xyzq, i)), t,
                                    aux_quad_values[i], equations)
        error_at_node = func(u_values[i], equations) - func(u_exact, equations)
        ref_index = mod(i - 1, rd.Nq) + 1
        node_weight = rd.wq[ref_index] * area_element(aux_quad_values[i], equations)
        component_l2_errors += node_weight * error_at_node .^ 2
        component_linf_errors = max.(component_linf_errors, abs.(error_at_node))
        total_volume += node_weight
    end
    return sqrt.(component_l2_errors ./ total_volume), component_linf_errors
end
end # @muladd