dg_1d_parabolic.jl

# By default, Julia/LLVM does not use fused multiply-add operations (FMAs).
# Since these FMAs can increase the performance of many numerical algorithms,
# we need to opt-in explicitly.
# See https://ranocha.de/blog/Optimizing_EC_Trixi for further details.
@muladd begin
#! format: noindent

# This method is called when a `SemidiscretizationHyperbolicParabolic` is constructed.
# It constructs the basic `cache` used throughout the simulation to compute
# the RHS etc.
function create_cache_parabolic(mesh::TreeMesh{1},
                                equations_hyperbolic::AbstractEquations,
                                dg::DG, n_elements, uEltype)
    parabolic_container = init_parabolic_container_1d(nvariables(equations_hyperbolic),
                                                      nnodes(dg), n_elements,
                                                      uEltype)

    cache_parabolic = (; parabolic_container)

    return cache_parabolic
end

# This file collects all methods that have been updated to work with parabolic systems of equations
#
# assumptions: parabolic terms are of the form div(f(u, grad(u))) and
# will be discretized first order form as follows:
#               1. compute grad(u)
#               2. compute f(u, grad(u))
#               3. compute div(f(u, grad(u))) (i.e., the "regular" rhs! call)
# boundary conditions will be applied to both grad(u) and div(f(u, grad(u))).
function rhs_parabolic!(du, u, t, mesh::TreeMesh{1},
                        equations_parabolic::AbstractEquationsParabolic,
                        boundary_conditions_parabolic, source_terms_parabolic,
                        dg::DG, parabolic_scheme, cache, cache_parabolic)
    @unpack parabolic_container = cache_parabolic
    @unpack u_transformed, gradients, flux_parabolic = parabolic_container

    # Convert conservative variables to a form more suitable for parabolic flux calculations
    @trixi_timeit timer() "transform variables" begin
        transform_variables!(u_transformed, u, mesh, equations_parabolic, dg,
                             cache)
    end

    # Compute the gradients of the transformed variables
    @trixi_timeit timer() "calculate gradient" begin
        calc_gradient!(gradients, u_transformed, t, mesh, equations_parabolic,
                       boundary_conditions_parabolic, dg,
                       parabolic_scheme, cache)
    end

    # Compute and store the parabolic fluxes
    @trixi_timeit timer() "calculate parabolic fluxes" begin
        calc_parabolic_fluxes!(flux_parabolic, gradients, u_transformed, mesh,
                               equations_parabolic, dg, cache)
    end

    # The remainder of this function is essentially a regular rhs! for
    # parabolic equations (i.e., it computes the divergence of the parabolic fluxes)
    #
    # OBS! In `calc_parabolic_fluxes!`, the parabolic flux values at the volume nodes of each element have
    # been computed and stored in `flux_parabolic`. In the following, we *reuse* (abuse) the
    # `interfaces` and `boundaries` containers in `cache` to interpolate and store the
    # *fluxes* at the element surfaces, as opposed to interpolating and storing the *solution* (as it
    # is done in the hyperbolic operator). That is, `interfaces.u`/`boundaries.u` store *parabolic flux values*
    # and *not the solution*.  The advantage is that a) we do not need to allocate more storage, b) we
    # do not need to recreate the existing data structure only with a different name, and c) we do not
    # need to interpolate solutions *and* gradients to the surfaces.

    # Reset du
    @trixi_timeit timer() "reset ∂u/∂t" set_zero!(du, dg, cache)

    # Calculate volume integral
    # This calls the specialized version for the parabolic flux.
    @trixi_timeit timer() "volume integral" begin
        calc_volume_integral!(du, flux_parabolic, mesh, equations_parabolic, dg, cache)
    end

    # Prolong solution to interfaces
    # This reuses `prolong2interfaces!` for the purely hyperbolic case.
    @trixi_timeit timer() "prolong2interfaces" begin
        prolong2interfaces!(cache, flux_parabolic, mesh, equations_parabolic, dg)
    end

    # Calculate interface fluxes.
    # This calls the specialized version for the parabolic flux.
    @trixi_timeit timer() "interface flux" begin
        calc_interface_flux!(cache.elements.surface_flux_values,
                             mesh, equations_parabolic, dg,
                             parabolic_scheme, cache)
    end

    # Prolong solution to boundaries.
    # This reuses `prolong2boundaries!` for the purely hyperbolic case.
    @trixi_timeit timer() "prolong2boundaries" begin
        prolong2boundaries!(cache, flux_parabolic, mesh, equations_parabolic, dg)
    end

    # Calculate boundary fluxes.
    # This calls the specialized version for parabolic equations.
    @trixi_timeit timer() "boundary flux" begin
        calc_boundary_flux_divergence!(cache, t,
                                       boundary_conditions_parabolic, mesh,
                                       equations_parabolic,
                                       dg.surface_integral, dg)
    end

    # Calculate surface integrals.
    # This reuses `calc_surface_integral!` for the purely hyperbolic case.
    @trixi_timeit timer() "surface integral" begin
        calc_surface_integral!(nothing, du, u, mesh, equations_parabolic,
                               dg.surface_integral, dg, cache)
    end

    # Apply Jacobian from mapping to reference element
    @trixi_timeit timer() "Jacobian" begin
        apply_jacobian_parabolic!(du, mesh, equations_parabolic, dg, cache)
    end

    # Calculate source terms
    @trixi_timeit timer() "source terms parabolic" begin
        calc_sources_parabolic!(du, u, gradients, t, source_terms_parabolic,
                                equations_parabolic, dg, cache)
    end

    return nothing
