894 lines
40 KiB
C++
894 lines
40 KiB
C++
module;
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <concepts>
|
|
#include <cstdint>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <span>
|
|
#include <stdexcept>
|
|
#include <type_traits>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include <mfem.hpp>
|
|
#include <mpi.h>
|
|
|
|
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 {
|
|
/* The maps are local true-dof maps. Gravity and retained coordinates need
|
|
* not be contiguous or share the residual ordering. Only the discretized
|
|
* sizes are runtime data; the production factory selects block identities
|
|
* from the compiled equation form. */
|
|
class GravityBlockPartition final {
|
|
public:
|
|
GravityBlockPartition(
|
|
const int rootValueSize,
|
|
const int rootResidualSize,
|
|
std::vector<int> gravityValues,
|
|
std::vector<int> gravityResiduals
|
|
)
|
|
: m_rootValueSize(rootValueSize),
|
|
m_rootResidualSize(rootResidualSize),
|
|
m_gravityValues(std::move(gravityValues)),
|
|
m_gravityResiduals(std::move(gravityResiduals)),
|
|
m_retainedValues(Complement(
|
|
rootValueSize,
|
|
m_gravityValues
|
|
)),
|
|
m_retainedResiduals(Complement(
|
|
rootResidualSize,
|
|
m_gravityResiduals
|
|
)) {
|
|
if (m_gravityValues.size() != m_gravityResiduals.size() ||
|
|
m_retainedValues.size() != m_retainedResiduals.size()) {
|
|
throw std::invalid_argument("Gravity elimination requires square gravity and retained blocks.");
|
|
}
|
|
}
|
|
|
|
[[nodiscard]] int RootValueSize() const noexcept {
|
|
return m_rootValueSize;
|
|
}
|
|
[[nodiscard]] int RootResidualSize() const noexcept {
|
|
return m_rootResidualSize;
|
|
}
|
|
[[nodiscard]] int GravityValueSize() const noexcept {
|
|
return static_cast<int>(m_gravityValues.size());
|
|
}
|
|
[[nodiscard]] int GravityResidualSize() const noexcept {
|
|
return static_cast<int>(m_gravityResiduals.size());
|
|
}
|
|
[[nodiscard]] int RetainedValueSize() const noexcept {
|
|
return static_cast<int>(m_retainedValues.size());
|
|
}
|
|
[[nodiscard]] int RetainedResidualSize() const noexcept {
|
|
return static_cast<int>(m_retainedResiduals.size());
|
|
}
|
|
|
|
[[nodiscard]] bool operator==(const GravityBlockPartition &) const = default;
|
|
|
|
void GatherGravityValues(
|
|
const mfem::Vector &full,
|
|
mfem::Vector &compact
|
|
) const {
|
|
Gather(m_gravityValues, m_rootValueSize, full, compact);
|
|
}
|
|
void GatherGravityResiduals(
|
|
const mfem::Vector &full,
|
|
mfem::Vector &compact
|
|
) const {
|
|
Gather(m_gravityResiduals, m_rootResidualSize, full, compact);
|
|
}
|
|
void GatherRetainedValues(
|
|
const mfem::Vector &full,
|
|
mfem::Vector &compact
|
|
) const {
|
|
Gather(m_retainedValues, m_rootValueSize, full, compact);
|
|
}
|
|
void GatherRetainedResiduals(
|
|
const mfem::Vector &full,
|
|
mfem::Vector &compact
|
|
) const {
|
|
Gather(m_retainedResiduals, m_rootResidualSize, full, compact);
|
|
}
|
|
void InjectGravityValues(
|
|
const mfem::Vector &compact,
|
|
mfem::Vector &full
|
|
) const {
|
|
Inject(m_gravityValues, m_rootValueSize, compact, full);
|
|
}
|
|
void InjectGravityResiduals(
|
|
const mfem::Vector &compact,
|
|
mfem::Vector &full
|
|
) const {
|
|
Inject(m_gravityResiduals, m_rootResidualSize, compact, full);
|
|
}
|
|
void InjectRetainedValues(
|
|
const mfem::Vector &compact,
|
|
mfem::Vector &full
|
|
) const {
|
|
Inject(m_retainedValues, m_rootValueSize, compact, full);
|
|
}
|
|
void InjectRetainedResiduals(
|
|
const mfem::Vector &compact,
|
|
mfem::Vector &full
|
|
) const {
|
|
Inject(m_retainedResiduals, m_rootResidualSize, compact, full);
|
|
}
|
|
void AddGravityValues(
|
|
const mfem::Vector &compact,
|
|
mfem::Vector &full
|
|
) const {
|
|
VerifySizes(m_gravityValues, m_rootValueSize, full, compact);
|
|
for (std::size_t index = 0; index < m_gravityValues.size(); ++index) {
|
|
full(m_gravityValues[index]) += compact(static_cast<int>(index));
|
|
}
|
|
}
|
|
|
|
private:
|
|
[[nodiscard]] static std::vector<int> Complement(
|
|
const int size,
|
|
const std::vector<int> &selected
|
|
) {
|
|
if (size < 0) {
|
|
throw std::invalid_argument("A gravity partition requires non-negative root dimensions.");
|
|
}
|
|
std::vector<bool> used(static_cast<std::size_t>(size), false);
|
|
for (const int index : selected) {
|
|
if (index < 0 || index >= size || used[static_cast<std::size_t>(index)]) {
|
|
throw std::invalid_argument("Gravity partition indices must be unique and in range.");
|
|
}
|
|
used[static_cast<std::size_t>(index)] = true;
|
|
}
|
|
std::vector<int> complement;
|
|
complement.reserve(static_cast<std::size_t>(size) - selected.size());
|
|
for (int index = 0; index < size; ++index) {
|
|
if (!used[static_cast<std::size_t>(index)]) {
|
|
complement.push_back(index);
|
|
}
|
|
}
|
|
return complement;
|
|
}
|
|
|
|
static void VerifySizes(
|
|
const std::vector<int> &indices,
|
|
const int fullSize,
|
|
const mfem::Vector &full,
|
|
const mfem::Vector &compact
|
|
) {
|
|
if (full.Size() != fullSize || compact.Size() != static_cast<int>(indices.size())) {
|
|
throw std::invalid_argument("Gravity partition vectors must have their preallocated sizes.");
|
|
}
|
|
const std::less<const mfem::real_t *> before;
|
|
if (compact.Size() != 0 && full.Size() != 0 && before(full.GetData(), compact.GetData() + compact.Size()) &&
|
|
before(compact.GetData(), full.GetData() + full.Size())) {
|
|
throw std::invalid_argument("Gravity gather and injection require distinct vector storage.");
|
|
}
|
|
}
|
|
|
|
static void Gather(
|
|
const std::vector<int> &indices,
|
|
const int fullSize,
|
|
const mfem::Vector &full,
|
|
mfem::Vector &compact
|
|
) {
|
|
VerifySizes(indices, fullSize, full, compact);
|
|
for (std::size_t index = 0; index < indices.size(); ++index) {
|
|
compact(static_cast<int>(index)) = full(indices[index]);
|
|
}
|
|
}
|
|
|
|
static void Inject(
|
|
const std::vector<int> &indices,
|
|
const int fullSize,
|
|
const mfem::Vector &compact,
|
|
mfem::Vector &full
|
|
) {
|
|
VerifySizes(indices, fullSize, full, compact);
|
|
full = 0.0;
|
|
for (std::size_t index = 0; index < indices.size(); ++index) {
|
|
full(indices[index]) = compact(static_cast<int>(index));
|
|
}
|
|
}
|
|
|
|
int m_rootValueSize;
|
|
int m_rootResidualSize;
|
|
std::vector<int> m_gravityValues;
|
|
std::vector<int> m_gravityResiduals;
|
|
std::vector<int> m_retainedValues;
|
|
std::vector<int> m_retainedResiduals;
|
|
};
|
|
|
|
template <typename Form>
|
|
concept GravityPartitionForm = requires {
|
|
typename Form::value_blocks;
|
|
typename Form::residual_blocks;
|
|
requires utils::blocks::block_form_is_valid_v<Form>;
|
|
requires utils::blocks::contains_type_v<utils::blocks::gravity::gradient::value, typename Form::value_blocks>;
|
|
requires utils::blocks::contains_type_v<utils::blocks::gravity::poisson::value, typename Form::value_blocks>;
|
|
requires utils::blocks::contains_type_v<
|
|
utils::blocks::gravity::gradient::residual, typename Form::residual_blocks>;
|
|
requires utils::blocks::contains_type_v<
|
|
utils::blocks::gravity::poisson::residual, typename Form::residual_blocks>;
|
|
};
|
|
|
|
template <GravityPartitionForm Form>
|
|
[[nodiscard]] GravityBlockPartition MakeGravityBlockPartition(const utils::blocks::form_layout<Form> &layout) {
|
|
using namespace utils::blocks;
|
|
std::vector<int> values;
|
|
std::vector<int> residuals;
|
|
const auto append = [](std::vector<int> &indices, const int offset, const int size) {
|
|
for (int index = 0; index < size; ++index) {
|
|
indices.push_back(offset + index);
|
|
}
|
|
};
|
|
const auto value = [&]<typename Block>() {
|
|
constexpr auto block = value_block<type_index_v<Block, typename Form::value_blocks>>{};
|
|
append(values, layout.offset(block), layout.size(block));
|
|
};
|
|
const auto residual = [&]<typename Block>() {
|
|
constexpr auto block = residual_block<type_index_v<Block, typename Form::residual_blocks>>{};
|
|
append(residuals, layout.offset(block), layout.size(block));
|
|
};
|
|
value.template operator()<gravity::gradient::value>();
|
|
value.template operator()<gravity::poisson::value>();
|
|
residual.template operator()<gravity::gradient::residual>();
|
|
residual.template operator()<gravity::poisson::residual>();
|
|
return {
|
|
layout.value_offsets().Last(), layout.residual_offsets().Last(), std::move(values), std::move(residuals)
|
|
};
|
|
}
|
|
|
|
template <typename Operation>
|
|
concept GravitySubsystemOperation =
|
|
std::derived_from<Operation, mfem::Operator> && requires(const Operation &operation) {
|
|
{ operation.IsPrepared() } -> std::same_as<bool>;
|
|
{ operation.GetProblem().GetPreparationGeneration() } -> std::convertible_to<std::uint64_t>;
|
|
};
|
|
|
|
template <typename Preconditioner>
|
|
concept GravitySubsystemPreconditioner =
|
|
std::derived_from<Preconditioner, mfem::Solver> &&
|
|
LinearPreconditionerApplicationContractAvailable<Preconditioner> && requires(const Preconditioner &inverse) {
|
|
{ inverse.IsCurrent() } -> std::same_as<bool>;
|
|
};
|
|
} // 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 <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,
|
|
MPI_Comm communicator
|
|
)
|
|
: mfem::Operator(
|
|
partition.GravityResidualSize(),
|
|
partition.GravityValueSize()
|
|
),
|
|
m_operation(std::addressof(operation)),
|
|
m_partition(std::addressof(partition)),
|
|
m_fullDirection(partition.RootValueSize()),
|
|
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();
|
|
}
|
|
|
|
void Refresh() {
|
|
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,
|
|
mfem::Vector &action
|
|
) const override {
|
|
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);
|
|
}
|
|
|
|
private:
|
|
const Operation *m_operation;
|
|
const GravityBlockPartition *m_partition;
|
|
std::uint64_t m_generation{0};
|
|
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,
|
|
GravityApplicationPolicy Applications>
|
|
class GravityRestrictedInverse final : public mfem::Solver {
|
|
public:
|
|
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 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_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.");
|
|
}
|
|
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())
|
|
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);
|
|
m_partition->GatherGravityValues(m_fullCorrection, correction);
|
|
}
|
|
|
|
private:
|
|
const Preconditioner *m_inverse;
|
|
const RestrictedOperator *m_operation;
|
|
const GravityBlockPartition *m_partition;
|
|
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. 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,
|
|
GravityApplicationPolicy Applications = AutomaticGravityApplications>
|
|
requires LinearBackendRuntimeAvailableFor<
|
|
Backend,
|
|
detail::GravityRestrictedOperator<Operation, Applications>,
|
|
detail::GravityRestrictedInverse<Operation, Preconditioner, Applications>>
|
|
class PreparedGravitySubsystem final {
|
|
public:
|
|
using OperatorType = Operation;
|
|
using PreconditionerType = Preconditioner;
|
|
using RestrictedOperator = detail::GravityRestrictedOperator<Operation, Applications>;
|
|
using RestrictedInverse = detail::GravityRestrictedInverse<Operation, Preconditioner, Applications>;
|
|
using BackendType = PreparedLinearBackendType<Backend, RestrictedOperator, RestrictedInverse>;
|
|
|
|
PreparedGravitySubsystem(
|
|
const Operation &operation,
|
|
const Preconditioner &preconditioner,
|
|
GravityBlockPartition partition,
|
|
const MPI_Comm communicator,
|
|
Backend backend = {},
|
|
Applications = {}
|
|
)
|
|
: m_fullOperator(std::addressof(operation)),
|
|
m_fullPreconditioner(std::addressof(preconditioner)),
|
|
m_partition(std::move(partition)),
|
|
m_operator(
|
|
operation,
|
|
m_partition,
|
|
communicator
|
|
),
|
|
m_inverse(
|
|
preconditioner,
|
|
m_operator,
|
|
m_partition,
|
|
communicator
|
|
),
|
|
m_backend(prepareLinearBackend(
|
|
std::move(backend),
|
|
m_operator,
|
|
m_inverse,
|
|
communicator
|
|
)) {
|
|
if (!IsCurrent()) {
|
|
throw std::logic_error("Gravity subsystem preparation requires a current inverse.");
|
|
}
|
|
}
|
|
|
|
PreparedGravitySubsystem(const PreparedGravitySubsystem &) = delete;
|
|
PreparedGravitySubsystem &operator=(const PreparedGravitySubsystem &) = delete;
|
|
PreparedGravitySubsystem(PreparedGravitySubsystem &&) = delete;
|
|
PreparedGravitySubsystem &operator=(PreparedGravitySubsystem &&) = delete;
|
|
|
|
void Refresh() {
|
|
if (!m_fullPreconditioner->IsCurrent()) {
|
|
throw std::logic_error("Refresh the root inverse before refreshing its gravity subsystem.");
|
|
}
|
|
m_operator.Refresh();
|
|
}
|
|
|
|
[[nodiscard]] bool IsCurrent() const {
|
|
return m_operator.IsCurrent() && m_inverse.IsCurrent() && m_backend.IsReady();
|
|
}
|
|
[[nodiscard]] const GravityBlockPartition &GetPartition() const noexcept {
|
|
return m_partition;
|
|
}
|
|
[[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;
|
|
}
|
|
[[nodiscard]] const Preconditioner &GetFullPreconditioner() const noexcept {
|
|
return *m_fullPreconditioner;
|
|
}
|
|
[[nodiscard]] MPI_Comm GetCommunicator() const noexcept {
|
|
return m_backend.GetCommunicator();
|
|
}
|
|
[[nodiscard]] std::uint64_t GetSolveCount() const noexcept {
|
|
return m_solveCount;
|
|
}
|
|
[[nodiscard]] std::uint64_t GetIterationCount() const noexcept {
|
|
return m_iterationCount;
|
|
}
|
|
[[nodiscard]] std::uint64_t GetFailureCount() const noexcept {
|
|
return m_failureCount;
|
|
}
|
|
[[nodiscard]] double GetSolveSeconds() const noexcept {
|
|
return m_solveSeconds;
|
|
}
|
|
[[nodiscard]] const std::optional<LinearSolveReport> &GetLastSolveReport() const noexcept {
|
|
return m_lastSolve;
|
|
}
|
|
|
|
[[nodiscard]] LinearSolveReport Solve(
|
|
const mfem::Vector &rightHandSide,
|
|
mfem::Vector &correction,
|
|
const LinearSolveControl &control
|
|
) {
|
|
++m_solveCount;
|
|
try {
|
|
const auto report = m_backend.Solve(rightHandSide, correction, control);
|
|
m_iterationCount += static_cast<std::uint64_t>(report.iterations);
|
|
m_solveSeconds += report.solveSeconds;
|
|
if (!report.Converged()) {
|
|
++m_failureCount;
|
|
}
|
|
m_lastSolve = report;
|
|
return report;
|
|
} catch (...) {
|
|
++m_failureCount;
|
|
throw;
|
|
}
|
|
}
|
|
|
|
private:
|
|
const Operation *m_fullOperator;
|
|
const Preconditioner *m_fullPreconditioner;
|
|
GravityBlockPartition m_partition;
|
|
RestrictedOperator m_operator;
|
|
RestrictedInverse m_inverse;
|
|
BackendType m_backend;
|
|
// Times and iteration counts cover completed calls; failed contracts
|
|
// are counted without attempting another collective during unwinding.
|
|
std::uint64_t m_solveCount{0};
|
|
std::uint64_t m_iterationCount{0};
|
|
std::uint64_t m_failureCount{0};
|
|
double m_solveSeconds{0.0};
|
|
std::optional<LinearSolveReport> m_lastSolve;
|
|
};
|
|
|
|
template <typename Subsystem>
|
|
concept GravitySubsystemRuntime = requires(
|
|
Subsystem &subsystem,
|
|
const Subsystem &constant,
|
|
const mfem::Vector &rhs,
|
|
mfem::Vector &correction,
|
|
const LinearSolveControl &control
|
|
) {
|
|
{ constant.IsCurrent() } -> std::same_as<bool>;
|
|
{ constant.GetPartition() } -> std::same_as<const GravityBlockPartition &>;
|
|
{ constant.GetCommunicator() } -> std::same_as<MPI_Comm>;
|
|
{ subsystem.Solve(rhs, correction, control) } -> std::same_as<LinearSolveReport>;
|
|
{ constant.GetFullOperator().Mult(rhs, correction) } -> std::same_as<void>;
|
|
{ constant.GetFullOperator().GetProblem().GetPreparationGeneration() } -> std::convertible_to<std::uint64_t>;
|
|
};
|
|
|
|
template <GravitySubsystemRuntime Subsystem> class PreparedGravityCurvatureCorrection final {
|
|
public:
|
|
explicit PreparedGravityCurvatureCorrection(Subsystem &subsystem)
|
|
: m_subsystem(std::addressof(subsystem)),
|
|
m_baseResidual(subsystem.GetPartition().GravityResidualSize()),
|
|
m_baseAction(subsystem.GetPartition().GravityResidualSize()),
|
|
m_fullAction(subsystem.GetPartition().RootResidualSize()),
|
|
m_rightHandSide(subsystem.GetPartition().GravityResidualSize()) {
|
|
}
|
|
|
|
PreparedGravityCurvatureCorrection(const PreparedGravityCurvatureCorrection &) = delete;
|
|
PreparedGravityCurvatureCorrection &operator=(const PreparedGravityCurvatureCorrection &) = delete;
|
|
PreparedGravityCurvatureCorrection(PreparedGravityCurvatureCorrection &&) = delete;
|
|
PreparedGravityCurvatureCorrection &operator=(PreparedGravityCurvatureCorrection &&) = delete;
|
|
|
|
// Capture must precede trial preparation: using a candidate Jacobian
|
|
// here would introduce a first-order change into the Newton tangent.
|
|
void CaptureBase(
|
|
const mfem::Vector &baseResidual,
|
|
const mfem::Vector &baseDirection
|
|
) {
|
|
m_hasBase = false;
|
|
const auto &partition = m_subsystem->GetPartition();
|
|
RequireCollectively(
|
|
m_subsystem->IsCurrent() && baseResidual.Size() == partition.RootResidualSize() &&
|
|
baseDirection.Size() == partition.RootValueSize(),
|
|
"Gravity curvature capture requires a current base and "
|
|
"complete vectors."
|
|
);
|
|
partition.GatherGravityResiduals(baseResidual, m_baseResidual);
|
|
m_subsystem->GetFullOperator().Mult(baseDirection, m_fullAction);
|
|
partition.GatherGravityResiduals(m_fullAction, m_baseAction);
|
|
m_baseGeneration = m_subsystem->GetFullOperator().GetProblem().GetPreparationGeneration();
|
|
m_hasBase = true;
|
|
}
|
|
|
|
void Invalidate() noexcept {
|
|
m_hasBase = false;
|
|
}
|
|
[[nodiscard]] bool HasBase() const noexcept {
|
|
return m_hasBase;
|
|
}
|
|
[[nodiscard]] std::uint64_t GetBaseGeneration() const noexcept {
|
|
return m_baseGeneration;
|
|
}
|
|
|
|
void BuildLinearPrediction(
|
|
const double alpha,
|
|
mfem::Vector &gravityResidual
|
|
) const {
|
|
RequireCollectively(
|
|
m_hasBase && gravityResidual.Size() == m_baseResidual.Size(),
|
|
"A gravity prediction requires a captured base and a preallocated gravity residual."
|
|
);
|
|
RequireStepLength(alpha);
|
|
bool finite = true;
|
|
for (int index = 0; index < gravityResidual.Size(); ++index) {
|
|
gravityResidual(index) = m_baseResidual(index) + alpha * m_baseAction(index);
|
|
finite = finite && std::isfinite(gravityResidual(index));
|
|
}
|
|
RequireCollectively(finite, "The captured gravity linear prediction is nonfinite.");
|
|
}
|
|
|
|
// The caller owns candidate preparation and commits only after a new
|
|
// complete residual evaluation. A failed inner solve never modifies
|
|
// an accepted state; its output remains a diagnostic correction only.
|
|
[[nodiscard]] LinearSolveReport SolveCorrection(
|
|
const double alpha,
|
|
const mfem::Vector &trialResidual,
|
|
mfem::Vector &gravityCorrection,
|
|
const LinearSolveControl &control
|
|
) {
|
|
const auto &partition = m_subsystem->GetPartition();
|
|
RequireCollectively(
|
|
m_hasBase && m_subsystem->IsCurrent() && trialResidual.Size() == partition.RootResidualSize() &&
|
|
gravityCorrection.Size() == partition.GravityValueSize(),
|
|
"Gravity curvature correction requires a captured base, a current "
|
|
"candidate, and valid vectors."
|
|
);
|
|
RequireStepLength(alpha);
|
|
partition.GatherGravityResiduals(trialResidual, m_rightHandSide);
|
|
for (int index = 0; index < m_rightHandSide.Size(); ++index) {
|
|
m_rightHandSide(index) = m_baseResidual(index) + alpha * m_baseAction(index) - m_rightHandSide(index);
|
|
}
|
|
gravityCorrection = 0.0;
|
|
return m_subsystem->Solve(m_rightHandSide, gravityCorrection, control);
|
|
}
|
|
|
|
private:
|
|
void RequireStepLength(const double alpha) const {
|
|
RequireCollectively(
|
|
std::isfinite(alpha) && alpha >= 0.0, "Gravity curvature requires a finite, non-negative step length."
|
|
);
|
|
double minimumAlpha = 0.0;
|
|
double maximumAlpha = 0.0;
|
|
if (MPI_Allreduce(&alpha, &minimumAlpha, 1, MPI_DOUBLE, MPI_MIN, m_subsystem->GetCommunicator()) !=
|
|
MPI_SUCCESS ||
|
|
MPI_Allreduce(&alpha, &maximumAlpha, 1, MPI_DOUBLE, MPI_MAX, m_subsystem->GetCommunicator()) !=
|
|
MPI_SUCCESS) {
|
|
throw std::runtime_error("Gravity curvature could not verify its distributed step length.");
|
|
}
|
|
if (minimumAlpha != maximumAlpha) {
|
|
throw std::invalid_argument("Gravity curvature requires the same step length on every rank.");
|
|
}
|
|
}
|
|
|
|
void RequireCollectively(
|
|
const bool valid,
|
|
const char *message
|
|
) const {
|
|
int local = valid ? 1 : 0;
|
|
int global = 0;
|
|
if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, m_subsystem->GetCommunicator()) != MPI_SUCCESS) {
|
|
throw std::runtime_error(
|
|
"Gravity curvature correction could not verify "
|
|
"distributed validity."
|
|
);
|
|
}
|
|
if (global == 0) {
|
|
throw std::logic_error(message);
|
|
}
|
|
}
|
|
|
|
Subsystem *m_subsystem;
|
|
mfem::Vector m_baseResidual;
|
|
mfem::Vector m_baseAction;
|
|
mfem::Vector m_fullAction;
|
|
mfem::Vector m_rightHandSide;
|
|
std::uint64_t m_baseGeneration{0};
|
|
bool m_hasBase{false};
|
|
};
|
|
} // namespace mean_field::solver
|