633 lines
26 KiB
Plaintext
633 lines
26 KiB
Plaintext
|
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;
|
||
|
|
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());
|
||
|
|
}
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
namespace mean_field::solver::detail {
|
||
|
|
template <GravitySubsystemOperation Operation> class GravityRestrictedOperator final : public mfem::Operator {
|
||
|
|
public:
|
||
|
|
GravityRestrictedOperator(
|
||
|
|
const Operation &operation,
|
||
|
|
const GravityBlockPartition &partition
|
||
|
|
)
|
||
|
|
: mfem::Operator(
|
||
|
|
partition.GravityResidualSize(),
|
||
|
|
partition.GravityValueSize()
|
||
|
|
),
|
||
|
|
m_operation(std::addressof(operation)),
|
||
|
|
m_partition(std::addressof(partition)),
|
||
|
|
m_fullDirection(partition.RootValueSize()),
|
||
|
|
m_fullAction(partition.RootResidualSize()) {
|
||
|
|
if (operation.Width() != partition.RootValueSize() || operation.Height() != partition.RootResidualSize()) {
|
||
|
|
throw std::invalid_argument("The gravity partition does not match its root operator.");
|
||
|
|
}
|
||
|
|
Refresh();
|
||
|
|
}
|
||
|
|
|
||
|
|
void Refresh() {
|
||
|
|
if (!m_operation->IsPrepared()) {
|
||
|
|
throw std::logic_error("Gravity requires a prepared root linearization before refresh.");
|
||
|
|
}
|
||
|
|
m_generation = m_operation->GetProblem().GetPreparationGeneration();
|
||
|
|
}
|
||
|
|
|
||
|
|
[[nodiscard]] bool IsCurrent() const {
|
||
|
|
return m_operation->IsPrepared() && m_generation == m_operation->GetProblem().GetPreparationGeneration();
|
||
|
|
}
|
||
|
|
|
||
|
|
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.");
|
||
|
|
}
|
||
|
|
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};
|
||
|
|
mutable mfem::Vector m_fullDirection;
|
||
|
|
mutable mfem::Vector m_fullAction;
|
||
|
|
};
|
||
|
|
|
||
|
|
template <GravitySubsystemOperation Operation, GravitySubsystemPreconditioner Preconditioner>
|
||
|
|
class GravityRestrictedInverse final : public mfem::Solver {
|
||
|
|
public:
|
||
|
|
static constexpr auto applicationContract = linearPreconditionerApplicationContract<Preconditioner>;
|
||
|
|
|
||
|
|
GravityRestrictedInverse(
|
||
|
|
const Preconditioner &inverse,
|
||
|
|
const GravityRestrictedOperator<Operation> &operation,
|
||
|
|
const GravityBlockPartition &partition
|
||
|
|
)
|
||
|
|
: 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()) {
|
||
|
|
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."
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
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.");
|
||
|
|
}
|
||
|
|
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 GravityRestrictedOperator<Operation> *m_operation;
|
||
|
|
const GravityBlockPartition *m_partition;
|
||
|
|
mutable mfem::Vector m_fullResidual;
|
||
|
|
mutable mfem::Vector m_fullCorrection;
|
||
|
|
};
|
||
|
|
} // 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. */
|
||
|
|
template <
|
||
|
|
GravitySubsystemOperation Operation,
|
||
|
|
GravitySubsystemPreconditioner Preconditioner,
|
||
|
|
LinearBackendConfiguration Backend = linear::FGMRES>
|
||
|
|
requires LinearBackendRuntimeAvailableFor<
|
||
|
|
Backend,
|
||
|
|
detail::GravityRestrictedOperator<Operation>,
|
||
|
|
detail::GravityRestrictedInverse<Operation, Preconditioner>>
|
||
|
|
class PreparedGravitySubsystem final {
|
||
|
|
public:
|
||
|
|
using OperatorType = Operation;
|
||
|
|
using PreconditionerType = Preconditioner;
|
||
|
|
using RestrictedOperator = detail::GravityRestrictedOperator<Operation>;
|
||
|
|
using RestrictedInverse = detail::GravityRestrictedInverse<Operation, Preconditioner>;
|
||
|
|
using BackendType = PreparedLinearBackendType<Backend, RestrictedOperator, RestrictedInverse>;
|
||
|
|
|
||
|
|
PreparedGravitySubsystem(
|
||
|
|
const Operation &operation,
|
||
|
|
const Preconditioner &preconditioner,
|
||
|
|
GravityBlockPartition partition,
|
||
|
|
const MPI_Comm communicator,
|
||
|
|
Backend backend = {}
|
||
|
|
)
|
||
|
|
: m_fullOperator(std::addressof(operation)),
|
||
|
|
m_fullPreconditioner(std::addressof(preconditioner)),
|
||
|
|
m_partition(std::move(partition)),
|
||
|
|
m_operator(
|
||
|
|
operation,
|
||
|
|
m_partition
|
||
|
|
),
|
||
|
|
m_inverse(
|
||
|
|
preconditioner,
|
||
|
|
m_operator,
|
||
|
|
m_partition
|
||
|
|
),
|
||
|
|
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 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;
|
||
|
|
}
|
||
|
|
|
||
|
|
// 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() && std::isfinite(alpha) && alpha >= 0.0 &&
|
||
|
|
trialResidual.Size() == partition.RootResidualSize() &&
|
||
|
|
gravityCorrection.Size() == partition.GravityValueSize(),
|
||
|
|
"Gravity curvature correction requires a captured base, a current "
|
||
|
|
"candidate, and valid vectors."
|
||
|
|
);
|
||
|
|
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 correction could not verify "
|
||
|
|
"its distributed step length."
|
||
|
|
);
|
||
|
|
}
|
||
|
|
if (minimumAlpha != maximumAlpha) {
|
||
|
|
throw std::invalid_argument(
|
||
|
|
"Gravity curvature correction requires the "
|
||
|
|
"same step length on every rank."
|
||
|
|
);
|
||
|
|
}
|
||
|
|
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 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
|