MeanField/libmeanfield/impl/operators/kernels/hydrostatic_equilibrium_kernels.cpp

module;

#include <algorithm>
#include <array>
#include <cmath>
#include <mfem.hpp>
#include <optional>

module mean_field;

import :operators.kernels.hydrostatic_equilibrium;

namespace {
    void true_to_local(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &trueVector,
        mfem::Vector &localVector
    ) {
        MFEM_VERIFY(
            trueVector.Size() == finiteElementSpace.GetTrueVSize(),
            "True vector has the wrong size."
        );

        localVector.SetSize(finiteElementSpace.GetVSize());

        const mfem::Operator *prolongation =
            finiteElementSpace.GetProlongationMatrix();

        if (prolongation != nullptr) {
            prolongation->Mult(trueVector, localVector);
        } else {
            localVector = trueVector;
        }
    }

    void local_to_true(
        const mfem::ParFiniteElementSpace &finiteElementSpace,
        const mfem::Vector &localVector,
        mfem::Vector &trueVector
    ) {
        MFEM_VERIFY(
            localVector.Size() == finiteElementSpace.GetVSize(),
            "Local vector has the wrong size."
        );

        trueVector.SetSize(finiteElementSpace.GetTrueVSize());

        trueVector = 0.0;

        const mfem::Operator *prolongation =
            finiteElementSpace.GetProlongationMatrix();

        if (prolongation != nullptr) {
            prolongation->MultTranspose(localVector, trueVector);
        } else {
            trueVector = localVector;
        }
    }

    void validate_fem(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapperStateless &domainMapper
    ) {
        MFEM_VERIFY(
            f.mesh != nullptr, "The hydrostatic kernel requires a mesh."
        );

        MFEM_VERIFY(
            f.enthalpyFes != nullptr, "The hydrostatic kernel requires the "
                                      "enthalpy finite-element space."
        );

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

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

        MFEM_VERIFY(
            f.compactificationFes != nullptr,
            "The hydrostatic kernel requires the "
            "compactification finite-element space."
        );

        MFEM_VERIFY(
            f.compactificationCoordinate != nullptr,
            "The hydrostatic kernel requires the "
            "compactification coordinate."
        );

        MFEM_VERIFY(
            f.quadratureFactory != nullptr,
            "The hydrostatic kernel requires the "
            "quadrature-rule factory."
        );

        MFEM_VERIFY(
            f.mesh->Dimension() == 3,
            "The rigid-rotation hydrostatic kernel "
            "currently requires a three-dimensional mesh."
        );

        MFEM_VERIFY(
            domainMapper.GetDimension() == f.mesh->Dimension(),
            "The domain-mapper dimension does not match "
            "the mesh dimension."
        );
    }

    const mfem::IntegrationRule &get_hydrostatic_rule(
        const mean_field::fem::FEM &f,
        const mfem::FiniteElement &enthalpyElement,
        const mfem::FiniteElement &potentialElement,
        const mfem::ElementTransformation &transformation
    ) {
        using EnthalpyField =
            mean_field::field::Field<mean_field::field::Enthalpy>;

        MFEM_VERIFY(
            enthalpyElement.GetOrder() ==
                mean_field::field::Enthalpy::Scalar::familyOrder,
            "The hydrostatic test element does not match "
            "the registered enthalpy field."
        );

        MFEM_VERIFY(
            potentialElement.GetOrder() ==
                mean_field::field::Gravity::Potential::familyOrder,
            "The hydrostatic potential element does not "
            "match the registered gravity-potential field."
        );

        const auto enthalpyQuery = EnthalpyField::make_query<
            mean_field::field::Enthalpy::Form::EquilibriumEnthalpy>(
            mean_field::quadrature::QuadratureRole::discretization,
            transformation.OrderW(), {}, mean_field::utils::DOMAINS::STELLAR,
            mean_field::quadrature::MappingKind::general
        );

        const auto gravityQuery = EnthalpyField::make_query<
            mean_field::field::Enthalpy::Form::EquilibriumGravity>(
            mean_field::quadrature::QuadratureRole::discretization,
            transformation.OrderW(), {}, mean_field::utils::DOMAINS::STELLAR,
            mean_field::quadrature::MappingKind::general
        );

        const auto rotationQuery = EnthalpyField::make_query<
            mean_field::field::Enthalpy::Form::EquilibriumRotation>(
            mean_field::quadrature::QuadratureRole::discretization,
            transformation.OrderW(), std::array<int, 1>{2},
            mean_field::utils::DOMAINS::STELLAR,
            mean_field::quadrature::MappingKind::general
        );

        const auto constantQuery = EnthalpyField::make_query<
            mean_field::field::Enthalpy::Form::EquilibriumConstant>(
            mean_field::quadrature::QuadratureRole::discretization,
            transformation.OrderW(), {}, mean_field::utils::DOMAINS::STELLAR,
            mean_field::quadrature::MappingKind::general
        );

        int integrationOrder    = 0;

        const auto update_order = [&f, &transformation, &integrationOrder](
                                      const mean_field::quadrature::Query &query
                                  ) {
            const auto rule = f.quadratureFactory->get(
                query, transformation.GetGeometryType()
            );

            MFEM_VERIFY(
                rule.integration_rule != nullptr,
                "The quadrature policy did not return "
                "a hydrostatic-equilibrium rule."
            );

            integrationOrder =
                std::max(integrationOrder, rule.resolution.order);
        };

        update_order(enthalpyQuery);
        update_order(gravityQuery);
        update_order(rotationQuery);
        update_order(constantQuery);

        return mfem::IntRules.Get(
            transformation.GetGeometryType(), integrationOrder
        );
    }

