MeanField/libmeanfield/impl/mapping/compactification/kelvin.cpp

module;

#include <cmath>
#include <mfem.hpp>
#include <stdexcept>

module mean_field;

namespace {
    bool vector_is_finite(const mfem::Vector &vector) {
        for (int i = 0; i < vector.Size(); ++i) {
            if (!std::isfinite(vector(i)))
                return false;
        }
        return true;
    }

    bool matrix_is_finite(const mfem::DenseMatrix &matrix) {
        for (int i = 0; i < matrix.Height(); ++i) {
            for (int j = 0; j < matrix.Width(); ++j) {
                if (!std::isfinite(matrix(i, j)))
                    return false;
            }
        }
        return true;
    }
} // namespace

namespace mean_field::mapping::compactification {
    KelvinCompactification::KelvinCompactification(options::KelvinCompactificationOptions options)
        : m_options(options) {
        if (!std::isfinite(m_options.r_star_ref) || !std::isfinite(m_options.r_inf_ref)) {
            throw std::invalid_argument("Kelvin compactification radii must be finite.");
        }

        if (m_options.r_star_ref <= 0.0 || m_options.r_inf_ref <= m_options.r_star_ref) {
            throw std::invalid_argument("Kelvin compactification requires 0 < r_star_ref < r_inf_ref.");
        }

        if (!std::isfinite(m_options.coordinate_tolerance) || m_options.coordinate_tolerance < 0.0 ||
            m_options.coordinate_tolerance >= 1.0) {
            throw std::invalid_argument(
                "Kelvin compactification coordinate tolerance must be finite "
                "and lie "
                "in [0, 1)."
            );
        }
    }

    MappingStatus KelvinCompactification::ComputeRadialFactors(
        const double compactification_coordinate,
        RadialFactors &factors
    ) const {
        if (!std::isfinite(compactification_coordinate))
            return MappingStatus::non_finite_input;

        // How close we will allow the code to get to compactified infinity
        const double tolerance = m_options.coordinate_tolerance;

        if (compactification_coordinate < -tolerance || compactification_coordinate > 1.0 + tolerance) {
            return MappingStatus::outside_reference_domain;
        }

        double coordinate = compactification_coordinate;
        if (coordinate < 0.0)
            coordinate = 0.0;

        if (coordinate >= 1.0 - tolerance) {
            return MappingStatus::at_compactified_infinity;
        }

        // Here we need to do some transformations from the options defined on the mesh to useful computational coordinates
        // r_inf_ref is the computational / reference radius of the infinity surface (the edge of the entire domain) and r_star_ref is the radius of the spherical
        // stellar model inscribed within. Therefore radial extent is the computational radial distance between the stellar surface and the infinity surface.
        // Note that this is separate from the parameterize compactification coordinate.
        const double radial_extent        = m_options.r_inf_ref - m_options.r_star_ref;

        // This places us at the correct spot in computational space given the current compactification coordinate. Say you have compactification = 0.5,
        // an r_star_ref of 2 and a radial extent of 3, this this will place you at 2 + 0.5 * 3 = 3.5 in computational space, which is half way between the stellar surface and the infinity surface.
        const double computational_radius = m_options.r_star_ref + coordinate * radial_extent;

        if (!std::isfinite(computational_radius) || computational_radius <= 0.0) {
            return MappingStatus::invalid_reference_radius;
        }

        // Invert the exterior coordinate so it runs from 0 at the star to 1 at compactified infinity
        const double one_minus_coordinate = 1.0 - coordinate;
        const double denominator          = computational_radius * one_minus_coordinate;

        if (!std::isfinite(denominator) || denominator <= 0.0) {
            return MappingStatus::non_finite_result;
        }

        // The scale here is the factor which stretches the finite computational domain into the infinite physical domain. Properties we need this to have
        // include that it should go to 1 at the stellar surface and go to infinity at the compactified infinity.
        // Mathematically this is scale = |r|/|x| where r is the physical radius and x is the computational radius.
        // Put another way, scale is the ratio of the target physical radius for the current compactification coordinate
        // to the current mesh radius.
        const double scale            = m_options.r_star_ref / denominator;
        const double scale_derivative = scale * (1.0 / one_minus_coordinate - radial_extent / computational_radius);

        if (!std::isfinite(scale) || !std::isfinite(scale_derivative)) {
            return MappingStatus::non_finite_result;
        }

        factors.coordinate           = coordinate;
        factors.computational_radius = computational_radius;
        factors.scale                = scale;
        factors.scale_derivative     = scale_derivative;

        return MappingStatus::valid;
    }

