MeanField/libmeanfield/impl/operators/contexts/gravity_field_context.cpp

module;
#include <cmath>
#include <memory>
#include <mfem.hpp>

module mean_field;
import :operators.context.gravity_field;

namespace {
    void validate_displacement(
        const mean_field::fem::FEM &f,
        const mfem::Vector &displacement_true
    ) {
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldGeometryContext requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            displacement_true.Size() == f.displacementFes->GetTrueVSize(),
            "GravityFieldGeometryContext received a displacement vector with "
            "the "
            "wrong size."
        );

        for (int i = 0; i < displacement_true.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(displacement_true(i)), "GravityFieldGeometryContext received a non-finite "
                                                     "displacement "
                                                     "value."
            );
        }
    }

    void validate_linearization_state(
        const mean_field::fem::FEM &f,
        const mean_field::operators::context::gravity_field::GravityFieldStateView &state
    ) {
        MFEM_VERIFY(
            f.densityFes != nullptr, "GravityFieldLinearizationContext "
                                     "requires the density finite-element "
                                     "space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "GravityFieldLinearizationContext requires the gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "GravityFieldLinearizationContext requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldLinearizationContext requires "
                                          "the displacement finite-element space."
        );

        MFEM_VERIFY(
            state.density.Size() == f.densityFes->GetTrueVSize(),
            "GravityFieldLinearizationContext received a density vector with "
            "the "
            "wrong size."
        );
        MFEM_VERIFY(
            state.displacement.Size() == f.displacementFes->GetTrueVSize(),
            "GravityFieldLinearizationContext received a displacement vector "
            "with "
            "the wrong size."
        );
        MFEM_VERIFY(
            state.gravity_gradient.Size() == f.gravityFluxFes->GetTrueVSize(),
            "GravityFieldLinearizationContext received a gravity-gradient "
            "vector "
            "with the wrong size."
        );
        MFEM_VERIFY(
            state.gravity_potential.Size() == f.gravityPotentialFes->GetTrueVSize(),
            "GravityFieldLinearizationContext received a gravity-potential "
            "vector "
            "with the wrong size."
        );

        for (int i = 0; i < state.density.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(state.density(i)), "GravityFieldLinearizationContext received a non-finite "
                                                 "density "
                                                 "value."
            );
        }

        for (int i = 0; i < state.displacement.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(state.displacement(i)), "GravityFieldLinearizationContext received a non-finite "
                                                      "displacement "
                                                      "value."
            );
        }

        for (int i = 0; i < state.gravity_gradient.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(state.gravity_gradient(i)), "GravityFieldLinearizationContext received a non-finite "
                                                          "gravity-gradient value."
            );
        }

        for (int i = 0; i < state.gravity_potential.Size(); ++i) {
            MFEM_VERIFY(
                std::isfinite(state.gravity_potential(i)), "GravityFieldLinearizationContext received a non-finite "
                                                           "gravity-potential value."
            );
        }
    }
} // namespace

