MeanField/libmeanfield/interface/operators/prepared_mass_normalization.cppm

module;

#include <compare>
#include <cstdint>
#include <mfem.hpp>
#include <vector>

export module mean_field:operators.prepared_mass_normalization;

export import :fem;
export import :mapping.domain_mapper;
export import :operators.context.gravity_field;
export import :utils.blocks;

export namespace mean_field::operators {
    struct MassNormalizationDependencyStamp final {
        std::uint64_t identity{0};
        std::uint64_t revision{0};

        constexpr auto operator<=>(const MassNormalizationDependencyStamp &) const = default;
    };

    struct MassNormalizationDependencies final {
        MassNormalizationDependencyStamp discretization;
        MassNormalizationDependencyStamp density;
        MassNormalizationDependencyStamp displacement;
        MassNormalizationDependencyStamp targetMass;

        constexpr auto operator<=>(const MassNormalizationDependencies &) const = default;
    };

    struct MassNormalizationStateView final {
        double targetMass{0.0};
    };

    struct PreparedMassNormalizationReport final {
        bool rebuiltStaticPlan{false};
        bool refreshedGeometry{false};
        bool refreshedDensity{false};
        bool updatedTargetMass{false};
        bool assembledResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return rebuiltStaticPlan || refreshedGeometry || refreshedDensity || updatedTargetMass || assembledResidual;
        }
    };

    struct PreparedMassNormalizationActionStatistics final {
        std::uint64_t densityApplications{0};
        std::uint64_t displacementApplications{0};
        std::uint64_t completeApplications{0};

        constexpr auto operator<=>(const PreparedMassNormalizationActionStatistics &) const = default;
    };

    /*
     * Prepared scalar row
     *
     *   R_M(rho, d) = integral_{Omega_star(d)} rho dV - M_target.
     *
     * Density and displacement are borrowed from the shared gravity-field
     * linearization context. This keeps the mass row on exactly the same
     * frozen state and geometry as the gravity and mechanical rows.
     * Jacobian directions use the corresponding registered FieldDof
     * coordinates; expansion to MFEM true-DOF vectors is private.
     */
    class PreparedMassNormalizationOperator final {
    public:
        PreparedMassNormalizationOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const context::gravity_field::GravityFieldLinearizationContext &gravityContext
        );

        PreparedMassNormalizationOperator(const PreparedMassNormalizationOperator &)            = delete;
        PreparedMassNormalizationOperator &operator=(const PreparedMassNormalizationOperator &) = delete;
        PreparedMassNormalizationOperator(PreparedMassNormalizationOperator &&)                 = delete;
        PreparedMassNormalizationOperator &operator=(PreparedMassNormalizationOperator &&)      = delete;

        PreparedMassNormalizationReport Prepare(
            const MassNormalizationStateView &state,
            const MassNormalizationDependencies &dependencies
        );

        void BuildResidual(mfem::Vector &residual) const;

        void ApplyDensityJacobianAction(
            const mfem::Vector &densityVariation,
            mfem::Vector &action
        ) const;

        void ApplyDisplacementJacobianAction(
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

        void ApplyCompleteJacobianAction(
            const mfem::Vector &densityVariation,
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] double GetCurrentMass() const;
        [[nodiscard]] double GetTargetMass() const;
        [[nodiscard]] std::uint64_t GetPreparationCount() const noexcept;
        [[nodiscard]] std::uint64_t GetResidualApplicationCount() const noexcept;
        [[nodiscard]] const PreparedMassNormalizationActionStatistics &GetActionStatistics() const noexcept;
        [[nodiscard]] const fem::FEM &GetFEM() const noexcept;
        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetGravityContext() const noexcept;

    private:
        struct QuadraturePointData final {
            mfem::IntegrationPoint integrationPoint;
            mfem::Vector densityShape;
            mapping::VolumeMappingContext mappingContext;
            double density{0.0};
        };

        struct ElementPAData final {
            int elementId{-1};

            mfem::Array<int> densityDofs;
            mfem::Array<int> displacementDofs;
            mfem::Array<int> compactificationDofs;

            mfem::DofTransformation *densityDofTransformation{nullptr};
            mfem::DofTransformation *displacementDofTransformation{nullptr};
            mfem::DofTransformation *compactificationDofTransformation{nullptr};

            mfem::Vector baseDisplacement;
            mfem::Vector compactification;

            std::vector<QuadraturePointData> quadraturePoints;
        };

        void BuildStaticPlan();
        void RefreshGeometry(const mfem::Vector &displacement);
        void RefreshDensity(const mfem::Vector &density);
        void AssembleResidual();
        void VerifyPrepared() const;

        [[nodiscard]] double EvaluateDensityActionLocal(const mfem::Vector &densityVariation) const;

        [[nodiscard]] double EvaluateDisplacementActionLocal(const mfem::Vector &displacementVariation) const;

        [[nodiscard]] double GlobalSum(double localValue) const;

        const fem::FEM &m_fem;
        const mapping::DomainMapper &m_domainMapper;
        const context::gravity_field::GravityFieldLinearizationContext &m_gravityContext;

        std::vector<ElementPAData> m_elements;
        MassNormalizationDependencies m_preparedDependencies;
        mfem::Vector m_cachedResidual;

        mutable mfem::Vector m_densityVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;

        double m_currentMass{0.0};
        double m_targetMass{0.0};

        std::uint64_t m_preparationCount{0};
        mutable std::uint64_t m_residualApplicationCount{0};
        mutable PreparedMassNormalizationActionStatistics m_actionStatistics;
        bool m_isPrepared{false};
    };

    using MassNormalizationLayout = utils::blocks::form_layout<utils::blocks::barotropic_equilibrium_form>;

    class PreparedMassNormalizationJacobianOperator final : public mfem::Operator {
    public:
        PreparedMassNormalizationJacobianOperator(
            const MassNormalizationLayout &layout,
            const PreparedMassNormalizationOperator &preparedOperator
        );

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override;

        [[nodiscard]] const MassNormalizationLayout &GetLayout() const noexcept;

    private:
        MassNormalizationLayout m_layout;
        const PreparedMassNormalizationOperator &m_preparedOperator;
    };
} // namespace mean_field::operators