stepsize_dg2d.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

function max_dt(u, t, mesh::TreeMesh{2},
                constant_speed::False, equations, dg::DG, cache)
    # Avoid division by zero if the speed vanishes everywhere,
    # e.g. for steady-state linear advection
    max_scaled_speed = nextfloat(zero(t))

    @batch reduction=(max, max_scaled_speed) for element in eachelement(dg, cache)
        max_lambda1 = max_lambda2 = zero(max_scaled_speed)
        for j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, element)
            lambda1, lambda2 = max_abs_speeds(u_node, equations)
            max_lambda1 = Base.max(max_lambda1, lambda1)
            max_lambda2 = Base.max(max_lambda2, lambda2)
        end
        inv_jacobian = cache.elements.inverse_jacobian[element]
        # Use `Base.max` to prevent silent failures, as `max` from `@fastmath` doesn't propagate
        # `NaN`s properly. See https://github.com/trixi-framework/Trixi.jl/pull/2445#discussion_r2336812323
        max_scaled_speed = Base.max(max_scaled_speed,
                                    inv_jacobian * (max_lambda1 + max_lambda2))
    end

    return 2 / (nnodes(dg) * max_scaled_speed)
end

function max_dt(u, t, mesh::TreeMesh{2},
                constant_diffusivity::False, equations,
                equations_parabolic::AbstractEquationsParabolic,
                dg::DG, cache)
    # Avoid division by zero if the diffusivity vanishes everywhere
    max_scaled_speed = nextfloat(zero(t))

    @batch reduction=(max, max_scaled_speed) for element in eachelement(dg, cache)
        max_diffusivity_ = zero(max_scaled_speed)
        for j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, element)
            # Note: For the currently supported parabolic equations
            # Diffusion & Navier-Stokes, we only have one diffusivity.
            diffusivity = max_diffusivity(u_node, equations_parabolic)
            max_diffusivity_ = max(max_diffusivity_, diffusivity)
        end
        inv_jacobian = cache.elements.inverse_jacobian[element] # 2 / Δx
        max_scaled_speed = max(max_scaled_speed, inv_jacobian^2 * max_diffusivity_)
    end

    # Factor 4 cancels with 2^2 from `inv_jacobian^2`, resulting in Δx^2
    return 4 / (nnodes(dg) * max_scaled_speed)
end

function max_dt(u, t, mesh::TreeMesh{2},
                constant_speed::True, equations, dg::DG, cache)
    # Avoid division by zero if the speed vanishes everywhere,
    # e.g. for steady-state linear advection
    max_scaled_speed = nextfloat(zero(t))

    max_lambda1, max_lambda2 = max_abs_speeds(equations)

    @batch reduction=(max, max_scaled_speed) for element in eachelement(dg, cache)
        inv_jacobian = cache.elements.inverse_jacobian[element]
        # Use `Base.max` to prevent silent failures, as `max` from `@fastmath` doesn't propagate
        # `NaN`s properly. See https://github.com/trixi-framework/Trixi.jl/pull/2445#discussion_r2336812323
        max_scaled_speed = Base.max(max_scaled_speed,
                                    inv_jacobian * (max_lambda1 + max_lambda2))
    end

    return 2 / (nnodes(dg) * max_scaled_speed)
end

function max_dt(u, t, mesh::TreeMeshParallel{2},
                constant_speed::False, equations, dg::DG, cache)
    # call the method accepting a general `mesh::TreeMesh{2}`
    # TODO: MPI, we should improve this; maybe we should dispatch on `u`
    #       and create some MPI array type, overloading broadcasting and mapreduce etc.
    #       Then, this specific array type should also work well with DiffEq etc.
    dt = invoke(max_dt,
                Tuple{typeof(u), typeof(t), TreeMesh{2},
                      typeof(constant_speed), typeof(equations), typeof(dg),
                      typeof(cache)},
                u, t, mesh, constant_speed, equations, dg, cache)
    # Base.min instead of min needed, see comment in src/auxiliary/math.jl
    dt = MPI.Allreduce!(Ref(dt), Base.min, mpi_comm())[]

    return dt
end