    struct HydrostaticAssemblyRequest {
        const mean_field::physics::RigidRotation *rotation{nullptr};

        const mfem::Vector *baseEnthalpyTrue{nullptr};
        const mfem::Vector *basePotentialTrue{nullptr};

        const mfem::Vector *enthalpyVariationTrue{nullptr};
        const mfem::Vector *potentialVariationTrue{nullptr};
        const mfem::Vector *displacementVariationTrue{nullptr};

        double bernoulliConstant{0.0};
        double constantVariation{0.0};

        bool buildResidual{false};
    };

    void assemble_hydrostatic_form(
        const mean_field::fem::FEM &f,
        const mean_field::mapping::DomainMapperStateless &domainMapper,
        const mfem::Vector &displacementTrue,
        const HydrostaticAssemblyRequest &request,
        mfem::Vector &result
    ) {
        validate_fem(f, domainMapper);

        MFEM_VERIFY(
            displacementTrue.Size() == f.displacementFes->GetTrueVSize(),
            "The hydrostatic displacement vector has "
            "the wrong size."
        );

        MFEM_VERIFY(
            std::isfinite(request.bernoulliConstant),
            "The Bernoulli constant is non-finite."
        );

        MFEM_VERIFY(
            std::isfinite(request.constantVariation),
            "The Bernoulli-constant variation is non-finite."
        );

        const bool requiresBaseState =
            request.buildResidual ||
            request.displacementVariationTrue != nullptr;

        if (requiresBaseState) {
            MFEM_VERIFY(
                request.rotation != nullptr,
                "The hydrostatic residual or geometry "
                "action requires the rotation model."
            );

            MFEM_VERIFY(
                request.baseEnthalpyTrue != nullptr,
                "The hydrostatic residual or geometry "
                "action requires the base enthalpy."
            );

            MFEM_VERIFY(
                request.basePotentialTrue != nullptr,
                "The hydrostatic residual or geometry "
                "action requires the base potential."
            );
        }

        if (request.baseEnthalpyTrue != nullptr) {
            MFEM_VERIFY(
                request.baseEnthalpyTrue->Size() ==
                    f.enthalpyFes->GetTrueVSize(),
                "The base enthalpy vector has the wrong size."
            );
        }

        if (request.basePotentialTrue != nullptr) {
            MFEM_VERIFY(
                request.basePotentialTrue->Size() ==
                    f.gravityPotentialFes->GetTrueVSize(),
                "The base potential vector has the wrong size."
            );
        }

        if (request.enthalpyVariationTrue != nullptr) {
            MFEM_VERIFY(
                request.enthalpyVariationTrue->Size() ==
                    f.enthalpyFes->GetTrueVSize(),
                "The enthalpy variation has the wrong size."
            );
        }

        if (request.potentialVariationTrue != nullptr) {
            MFEM_VERIFY(
                request.potentialVariationTrue->Size() ==
                    f.gravityPotentialFes->GetTrueVSize(),
                "The potential variation has the wrong size."
            );
        }

        if (request.displacementVariationTrue != nullptr) {
            MFEM_VERIFY(
                request.displacementVariationTrue->Size() ==
                    f.displacementFes->GetTrueVSize(),
                "The displacement variation has the wrong size."
            );
        }

        mfem::Vector displacementLocal;
        true_to_local(*f.displacementFes, displacementTrue, displacementLocal);

        mfem::Vector baseEnthalpyLocal;
        mfem::Vector basePotentialLocal;
        mfem::Vector enthalpyVariationLocal;
        mfem::Vector potentialVariationLocal;
        mfem::Vector displacementVariationLocal;

        if (request.baseEnthalpyTrue != nullptr) {
            true_to_local(
                *f.enthalpyFes, *request.baseEnthalpyTrue, baseEnthalpyLocal
            );
        }

        if (request.basePotentialTrue != nullptr) {
            true_to_local(
                *f.gravityPotentialFes, *request.basePotentialTrue,
                basePotentialLocal
            );
        }

        if (request.enthalpyVariationTrue != nullptr) {
            true_to_local(
                *f.enthalpyFes, *request.enthalpyVariationTrue,
                enthalpyVariationLocal
            );
        }

        if (request.potentialVariationTrue != nullptr) {
            true_to_local(
                *f.gravityPotentialFes, *request.potentialVariationTrue,
                potentialVariationLocal
            );
        }

        if (request.displacementVariationTrue != nullptr) {
            true_to_local(
                *f.displacementFes, *request.displacementVariationTrue,
                displacementVariationLocal
            );
        }

        mfem::Vector localResult(f.enthalpyFes->GetVSize());

        localResult = 0.0;

        mean_field::mapping::DomainMapperStateless::Workspace workspace(
            f.mesh->Dimension()
        );

        mfem::Array<int> enthalpyDofs;
        mfem::Array<int> potentialDofs;
        mfem::Array<int> displacementDofs;
        mfem::Array<int> compactificationDofs;

        mfem::Vector elementBaseEnthalpy;
        mfem::Vector elementBasePotential;
        mfem::Vector elementEnthalpyVariation;
        mfem::Vector elementPotentialVariation;
        mfem::Vector elementDisplacement;
        mfem::Vector elementDisplacementVariation;
        mfem::Vector elementCompactification;
        mfem::Vector elementResult;

        mfem::Vector enthalpyShape;
        mfem::Vector potentialShape;

        const int vacuumAttribute = domainMapper.GetVacuumElementAttribute();

        for (int elementId = 0; elementId < f.mesh->GetNE(); ++elementId) {
            mfem::ElementTransformation *transformation =
                f.mesh->GetElementTransformation(elementId);

            MFEM_VERIFY(
                transformation != nullptr,
                "The hydrostatic kernel received a null "
                "element transformation."
            );

            if (transformation->Attribute == vacuumAttribute) {
                continue;
            }

            const mfem::FiniteElement &enthalpyElement =
                *f.enthalpyFes->GetFE(elementId);

            const mfem::FiniteElement &potentialElement =
                *f.gravityPotentialFes->GetFE(elementId);

            const mfem::FiniteElement &displacementElement =
                *f.displacementFes->GetFE(elementId);

            const mfem::FiniteElement &compactificationElement =
                *f.compactificationFes->GetFE(elementId);

            mfem::DofTransformation *enthalpyDofTransformation =
                f.enthalpyFes->GetElementDofs(elementId, enthalpyDofs);

            mfem::DofTransformation *potentialDofTransformation =
                f.gravityPotentialFes->GetElementDofs(elementId, potentialDofs);

            mfem::DofTransformation *displacementDofTransformation =
                f.displacementFes->GetElementVDofs(elementId, displacementDofs);

            mfem::DofTransformation *compactificationDofTransformation =
                f.compactificationFes->GetElementDofs(
                    elementId, compactificationDofs
                );

            displacementLocal.GetSubVector(
                displacementDofs, elementDisplacement
            );

            f.compactificationCoordinate->GetSubVector(
                compactificationDofs, elementCompactification
            );

            if (request.baseEnthalpyTrue != nullptr) {
                baseEnthalpyLocal.GetSubVector(
                    enthalpyDofs, elementBaseEnthalpy
                );
            }

            if (request.basePotentialTrue != nullptr) {
                basePotentialLocal.GetSubVector(
                    potentialDofs, elementBasePotential
                );
            }

            if (request.enthalpyVariationTrue != nullptr) {
                enthalpyVariationLocal.GetSubVector(
                    enthalpyDofs, elementEnthalpyVariation
                );
            }

            if (request.potentialVariationTrue != nullptr) {
                potentialVariationLocal.GetSubVector(
                    potentialDofs, elementPotentialVariation
                );
            }

            if (request.displacementVariationTrue != nullptr) {
                displacementVariationLocal.GetSubVector(
                    displacementDofs, elementDisplacementVariation
                );
            }

            if (enthalpyDofTransformation != nullptr) {
                if (request.baseEnthalpyTrue != nullptr) {
                    enthalpyDofTransformation->InvTransformPrimal(
                        elementBaseEnthalpy
                    );
                }

                if (request.enthalpyVariationTrue != nullptr) {
                    enthalpyDofTransformation->InvTransformPrimal(
                        elementEnthalpyVariation
                    );
                }
            }

            if (potentialDofTransformation != nullptr) {
                if (request.basePotentialTrue != nullptr) {
                    potentialDofTransformation->InvTransformPrimal(
                        elementBasePotential
                    );
                }

                if (request.potentialVariationTrue != nullptr) {
                    potentialDofTransformation->InvTransformPrimal(
                        elementPotentialVariation
                    );
                }
            }

            if (displacementDofTransformation != nullptr) {
                displacementDofTransformation->InvTransformPrimal(
                    elementDisplacement
                );

                if (request.displacementVariationTrue != nullptr) {
                    displacementDofTransformation->InvTransformPrimal(
                        elementDisplacementVariation
                    );
                }
            }

            if (compactificationDofTransformation != nullptr) {
                compactificationDofTransformation->InvTransformPrimal(
                    elementCompactification
                );
            }

            const mean_field::mapping::ElementDisplacementData
                displacementData = mean_field::mapping::
                    ElementDisplacementDataFromElementVDofs(
                        displacementElement, elementDisplacement
                    );

            const mean_field::mapping::ElementCompactificationData
                compactificationData(
                    compactificationElement, elementCompactification
                );

            const mean_field::mapping::ElementMappingData mappingData{
                .displacement     = displacementData,
                .compactification = compactificationData
            };

            std::optional<mean_field::mapping::ElementDisplacementData>
                displacementVariationData;

            if (request.displacementVariationTrue != nullptr) {
                displacementVariationData.emplace(
                    mean_field::mapping::
                        ElementDisplacementDataFromElementVDofs(
                            displacementElement, elementDisplacementVariation
                        )
                );
            }

            elementResult.SetSize(enthalpyElement.GetDof());

            elementResult = 0.0;

            enthalpyShape.SetSize(enthalpyElement.GetDof());

            potentialShape.SetSize(potentialElement.GetDof());

            const mfem::IntegrationRule &integrationRule = get_hydrostatic_rule(
                f, enthalpyElement, potentialElement, *transformation
            );

            for (int quadraturePoint = 0;
                 quadraturePoint < integrationRule.GetNPoints();
                 ++quadraturePoint) {
                const mfem::IntegrationPoint &integrationPoint =
                    integrationRule.IntPoint(quadraturePoint);

                transformation->SetIntPoint(&integrationPoint);

                mean_field::mapping::VolumeMappingContext mappingContext;

                const mean_field::mapping::MappingStatus mappingStatus =
                    domainMapper.EvaluateVolume(
                        mappingData, *transformation, integrationPoint,
                        workspace, mappingContext
                    );

                MFEM_VERIFY(
                    mappingStatus == mean_field::mapping::MappingStatus::valid,
                    "The base mapping is invalid in the "
                    "hydrostatic kernel. Element: "
                        << elementId
                        << ", quadrature point: " << quadraturePoint
                        << ", status: " << static_cast<int>(mappingStatus)
                );

                enthalpyElement.CalcShape(integrationPoint, enthalpyShape);

                potentialElement.CalcShape(integrationPoint, potentialShape);

                double baseIntegrand = 0.0;

                if (requiresBaseState) {
                    const double enthalpyValue =
                        elementBaseEnthalpy * enthalpyShape;

                    const double potentialValue =
                        elementBasePotential * potentialShape;

                    const double rotationPotential =
                        request.rotation->potential(
                            mappingContext.mapping.physical_position
                        );

                    baseIntegrand = enthalpyValue + potentialValue -
                                    rotationPotential -
                                    request.bernoulliConstant;
                }

                if (request.buildResidual) {
                    elementResult.Add(
                        mappingContext.quadrature.weight * baseIntegrand,
                        enthalpyShape
                    );

                    continue;
                }

                double materialVariation = -request.constantVariation;

                if (request.enthalpyVariationTrue != nullptr) {
                    materialVariation +=
                        elementEnthalpyVariation * enthalpyShape;
                }

                if (request.potentialVariationTrue != nullptr) {
                    materialVariation +=
                        elementPotentialVariation * potentialShape;
                }

                double weightedVariation =
                    mappingContext.quadrature.weight * materialVariation;

                if (request.displacementVariationTrue != nullptr) {
                    mean_field::mapping::VolumeMappingVariation
                        mappingVariation;

                    const mean_field::mapping::MappingStatus variationStatus =
                        domainMapper.EvaluateVolumeVariation(
                            mappingData, *displacementVariationData,
                            *transformation, integrationPoint, mappingContext,
                            workspace, mappingVariation
                        );

                    MFEM_VERIFY(
                        variationStatus ==
                            mean_field::mapping::MappingStatus::valid,
                        "The mapping variation is invalid "
                        "in the hydrostatic kernel."
                    );

                    const double rotationVariation =
                        request.rotation->potential_directional_derivative(
                            mappingContext.mapping.physical_position,
                            mappingVariation.mapping.physical_position_variation
                        );

                    weightedVariation +=
                        baseIntegrand * mappingVariation.weight_variation -
                        rotationVariation * mappingContext.quadrature.weight;
                }

                elementResult.Add(weightedVariation, enthalpyShape);
            }

            if (enthalpyDofTransformation != nullptr) {
                enthalpyDofTransformation->TransformDual(elementResult);
            }

            localResult.AddElementVector(enthalpyDofs, elementResult);
        }

        local_to_true(*f.enthalpyFes, localResult, result);
    }
} // namespace

