MeanField/libmeanfield/impl/mapping/coefficients.cpp

module;
#include <mfem.hpp>

module mean_field;
import :mapping.types;

namespace mean_field::mapping {
    ///////////////////////////////
    /// MappedScalarCoefficient ///
    //////////////////////////////
    MappedScalarCoefficient::MappedScalarCoefficient(
        const DomainMapper &map,
        mfem::Coefficient &coeff,
        const COORDINATE_SPACE coord_space
    ) : m_map(map),
        m_coeff(coeff),
        m_coord_space(coord_space) {};

    double MappedScalarCoefficient::Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) {
        T.SetIntPoint(&ip);
        double f_val = 0.0;

        switch (m_coord_space) {
            case COORDINATE_SPACE::PHYSICAL: {
                f_val = eval_at_point(m_coeff, T, ip);
                const double detJ = m_map.ComputeDetJ(T, ip);
                return f_val * fabs(detJ);
            }
            case COORDINATE_SPACE::REFERENCE: {
                f_val = m_coeff.Eval(T, ip);
                return f_val;
            }
        }
    }

    double MappedScalarCoefficient::eval_at_point(mfem::Coefficient &c, mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) {
        return c.Eval(T, ip);
    }

    //////////////////////////////////
    /// MappedDiffusionCoefficient ///
    //////////////////////////////////

    MappedDiffusionCoefficient::MappedDiffusionCoefficient(
        const DomainMapper &map,
        mfem::Coefficient &sigma,
        const int dim
    ) : mfem::MatrixCoefficient(dim),
        m_map(map),
        m_scalar(&sigma),
        m_tensor(nullptr) {
    };

    MappedDiffusionCoefficient::MappedDiffusionCoefficient(
        const DomainMapper &map,
        mfem::MatrixCoefficient &sigma
    ) : mfem::MatrixCoefficient(sigma.GetHeight()),
        m_map(map),
        m_scalar(nullptr),
        m_tensor(&sigma) {
    };

    void MappedDiffusionCoefficient::Eval(mfem::DenseMatrix &K, mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) {
        const int dim = height;
        T.SetIntPoint(&ip);

        mfem::DenseMatrix J(dim, dim), JInv(dim, dim);
        m_map.ComputeJacobian(T, J);
        const double detJ = J.Det();
        mfem::CalcInverse(J, JInv);

        if (m_scalar) {
            const double sig_val = m_scalar->Eval(T, ip);
            mfem::MultABt(JInv, JInv, K);
            K *= sig_val * fabs(detJ);
        } else {
            mfem::DenseMatrix sig_mat(dim, dim);
            m_tensor->Eval(sig_mat, T, ip);

            mfem::DenseMatrix temp(dim, dim);
            Mult(JInv, sig_mat, temp);

            MultABt(temp, JInv, K);
            K *= fabs(detJ);
        }
    }

    ///////////////////////////////
    /// MappedVectorCoefficient ///
    ///////////////////////////////
    MappedVectorCoefficient::MappedVectorCoefficient(
        const DomainMapper &map,
        mfem::VectorCoefficient &coeff
    ) : mfem::VectorCoefficient(coeff.GetVDim()),
        m_map(map),
        m_coeff(coeff) {
    };

    void MappedVectorCoefficient::Eval(mfem::Vector &V, mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) {
        const int dim = vdim;
        T.SetIntPoint(&ip);

        mfem::DenseMatrix JInv(dim, dim);
        m_map.ComputeInverseJacobian(T, JInv);
        double detJ = m_map.ComputeDetJ(T, ip);

        mfem::Vector C_phys(dim);
        m_coeff.Eval(C_phys, T, ip);

        V.SetSize(dim);
        JInv.MultTranspose(C_phys, V);
        V *= fabs(detJ);
    }

    ///////////////////////////////////////////
    /// PhysicalPositionFunctionCoefficient ///
    ///////////////////////////////////////////
    PhysicalPositionFunctionCoefficient::PhysicalPositionFunctionCoefficient(
        const DomainMapper &map,
        Func f // std::function<double(const mfem::Vector&)>
    ) : m_f(std::move(f)),
        m_map(map) {};

    double PhysicalPositionFunctionCoefficient::Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) {
        T.SetIntPoint(&ip);
        mfem::Vector x;
        m_map.GetPhysicalPoint(T, ip, x);
        return m_f(x);
    }

    MappedHDivMassCoefficient::MappedHDivMassCoefficient(const DomainMapper& map, const int dim)
    : mfem::MatrixCoefficient(dim),
      m_map(map) {}

    void MappedHDivMassCoefficient::Eval(
        mfem::DenseMatrix& matrix,
        mfem::ElementTransformation& transformation,
        const mfem::IntegrationPoint& integration_point
    ) {
        transformation.SetIntPoint(&integration_point);

        mfem::DenseMatrix map_jacobian(height, height);
        m_map.ComputeJacobian(transformation, map_jacobian);

        const double map_determinant = map_jacobian.Det();

        MFEM_VERIFY(map_determinant > 0.0, "Domain mapping has a non-positive Jacobian determinant.");

        mfem::MultAtB(map_jacobian, map_jacobian, matrix);
        matrix *= 1.0 / std::abs(map_determinant);
    }
}