function max_dt(u, t, mesh::TreeMeshParallel{2},
                constant_speed::True, equations, dg::DG, cache)
    # call the method accepting a general `mesh::TreeMesh{2}`
    # TODO: MPI, we should improve this; maybe we should dispatch on `u`
    #       and create some MPI array type, overloading broadcasting and mapreduce etc.
    #       Then, this specific array type should also work well with DiffEq etc.
    dt = invoke(max_dt,
                Tuple{typeof(u), typeof(t), TreeMesh{2},
                      typeof(constant_speed), typeof(equations), typeof(dg),
                      typeof(cache)},
                u, t, mesh, constant_speed, equations, dg, cache)
    # Base.min instead of min needed, see comment in src/auxiliary/math.jl
    dt = MPI.Allreduce!(Ref(dt), Base.min, mpi_comm())[]

    return dt
end

# Hyperbolic-parabolic simulations are not yet supported on MPI-parallel meshes.
# Thus, there is no `max_dt` function for `TreeMeshParallel{2}` and
# `equations_parabolic::AbstractEquationsParabolic` implemented.

function max_dt(u, t,
                mesh::Union{StructuredMesh{2}, UnstructuredMesh2D, P4estMesh{2},
                            P4estMeshView{2}, T8codeMesh{2}, StructuredMeshView{2}},
                constant_speed, equations, dg::DG, cache)
    backend = trixi_backend(u)

    max_lambda = calc_max_scaled_speed(backend, u, mesh, constant_speed, equations, dg,
                                       cache)

    # Avoid division by zero if the speed vanishes everywhere,
    # e.g. for steady-state linear advection
    max_scaled_speed = Base.max(nextfloat(zero(t)), max_lambda)

    return 2 / (nnodes(dg) * max_scaled_speed)
end

@inline function max_scaled_speed_per_element(u,
                                              ::Type{<:Union{StructuredMesh{2},
                                                             UnstructuredMesh2D,
                                                             P4estMesh{2},
                                                             T8codeMesh{2},
                                                             StructuredMeshView{2}}},
                                              constant_speed::False, equations, dg::DG,
                                              contravariant_vectors, inverse_jacobian,
                                              element)
    max_lambda1 = max_lambda2 = zero(eltype(u))
    for j in eachnode(dg), i in eachnode(dg)
        u_node = get_node_vars(u, equations, dg, i, j, element)
        lambda1, lambda2 = max_abs_speeds(u_node, equations)

        # Local speeds transformed to the reference element
        Ja11, Ja12 = get_contravariant_vector(1, contravariant_vectors,
                                              i, j, element)
        lambda1_transformed = abs(Ja11 * lambda1 + Ja12 * lambda2)
        Ja21, Ja22 = get_contravariant_vector(2, contravariant_vectors,
                                              i, j, element)
        lambda2_transformed = abs(Ja21 * lambda1 + Ja22 * lambda2)

        inv_jacobian = abs(inverse_jacobian[i, j, element])

        max_lambda1 = Base.max(max_lambda1, lambda1_transformed * inv_jacobian)
        max_lambda2 = Base.max(max_lambda2, lambda2_transformed * inv_jacobian)
    end
    return max_lambda1 + max_lambda2
end

function max_dt(u, t,
                mesh::P4estMesh{2}, # Parabolic terms currently only for `TreeMesh` and `P4estMesh`
                constant_diffusivity::False, equations,
                equations_parabolic::AbstractEquationsParabolic,
                dg::DG, cache)
    # Avoid division by zero if the diffusivity vanishes everywhere
    max_scaled_diffusivity = nextfloat(zero(t))

    @unpack contravariant_vectors, inverse_jacobian = cache.elements

    @batch reduction=(max, max_scaled_diffusivity) for element in eachelement(dg, cache)
        max_diffusivity1 = max_diffusivity2 = zero(max_scaled_diffusivity)
        for j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, element)
            diffusivity = max_diffusivity(u_node, equations_parabolic)

            # Local diffusivity transformed to the reference element
            Ja11, Ja12 = get_contravariant_vector(1, contravariant_vectors,
                                                  i, j, element)
            # The metric terms are squared due to the second derivatives in the parabolic terms,
            # which lead to application of the chain rule (coordinate transformation) twice.
            # See for instance also the implementation in FLEXI
            # https://github.com/flexi-framework/flexi/blob/e980e8635e6605daf906fc176c9897314a9cd9b1/src/equations/navierstokes/calctimestep.f90#L75-L77
            # or FLUXO:
            # https://github.com/project-fluxo/fluxo/blob/c7e0cc9b7fd4569dcab67bbb6e5a25c0a84859f1/src/equation/navierstokes/calctimestep.f90#L130-L132
            diffusivity1_transformed = diffusivity * (Ja11^2 + Ja12^2)
            Ja21, Ja22 = get_contravariant_vector(2, contravariant_vectors,
                                                  i, j, element)
            diffusivity2_transformed = diffusivity * (Ja21^2 + Ja22^2)

            inv_jacobian = abs(inverse_jacobian[i, j, element])

            max_diffusivity1 = Base.max(max_diffusivity1,
                                        diffusivity1_transformed * inv_jacobian^2)
            max_diffusivity2 = Base.max(max_diffusivity2,
                                        diffusivity2_transformed * inv_jacobian^2)
        end
        # Use `Base.max` to prevent silent failures, as `max` from `@fastmath` doesn't propagate
        # `NaN`s properly. See https://github.com/trixi-framework/Trixi.jl/pull/2445#discussion_r2336812323
        max_scaled_diffusivity = Base.max(max_scaled_diffusivity,
                                          max_diffusivity1 + max_diffusivity2)
    end

    return 4 / (nnodes(dg) * max_scaled_diffusivity)