namespace mean_field::operators::kernels {
    void apply_hydrostatic_equilibrium(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const physics::RigidRotation &rotation,
        const mfem::Vector &enthalpyTrue,
        const mfem::Vector &potentialTrue,
        const mfem::Vector &displacementTrue,
        const double bernoulliConstant,
        mfem::Vector &residual
    ) {
        HydrostaticAssemblyRequest request;

        request.rotation          = &rotation;
        request.baseEnthalpyTrue  = &enthalpyTrue;
        request.basePotentialTrue = &potentialTrue;
        request.bernoulliConstant = bernoulliConstant;
        request.buildResidual     = true;

        assemble_hydrostatic_form(
            f, domainMapper, displacementTrue, request, residual
        );
    }

    void apply_hydrostatic_equilibrium_enthalpy_action(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const mfem::Vector &enthalpyVariationTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.enthalpyVariationTrue = &enthalpyVariationTrue;

        assemble_hydrostatic_form(
            f, domainMapper, displacementTrue, request, action
        );
    }

    void apply_hydrostatic_equilibrium_potential_action(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const mfem::Vector &potentialVariationTrue,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.potentialVariationTrue = &potentialVariationTrue;

        assemble_hydrostatic_form(
            f, domainMapper, displacementTrue, request, action
        );
    }