end

# Transform solution variables prior to taking the gradient
# (e.g., conservative to primitive variables). Defaults to doing nothing.
# TODO: can we avoid copying data?
function transform_variables!(u_transformed, u, mesh::TreeMesh{1},
                              equations_parabolic::AbstractEquationsParabolic,
                              dg::DG, cache)
    transformation = gradient_variable_transformation(equations_parabolic)

    @threaded for element in eachelement(dg, cache)
        # Calculate volume terms in one element
        for i in eachnode(dg)
            u_node = get_node_vars(u, equations_parabolic, dg, i, element)
            u_transformed_node = transformation(u_node, equations_parabolic)
            set_node_vars!(u_transformed, u_transformed_node, equations_parabolic, dg,
                           i, element)
        end
    end

    return nothing
end

# This is the version used when calculating the divergence of the parabolic fluxes.
# Identical to weak-form volume integral/kernel for the purely hyperbolic case,
# except that the fluxes are here already precomputed in `calc_parabolic_fluxes!`
function calc_volume_integral!(du, flux_parabolic, mesh::TreeMesh{1},
                               equations_parabolic::AbstractEquationsParabolic,
                               dg::DGSEM, cache)
    @unpack derivative_hat = dg.basis

    @threaded for element in eachelement(dg, cache)
        # Calculate volume terms in one element
        for i in eachnode(dg)
            flux_1_node = get_node_vars(flux_parabolic, equations_parabolic, dg, i,
                                        element)

            for ii in eachnode(dg)
                multiply_add_to_node_vars!(du, derivative_hat[ii, i], flux_1_node,
                                           equations_parabolic, dg, ii, element)
            end
        end
    end

    return nothing
end

# This is the version used when calculating the divergence of the parabolic fluxes
function calc_interface_flux!(surface_flux_values, mesh::TreeMesh{1},
                              equations_parabolic, dg::DG, parabolic_scheme,
                              cache)
    @unpack neighbor_ids, orientations = cache.interfaces

    @threaded for interface in eachinterface(dg, cache)
        # Get neighboring elements
        left_id = neighbor_ids[1, interface]
        right_id = neighbor_ids[2, interface]

        # Determine interface direction with respect to elements:
        # orientation = 1: left -> 2, right -> 1
        left_direction = 2 * orientations[interface]
        right_direction = 2 * orientations[interface] - 1

        # Get precomputed fluxes at interfaces
        flux_ll, flux_rr = get_surface_node_vars(cache.interfaces.u,
                                                 equations_parabolic,
                                                 dg, interface)

        # compute interface flux for the DG divergence 
        flux = flux_parabolic(flux_ll, flux_rr, Divergence(),
                              equations_parabolic, parabolic_scheme)

        # Copy flux to left and right element storage
        for v in eachvariable(equations_parabolic)
            surface_flux_values[v, left_direction, left_id] = flux[v]
            surface_flux_values[v, right_direction, right_id] = flux[v]
        end
    end

    return nothing