end

@inline function max_scaled_speed_per_element(u,
                                              ::Type{<:Union{StructuredMesh{2},
                                                             UnstructuredMesh2D,
                                                             P4estMesh{2},
                                                             P4estMeshView{2},
                                                             T8codeMesh{2},
                                                             StructuredMeshView{2}}},
                                              constant_speed::True, equations, dg::DG,
                                              contravariant_vectors, inverse_jacobian,
                                              element)
    max_scaled_speed = zero(eltype(u))
    max_lambda1, max_lambda2 = max_abs_speeds(equations)
    for j in eachnode(dg), i in eachnode(dg)
        # Local speeds transformed to the reference element
        Ja11, Ja12 = get_contravariant_vector(1, contravariant_vectors,
                                              i, j, element)
        lambda1_transformed = abs(Ja11 * max_lambda1 + Ja12 * max_lambda2)
        Ja21, Ja22 = get_contravariant_vector(2, contravariant_vectors,
                                              i, j, element)
        lambda2_transformed = abs(Ja21 * max_lambda1 + Ja22 * max_lambda2)

        inv_jacobian = abs(inverse_jacobian[i, j, element])

        max_scaled_speed = Base.max(max_scaled_speed,
                                    inv_jacobian *
                                    (lambda1_transformed + lambda2_transformed))
    end

    return max_scaled_speed
end

function max_dt(u, t,
                mesh::P4estMesh{2}, # Parabolic terms currently only for `TreeMesh` and `P4estMesh`
                constant_diffusivity::True, equations,
                equations_parabolic::AbstractEquationsParabolic,
                dg::DG, cache)
    # Avoid division by zero if the diffusivity vanishes everywhere
    max_scaled_diffusivity = nextfloat(zero(t))

    diffusivity = max_diffusivity(equations_parabolic)

    @unpack contravariant_vectors, inverse_jacobian = cache.elements

    @batch reduction=(max, max_scaled_diffusivity) for element in eachelement(dg, cache)
        max_diffusivity1 = max_diffusivity2 = zero(max_scaled_diffusivity)
        for j in eachnode(dg), i in eachnode(dg)
            # Local diffusivity transformed to the reference element
            Ja11, Ja12 = get_contravariant_vector(1, contravariant_vectors,
                                                  i, j, element)
            # The metric terms are squared due to the second derivatives in the parabolic terms,
            # which lead to application of the chain rule (coordinate transformation) twice.
            # See for instance also the implementation in FLEXI
            # https://github.com/flexi-framework/flexi/blob/e980e8635e6605daf906fc176c9897314a9cd9b1/src/equations/navierstokes/calctimestep.f90#L75-L77
            # or FLUXO:
            # https://github.com/project-fluxo/fluxo/blob/c7e0cc9b7fd4569dcab67bbb6e5a25c0a84859f1/src/equation/navierstokes/calctimestep.f90#L130-L132
            diffusivity1_transformed = diffusivity * (Ja11^2 + Ja12^2)
            Ja21, Ja22 = get_contravariant_vector(2, contravariant_vectors,
                                                  i, j, element)
            diffusivity2_transformed = diffusivity * (Ja21^2 + Ja22^2)

            inv_jacobian = abs(inverse_jacobian[i, j, element])

            max_diffusivity1 = Base.max(max_diffusivity1,
                                        diffusivity1_transformed * inv_jacobian^2)
            max_diffusivity2 = Base.max(max_diffusivity2,
                                        diffusivity2_transformed * inv_jacobian^2)
        end
        # Use `Base.max` to prevent silent failures, as `max` from `@fastmath` doesn't propagate
        # `NaN`s properly. See https://github.com/trixi-framework/Trixi.jl/pull/2445#discussion_r2336812323
        max_scaled_diffusivity = Base.max(max_scaled_diffusivity,
                                          max_diffusivity1 + max_diffusivity2)
    end

    return 4 / (nnodes(dg) * max_scaled_diffusivity)