    void apply_hydrostatic_equilibrium_constant_action(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const double constantVariation,
        const mfem::Vector &displacementTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.constantVariation = constantVariation;

        assemble_hydrostatic_form(
            f, domainMapper, displacementTrue, request, action
        );
    }

    void apply_hydrostatic_equilibrium_displacement_action(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const physics::RigidRotation &rotation,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &basePotentialTrue,
        const mfem::Vector &baseDisplacementTrue,
        const double baseBernoulliConstant,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.rotation                  = &rotation;
        request.baseEnthalpyTrue          = &baseEnthalpyTrue;
        request.basePotentialTrue         = &basePotentialTrue;
        request.displacementVariationTrue = &displacementVariationTrue;
        request.bernoulliConstant         = baseBernoulliConstant;

        assemble_hydrostatic_form(
            f, domainMapper, baseDisplacementTrue, request, action
        );
    }

    void apply_hydrostatic_equilibrium_action(
        const fem::FEM &f,
        const mapping::DomainMapperStateless &domainMapper,
        const physics::RigidRotation &rotation,
        const mfem::Vector &baseEnthalpyTrue,
        const mfem::Vector &basePotentialTrue,
        const mfem::Vector &baseDisplacementTrue,
        const double baseBernoulliConstant,
        const mfem::Vector &enthalpyVariationTrue,
        const mfem::Vector &potentialVariationTrue,
        const double constantVariation,
        const mfem::Vector &displacementVariationTrue,
        mfem::Vector &action
    ) {
        HydrostaticAssemblyRequest request;

        request.rotation                  = &rotation;
        request.baseEnthalpyTrue          = &baseEnthalpyTrue;
        request.basePotentialTrue         = &basePotentialTrue;
        request.enthalpyVariationTrue     = &enthalpyVariationTrue;
        request.potentialVariationTrue    = &potentialVariationTrue;
        request.displacementVariationTrue = &displacementVariationTrue;
        request.bernoulliConstant         = baseBernoulliConstant;
        request.constantVariation         = constantVariation;

        assemble_hydrostatic_form(
            f, domainMapper, baseDisplacementTrue, request, action
        );
    }
} // namespace mean_field::operators::kernels