combined_mapping.hpp

// SPDX-License-Identifier: MIT
#pragma once
 
#include "sll/math_tools.hpp"
 
#include "inv_jacobian_o_point.hpp"
#include "inverse_jacobian_matrix.hpp"
#include "mapping_tools.hpp"
 
template <class Mapping1, class Mapping2>
class CombinedMapping
{
    static_assert(is_mapping_v<Mapping1>);
    static_assert(is_mapping_v<Mapping2>);
    static_assert(std::is_same_v<typename Mapping2::CoordResult, typename Mapping1::CoordArg>);
    static_assert(
            is_analytical_mapping_v<Mapping2>,
            "The second mapping must be analytical to evaluate the jacobian");
 
public:
    using CoordArg = typename Mapping2::CoordArg;
    using CoordResult = typename Mapping1::CoordResult;
    using CoordJacobian = typename Mapping2::CoordResult;
 
private:
    using InverseMapping2 = inverse_mapping_t<Mapping2>;
 
    static_assert(has_2d_jacobian_v<Mapping1, CoordJacobian>);
    static_assert(has_2d_jacobian_v<InverseMapping2, CoordJacobian>);
 
private:
    Mapping1 m_mapping_1;
    Mapping2 m_mapping_2;
    InverseMapping2 m_inv_mapping_2;
    // The Jacobian defined on CoordJacobian is the inverse of the inverse mapping
    InverseJacobianMatrix<InverseMapping2, CoordJacobian> m_jacobian_mapping_2;
    double m_epsilon;
 
public:
    template <
            class Map1,
            std::enable_if_t<
                    (has_singular_o_point_inv_jacobian_v<Map1>)
                            || (has_singular_o_point_inv_jacobian_v<InverseMapping2>),
                    bool> = true>
    CombinedMapping(Map1 mapping_1, Mapping2 mapping_2, double epsilon)
        : m_mapping_1(mapping_1)
        , m_mapping_2(mapping_2)
        , m_inv_mapping_2(mapping_2.get_inverse_mapping())
        , m_jacobian_mapping_2(mapping_2.get_inverse_mapping())
        , m_epsilon(epsilon)
    {
        static_assert(std::is_same_v<Mapping1, Map1>);
    }
 
    template <
            class Map1,
            std::enable_if_t<
                    !((has_singular_o_point_inv_jacobian_v<Map1>)
                      || (has_singular_o_point_inv_jacobian_v<InverseMapping2>)),
                    bool> = true>
    CombinedMapping(Map1 mapping_1, Mapping2 mapping_2)
        : m_mapping_1(mapping_1)
        , m_mapping_2(mapping_2)
        , m_inv_mapping_2(mapping_2.get_inverse_mapping())
        , m_jacobian_mapping_2(mapping_2.get_inverse_mapping())
        , m_epsilon(0.0)
    {
        static_assert(std::is_same_v<Mapping1, Map1>);
    }
 
    CoordResult operator()(CoordArg coord)
    {
        return m_mapping_1(m_mapping_2(coord));
    }
 
    KOKKOS_INLINE_FUNCTION void jacobian_matrix(CoordJacobian const& coord, Matrix_2x2& matrix)
            const
    {
        Matrix_2x2 jacobian_mapping_1;
        m_mapping_1.jacobian_matrix(coord, jacobian_mapping_1);
 
        Matrix_2x2 m_jacobian_mapping_2 = m_jacobian_mapping_2(coord);
 
        matrix = mat_mul(jacobian_mapping_1, m_jacobian_mapping_2);
    }
 
    KOKKOS_INLINE_FUNCTION double jacobian_11(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        jacobian_matrix(coord_rtheta, J);
        return J[0][0];
    }
 
    KOKKOS_INLINE_FUNCTION double jacobian_12(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        jacobian_matrix(coord_rtheta, J);
        return J[0][1];
    }
 
    KOKKOS_INLINE_FUNCTION double jacobian_21(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        jacobian_matrix(coord_rtheta, J);
        return J[1][0];
    }
 
    KOKKOS_INLINE_FUNCTION double jacobian_22(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        jacobian_matrix(coord_rtheta, J);
        return J[1][1];
    }
 
    KOKKOS_INLINE_FUNCTION double jacobian(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        jacobian_matrix(coord_rtheta, J);
        return determinant(J);
    }
 
