feat(libmeanfield): centrifugal + pressure

libmeanfield/interface/operators/prepared_hydrostatic_equilibrium_operator.cppm

@@ -0,0 +1,278 @@
+module;
+
+#include <compare>
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+#include <mfem.hpp>
+
+export module mean_field:operators.prepared_hydrostatic_equilibrium;
+
+export import :fem;
+export import :mapping.domain_mapper;
+export import :operators.context.hydrostatic_equilibrium;
+export import :physics.rigid_rotation;
+
+export namespace mean_field::operators {
+    struct PreparedHydrostaticEquilibriumReport {
+        context::hydrostatic::HydrostaticPreparationReport contextReport;
+
+        bool updatedRotation{false};
+        bool preparedAlgebraicJacobianBlocks{false};
+        bool preparedDisplacementJacobianData{false};
+        bool preparedResidual{false};
+
+        [[nodiscard]] bool DidAnyWork() const noexcept {
+            return contextReport.DidAnyWork() || updatedRotation ||
+                   preparedAlgebraicJacobianBlocks ||
+                   preparedDisplacementJacobianData || preparedResidual;
+        }
+    };
+
+    struct PreparedHydrostaticAlgebraicJacobianStatistics {
+        std::uint64_t preparations{0};
+        std::uint64_t enthalpyApplications{0};
+        std::uint64_t gravityPotentialApplications{0};
+        std::uint64_t bernoulliConstantApplications{0};
+        std::uint64_t combinedApplications{0};
+
+        constexpr auto operator<=>(
+            const PreparedHydrostaticAlgebraicJacobianStatistics &
+        ) const = default;
+    };
+
+    struct PreparedHydrostaticDisplacementJacobianStatistics {
+        std::uint64_t preparations{0};
+        std::uint64_t applications{0};
+
+        constexpr auto operator<=>(
+            const PreparedHydrostaticDisplacementJacobianStatistics &
+        ) const = default;
+    };
+
+    struct PreparedHydrostaticCompleteJacobianStatistics {
+        std::uint64_t applications{0};
+
+        constexpr auto operator<=>(
+            const PreparedHydrostaticCompleteJacobianStatistics &
+        ) const = default;
+    };
+
+    enum class HydrostaticJacobianInputBlock : int {
+        enthalpy          = 0,
+        gravityPotential  = 1,
+        bernoulliConstant = 2,
+        displacement      = 3
+    };
+
+    class HydrostaticJacobianBlockLayout final {
+    public:
+        explicit HydrostaticJacobianBlockLayout(const fem::FEM &f);
+
+        [[nodiscard]] int Offset(HydrostaticJacobianInputBlock block) const;
+
+        [[nodiscard]] int Size(HydrostaticJacobianInputBlock block) const;
+
+        [[nodiscard]] int GetTotalSize() const noexcept;
+
+        [[nodiscard]] int GetResidualSize() const noexcept;
+
+    private:
+        int m_enthalpySize{0};
+        int m_gravityPotentialSize{0};
+        int m_displacementSize{0};
+        int m_totalSize{0};
+        int m_residualSize{0};
+    };
+
+    class PreparedHydrostaticEquilibriumOperator final {
+    public:
+        PreparedHydrostaticEquilibriumOperator(
+            const fem::FEM &f,
+            const mapping::DomainMapperStateless &domainMapper
+        );
+
+        PreparedHydrostaticEquilibriumOperator(
+            const PreparedHydrostaticEquilibriumOperator &
+        ) = delete;
+
+        PreparedHydrostaticEquilibriumOperator &
+        operator=(const PreparedHydrostaticEquilibriumOperator &) = delete;
+
+        PreparedHydrostaticEquilibriumOperator(
+            PreparedHydrostaticEquilibriumOperator &&
+        ) = delete;
+
+        PreparedHydrostaticEquilibriumOperator &
+        operator=(PreparedHydrostaticEquilibriumOperator &&) = delete;
+
+        PreparedHydrostaticEquilibriumReport Prepare(
+            const context::hydrostatic::HydrostaticEquilibriumStateView &state,
+            const context::hydrostatic::HydrostaticEquilibriumDependencies
+                &dependencies,
+            const physics::RigidRotation &rotation
+        );
+
+        void BuildResidual(mfem::Vector &residual) const;
+
+        void ApplyEnthalpyJacobianAction(
+            const mfem::Vector &enthalpyVariation,
+            mfem::Vector &action
+        ) const;
+
+        void ApplyGravityPotentialJacobianAction(
+            const mfem::Vector &gravityPotentialVariation,
+            mfem::Vector &action
+        ) const;
+
+        void ApplyBernoulliConstantJacobianAction(
+            double bernoulliConstantVariation,
+            mfem::Vector &action
+        ) const;
+
+        void ApplyAlgebraicJacobianAction(
+            const mfem::Vector &enthalpyVariation,
+            const mfem::Vector &gravityPotentialVariation,
+            double bernoulliConstantVariation,
+            mfem::Vector &action
+        ) const;
+
+        void ApplyDisplacementJacobianAction(
+            const mfem::Vector &displacementVariation,
+            mfem::Vector &action
+        ) const;
+
+        void ApplyCompleteJacobianAction(
+            const mfem::Vector &enthalpyVariation,
+            const mfem::Vector &gravityPotentialVariation,
+            double bernoulliConstantVariation,
+            const mfem::Vector &displacementVariation,
+            mfem::Vector &action
+        ) const;
+
+        [[nodiscard]] bool IsPrepared() const noexcept;
+
+        [[nodiscard]] const context::hydrostatic::
+            HydrostaticPreparationStatistics &
+            GetContextPreparationStatistics() const noexcept;
+
+        [[nodiscard]] std::uint64_t
+        GetResidualPreparationCount() const noexcept;
+
+        [[nodiscard]] std::uint64_t
+        GetResidualApplicationCount() const noexcept;
+
+        [[nodiscard]] const PreparedHydrostaticAlgebraicJacobianStatistics &
+        GetAlgebraicJacobianStatistics() const noexcept;
+
+        [[nodiscard]] const PreparedHydrostaticDisplacementJacobianStatistics &
+        GetDisplacementJacobianStatistics() const noexcept;
+
+        [[nodiscard]] const PreparedHydrostaticCompleteJacobianStatistics &
+        GetCompleteJacobianStatistics() const noexcept;
+
+        [[nodiscard]] std::size_t GetStellarElementCount() const noexcept;
+
+        [[nodiscard]] const fem::FEM &GetFEM() const noexcept;
+
+    private:
+        struct ElementPAData {
+            int elementId{-1};
+
+            mfem::Array<int> enthalpyDofs;
+            mfem::Array<int> gravityPotentialDofs;
+            mfem::Array<int> displacementDofs;
+
+            mfem::DofTransformation *enthalpyDofTransformation{nullptr};
+
+            mfem::DofTransformation *gravityPotentialDofTransformation{nullptr};
+
+            mfem::DofTransformation *displacementDofTransformation{nullptr};
+
+            const mfem::IntegrationRule *integrationRule{nullptr};
+
+            // Rows are quadrature points and columns are element DOFs.
+            mfem::DenseMatrix enthalpyBasis;
+            mfem::DenseMatrix gravityPotentialBasis;
+
+            // Rows are quadrature points and columns are physical components.
+            mfem::DenseMatrix physicalPositions;
+
+            mfem::Vector quadratureWeights;
+            std::vector<mapping::VolumeMappingContext> baseMappingContexts;
+
+            std::optional<mapping::ElementDisplacementData>
+                baseDisplacementData;
+
+            std::optional<mapping::ElementCompactificationData>
+                compactificationData;
+
+            mfem::Vector rotationPotential;
+            mfem::DenseMatrix rotationGradient;
+            mfem::Vector hydrostaticImbalance;
+            mfem::Vector weightedResidual;
+
+            // Geometry-dependent algebraic Jacobian blocks.
+            mfem::DenseMatrix enthalpyJacobian;
+            mfem::DenseMatrix gravityPotentialJacobian;
+            mfem::Vector bernoulliConstantJacobian;
+        };
+
+        void PrepareStaticPlan();
+        void PrepareGeometry();
+        void PrepareAlgebraicJacobianBlocks();
+        void PrepareRotation();
+        void PrepareBaseState();
+        void FinalizeDisplacementJacobianPreparation();
+        void AssembleCachedResidual();
+        void VerifyPrepared() const;
+
+        const fem::FEM &m_fem;
+        const mapping::DomainMapperStateless &m_domainMapper;
+
+        context::hydrostatic::HydrostaticEquilibriumContext m_context;
+
+        std::optional<physics::RigidRotation> m_rotation;
+        std::vector<ElementPAData> m_elements;
+        mfem::Vector m_cachedResidual;
+
+        std::uint64_t m_residualPreparationCount{0};
+
+        mutable std::uint64_t m_residualApplicationCount{0};
+
+        mutable PreparedHydrostaticAlgebraicJacobianStatistics
+            m_algebraicJacobianStatistics;
+
+        mutable PreparedHydrostaticDisplacementJacobianStatistics
+            m_displacementJacobianStatistics;
+
+        mutable PreparedHydrostaticCompleteJacobianStatistics
+            m_completeJacobianStatistics;
+
+        bool m_isPrepared{false};
+    };
+
+    class PreparedHydrostaticEquilibriumJacobianOperator final
+        : public mfem::Operator {
+    public:
+        PreparedHydrostaticEquilibriumJacobianOperator(
+            const fem::FEM &f,
+            const PreparedHydrostaticEquilibriumOperator &preparedOperator
+        );
+
+        void Mult(
+            const mfem::Vector &direction,
+            mfem::Vector &action
+        ) const override;
+
+        [[nodiscard]] const HydrostaticJacobianBlockLayout &
+        GetLayout() const noexcept;
+
+    private:
+        HydrostaticJacobianBlockLayout m_layout;
+
+        const PreparedHydrostaticEquilibriumOperator &m_preparedOperator;
+    };
+} // namespace mean_field::operators