end

function calc_parabolic_fluxes!(flux_parabolic, gradients, u_transformed,
                                mesh::TreeMesh{1},
                                equations_parabolic::AbstractEquationsParabolic,
                                dg::DG, cache)
    @threaded for element in eachelement(dg, cache)
        for i in eachnode(dg)
            # Get solution and gradients
            u_node = get_node_vars(u_transformed, equations_parabolic, dg, i, element)
            gradients_1_node = get_node_vars(gradients, equations_parabolic, dg,
                                             i, element)

            # Calculate parabolic flux and store each component for later use
            flux_parabolic_node = flux(u_node, (gradients_1_node,), 1,
                                       equations_parabolic)
            set_node_vars!(flux_parabolic, flux_parabolic_node, equations_parabolic, dg,
                           i, element)
        end
    end

    return nothing
end

function calc_boundary_flux_gradient!(cache, t,
                                      boundary_conditions_parabolic::BoundaryConditionPeriodic,
                                      mesh::TreeMesh{1},
                                      equations_parabolic::AbstractEquationsParabolic,
                                      surface_integral, dg::DG)
    return nothing
end

function calc_boundary_flux_divergence!(cache, t,
                                        boundary_conditions_parabolic::BoundaryConditionPeriodic,
                                        mesh::TreeMesh{1},
                                        equations_parabolic::AbstractEquationsParabolic,
                                        surface_integral, dg::DG)
    return nothing
end

function calc_boundary_flux_gradient!(cache, t,
                                      boundary_conditions_parabolic::NamedTuple,
                                      mesh::TreeMesh{1}, # for dispatch only
                                      equations_parabolic::AbstractEquationsParabolic,
                                      surface_integral, dg::DG)
    @unpack surface_flux_values = cache.elements
    @unpack n_boundaries_per_direction = cache.boundaries

    # Calculate indices
    lasts = accumulate(+, n_boundaries_per_direction)
    firsts = lasts - n_boundaries_per_direction .+ 1

    # Calc boundary fluxes in each direction
    calc_boundary_flux_by_direction_gradient!(surface_flux_values, t,
                                              boundary_conditions_parabolic[1],
                                              equations_parabolic, surface_integral,
                                              dg, cache,
                                              1, firsts[1], lasts[1])
    calc_boundary_flux_by_direction_gradient!(surface_flux_values, t,
                                              boundary_conditions_parabolic[2],
                                              equations_parabolic, surface_integral,
                                              dg, cache,
                                              2, firsts[2], lasts[2])

    return nothing
end

function calc_boundary_flux_by_direction_gradient!(surface_flux_values::AbstractArray{<:Any,
                                                                                      3},
                                                   t, boundary_condition,
                                                   equations_parabolic::AbstractEquationsParabolic,
                                                   surface_integral, dg::DG, cache,
                                                   direction, first_boundary,
                                                   last_boundary)
    @unpack surface_flux = surface_integral
    @unpack u, neighbor_ids, neighbor_sides, node_coordinates, orientations = cache.boundaries

    @threaded for boundary in first_boundary:last_boundary
        # Get neighboring element
        neighbor = neighbor_ids[boundary]

        # Get boundary flux
        u_ll, u_rr = get_surface_node_vars(u, equations_parabolic, dg, boundary)
        if neighbor_sides[boundary] == 1 # Element is on the left, boundary on the right
            u_inner = u_ll
        else # Element is on the right, boundary on the left
            u_inner = u_rr
        end

        # TODO: revisit if we want more general boundary treatments.
        # This assumes the gradient numerical flux at the boundary is the gradient variable,
        # which is consistent with BR1, LDG.
        flux_inner = u_inner

        x = get_node_coords(node_coordinates, equations_parabolic, dg, boundary)
        flux = boundary_condition(flux_inner, u_inner, orientations[boundary],
                                  direction,
                                  x, t, Gradient(), equations_parabolic)

        # Copy flux to left and right element storage
        for v in eachvariable(equations_parabolic)
            surface_flux_values[v, direction, neighbor] = flux[v]
        end
    end

    return nothing
