MeanField/libmeanfield/interface/operators/prepared_barotropic_closure_operator.cppm

module;

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

export module mean_field:operators.prepared_barotropic_closure;

export import :eos.polytrope;
export import :fem;
export import :field.mfem;
export import :mapping.domain_mapper;
export import :operators.context.barotropic_closure_linearization;

export namespace mean_field::operators {
    struct PreparedBarotropicClosureReport final {
        context::barotropic::BarotropicClosurePreparationReport contextReport;
        bool preparedElementData{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return contextReport.DidAnyWork() || preparedElementData;
        }
    };

    class PreparedBarotropicClosureOperator final : public mfem::Operator {
    public:
        PreparedBarotropicClosureOperator(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const eos::Polytrope &equationOfState
        );

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

        PreparedBarotropicClosureReport Prepare(
            const context::barotropic::BarotropicClosureStateView &state,
            const context::barotropic::BarotropicClosureDependencies &dependencies
        );

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

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

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

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] std::uint64_t GetPreparationCount() const noexcept;
        [[nodiscard]] int GetDensitySize() const noexcept;
        [[nodiscard]] int GetEnthalpySize() const noexcept;
        [[nodiscard]] int GetDisplacementSize() const noexcept;

        [[nodiscard]] const context::barotropic::BarotropicClosureLinearizationContext &GetContext() const noexcept;
        [[nodiscard]] const context::barotropic::BarotropicClosurePreparationStatistics &
        GetContextPreparationStatistics() const noexcept;

    private:
        struct ConstructionData;

        [[nodiscard]] static ConstructionData MakeConstructionData(const fem::FEM &f);

        PreparedBarotropicClosureOperator(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const eos::Polytrope &equationOfState,
            ConstructionData constructionData
        );

        void VerifyPrepared() const;

        void ApplyThermodynamicActionFull(
            const mfem::Vector &densityVariationTrue,
            const mfem::Vector &enthalpyVariationTrue,
            mfem::Vector &actionTrue
        ) const;

        struct ElementPAData {
            mfem::Array<int> densityDofs;
            mfem::Array<int> enthalpyDofs;

            mfem::DofTransformation *densityDofTransformation{nullptr};
            mfem::DofTransformation *enthalpyDofTransformation{nullptr};

            mfem::DenseMatrix densityBasis;
            mfem::DenseMatrix enthalpyBasis;

            mfem::Vector weightedResidual;
            mfem::Vector quadratureWeights;
            mfem::Vector weightedEnthalpyDerivative;
        };

        const fem::FEM &m_fem;
        const mapping::DomainMapperStateless &m_domainMapper;
        const eos::Polytrope &m_equationOfState;

        field::FieldDofMap m_densityMap;
        field::FieldDofMap m_enthalpyMap;
        field::FieldDofMap m_displacementMap;

        context::barotropic::BarotropicClosureLinearizationContext m_context;

        std::vector<ElementPAData> m_elements;

        mfem::Vector m_baseDensityTrue;
        mfem::Vector m_baseEnthalpyTrue;
        mfem::Vector m_baseDisplacementTrue;

        mutable mfem::Vector m_densityVariationTrue;
        mutable mfem::Vector m_enthalpyVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;
        mutable mfem::Vector m_fullThermodynamicAction;
        mutable mfem::Vector m_fullDisplacementAction;
        mutable mfem::Vector m_fullResidual;

        std::uint64_t m_preparationCount{0};
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators