Add fixed-work mass-consistent gravity Schur refinement

This commit is contained in:
2026-09-17 16:19:22 -04:00
parent 018fc1c992
commit b18f41f88f
3 changed files with 531 additions and 7 deletions

View File

@@ -1,6 +1,7 @@
module;
#include <algorithm>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <memory>
@@ -9,6 +10,7 @@ module;
#include <utility>
#include <mfem.hpp>
#include <mpi.h>
export module mean_field:preconditioning.gravity_field;
@@ -23,12 +25,21 @@ export namespace mean_field::preconditioning {
struct GravityUpperTriangular final { };
struct GravityApproximateLDU final { };
// A fixed number of corrections preserves a stationary linear application
// when both child inverses are stationary. Relaxation is a numerical choice;
// positivity alone does not establish Richardson contraction.
struct GravityConsistentSchurLDU final {
int correctionSweeps{1};
double relaxation{1.0};
};
template <typename Candidate> struct IsGravityFactorizationPolicy : std::false_type { };
template <> struct IsGravityFactorizationPolicy<GravityBlockDiagonal> : std::true_type { };
template <> struct IsGravityFactorizationPolicy<GravityLowerTriangular> : std::true_type { };
template <> struct IsGravityFactorizationPolicy<GravityUpperTriangular> : std::true_type { };
template <> struct IsGravityFactorizationPolicy<GravityApproximateLDU> : std::true_type { };
template <> struct IsGravityFactorizationPolicy<GravityConsistentSchurLDU> : std::true_type { };
template <typename Candidate>
concept GravityFactorizationPolicy = IsGravityFactorizationPolicy<std::remove_cvref_t<Candidate>>::value;
@@ -179,6 +190,7 @@ export namespace mean_field::preconditioning {
std::uint64_t potentialSchurApplications{0};
std::uint64_t divergenceApplications{0};
std::uint64_t transposeDivergenceApplications{0};
std::uint64_t schurRefinementSweeps{0};
};
template <GravityFactorizationPolicy Policy> class GravityFactorizationOperator final : public mfem::Solver {
@@ -197,7 +209,17 @@ export namespace mean_field::preconditioning {
m_offsets(3),
m_potentialWorkspace(potentialSchurInverse.Height()),
m_gradientWorkspace(massInverse.Height()),
m_massCorrection(massInverse.Height()) {
m_massCorrection(massInverse.Height()),
m_schurResidual(RefinementWorkspaceSize(m_policy, potentialSchurInverse.Height())),
m_potentialCorrection(RefinementWorkspaceSize(m_policy, potentialSchurInverse.Height())) {
if constexpr (std::same_as<Policy, GravityConsistentSchurLDU>) {
if (m_policy.correctionSweeps < 0 || !std::isfinite(m_policy.relaxation) ||
m_policy.relaxation <= 0.0) {
throw std::invalid_argument(
"Gravity Schur refinement requires nonnegative sweeps and finite positive relaxation."
);
}
}
if (massInverse.Height() <= 0 || massInverse.Height() != massInverse.Width()) {
throw std::invalid_argument("The gravity factorization requires a square gradient-mass inverse.");
}
@@ -269,11 +291,36 @@ export namespace mean_field::preconditioning {
++m_statistics.transposeDivergenceApplications;
++m_statistics.massInverseApplications;
} else {
static_assert(std::same_as<Policy, GravityApproximateLDU>);
static_assert(
std::same_as<Policy, GravityApproximateLDU> || std::same_as<Policy, GravityConsistentSchurLDU>
);
m_massInverse->Mult(gradientRightHandSide, gradientAction);
m_divergence->Mult(gradientAction, m_potentialWorkspace);
m_potentialWorkspace -= potentialRightHandSide;
m_potentialSchurInverse->Mult(m_potentialWorkspace, potentialAction);
if constexpr (std::same_as<Policy, GravityConsistentSchurLDU>) {
// For [M B^T; B 0], flux recovery uses R(f-B^T phi), so its
// consistent positive Schur operator is B R B^T and its RHS
// is B R f-g. The assembled diagonal-mass Schur inverse Q
// remains a cheap correction, rather than the target operator.
// Zero sweeps deliberately retain the old LDU arithmetic.
for (int sweep = 0; sweep < m_policy.correctionSweeps; ++sweep) {
m_divergence->MultTranspose(potentialAction, m_gradientWorkspace);
m_massInverse->Mult(m_gradientWorkspace, m_massCorrection);
m_divergence->Mult(m_massCorrection, m_schurResidual);
m_schurResidual *= -1.0;
m_schurResidual += m_potentialWorkspace;
m_potentialSchurInverse->Mult(m_schurResidual, m_potentialCorrection);
potentialAction.Add(m_policy.relaxation, m_potentialCorrection);
++m_statistics.massInverseApplications;
++m_statistics.potentialSchurApplications;
++m_statistics.divergenceApplications;
++m_statistics.transposeDivergenceApplications;
++m_statistics.schurRefinementSweeps;
}
}
m_divergence->MultTranspose(potentialAction, m_gradientWorkspace);
m_massInverse->Mult(m_gradientWorkspace, m_massCorrection);
gradientAction -= m_massCorrection;
@@ -295,6 +342,14 @@ export namespace mean_field::preconditioning {
}
private:
[[nodiscard]] static int RefinementWorkspaceSize(const Policy &policy, const int potentialSize) noexcept {
if constexpr (std::same_as<Policy, GravityConsistentSchurLDU>) {
return policy.correctionSweeps > 0 ? potentialSize : 0;
} else {
return 0;
}
}
Policy m_policy;
const mfem::Solver *m_massInverse;
const mfem::Solver *m_potentialSchurInverse;
@@ -303,6 +358,8 @@ export namespace mean_field::preconditioning {
mutable mfem::Vector m_potentialWorkspace;
mutable mfem::Vector m_gradientWorkspace;
mutable mfem::Vector m_massCorrection;
mutable mfem::Vector m_schurResidual;
mutable mfem::Vector m_potentialCorrection;
mutable GravityFactorizationStatistics m_statistics;
};
@@ -474,12 +531,12 @@ export namespace mean_field::preconditioning {
Block block
)
: mfem::Solver(GravitySize(geometryContext)),
m_block(std::move(block)),
m_block(ValidateConfiguration(f, std::move(block))),
m_geometryContext(std::addressof(geometryContext)),
m_gradientMap(geometryContext.GetMassOperator().GetFluxMap()),
m_potentialMap(geometryContext.GetSourceOperator().GetPotentialMap()),
m_divergence(
geometryContext.GetDivergenceOperator(),
DivergenceForFactorization(f, geometryContext),
m_gradientMap,
m_potentialMap
),
@@ -560,13 +617,13 @@ export namespace mean_field::preconditioning {
return report;
}
m_divergence.Rebind(geometryContext.GetDivergenceOperator());
if (report.discretizationChanged) m_fixedDivergence.Clear();
m_divergence.Rebind(DivergenceForFactorization(f, geometryContext));
report.rebuiltDivergenceBinding = report.discretizationChanged;
RefreshMassInverse(geometryContext);
report.rebuiltMassInverse = true;
if (report.discretizationChanged) m_fixedDivergence.Clear();
auto potentialSchur = AssemblePotentialSchur(f, geometryContext);
m_potentialSchurInverse.Refresh(*potentialSchur);
m_potentialSchurSurrogate = std::move(potentialSchur);
@@ -614,6 +671,56 @@ export namespace mean_field::preconditioning {
}
private:
[[nodiscard]] static Block ValidateConfiguration(const fem::FEM &f, Block block) {
if constexpr (std::same_as<Policy, GravityConsistentSchurLDU>) {
if (f.mesh == nullptr) {
throw std::invalid_argument("Gravity Schur refinement requires a distributed mesh.");
}
const auto &policy = block.factorizationPolicy();
const int localValid = policy.correctionSweeps >= 0 && std::isfinite(policy.relaxation) &&
policy.relaxation > 0.0;
int globallyValid = 0;
const MPI_Comm communicator = f.mesh->GetComm();
if (MPI_Allreduce(&localValid, &globallyValid, 1, MPI_INT, MPI_MIN, communicator) != MPI_SUCCESS) {
throw std::runtime_error("Could not synchronize gravity Schur refinement validity.");
}
if (!globallyValid) {
throw std::invalid_argument(
"Gravity Schur refinement requires nonnegative sweeps and finite positive relaxation."
);
}
// Different work counts would enter collective child inverses in
// different orders; different damping would change the global
// operator. Check once before constructing either child inverse.
const double localOptions[2]{static_cast<double>(policy.correctionSweeps), policy.relaxation};
double minimumOptions[2], maximumOptions[2];
if (MPI_Allreduce(localOptions, minimumOptions, 2, MPI_DOUBLE, MPI_MIN, communicator) != MPI_SUCCESS ||
MPI_Allreduce(localOptions, maximumOptions, 2, MPI_DOUBLE, MPI_MAX, communicator) != MPI_SUCCESS) {
throw std::runtime_error("Could not synchronize gravity Schur refinement options.");
}
if (minimumOptions[0] != maximumOptions[0] || minimumOptions[1] != maximumOptions[1]) {
throw std::invalid_argument("Gravity Schur refinement options must agree on every rank.");
}
}
return block;
}
[[nodiscard]] const mfem::Operator &DivergenceForFactorization(
const fem::FEM &f,
const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext
) {
if constexpr (std::same_as<Policy, GravityConsistentSchurLDU>) {
if (m_block.factorizationPolicy().correctionSweeps > 0) {
// The Schur surrogate already uses this cached sparse
// computational-mesh divergence. Reuse it for the extra
// applications; zero sweeps retain the original arithmetic.
return m_fixedDivergence.Get(f);
}
}
return geometryContext.GetDivergenceOperator();
}
[[nodiscard]] static int
GravitySize(const operators::context::gravity_field::GravityFieldGeometryContext &geometryContext) {
if (!geometryContext.IsPrepared()) {
@@ -668,9 +775,9 @@ export namespace mean_field::preconditioning {
const operators::context::gravity_field::GravityFieldGeometryContext *m_geometryContext;
field::FieldDofMap m_gradientMap;
field::FieldDofMap m_potentialMap;
PreparedGravityDivergenceCache m_fixedDivergence;
ReducedGravityDivergenceOperator m_divergence;
PreparedMassInverse m_massInverse;
PreparedGravityDivergenceCache m_fixedDivergence;
std::unique_ptr<mfem::HypreParMatrix> m_potentialSchurSurrogate;
backend::PreparedHypreBoomerAMG<Mode> m_potentialSchurInverse;
ReducedFieldSolverAdapter m_reducedPotentialSchurInverse;