end

function calc_boundary_flux_divergence!(cache, t,
                                        boundary_conditions_parabolic::NamedTuple,
                                        mesh::TreeMesh{1},
                                        equations_parabolic::AbstractEquationsParabolic,
                                        surface_integral, dg::DG)
    @unpack surface_flux_values = cache.elements
    @unpack n_boundaries_per_direction = cache.boundaries

    # Calculate indices
    lasts = accumulate(+, n_boundaries_per_direction)
    firsts = lasts - n_boundaries_per_direction .+ 1

    # Calc boundary fluxes in each direction
    calc_boundary_flux_by_direction_divergence!(surface_flux_values, t,
                                                boundary_conditions_parabolic[1],
                                                equations_parabolic, surface_integral,
                                                dg, cache,
                                                1, firsts[1], lasts[1])
    calc_boundary_flux_by_direction_divergence!(surface_flux_values, t,
                                                boundary_conditions_parabolic[2],
                                                equations_parabolic, surface_integral,
                                                dg, cache,
                                                2, firsts[2], lasts[2])

    return nothing
end

function calc_boundary_flux_by_direction_divergence!(surface_flux_values::AbstractArray{<:Any,
                                                                                        3},
                                                     t,
                                                     boundary_condition,
                                                     equations_parabolic::AbstractEquationsParabolic,
                                                     surface_integral, dg::DG, cache,
                                                     direction, first_boundary,
                                                     last_boundary)
    @unpack surface_flux = surface_integral

    # Note: cache.boundaries.u contains the unsigned normal component (using "orientation", not "direction")
    # of the parabolic flux, as computed in `prolong2boundaries!`
    @unpack u, neighbor_ids, neighbor_sides, node_coordinates, orientations = cache.boundaries

    @threaded for boundary in first_boundary:last_boundary
        # Get neighboring element
        neighbor = neighbor_ids[boundary]

        # Get parabolic boundary fluxes
        flux_ll, flux_rr = get_surface_node_vars(u, equations_parabolic, dg, boundary)
        if neighbor_sides[boundary] == 1 # Element is on the left, boundary on the right
            flux_inner = flux_ll
        else # Element is on the right, boundary on the left
            flux_inner = flux_rr
        end

        x = get_node_coords(node_coordinates, equations_parabolic, dg, boundary)

        # TODO: add a field in `cache.boundaries` for gradient information.
        # Here, we pass in `u_inner = nothing` since we overwrite cache.boundaries.u with gradient information.
        # This currently works with Dirichlet/Neuman boundary conditions for LaplaceDiffusion2D and
        # NoSlipWall/Adiabatic boundary conditions for CompressibleNavierStokesDiffusion2D as of 2022-6-27.
        # It will not work with implementations which utilize `u_inner` to impose boundary conditions.
        flux = boundary_condition(flux_inner, nothing, orientations[boundary],
                                  direction,
                                  x, t, Divergence(), equations_parabolic)

        # Copy flux to left and right element storage
        for v in eachvariable(equations_parabolic)
            surface_flux_values[v, direction, neighbor] = flux[v]
        end
    end

    return nothing
end

