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

module;

#include <cmath>

#include <mfem.hpp>

module mean_field;

import :operators.context.hydrostatic_equilibrium;

namespace {
    void validate_finite_vector(
        const mfem::Vector &vector,
        const char *message
    ) {
        for (int i = 0; i < vector.Size(); ++i) {
            MFEM_VERIFY(std::isfinite(vector(i)), message);
        }
    }

    void validate_state(
        const mean_field::fem::FEM &f,
        const mean_field::operators::context::hydrostatic::
            HydrostaticEquilibriumStateView &state
    ) {
        MFEM_VERIFY(
            f.enthalpyFes != nullptr,
            "HydrostaticEquilibriumContext requires the "
            "enthalpy finite-element space."
        );

        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr,
            "HydrostaticEquilibriumContext requires the "
            "gravity-potential finite-element space."
        );

        MFEM_VERIFY(
            f.displacementFes != nullptr,
            "HydrostaticEquilibriumContext requires the "
            "displacement finite-element space."
        );

        MFEM_VERIFY(
            state.enthalpy.Size() == f.enthalpyFes->GetTrueVSize(),
            "HydrostaticEquilibriumContext received an "
            "enthalpy vector with the wrong size."
        );

        MFEM_VERIFY(
            state.gravityPotential.Size() ==
                f.gravityPotentialFes->GetTrueVSize(),
            "HydrostaticEquilibriumContext received a "
            "gravity-potential vector with the wrong size."
        );

        MFEM_VERIFY(
            state.displacement.Size() == f.displacementFes->GetTrueVSize(),
            "HydrostaticEquilibriumContext received a "
            "displacement vector with the wrong size."
        );

        validate_finite_vector(
            state.enthalpy, "HydrostaticEquilibriumContext received a "
                            "non-finite enthalpy value."
        );

        validate_finite_vector(
            state.gravityPotential, "HydrostaticEquilibriumContext received a "
                                    "non-finite gravity-potential value."
        );

        validate_finite_vector(
            state.displacement, "HydrostaticEquilibriumContext received a "
                                "non-finite displacement value."
        );

        MFEM_VERIFY(
            std::isfinite(state.bernoulliConstant),
            "HydrostaticEquilibriumContext received a "
            "non-finite Bernoulli constant."
        );
    }

    template <typename Stamp>
    void validate_dependency_transition(
        const Stamp &prepared,
        const Stamp &requested,
        const char *message
    ) {
        MFEM_VERIFY(requested.CanFollow(prepared), message);
    }
} // namespace

namespace mean_field::operators::context::hydrostatic {
    HydrostaticEquilibriumContext::HydrostaticEquilibriumContext(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper
    )
        : m_f(f),
          m_domainMapper(domainMapper) {
        MFEM_VERIFY(
            m_f.mesh != nullptr,
            "HydrostaticEquilibriumContext requires a mesh."
        );

        MFEM_VERIFY(
            m_f.enthalpyFes != nullptr,
            "HydrostaticEquilibriumContext requires the "
            "enthalpy finite-element space."
        );

        MFEM_VERIFY(
            m_f.gravityPotentialFes != nullptr,
            "HydrostaticEquilibriumContext requires the "
            "gravity-potential finite-element space."
        );

        MFEM_VERIFY(
            m_f.displacementFes != nullptr,
            "HydrostaticEquilibriumContext requires the "
            "displacement finite-element space."
        );

        MFEM_VERIFY(
            m_domainMapper.GetDimension() == m_f.mesh->Dimension(),
            "The hydrostatic context's stateless "
            "domain-mapper dimension does not match the mesh "
            "dimension."
        );
    }