namespace mean_field::operators::context::gravity_field {
    GravityFieldGeometryContext::GravityFieldGeometryContext(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domain_mapper
    )
        : m_fem(f),
          m_domain_mapper(domain_mapper) {
        MFEM_VERIFY(f.mesh != nullptr, "GravityFieldGeometryContext requires a mesh.");
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "GravityFieldGeometryContext requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.densityFes != nullptr, "GravityFieldGeometryContext requires the density finite-element "
                                     "space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "GravityFieldGeometryContext requires the gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldGeometryContext requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationFes != nullptr, "GravityFieldGeometryContext requires the compactification "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr, "GravityFieldGeometryContext requires the compactification "
                                                     "coordinate."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "GravityFieldGeometryContext requires the quadrature-rule factory."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "The stateless domain-mapper dimension does not match the mesh "
            "dimension."
        );
    }

    GravityFieldGeometryPreparation GravityFieldGeometryContext::Prepare(
        const mfem::Vector &displacement_true,
        const DiscretizationRevision discretization_revision,
        const DisplacementRevision displacement_revision
    ) {
        validate_displacement(m_fem, displacement_true);

        if (m_is_prepared) {
            MFEM_VERIFY(
                discretization_revision >= m_discretization_revision,
                "GravityFieldGeometryContext received an older discretization "
                "revision."
            );
            MFEM_VERIFY(
                displacement_revision >= m_displacement_revision,
                "GravityFieldGeometryContext received an older displacement "
                "revision."
            );
        }

        const bool discretization_changed = !m_is_prepared || discretization_revision != m_discretization_revision;
        const bool displacement_changed   = !m_is_prepared || displacement_revision != m_displacement_revision;

        GravityFieldGeometryPreparation preparation;

        if (!discretization_changed && !displacement_changed) {
            return preparation;
        }

        if (discretization_changed) {
            auto mass_operator   = std::make_unique<PreparedMappedHDivMassOperator>(m_fem, m_domain_mapper);
            auto source_operator = std::make_unique<PreparedMappedGravitySourceOperator>(m_fem, m_domain_mapper);

            mass_operator->Prepare(displacement_true);
            source_operator->Prepare(displacement_true);

            m_mass_operator                     = std::move(mass_operator);
            m_source_operator                   = std::move(source_operator);

            preparation.reconstructed_operators = true;
            preparation.rebuilt_mass_operator   = true;
            preparation.rebuilt_source_operator = true;
        } else {
            MFEM_VERIFY(
                m_mass_operator != nullptr, "GravityFieldGeometryContext has "
                                            "no prepared H(div) mass operator."
            );
            MFEM_VERIFY(
                m_source_operator != nullptr, "GravityFieldGeometryContext has no prepared gravity source "
                                              "operator."
            );

            m_mass_operator->Prepare(displacement_true);
            m_source_operator->Prepare(displacement_true);

            preparation.rebuilt_mass_operator   = true;
            preparation.rebuilt_source_operator = true;
        }

        m_displacement_true                   = displacement_true;
        m_discretization_revision             = discretization_revision;
        m_displacement_revision               = displacement_revision;
        m_is_prepared                         = true;

        preparation.refreshed_variation_state = true;

        return preparation;
    }

    const PreparedMappedHDivMassOperator &GravityFieldGeometryContext::GetMassOperator() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldGeometryContext must be prepared before "
                           "accessing its mass operator."
        );
        MFEM_VERIFY(m_mass_operator != nullptr, "GravityFieldGeometryContext has no prepared H(div) mass operator.");
        return *m_mass_operator;
    }

    const PreparedMappedGravitySourceOperator &GravityFieldGeometryContext::GetSourceOperator() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldGeometryContext must be prepared before "
                           "accessing its source operator."
        );
        MFEM_VERIFY(
            m_source_operator != nullptr, "GravityFieldGeometryContext has no "
                                          "prepared gravity source operator."
        );
        return *m_source_operator;
    }

    const mfem::Vector &GravityFieldGeometryContext::GetDisplacement() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldGeometryContext must be prepared before "
                           "accessing its displacement."
        );
        return m_displacement_true;
    }

    DiscretizationRevision GravityFieldGeometryContext::GetDiscretizationRevision() const noexcept {
        return m_discretization_revision;
    }

    DisplacementRevision GravityFieldGeometryContext::GetDisplacementRevision() const noexcept {
        return m_displacement_revision;
    }

    bool GravityFieldGeometryContext::IsPrepared() const noexcept {
        return m_is_prepared;
    }

    GravityFieldLinearizationContext::GravityFieldLinearizationContext(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domain_mapper
    )
        : m_fem(f),
          m_geometry_context(
              f,
              domain_mapper
          ) {
        MFEM_VERIFY(
            f.densityFes != nullptr, "GravityFieldLinearizationContext "
                                     "requires the density finite-element "
                                     "space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "GravityFieldLinearizationContext requires the gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "GravityFieldLinearizationContext requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldLinearizationContext requires "
                                          "the displacement finite-element space."
        );
    }

    GravityFieldPreparationReport GravityFieldLinearizationContext::Prepare(
        const GravityFieldStateView &state,
        const GravityFieldRevisions &revisions
    ) {
        validate_linearization_state(m_fem, state);

        if (m_is_prepared) {
            MFEM_VERIFY(
                revisions.discretization >= m_revisions.discretization,
                "GravityFieldLinearizationContext received an older "
                "discretization "
                "revision."
            );
            MFEM_VERIFY(
                revisions.displacement >= m_revisions.displacement,
                "GravityFieldLinearizationContext received an older "
                "displacement "
                "revision."
            );
            MFEM_VERIFY(
                revisions.density >= m_revisions.density, "GravityFieldLinearizationContext received an older density "
                                                          "revision."
            );
            MFEM_VERIFY(
                revisions.gravity_gradient >= m_revisions.gravity_gradient,
                "GravityFieldLinearizationContext received an older "
                "gravity-gradient "
                "revision."
            );
            MFEM_VERIFY(
                revisions.gravity_potential >= m_revisions.gravity_potential,
                "GravityFieldLinearizationContext received an older "
                "gravity-potential revision."
            );
        }

        const bool discretization_changed = !m_is_prepared || revisions.discretization != m_revisions.discretization;
        const bool density_changed =
            !m_is_prepared || discretization_changed || revisions.density != m_revisions.density;
        const bool gravity_gradient_changed =
            !m_is_prepared || discretization_changed || revisions.gravity_gradient != m_revisions.gravity_gradient;

        GravityFieldPreparationReport report;

        report.geometry =
            m_geometry_context.Prepare(state.displacement, revisions.discretization, revisions.displacement);

        if (density_changed) {
            m_density_true         = state.density;
            report.updated_density = true;
        }

        if (gravity_gradient_changed) {
            m_gravity_gradient_true         = state.gravity_gradient;
            report.updated_gravity_gradient = true;
        }

        m_revisions   = revisions;
        m_is_prepared = true;

        return report;
    }

    const GravityFieldGeometryContext &GravityFieldLinearizationContext::GetGeometryContext() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its geometry context."
        );
        return m_geometry_context;
    }

    const mfem::Vector &GravityFieldLinearizationContext::GetDensity() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its density."
        );
        return m_density_true;
    }

    const mfem::Vector &GravityFieldLinearizationContext::GetGravityGradient() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its gravity gradient."
        );
        return m_gravity_gradient_true;
    }

    const GravityFieldRevisions &GravityFieldLinearizationContext::GetRevisions() const {
        MFEM_VERIFY(
            m_is_prepared, "GravityFieldLinearizationContext must be prepared "
                           "before accessing its revisions."
        );
        return m_revisions;
    }

    bool GravityFieldLinearizationContext::IsPrepared() const noexcept {
        return m_is_prepared;
    }
} // namespace mean_field::operators::context::gravity_field