MeanField/libmeanfield/interface/operators/prepared_gravity_displacement_force.cppm

module;

#include <compare>
#include <cstdint>
#include <expected>
#include <vector>

#include <mfem.hpp>

export module mean_field:operators.prepared_gravity_displacement_force;

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

export namespace mean_field::operators {
    struct PreparedGravityDisplacementForceReport final {
        bool preparedResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return preparedResidual;
        }
    };

    struct PreparedGravityDisplacementForceColumnStatistics final {
        std::uint64_t applications{0};

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

    struct PreparedGravityDisplacementForceCompleteStatistics final {
        std::uint64_t applications{0};

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

    class PreparedGravityDisplacementForceOperator final {
    public:
        PreparedGravityDisplacementForceOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const context::gravity_field::GravityFieldLinearizationContext &gravityContext
        );

        PreparedGravityDisplacementForceOperator(const PreparedGravityDisplacementForceOperator &)            = delete;

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

        PreparedGravityDisplacementForceOperator(PreparedGravityDisplacementForceOperator &&)                 = delete;

        PreparedGravityDisplacementForceOperator &operator=(PreparedGravityDisplacementForceOperator &&)      = delete;

        /*
         * Freeze the already-prepared shared gravity context. The caller must
         * first prepare GravityFieldLinearizationContext for the desired
         * state. Repeated calls with unchanged revisions do no work.
         */
        PreparedGravityDisplacementForceReport Prepare();

        [[nodiscard]] std::expected<
            PreparedGravityDisplacementForceReport,
            kernels::GravityDisplacementForceRejection>
        TryPrepare();

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

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

        void ApplyGravityGradientJacobianAction(
            const mfem::Vector &gravityGradientVariation,
            mfem::Vector &action
        ) const;

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

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

        [[nodiscard]] bool IsPrepared() const noexcept;

        [[nodiscard]] std::uint64_t GetResidualPreparationCount() const noexcept;

        [[nodiscard]] std::uint64_t GetResidualApplicationCount() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceColumnStatistics &
        GetDensityJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceColumnStatistics &
        GetGravityGradientJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceColumnStatistics &
        GetDisplacementJacobianStatistics() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceCompleteStatistics &
        GetCompleteJacobianStatistics() const noexcept;

        [[nodiscard]] const fem::FEM &GetFEM() const noexcept;

        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetGravityContext() const noexcept;

    private:
        void VerifyPrepared() const;
        [[nodiscard]] std::expected<
            void,
            kernels::GravityDisplacementForceRejection>
        TryPrepareElementData();
        void ApplyPreparedCompleteJacobianActionTrue(
            const mfem::Vector &densityVariationTrue,
            const mfem::Vector &displacementVariationTrue,
            const mfem::Vector &gravityGradientVariationTrue,
            mfem::Vector &actionTrue
        ) const;

        struct ElementPAData {
            int elementId{-1};
            mfem::Array<int> densityDofs;
            mfem::Array<int> gravityGradientDofs;
            mfem::Array<int> displacementDofs;
            mfem::DofTransformation *densityDofTransformation{nullptr};
            mfem::DofTransformation *gravityGradientDofTransformation{nullptr};
            mfem::DofTransformation *displacementDofTransformation{nullptr};
            const mfem::IntegrationRule *integrationRule{nullptr};
            mfem::DenseMatrix mappingJacobians;
            mfem::DenseMatrix inverseMeshJacobians;
            mfem::DenseMatrix baseGravityReferenceValues;
            mfem::Vector baseDensityValues;
            mfem::Vector referenceWeights;
        };

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

        context::gravity_field::GravityFieldRevisions m_preparedRevisions;
        mfem::Vector m_cachedResidual;
        std::vector<ElementPAData> m_elements;

        mutable mfem::Vector m_densityVariationTrue;
        mutable mfem::Vector m_gravityGradientVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;
        mutable mfem::Vector m_actionTrue;
        mutable mfem::Vector m_densityVariationLocal;
        mutable mfem::Vector m_gravityGradientVariationLocal;
        mutable mfem::Vector m_displacementVariationLocal;
        mutable mfem::Vector m_localAction;
        mutable mfem::Vector m_elementDensityVariation;
        mutable mfem::Vector m_elementGravityGradientVariation;
        mutable mfem::Vector m_elementDisplacementVariation;
        mutable mfem::Vector m_elementAction;
        mutable mfem::Vector m_densityShape;
        mutable mfem::Vector m_displacementShape;
        mutable mfem::Vector m_baseGravityReferenceValue;
        mutable mfem::Vector m_gravityVariationReferenceValue;
        mutable mfem::Vector m_mappedBaseGravity;
        mutable mfem::Vector m_mappedGravityVariation;
        mutable mfem::Vector m_mappedGeometryVariation;
        mutable mfem::Vector m_forceValue;
        mutable mfem::DenseMatrix m_gravityGradientShape;
        mutable mfem::DenseMatrix m_referenceDisplacementDShape;
        mutable mfem::DenseMatrix m_referenceDisplacementJacobian;
        mutable mfem::DenseMatrix m_displacementJacobianVariation;
        mutable mfem::DenseMatrix m_mappingJacobian;
        mutable mfem::DenseMatrix m_inverseMeshJacobian;

        std::uint64_t m_residualPreparationCount{0};
        mutable std::uint64_t m_residualApplicationCount{0};

        mutable PreparedGravityDisplacementForceColumnStatistics m_densityJacobianStatistics;

        mutable PreparedGravityDisplacementForceColumnStatistics m_gravityGradientJacobianStatistics;

        mutable PreparedGravityDisplacementForceColumnStatistics m_displacementJacobianStatistics;

        mutable PreparedGravityDisplacementForceCompleteStatistics m_completeJacobianStatistics;

        bool m_isPrepared{false};
    };

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

    class PreparedGravityDisplacementForceJacobianOperator final : public mfem::Operator {
    public:
        PreparedGravityDisplacementForceJacobianOperator(
            const GravityDisplacementForceLayout &layout,
            const PreparedGravityDisplacementForceOperator &preparedOperator
        );

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

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

    private:
        GravityDisplacementForceLayout m_layout;
        const PreparedGravityDisplacementForceOperator &m_preparedOperator;
    };
} // namespace mean_field::operators