end

function max_dt(u, t, mesh::P4estMeshParallel{2},
                constant_speed::False, equations, dg::DG, cache)
    # call the method accepting a general `mesh::P4estMesh{2}`
    # TODO: MPI, we should improve this; maybe we should dispatch on `u`
    #       and create some MPI array type, overloading broadcasting and mapreduce etc.
    #       Then, this specific array type should also work well with DiffEq etc.
    dt = invoke(max_dt,
                Tuple{typeof(u), typeof(t), P4estMesh{2},
                      typeof(constant_speed), typeof(equations), typeof(dg),
                      typeof(cache)},
                u, t, mesh, constant_speed, equations, dg, cache)
    # Base.min instead of min needed, see comment in src/auxiliary/math.jl
    dt = MPI.Allreduce!(Ref(dt), Base.min, mpi_comm())[]

    return dt
end

function max_dt(u, t, mesh::P4estMeshParallel{2},
                constant_speed::True, equations, dg::DG, cache)
    # call the method accepting a general `mesh::P4estMesh{2}`
    # TODO: MPI, we should improve this; maybe we should dispatch on `u`
    #       and create some MPI array type, overloading broadcasting and mapreduce etc.
    #       Then, this specific array type should also work well with DiffEq etc.
    dt = invoke(max_dt,
                Tuple{typeof(u), typeof(t), P4estMesh{2},
                      typeof(constant_speed), typeof(equations), typeof(dg),
                      typeof(cache)},
                u, t, mesh, constant_speed, equations, dg, cache)
    # Base.min instead of min needed, see comment in src/auxiliary/math.jl
    dt = MPI.Allreduce!(Ref(dt), Base.min, mpi_comm())[]

    return dt
end

function max_dt(u, t, mesh::T8codeMeshParallel{2},
                constant_speed::False, equations, dg::DG, cache)
    # call the method accepting a general `mesh::T8codeMesh{2}`
    # TODO: MPI, we should improve this; maybe we should dispatch on `u`
    #       and create some MPI array type, overloading broadcasting and mapreduce etc.
    #       Then, this specific array type should also work well with DiffEq etc.
    dt = invoke(max_dt,
                Tuple{typeof(u), typeof(t), T8codeMesh{2},
                      typeof(constant_speed), typeof(equations), typeof(dg),
                      typeof(cache)},
                u, t, mesh, constant_speed, equations, dg, cache)
    # Base.min instead of min needed, see comment in src/auxiliary/math.jl
    dt = MPI.Allreduce!(Ref(dt), Base.min, mpi_comm())[]

    return dt
end

function max_dt(u, t, mesh::T8codeMeshParallel{2},
                constant_speed::True, equations, dg::DG, cache)
    # call the method accepting a general `mesh::T8codeMesh{2}`
    # TODO: MPI, we should improve this; maybe we should dispatch on `u`
    #       and create some MPI array type, overloading broadcasting and mapreduce etc.
    #       Then, this specific array type should also work well with DiffEq etc.
    dt = invoke(max_dt,
                Tuple{typeof(u), typeof(t), T8codeMesh{2},
                      typeof(constant_speed), typeof(equations), typeof(dg),
                      typeof(cache)},
                u, t, mesh, constant_speed, equations, dg, cache)
    # Base.min instead of min needed, see comment in src/auxiliary/math.jl
    dt = MPI.Allreduce!(Ref(dt), Base.min, mpi_comm())[]

    return dt
end
end # @muladd