stepsize_dg3d.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{3},
                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 = max_lambda3 = zero(max_scaled_speed)
        for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, k, element)
            lambda1, lambda2, lambda3 = max_abs_speeds(u_node, equations)
            max_lambda1 = Base.max(max_lambda1, lambda1)
            max_lambda2 = Base.max(max_lambda2, lambda2)
            max_lambda3 = Base.max(max_lambda3, lambda3)
        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 + max_lambda3))
    end

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

function max_dt(u, t, mesh::TreeMesh{3},
                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))

    @batch reduction=(max, max_scaled_diffusivity) for element in eachelement(dg, cache)
        max_diffusivity_ = zero(max_scaled_diffusivity)
        for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, k, 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_diffusivity = max(max_scaled_diffusivity,
                                     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_diffusivity)
end

function max_dt(u, t, mesh::TreeMesh{3},
                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_lambda3 = 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 + max_lambda3))
    end

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

function max_dt(u, t,
                mesh::Union{StructuredMesh{3}, P4estMesh{3}, T8codeMesh{3}},
                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{3},
                                                             P4estMesh{3},
                                                             T8codeMesh{3}}},
                                              constant_speed::False, equations, dg,
                                              contravariant_vectors, inverse_jacobian,
                                              element)
    max_lambda1 = max_lambda2 = max_lambda3 = zero(eltype(u))
    for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
        u_node = get_node_vars(u, equations, dg, i, j, k, element)
        lambda1, lambda2, lambda3 = max_abs_speeds(u_node, equations)

        Ja11, Ja12, Ja13 = get_contravariant_vector(1, contravariant_vectors,
                                                    i, j, k, element)
        lambda1_transformed = abs(Ja11 * lambda1 + Ja12 * lambda2 + Ja13 * lambda3)
        Ja21, Ja22, Ja23 = get_contravariant_vector(2, contravariant_vectors,
                                                    i, j, k, element)
        lambda2_transformed = abs(Ja21 * lambda1 + Ja22 * lambda2 + Ja23 * lambda3)
        Ja31, Ja32, Ja33 = get_contravariant_vector(3, contravariant_vectors,
                                                    i, j, k, element)
        lambda3_transformed = abs(Ja31 * lambda1 + Ja32 * lambda2 + Ja33 * lambda3)

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

        max_lambda1 = max(max_lambda1, inv_jacobian * lambda1_transformed)
        max_lambda2 = max(max_lambda2, inv_jacobian * lambda2_transformed)
        max_lambda3 = max(max_lambda3, inv_jacobian * lambda3_transformed)
    end
    return max_lambda1 + max_lambda2 + max_lambda3
end

function max_dt(u, t,
                mesh::P4estMesh{3}, # 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 = max_diffusivity3 = zero(max_scaled_diffusivity)
        for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
            u_node = get_node_vars(u, equations, dg, i, j, k, element)
            diffusivity = max_diffusivity(u_node, equations_parabolic)

            # Local diffusivity transformed to the reference element
            Ja11, Ja12, Ja13 = get_contravariant_vector(1, contravariant_vectors,
                                                        i, j, k, 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 + Ja13^2)
            Ja21, Ja22, Ja23 = get_contravariant_vector(2, contravariant_vectors,
                                                        i, j, k, element)
            diffusivity2_transformed = diffusivity * (Ja21^2 + Ja22^2 + Ja23^2)
            Ja31, Ja32, Ja33 = get_contravariant_vector(3, contravariant_vectors,
                                                        i, j, k, element)
            diffusivity3_transformed = diffusivity * (Ja31^2 + Ja32^2 + Ja33^2)

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

            max_diffusivity1 = Base.max(max_diffusivity1,
                                        diffusivity1_transformed * inv_jacobian^2)
            max_diffusivity2 = Base.max(max_diffusivity2,
                                        diffusivity2_transformed * inv_jacobian^2)
            max_diffusivity3 = Base.max(max_diffusivity3,
                                        diffusivity3_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 +
                                          max_diffusivity3)
    end

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

@inline function max_scaled_speed_per_element(u,
                                              ::Type{<:Union{StructuredMesh{3},
                                                             P4estMesh{3},
                                                             T8codeMesh{3}}},
                                              constant_speed::True, equations, dg::DG,
                                              contravariant_vectors, inverse_jacobian,
                                              element)
    max_scaled_speed = zero(eltype(u))
    max_lambda1, max_lambda2, max_lambda3 = max_abs_speeds(equations)
    for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
        Ja11, Ja12, Ja13 = get_contravariant_vector(1, contravariant_vectors,
                                                    i, j, k, element)
        lambda1_transformed = abs(Ja11 * max_lambda1 + Ja12 * max_lambda2 +
                                  Ja13 * max_lambda3)
        Ja21, Ja22, Ja23 = get_contravariant_vector(2, contravariant_vectors,
                                                    i, j, k, element)
        lambda2_transformed = abs(Ja21 * max_lambda1 + Ja22 * max_lambda2 +
                                  Ja23 * max_lambda3)
        Ja31, Ja32, Ja33 = get_contravariant_vector(3, contravariant_vectors,
                                                    i, j, k, element)
        lambda3_transformed = abs(Ja31 * max_lambda1 + Ja32 * max_lambda2 +
                                  Ja33 * max_lambda3)

        inv_jacobian = abs(inverse_jacobian[i, j, k, 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 *
                                    (lambda1_transformed + lambda2_transformed +
                                     lambda3_transformed))
    end

    return max_scaled_speed
end

function max_dt(u, t,
                mesh::P4estMesh{3}, # 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 = max_diffusivity3 = zero(max_scaled_diffusivity)
        for k in eachnode(dg), j in eachnode(dg), i in eachnode(dg)
            # Local diffusivity transformed to the reference element
            Ja11, Ja12, Ja13 = get_contravariant_vector(1, contravariant_vectors,
                                                        i, j, k, 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 + Ja13^2)
            Ja21, Ja22, Ja23 = get_contravariant_vector(2, contravariant_vectors,
                                                        i, j, k, element)
            diffusivity2_transformed = diffusivity * (Ja21^2 + Ja22^2 + Ja23^2)
            Ja31, Ja32, Ja33 = get_contravariant_vector(3, contravariant_vectors,
                                                        i, j, k, element)
            diffusivity3_transformed = diffusivity * (Ja31^2 + Ja32^2 + Ja33^2)

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

            max_diffusivity1 = Base.max(max_diffusivity1,
                                        diffusivity1_transformed * inv_jacobian^2)
            max_diffusivity2 = Base.max(max_diffusivity2,
                                        diffusivity2_transformed * inv_jacobian^2)
            max_diffusivity3 = Base.max(max_diffusivity3,
                                        diffusivity3_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 +
                                          max_diffusivity3)
    end

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

function max_dt(u, t, mesh::P4estMeshParallel{3},
                constant_speed::False, equations, dg::DG, cache)
    # call the method accepting a general `mesh::P4estMesh{3}`
    # 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{3},
                      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{3},
                constant_speed::True, equations, dg::DG, cache)
    # call the method accepting a general `mesh::P4estMesh{3}`
    # 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{3},
                      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{3},
                constant_speed::False, equations, dg::DG, cache)
    # call the method accepting a general `mesh::T8codeMesh{3}`
    # 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{3},
                      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{3},
                constant_speed::True, equations, dg::DG, cache)
    # call the method accepting a general `mesh::T8codeMesh{3}`
    # 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{3},
                      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