    KOKKOS_FUNCTION void inv_jacobian_matrix(CoordJacobian const& coord, Matrix_2x2& matrix) const
    {
        if constexpr (
                (has_singular_o_point_inv_jacobian_v<Mapping1>)
                || (has_singular_o_point_inv_jacobian_v<InverseMapping2>)) {
            using R = ddc::type_seq_element_t<0, ddc::to_type_seq_t<CoordJacobian>>;
            using Theta = ddc::type_seq_element_t<1, ddc::to_type_seq_t<CoordJacobian>>;
            double r = ddc::get<R>(coord);
            if (r < m_epsilon) {
                InvJacobianOPoint<CombinedMapping<Mapping1, Mapping2>, CoordJacobian> o_point_val(
                        *this);
                CoordJacobian coord_eps(m_epsilon, ddc::get<Theta>(coord));
                Matrix_2x2 J_0 = o_point_val();
                Matrix_2x2 J_eps;
                non_singular_inverse_jacobian_matrix(coord_eps, J_eps);
                matrix[0][0] = (1 - r / m_epsilon) * J_0[0][0] + r / m_epsilon * J_eps[0][0];
                matrix[0][1] = (1 - r / m_epsilon) * J_0[0][1] + r / m_epsilon * J_eps[0][1];
                matrix[1][0] = (1 - r / m_epsilon) * J_0[1][0] + r / m_epsilon * J_eps[1][0];
                matrix[1][1] = (1 - r / m_epsilon) * J_0[1][1] + r / m_epsilon * J_eps[1][1];
            } else {
                non_singular_inverse_jacobian_matrix(coord, matrix);
            }
        } else {
            non_singular_inverse_jacobian_matrix(coord, matrix);
        }
    }
 
    KOKKOS_INLINE_FUNCTION double inv_jacobian_11(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        inv_jacobian_matrix(coord_rtheta, J);
        return J[0][0];
    }
 
    KOKKOS_INLINE_FUNCTION double inv_jacobian_12(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        inv_jacobian_matrix(coord_rtheta, J);
        return J[0][1];
    }
 
    KOKKOS_INLINE_FUNCTION double inv_jacobian_21(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        inv_jacobian_matrix(coord_rtheta, J);
        return J[1][0];
    }
 
    KOKKOS_INLINE_FUNCTION double inv_jacobian_22(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        inv_jacobian_matrix(coord_rtheta, J);
        return J[1][1];
    }
 
    KOKKOS_INLINE_FUNCTION double inv_jacobian(CoordJacobian const& coord_rtheta) const
    {
        Matrix_2x2 J;
        inv_jacobian_matrix(coord_rtheta, J);
        return determinant(J);
    }
 
    template <class Mapping>
    KOKKOS_INLINE_FUNCTION Mapping const& get() const
    {
        static_assert(std::is_same_v<Mapping, Mapping1> || std::is_same_v<Mapping, Mapping2>);
        if constexpr (std::is_same_v<Mapping, Mapping1>) {
            return m_mapping_1;
        } else {
            return m_mapping_2;
        }
    }
 
private:
    KOKKOS_INLINE_FUNCTION void non_singular_inverse_jacobian_matrix(
            CoordJacobian const& coord,
            Matrix_2x2& matrix) const
    {
        InverseJacobianMatrix<Mapping1, CoordJacobian> inv_jacobian_matrix_1(m_mapping_1);
        Matrix_2x2 inv_jacobian_mapping_1 = inv_jacobian_matrix_1(coord);
 
        Matrix_2x2 inv_jacobian_mapping_2;
        m_inv_mapping_2.jacobian_matrix(coord, inv_jacobian_mapping_2);
 
        matrix = mat_mul(inv_jacobian_mapping_2, inv_jacobian_mapping_1);
    }
};
 
 
namespace mapping_detail {
template <class Mapping1, class Mapping2, class ExecSpace>
struct MappingAccessibility<ExecSpace, CombinedMapping<Mapping1, Mapping2>>
{
    static constexpr bool value = MappingAccessibility<ExecSpace, Mapping1>::value
                                  && MappingAccessibility<ExecSpace, Mapping2>::value;
};
} // namespace mapping_detail