tboudreaux/MeanField
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Apply canonical gravity blocks without full root work

Browse Source
This commit is contained in:
Emily Boudreaux
2026-09-17 10:15:35 -04:00
parent b7ea1bdcf8
commit 6e812aeca2
4 changed files with 725 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

318
libmeanfield/interface/solver/gravity_subsystem.cppm
View File

@@ -19,6 +19,8 @@ module;
export module mean_field:solver.gravity_subsystem;
export import :solver.linear_backend;
import :normalization.stellar_equilibrium;
import :preconditioning.equilibrium_coordinates;
export import :utils.blocks;
export namespace mean_field::solver {
@@ -71,6 +73,8 @@ export namespace mean_field::solver {
return static_cast<int>(m_retainedResiduals.size());
}
[[nodiscard]] bool operator==(const GravityBlockPartition &) const = default;
void GatherGravityValues(
const mfem::Vector &full,
mfem::Vector &compact
@@ -258,12 +262,158 @@ export namespace mean_field::solver {
};
} // namespace mean_field::solver
export namespace mean_field::solver {
struct AutomaticGravityApplications final { };
struct RootRestrictedGravityApplications final { };
template <typename Candidate>
concept GravityApplicationPolicy = std::same_as<Candidate, AutomaticGravityApplications> ||
std::same_as<Candidate, RootRestrictedGravityApplications>;
/* An opt-in capability promises the exact physical gravity principal block
* in canonical gradient/potential order. Merely exposing similarly named
* getters does not satisfy this contract. Custom operators keep the root
* restriction unless their own specialization makes this promise. */
template <typename Operation> struct GravityCompactOperatorAccess {
static constexpr bool available = false;
};
template <typename Operation, typename Preconditioner> struct GravityCompactInverseAccess {
static constexpr bool available = false;
};
} // namespace mean_field::solver
namespace mean_field::solver::detail {
template <GravitySubsystemOperation Operation> class GravityRestrictedOperator final : public mfem::Operator {
template <typename Couplings> struct SpecificationsLeaveGravityRows;
template <typename... Couplings>
struct SpecificationsLeaveGravityRows<utils::blocks::type_list<Couplings...>>
: std::bool_constant<
((!std::same_as<typename Couplings::Residual, utils::blocks::gravity::gradient::residual> &&
!std::same_as<typename Couplings::Residual, utils::blocks::gravity::poisson::residual>) &&
...)> { };
[[nodiscard]] inline bool CollectiveGravityApplicationChoice(
const bool local,
const MPI_Comm communicator
) {
const int localChoice = local ? 1 : 0;
int globalChoice = 0;
if (MPI_Allreduce(&localChoice, &globalChoice, 1, MPI_INT, MPI_MIN, communicator) != MPI_SUCCESS) {
throw std::runtime_error("Could not select a collective gravity application path.");
}
return globalChoice != 0;
}
} // namespace mean_field::solver::detail
export namespace mean_field::solver {
template <typename Problem>
requires std::same_as<typename Problem::PhysicalCoreType, operators::PreparedStellarEquilibriumOperator> &&
detail::SpecificationsLeaveGravityRows<typename operators::CompiledStellarEquilibriumSystem<
typename Problem::ModelType>::ContributionJacobianCouplings>::value
struct GravityCompactOperatorAccess<normalization::NormalizedStellarEquilibriumOperator<Problem>> {
using Operation = normalization::NormalizedStellarEquilibriumOperator<Problem>;
static constexpr bool available = true;
[[nodiscard]] static GravityBlockPartition CanonicalPartition(const Operation &operation) {
return MakeGravityBlockPartition(operation.GetProblem().GetManifest().layout());
}
[[nodiscard]] static const mfem::Vector &StateFactors(const Operation &operation) {
return operation.GetNormalization().StateFactors();
}
[[nodiscard]] static const mfem::Vector &ResidualFactors(const Operation &operation) {
return operation.GetNormalization().ResidualFactors();
}
static void ApplyPhysical(
const Operation &operation,
const mfem::Vector &direction,
mfem::Vector &action
) {
const auto &physical = operation.GetProblem().GetPhysicalOperator();
const auto &geometry = physical.GetGravityContext().GetGeometryContext();
const int gradientSize = physical.GetGravityContext().GetGravityGradientMap().reduced_size();
const int potentialSize = physical.GetGravityContext().GetGravityPotentialMap().reduced_size();
if (direction.Size() != gradientSize + potentialSize || action.Size() != direction.Size()) {
throw std::invalid_argument("Compact gravity does not match the physical field maps.");
}
const mfem::Vector gradient(const_cast<mfem::real_t *>(direction.GetData()), gradientSize);
auto *potentialData = const_cast<mfem::real_t *>(direction.GetData());
if (gradientSize != 0)
potentialData += gradientSize;
const mfem::Vector potential(potentialData, potentialSize);
// Gravity is affine in these unknowns at fixed geometry and density.
// Reusing its residual action retains boundary maps and signs while
// avoiding the zero material/geometry variations in the full Jacobian.
physical.GetGravityOperator().ApplyGravityUnknowns(gradient, potential, geometry, action);
}
};
template <typename Problem, typename Block>
requires GravityCompactOperatorAccess<
normalization::NormalizedStellarEquilibriumOperator<Problem>>::available &&
std::same_as<
typename Block::StructureComponent::Factorization,
preconditioning::IndependentStellarSubsystems> &&
std::same_as<
typename Block::StructureComponent::GravityComponent::CorrectionBlocks,
utils::blocks::
type_list<utils::blocks::gravity::gradient::value, utils::blocks::gravity::poisson::value>> &&
std::same_as<
typename Block::StructureComponent::GravityComponent::ResidualBlocks,
utils::blocks::type_list<
utils::blocks::gravity::gradient::residual,
utils::blocks::gravity::poisson::residual>>
struct GravityCompactInverseAccess<
normalization::NormalizedStellarEquilibriumOperator<Problem>,
normalization::
NormalizedStellarPreconditioner<Problem, preconditioning::PreparedStellarPreconditioner<Problem, Block>>> {
using Operation = normalization::NormalizedStellarEquilibriumOperator<Problem>;
using Preconditioner = normalization::
NormalizedStellarPreconditioner<Problem, preconditioning::PreparedStellarPreconditioner<Problem, Block>>;
using GravityBlock = typename Block::StructureComponent::GravityComponent;
static constexpr bool available = true;
static constexpr auto applicationContract =
preconditioning::backend::Traits<typename GravityBlock::BackendType>::applicationContract;
[[nodiscard]] static const auto &PhysicalInverse(const Preconditioner &inverse) {
return inverse.GetPhysicalInverse()
.GetGroupedPreconditioner()
.GetStructurePreconditioner()
.GetGravityPreconditioner();
}
[[nodiscard]] static bool Matches(
const Operation &operation,
const Preconditioner &inverse
) {
const auto &gravity = PhysicalInverse(inverse);
const auto canonical = GravityCompactOperatorAccess<Operation>::CanonicalPartition(operation);
return std::addressof(inverse.GetNormalizedJacobian()) == std::addressof(operation) &&
gravity.Height() == canonical.GravityValueSize() &&
gravity.Width() == canonical.GravityResidualSize();
}
[[nodiscard]] static bool IsCurrent(const Preconditioner &inverse) {
return PhysicalInverse(inverse).IsCurrent();
}
static void ApplyPhysical(
const Preconditioner &inverse,
const mfem::Vector &residual,
mfem::Vector &correction
) {
PhysicalInverse(inverse).Mult(residual, correction);
}
};
} // namespace mean_field::solver
namespace mean_field::solver::detail {
template <GravitySubsystemOperation Operation, GravityApplicationPolicy Applications>
class GravityRestrictedOperator final : public mfem::Operator {
public:
using Compact = GravityCompactOperatorAccess<Operation>;
static constexpr bool canCompact =
std::same_as<Applications, AutomaticGravityApplications> && Compact::available;
GravityRestrictedOperator(
const Operation &operation,
const GravityBlockPartition &partition
const GravityBlockPartition &partition,
MPI_Comm communicator
)
: mfem::Operator(
partition.GravityResidualSize(),
@@ -272,10 +422,21 @@ namespace mean_field::solver::detail {
m_operation(std::addressof(operation)),
m_partition(std::addressof(partition)),
m_fullDirection(partition.RootValueSize()),
m_fullAction(partition.RootResidualSize()) {
m_fullAction(partition.RootResidualSize()),
m_stateFactors(Width()),
m_residualFactors(Height()),
m_physicalDirection(Width()),
m_physicalAction(Height()) {
if (operation.Width() != partition.RootValueSize() || operation.Height() != partition.RootResidualSize()) {
throw std::invalid_argument("The gravity partition does not match its root operator.");
}
if constexpr (canCompact) {
// All ranks must enter the same child collectives, even if only
// one rank was given a noncanonical public partition.
m_compact = CollectiveGravityApplicationChoice(
Compact::CanonicalPartition(operation) == partition, communicator
);
}
Refresh();
}
@@ -283,12 +444,29 @@ namespace mean_field::solver::detail {
if (!m_operation->IsPrepared()) {
throw std::logic_error("Gravity requires a prepared root linearization before refresh.");
}
if constexpr (canCompact) {
if (m_compact) {
m_partition->GatherGravityValues(Compact::StateFactors(*m_operation), m_stateFactors);
m_partition->GatherGravityResiduals(Compact::ResidualFactors(*m_operation), m_residualFactors);
}
}
m_generation = m_operation->GetProblem().GetPreparationGeneration();
}
[[nodiscard]] bool IsCurrent() const {
return m_operation->IsPrepared() && m_generation == m_operation->GetProblem().GetPreparationGeneration();
}
[[nodiscard]] bool UsesCompactApplication() const noexcept {
return m_compact;
}
[[nodiscard]] const Operation &GetFullOperator() const noexcept {
return *m_operation;
}
[[nodiscard]] const mfem::Vector &StateFactors() const noexcept {
return m_stateFactors;
}
[[nodiscard]] const mfem::Vector &ResidualFactors() const noexcept {
return m_residualFactors;
}
void Mult(
const mfem::Vector &direction,
@@ -297,6 +475,19 @@ namespace mean_field::solver::detail {
if (!IsCurrent()) {
throw std::logic_error("The gravity linearization must be refreshed after preparation.");
}
if (direction.Size() != Width() || action.Size() != Height()) {
throw std::invalid_argument("Gravity actions require compatible preallocated vectors.");
}
if constexpr (canCompact) {
if (m_compact) {
for (int i = 0; i < Width(); ++i)
m_physicalDirection(i) = direction(i) / m_stateFactors(i);
Compact::ApplyPhysical(*m_operation, m_physicalDirection, m_physicalAction);
for (int i = 0; i < Height(); ++i)
action(i) = m_residualFactors(i) * m_physicalAction(i);
return;
}
}
m_partition->InjectGravityValues(direction, m_fullDirection);
m_operation->Mult(m_fullDirection, m_fullAction);
m_partition->GatherGravityResiduals(m_fullAction, action);
@@ -306,50 +497,89 @@ namespace mean_field::solver::detail {
const Operation *m_operation;
const GravityBlockPartition *m_partition;
std::uint64_t m_generation{0};
mutable mfem::Vector m_fullDirection;
mutable mfem::Vector m_fullAction;
bool m_compact{false};
mutable mfem::Vector m_fullDirection, m_fullAction;
mfem::Vector m_stateFactors, m_residualFactors;
mutable mfem::Vector m_physicalDirection, m_physicalAction;
};
template <GravitySubsystemOperation Operation, GravitySubsystemPreconditioner Preconditioner>
template <
GravitySubsystemOperation Operation,
GravitySubsystemPreconditioner Preconditioner,
GravityApplicationPolicy Applications>
class GravityRestrictedInverse final : public mfem::Solver {
public:
static constexpr auto applicationContract = linearPreconditionerApplicationContract<Preconditioner>;
using Compact = GravityCompactInverseAccess<Operation, Preconditioner>;
using RestrictedOperator = GravityRestrictedOperator<Operation, Applications>;
static constexpr bool canCompact =
std::same_as<Applications, AutomaticGravityApplications> && Compact::available;
static constexpr auto applicationContract = [] {
if constexpr (canCompact) {
if constexpr (Compact::applicationContract == preconditioning::ApplicationContract::flexible)
return preconditioning::ApplicationContract::flexible;
}
return linearPreconditionerApplicationContract<Preconditioner>;
}();
GravityRestrictedInverse(
const Preconditioner &inverse,
const GravityRestrictedOperator<Operation> &operation,
const GravityBlockPartition &partition
const RestrictedOperator &operation,
const GravityBlockPartition &partition,
MPI_Comm communicator
)
: mfem::Solver(partition.GravityValueSize()),
m_inverse(std::addressof(inverse)),
m_operation(std::addressof(operation)),
m_partition(std::addressof(partition)),
m_fullResidual(partition.RootResidualSize()),
m_fullCorrection(partition.RootValueSize()) {
m_fullCorrection(partition.RootValueSize()),
m_physicalResidual(Width()),
m_physicalCorrection(Height()) {
if (inverse.Width() != partition.RootResidualSize() || inverse.Height() != partition.RootValueSize()) {
throw std::invalid_argument("The gravity partition does not match its root inverse.");
}
}
[[nodiscard]] bool IsCurrent() const {
return m_operation->IsCurrent() && m_inverse->IsCurrent();
}
void SetOperator(const mfem::Operator &operation) override {
if (std::addressof(operation) != m_operation) {
throw std::invalid_argument(
"A prepared gravity inverse cannot be "
"rebound to a different operator."
if constexpr (canCompact) {
m_compact = CollectiveGravityApplicationChoice(
operation.UsesCompactApplication() && Compact::Matches(operation.GetFullOperator(), inverse),
communicator
);
}
}
[[nodiscard]] bool IsCurrent() const {
if constexpr (canCompact) {
if (m_compact && !Compact::IsCurrent(*m_inverse))
return false;
}
return m_operation->IsCurrent() && m_inverse->IsCurrent();
}
[[nodiscard]] bool UsesCompactApplication() const noexcept {
return m_compact;
}
void SetOperator(const mfem::Operator &operation) override {
if (std::addressof(operation) != m_operation) {
throw std::invalid_argument("A prepared gravity inverse cannot be rebound to a different operator.");
}
}
void Mult(
const mfem::Vector &residual,
mfem::Vector &correction
) const override {
if (!IsCurrent()) {
if (!IsCurrent())
throw std::logic_error("The gravity inverse requires a current root inverse.");
if (residual.Size() != Width() || correction.Size() != Height()) {
throw std::invalid_argument("Gravity inverse actions require compatible preallocated vectors.");
}
if constexpr (canCompact) {
if (m_compact) {
const auto &state = m_operation->StateFactors();
const auto &rows = m_operation->ResidualFactors();
for (int i = 0; i < Width(); ++i)
m_physicalResidual(i) = residual(i) / rows(i);
Compact::ApplyPhysical(*m_inverse, m_physicalResidual, m_physicalCorrection);
for (int i = 0; i < Height(); ++i)
correction(i) = state(i) * m_physicalCorrection(i);
return;
}
}
m_partition->InjectGravityResiduals(residual, m_fullResidual);
m_inverse->Mult(m_fullResidual, m_fullCorrection);
@@ -358,32 +588,34 @@ namespace mean_field::solver::detail {
private:
const Preconditioner *m_inverse;
const GravityRestrictedOperator<Operation> *m_operation;
const RestrictedOperator *m_operation;
const GravityBlockPartition *m_partition;
mutable mfem::Vector m_fullResidual;
mutable mfem::Vector m_fullCorrection;
bool m_compact{false};
mutable mfem::Vector m_fullResidual, m_fullCorrection, m_physicalResidual, m_physicalCorrection;
};
} // namespace mean_field::solver::detail
export namespace mean_field::solver {
/* Restricting the supplied inverse is a preconditioner, not an assertion
* that its principal block equals G_z^{-1}. The backend still verifies the
* true gravity defect. Keeping the exact root operator as the source of
* every action also retains its boundary rows and normalization. */
* true gravity defect. Registered compact applications retain the same
* boundary rows and diagonal coordinate maps; unknown recipes retain
* the complete root restriction. */
template <
GravitySubsystemOperation Operation,
GravitySubsystemPreconditioner Preconditioner,
LinearBackendConfiguration Backend = linear::FGMRES>
LinearBackendConfiguration Backend = linear::FGMRES,
GravityApplicationPolicy Applications = AutomaticGravityApplications>
requires LinearBackendRuntimeAvailableFor<
Backend,
detail::GravityRestrictedOperator<Operation>,
detail::GravityRestrictedInverse<Operation, Preconditioner>>
detail::GravityRestrictedOperator<Operation, Applications>,
detail::GravityRestrictedInverse<Operation, Preconditioner, Applications>>
class PreparedGravitySubsystem final {
public:
using OperatorType = Operation;
using PreconditionerType = Preconditioner;
using RestrictedOperator = detail::GravityRestrictedOperator<Operation>;
using RestrictedInverse = detail::GravityRestrictedInverse<Operation, Preconditioner>;
using RestrictedOperator = detail::GravityRestrictedOperator<Operation, Applications>;
using RestrictedInverse = detail::GravityRestrictedInverse<Operation, Preconditioner, Applications>;
using BackendType = PreparedLinearBackendType<Backend, RestrictedOperator, RestrictedInverse>;
PreparedGravitySubsystem(
@@ -391,19 +623,22 @@ export namespace mean_field::solver {
const Preconditioner &preconditioner,
GravityBlockPartition partition,
const MPI_Comm communicator,
Backend backend = {}
Backend backend = {},
Applications = {}
)
: m_fullOperator(std::addressof(operation)),
m_fullPreconditioner(std::addressof(preconditioner)),
m_partition(std::move(partition)),
m_operator(
operation,
m_partition
m_partition,
communicator
),
m_inverse(
preconditioner,
m_operator,
m_partition
m_partition,
communicator
),
m_backend(prepareLinearBackend(
std::move(backend),
@@ -437,6 +672,15 @@ export namespace mean_field::solver {
[[nodiscard]] const RestrictedOperator &GetOperator() const noexcept {
return m_operator;
}
[[nodiscard]] const RestrictedInverse &GetPreconditioner() const noexcept {
return m_inverse;
}
[[nodiscard]] bool UsesCompactOperator() const noexcept {
return m_operator.UsesCompactApplication();
}
[[nodiscard]] bool UsesCompactPreconditioner() const noexcept {
return m_inverse.UsesCompactApplication();
}
[[nodiscard]] const Operation &GetFullOperator() const noexcept {
return *m_fullOperator;
}