    HydrostaticPreparationReport HydrostaticEquilibriumContext::Prepare(
        const HydrostaticEquilibriumStateView &state,
        const HydrostaticEquilibriumDependencies &dependencies
    ) {
        validate_state(m_f, state);

        if (m_isPrepared) {
            validate_dependency_transition(
                m_dependencies.discretization, dependencies.discretization,
                "HydrostaticEquilibriumContext received an older "
                "discretization revision for the same identity."
            );

            validate_dependency_transition(
                m_dependencies.enthalpy, dependencies.enthalpy,
                "HydrostaticEquilibriumContext received an older "
                "enthalpy revision for the same identity."
            );

            validate_dependency_transition(
                m_dependencies.gravityPotential, dependencies.gravityPotential,
                "HydrostaticEquilibriumContext received an older "
                "gravity-potential revision for the same identity."
            );

            validate_dependency_transition(
                m_dependencies.displacement, dependencies.displacement,
                "HydrostaticEquilibriumContext received an older "
                "displacement revision for the same identity."
            );

            validate_dependency_transition(
                m_dependencies.rotation, dependencies.rotation,
                "HydrostaticEquilibriumContext received an older "
                "rotation revision for the same identity."
            );

            validate_dependency_transition(
                m_dependencies.bernoulliConstant,
                dependencies.bernoulliConstant,
                "HydrostaticEquilibriumContext received an older "
                "Bernoulli-constant revision for the same identity."
            );
        }

        const bool staticChanged =
            !m_isPrepared ||
            dependencies.discretization != m_dependencies.discretization;

        const bool enthalpyChanged =
            !m_isPrepared || dependencies.enthalpy != m_dependencies.enthalpy;

        const bool gravityPotentialChanged =
            !m_isPrepared ||
            dependencies.gravityPotential != m_dependencies.gravityPotential;

        const bool displacementChanged =
            !m_isPrepared ||
            dependencies.displacement != m_dependencies.displacement;

        const bool rotationChanged =
            !m_isPrepared || dependencies.rotation != m_dependencies.rotation;

        const bool bernoulliConstantChanged =
            !m_isPrepared ||
            dependencies.bernoulliConstant != m_dependencies.bernoulliConstant;

        const bool geometryPreparationRequired =
            staticChanged || displacementChanged;

        const bool rotationPreparationRequired =
            geometryPreparationRequired || rotationChanged;

        const bool baseStatePreparationRequired =
            rotationPreparationRequired || enthalpyChanged ||
            gravityPotentialChanged || bernoulliConstantChanged;

        HydrostaticPreparationReport report;

        report.preparedStaticDependencies   = staticChanged;
        report.preparedGeometryState        = geometryPreparationRequired;
        report.preparedRotationDependencies = rotationPreparationRequired;
        report.preparedBaseState            = baseStatePreparationRequired;

        if (staticChanged || enthalpyChanged) {
            m_baseEnthalpyTrue     = state.enthalpy;
            report.updatedEnthalpy = true;
        }

        if (staticChanged || gravityPotentialChanged) {
            m_baseGravityPotentialTrue     = state.gravityPotential;

            report.updatedGravityPotential = true;
        }

        if (geometryPreparationRequired) {
            m_displacementTrue         = state.displacement;
            report.updatedDisplacement = true;
        }

        if (staticChanged || bernoulliConstantChanged) {
            m_bernoulliConstant             = state.bernoulliConstant;
            report.updatedBernoulliConstant = true;
        }

        if (report.preparedStaticDependencies) {
            ++m_statistics.staticPreparations;
        }

        if (report.preparedGeometryState) {
            ++m_statistics.geometryPreparations;
        }

        if (report.preparedRotationDependencies) {
            ++m_statistics.rotationPreparations;
        }

        if (report.preparedBaseState) {
            ++m_statistics.baseStatePreparations;
        }

        m_dependencies = dependencies;
        m_isPrepared   = true;

        return report;
    }

    bool HydrostaticEquilibriumContext::IsPrepared() const noexcept {
        return m_isPrepared;
    }

    bool HydrostaticEquilibriumContext::MatchesDependencies(
        const HydrostaticEquilibriumDependencies &dependencies
    ) const noexcept {
        return m_isPrepared && dependencies == m_dependencies;
    }

    const HydrostaticEquilibriumDependencies &
    HydrostaticEquilibriumContext::GetDependencies() const {
        VerifyPrepared();
        return m_dependencies;
    }

    const HydrostaticPreparationStatistics &
    HydrostaticEquilibriumContext::GetPreparationStatistics() const noexcept {
        return m_statistics;
    }

    const mfem::Vector &
    HydrostaticEquilibriumContext::GetBaseEnthalpyTrue() const {
        VerifyPrepared();
        return m_baseEnthalpyTrue;
    }

    const mfem::Vector &
    HydrostaticEquilibriumContext::GetBaseGravityPotentialTrue() const {
        VerifyPrepared();
        return m_baseGravityPotentialTrue;
    }

    const mfem::Vector &
    HydrostaticEquilibriumContext::GetDisplacementTrue() const {
        VerifyPrepared();
        return m_displacementTrue;
    }

    double HydrostaticEquilibriumContext::GetBernoulliConstant() const {
        VerifyPrepared();
        return m_bernoulliConstant;
    }

    void HydrostaticEquilibriumContext::VerifyPrepared() const {
        MFEM_VERIFY(
            m_isPrepared, "HydrostaticEquilibriumContext has not been prepared."
        );
    }
} // namespace mean_field::operators::context::hydrostatic