function calc_volume_integral_gradient!(gradients, u_transformed,
                                        mesh::TreeMesh{1}, # for dispatch only
                                        equations_parabolic::AbstractEquationsParabolic,
                                        dg::DGSEM, cache)
    @unpack derivative_hat = dg.basis

    @threaded for element in eachelement(dg, cache)
        # Calculate volume terms in one element,
        # corresponds to `kernel` functions for the hyperbolic part of the flux
        for i in eachnode(dg)
            u_node = get_node_vars(u_transformed, equations_parabolic, dg,
                                   i, element)

            for ii in eachnode(dg)
                multiply_add_to_node_vars!(gradients, derivative_hat[ii, i],
                                           u_node, equations_parabolic, dg,
                                           ii, element)
            end
        end
    end

    return nothing
end

function calc_interface_flux_gradient!(surface_flux_values,
                                       mesh::TreeMesh{1},
                                       equations_parabolic,
                                       dg::DG, parabolic_scheme, cache)
    @unpack neighbor_ids, orientations = cache.interfaces

    @threaded for interface in eachinterface(dg, cache)
        # Get neighboring elements
        left_id = neighbor_ids[1, interface]
        right_id = neighbor_ids[2, interface]

        # Determine interface direction with respect to elements:
        # orientation = 1: left -> 2, right -> 1
        left_direction = 2 * orientations[interface]
        right_direction = 2 * orientations[interface] - 1

        # Call pointwise Riemann solver
        u_ll, u_rr = get_surface_node_vars(cache.interfaces.u,
                                           equations_parabolic, dg, interface)

        flux = flux_parabolic(u_ll, u_rr, Gradient(),
                              equations_parabolic, parabolic_scheme)

        # Copy flux to left and right element storage
        for v in eachvariable(equations_parabolic)
            surface_flux_values[v, left_direction, left_id] = flux[v]
            # No sign flip needed for gradient computation because for parabolic terms, 
            # the normals are not embedded in `flux_` for gradient computations.
            surface_flux_values[v, right_direction, right_id] = flux[v]
        end
    end

    return nothing
end

function calc_surface_integral_gradient!(gradients,
                                         mesh::TreeMesh{1}, # for dispatch only
                                         equations_parabolic::AbstractEquationsParabolic,
                                         dg::DGSEM, cache)
    @unpack inverse_weights = dg.basis
    @unpack surface_flux_values = cache.elements

    # Note that all fluxes have been computed with outward-pointing normal vectors.
    # We also use explicit assignments instead of `+=` to let `@muladd` turn these
    # into FMAs (see comment at the top of the file).
    factor = inverse_weights[1] # For LGL basis: Identical to weighted boundary interpolation at x = ±1
    @threaded for element in eachelement(dg, cache)
        for v in eachvariable(equations_parabolic)
            # surface at -x
            gradients[v, 1, element] = (gradients[v, 1, element] -
                                        surface_flux_values[v, 1, element] *
                                        factor)

            # surface at +x
            gradients[v, nnodes(dg), element] = (gradients[v, nnodes(dg), element] +
                                                 surface_flux_values[v, 2, element] *
                                                 factor)
        end
    end

    return nothing
end

function calc_surface_integral_gradient!(gradients,
                                         mesh::TreeMesh{1}, # for dispatch only
                                         equations_parabolic::AbstractEquationsParabolic,
                                         dg::DGSEM{<:GaussLegendreBasis}, cache)
    @unpack boundary_interpolation_inverse_weights = dg.basis
    @unpack surface_flux_values = cache.elements

    # Note that all fluxes have been computed with outward-pointing normal vectors.
    # We also use explicit assignments instead of `+=` to let `@muladd` turn these
    # into FMAs (see comment at the top of the file).
    @threaded for element in eachelement(dg, cache)
        for v in eachvariable(equations_parabolic)
            # Aliases for repeatedly accessed variables
            surface_flux_minus = surface_flux_values[v, 1, element]
            surface_flux_plus = surface_flux_values[v, 2, element]
            for ii in eachnode(dg)
                # surface at -x
                gradients[v, ii, element] = (gradients[v, ii, element] -
                                             surface_flux_minus *
                                             boundary_interpolation_inverse_weights[ii,
                                                                                    1])

                # surface at +x
                gradients[v, ii, element] = (gradients[v, ii, element] +
                                             surface_flux_plus *
                                             boundary_interpolation_inverse_weights[ii,
                                                                                    2])
            end
        end
    end

    return nothing