    MappingStatus KelvinCompactification::Evaluate(
        const ExteriorMapInput &input,
        ExteriorMapResult &result
    ) const {
        const int dimension = input.reference_position.Size();

        if (dimension <= 0 || input.displaced_position.Size() != dimension ||
            input.compactification_coordinate_gradient.Size() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (input.displacement_jacobian.Height() != dimension || input.displacement_jacobian.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (!vector_is_finite(input.reference_position) || !vector_is_finite(input.displaced_position) ||
            !vector_is_finite(input.compactification_coordinate_gradient) ||
            !matrix_is_finite(input.displacement_jacobian)) {
            return MappingStatus::non_finite_input;
        }

        RadialFactors factors;

        // The key radial factors we use are the scale (which stretches the finite computational domain into the infinite physical domain) and the scale derivative (which is used to compute the mapping jacobian).
        const MappingStatus factor_status = ComputeRadialFactors(input.compactification_coordinate, factors);
        if (factor_status != MappingStatus::valid)
            return factor_status;

        result.physical_position.SetSize(dimension);
        result.mapping_jacobian.SetSize(dimension, dimension);

        for (int i = 0; i < dimension; ++i) {
            // Note how the physical position is just the product of the displaced position and the scale factor.
            result.physical_position(i) = factors.scale * input.displaced_position(i);

            // The mapping jacobian comes from trivial application of the product rule
            // recall: r_i = scale * x_i where r is the physical position and x is the displaced position.
            // then we can differentiate wrt. X_j holding nothing fixed. Note the capital X here, this is the mesh coordinate not the displaced position.
            // Lets call this jacobian F
            // F = \frac{\partial r_i}{\partial X_{j}}
            // F then tells us how the physical position changes as we move along mesh coordinates
            // Lets then apply this to the function we have for the kelvin compactification
            // F = scale * \frac{\partial x_i}{\partial X_j} + x_i * \frac{\partial scale}{\partial X_j}
            // Below you can see the displacement jacobian (\frac{\partial x_i}{\partial X_j}) and the scale derivative (\frac{\partial scale}{\partial X_j}) being applied to compute the mapping jacobian.
            for (int j = 0; j < dimension; ++j) {
                const double scale_gradient = factors.scale_derivative * input.compactification_coordinate_gradient(j);
                result.mapping_jacobian(i, j) = factors.scale * input.displacement_jacobian(i, j) + input.displaced_position(i) * scale_gradient;
            }
        }

        if (!vector_is_finite(result.physical_position) || !matrix_is_finite(result.mapping_jacobian)) {
            return MappingStatus::non_finite_result;
        }

        const double mapping_determinant = result.mapping_jacobian.Det();

        if (!std::isfinite(mapping_determinant))
            return MappingStatus::non_finite_result;
        if (mapping_determinant <= 0.0)
            return MappingStatus::non_positive_determinant;

        return MappingStatus::valid;
    }

    MappingStatus KelvinCompactification::EvaluateVariation(
        const ExteriorMapInput &input,
        const ExteriorMapResult &result,
        const ExteriorMapDirection &direction,
        ExteriorMapVariation &variation
    ) const {
        const int dimension = input.reference_position.Size();

        if (dimension <= 0 || input.displaced_position.Size() != dimension ||
            input.compactification_coordinate_gradient.Size() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (input.displacement_jacobian.Height() != dimension || input.displacement_jacobian.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (result.physical_position.Size() != dimension || result.mapping_jacobian.Height() != dimension ||
            result.mapping_jacobian.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (direction.displaced_position_variation.Size() != dimension ||
            direction.displacement_jacobian_variation.Height() != dimension ||
            direction.displacement_jacobian_variation.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (!vector_is_finite(input.reference_position) || !vector_is_finite(input.displaced_position) ||
            !vector_is_finite(input.compactification_coordinate_gradient) ||
            !matrix_is_finite(input.displacement_jacobian)) {
            return MappingStatus::non_finite_input;
        }

        if (!vector_is_finite(result.physical_position) || !matrix_is_finite(result.mapping_jacobian) ||
            !vector_is_finite(direction.displaced_position_variation) ||
            !matrix_is_finite(direction.displacement_jacobian_variation)) {
            return MappingStatus::non_finite_input;
        }

        RadialFactors factors;
        const MappingStatus factor_status = ComputeRadialFactors(input.compactification_coordinate, factors);
        if (factor_status != MappingStatus::valid)
            return factor_status;

        variation.physical_position_variation.SetSize(dimension);
        variation.mapping_jacobian_variation.SetSize(dimension, dimension);

        for (int i = 0; i < dimension; ++i) {
            variation.physical_position_variation(i) = factors.scale * direction.displaced_position_variation(i);

            for (int j = 0; j < dimension; ++j) {
                const double scale_gradient = factors.scale_derivative * input.compactification_coordinate_gradient(j);
                variation.mapping_jacobian_variation(i, j) =
                    factors.scale * direction.displacement_jacobian_variation(i, j) +
                    direction.displaced_position_variation(i) * scale_gradient;
            }
        }

        if (!vector_is_finite(variation.physical_position_variation) ||
            !matrix_is_finite(variation.mapping_jacobian_variation)) {
            return MappingStatus::non_finite_result;
        }

        return MappingStatus::valid;
    }

    std::string_view KelvinCompactification::GetName() const noexcept {
        return "KelvinCompactification";
    }

    double KelvinCompactification::GetReferenceStellarRadius() const noexcept {
        return m_options.r_star_ref;
    }

    double KelvinCompactification::GetReferenceInfinityRadius() const noexcept {
        return m_options.r_inf_ref;
    }

    double KelvinCompactification::GetCoordinateTolerance() const noexcept {
        return m_options.coordinate_tolerance;
    }
} // namespace mean_field::mapping::compactification