end

# Calculate the gradient of the transformed variables
function calc_gradient!(gradients, u_transformed, t, mesh::TreeMesh{1},
                        equations_parabolic, boundary_conditions_parabolic,
                        dg::DG, parabolic_scheme, cache)

    # Reset gradients
    @trixi_timeit timer() "reset gradients" begin
        set_zero!(gradients, dg, cache)
    end

    # Calculate volume integral
    @trixi_timeit timer() "volume integral" begin
        calc_volume_integral_gradient!(gradients, u_transformed,
                                       mesh, equations_parabolic, dg, cache)
    end

    # Prolong solution to interfaces.
    # This reusues `prolong2interfaces!` for the purely hyperbolic case.
    @trixi_timeit timer() "prolong2interfaces" begin
        prolong2interfaces!(cache, u_transformed, mesh,
                            equations_parabolic, dg)
    end

    # Calculate interface fluxes
    @trixi_timeit timer() "interface flux" begin
        @unpack surface_flux_values = cache.elements
        calc_interface_flux_gradient!(surface_flux_values, mesh, equations_parabolic,
                                      dg, parabolic_scheme, cache)
    end

    # Prolong solution to boundaries.
    # This reuses `prolong2boundaries!` for the purely hyperbolic case.
    @trixi_timeit timer() "prolong2boundaries" begin
        prolong2boundaries!(cache, u_transformed, mesh, equations_parabolic, dg)
    end

    # Calculate boundary fluxes
    @trixi_timeit timer() "boundary flux" begin
        calc_boundary_flux_gradient!(cache, t,
                                     boundary_conditions_parabolic, mesh,
                                     equations_parabolic,
                                     dg.surface_integral, dg)
    end

    # Calculate surface integrals
    @trixi_timeit timer() "surface integral" begin
        calc_surface_integral_gradient!(gradients, mesh, equations_parabolic,
                                        dg, cache)
    end

    # Apply Jacobian from mapping to reference element
    @trixi_timeit timer() "Jacobian" begin
        apply_jacobian_parabolic!(gradients, mesh, equations_parabolic, dg,
                                  cache)
    end

    return nothing
end

# Needed to *not* flip the sign of the inverse Jacobian.
# This is because the parabolic fluxes are assumed to be of the form
#   `du/dt + df/dx = dg/dx + source(x,t)`,
# where f(u) is the inviscid flux and g(u) is the parabolic flux.
function apply_jacobian_parabolic!(du::AbstractArray, mesh::TreeMesh{1},
                                   equations_parabolic::AbstractEquationsParabolic,
                                   dg::DG, cache)
    @unpack inverse_jacobian = cache.elements

    @threaded for element in eachelement(dg, cache)
        factor = inverse_jacobian[element]

        for i in eachnode(dg)
            for v in eachvariable(equations_parabolic)
                du[v, i, element] *= factor
            end
        end
    end

    return nothing
end

# Need dimension specific version to avoid error at dispatching
function calc_sources_parabolic!(du, u, gradients, t, source_terms_parabolic::Nothing,
                                 equations_parabolic::AbstractEquations{1}, dg::DG,
                                 cache)
    return nothing
end

function calc_sources_parabolic!(du, u, gradients, t, source_terms_parabolic,
                                 equations_parabolic::AbstractEquations{1}, dg::DG,
                                 cache)
    @unpack node_coordinates = cache.elements

    @threaded for element in eachelement(dg, cache)
        for i in eachnode(dg)
            u_local = get_node_vars(u, equations_parabolic, dg, i, element)
            gradients_x_local = get_node_vars(gradients, equations_parabolic, dg, i,
                                              element)
            x_local = get_node_coords(node_coordinates, equations_parabolic, dg,
                                      i, element)
            du_local = source_terms_parabolic(u_local, (gradients_x_local,), x_local, t,
                                              equations_parabolic)
            add_to_node_vars!(du, du_local, equations_parabolic, dg, i, element)
        end
    end

    return nothing
end
end # @muladd