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

Add bounded coupled damping and prediction-based attempt selection

Browse Source
This commit is contained in:
Emily Boudreaux
2026-09-18 10:45:45 -04:00
parent 46249d0758
commit 422c46bcbf
11 changed files with 2296 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

296
libmeanfield/interface/solver/coupled_trial_correction.cppm
View File

@@ -12,6 +12,7 @@ module;
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>
#include <mfem.hpp>
#include <mpi.h>
@@ -21,6 +22,33 @@ export module mean_field:solver.coupled_trial_correction;
export import :solver.linear_backend;
export namespace mean_field::solver {
struct FullCoupledStep final { };
struct BacktrackedCoupledStep final {
int maximumSteps{4};
double reduction{.5};
void Validate() const {
if (maximumSteps <= 0 || !std::isfinite(reduction) || reduction <= 0.0 || reduction >= 1.0 ||
std::pow(reduction, maximumSteps - 1) == 0.0)
throw std::invalid_argument("Coupled damping requires positive resolved steps and a reduction strictly between zero and one.");
}
};
struct EveryRejectedCoupledAttempt final { };
struct NearAcceptanceCoupledAttempt final {
double maximumExcessRatio{1.0};
void Validate() const {
if (!std::isfinite(maximumExcessRatio) || maximumExcessRatio <= 0.0)
throw std::invalid_argument("A coupled prediction gate requires a finite positive excess ratio.");
}
};
template <typename Candidate>
concept CoupledCorrectionSearch = std::same_as<Candidate, FullCoupledStep> ||
std::same_as<Candidate, BacktrackedCoupledStep>;
template <typename Candidate>
concept CoupledAttemptCriterion = std::same_as<Candidate, EveryRejectedCoupledAttempt> ||
std::same_as<Candidate, NearAcceptanceCoupledAttempt>;
struct CoupledTrialCorrectionOptions final {
LinearSolveControl control{.relativeTolerance = .03, .absoluteTolerance = 1.e-14, .maximumIterations = 160};
linear::FGMRESOptions backend{.restartLength = 160, .printLevel = -1};
@@ -52,7 +80,16 @@ export namespace mean_field::solver {
enum class CoupledTrialAttemptStatus : std::uint8_t {
not_attempted, original_accepted, budget_exhausted, kept, linear_failure, norm_limit,
non_finite, inadmissible, not_improved
non_finite, inadmissible, not_improved, prediction_deferred
};
enum class CoupledTrialCandidateStatus : std::uint8_t { evaluated, inadmissible, non_finite };
struct CoupledTrialCandidateReport final {
double beta{1.0};
CoupledTrialCandidateStatus status{CoupledTrialCandidateStatus::inadmissible};
double merit{std::numeric_limits<double>::quiet_NaN()};
bool accepted{false};
double preparationSeconds{0.0};
};
struct CoupledTrialDiagnostics final {
@@ -64,6 +101,16 @@ export namespace mean_field::solver {
std::optional<CoupledTrialCorrectionReport> currentSolve{};
double originalMerit{std::numeric_limits<double>::quiet_NaN()};
double correctedMerit{std::numeric_limits<double>::quiet_NaN()};
std::uint64_t deferredAttempts{0};
int candidateCount{0};
double selectedBeta{std::numeric_limits<double>::quiet_NaN()};
double baseMerit{std::numeric_limits<double>::quiet_NaN()};
double linearPredictionMerit{std::numeric_limits<double>::quiet_NaN()};
double predictedDecrease{std::numeric_limits<double>::quiet_NaN()};
double predictionExcessRatio{std::numeric_limits<double>::quiet_NaN()};
// Search candidates exclude the possible final replay of an earlier
// best beta. Preparations/restorations still account for that work.
std::vector<CoupledTrialCandidateReport> candidates{};
};
}
@@ -194,6 +241,7 @@ export namespace mean_field::solver {
[[nodiscard]] double GetBaseDirectionNorm() const noexcept { return m_directionNorm; }
[[nodiscard]] MPI_Comm GetCommunicator() const noexcept { return m_communicator; }
[[nodiscard]] const mfem::Vector &GetRightHandSide() const noexcept { return m_rightHandSide; }
[[nodiscard]] const mfem::Vector &GetBaseResidual() const noexcept { return m_baseResidual; }
[[nodiscard]] const mfem::Vector &GetVerifiedAction() const noexcept { return m_verifiedAction; }
[[nodiscard]] const mfem::Vector &GetVerifiedDefect() const noexcept { return m_defect; }
[[nodiscard]] const auto &GetLastReport() const noexcept { return m_lastReport; }
@@ -302,23 +350,24 @@ export namespace mean_field::solver {
// execute the same transaction while their typed callbacks choose the metric
// and enthalpy policy. This helper never reads or writes accepted-state buffers.
export namespace mean_field::solver::detail {
template <typename State> class PreparedCoupledTrialTransaction final {
template <typename State, CoupledCorrectionSearch Search = FullCoupledStep,
CoupledAttemptCriterion Attempt = EveryRejectedCoupledAttempt>
class PreparedCoupledTrialTransaction final {
using Operation = std::remove_cvref_t<decltype(*std::declval<State>().normalizedOperator)>;
using Inverse = std::remove_cvref_t<decltype(*std::declval<State>().normalizedInverse)>;
using Clock = std::chrono::steady_clock;
public:
PreparedCoupledTrialTransaction(State &state, const CoupledTrialCorrectionOptions &options,
int maximumAttemptsPerIteration = 1)
int maximumAttemptsPerIteration = 1, Search search = {}, Attempt attempt = {})
: m_core(*state.normalizedOperator, *state.normalizedInverse, state.Problem().GetCommunicator(), options),
m_maximumAttempts(maximumAttemptsPerIteration), m_originalState(state.normalizedOperator->Width()),
m_originalResidual(state.normalizedOperator->Height()), m_originalPhysical(state.normalizedOperator->Width()),
m_correction(state.normalizedOperator->Width()) {
CoupledRequire(maximumAttemptsPerIteration > 0, m_core.GetCommunicator(), "A coupled trial attempt budget must be positive.");
int minimum = 0, maximum = 0;
if (MPI_Allreduce(&maximumAttemptsPerIteration, &minimum, 1, MPI_INT, MPI_MIN, m_core.GetCommunicator()) != MPI_SUCCESS ||
MPI_Allreduce(&maximumAttemptsPerIteration, &maximum, 1, MPI_INT, MPI_MAX, m_core.GetCommunicator()) != MPI_SUCCESS)
throw std::runtime_error("Coupled trial attempt budgets could not be compared.");
if (minimum != maximum) throw std::invalid_argument("Coupled trial attempt budgets must agree across ranks.");
m_maximumAttempts(maximumAttemptsPerIteration), m_search(search), m_attempt(attempt),
m_originalState(state.normalizedOperator->Width()), m_originalResidual(state.normalizedOperator->Height()),
m_originalPhysical(state.normalizedOperator->Width()), m_correction(state.normalizedOperator->Width()),
m_selectedCorrection(state.normalizedOperator->Width()),
m_prediction(std::same_as<Attempt, NearAcceptanceCoupledAttempt> ? state.normalizedOperator->Height() : 0) {
ValidatePoliciesCollectively();
m_selectedCorrection = 0.0;
m_diagnostics.candidates.reserve(MaximumSteps());
}
void CaptureBase(const mfem::Vector &residual, const mfem::Vector &direction) {
m_core.CaptureBase(residual, direction);
@@ -327,11 +376,19 @@ export namespace mean_field::solver::detail {
void BeginTrial() noexcept {
m_diagnostics.status = CoupledTrialAttemptStatus::not_attempted;
m_diagnostics.currentSolve.reset();
m_diagnostics.originalMerit = m_diagnostics.correctedMerit = std::numeric_limits<double>::quiet_NaN();
m_diagnostics.originalMerit = m_diagnostics.correctedMerit = m_diagnostics.selectedBeta =
m_diagnostics.baseMerit = m_diagnostics.linearPredictionMerit = m_diagnostics.predictedDecrease =
m_diagnostics.predictionExcessRatio = std::numeric_limits<double>::quiet_NaN();
m_diagnostics.candidateCount = 0;
m_diagnostics.candidates.clear();
m_selectedCorrection = 0.0;
}
[[nodiscard]] const CoupledTrialDiagnostics &GetDiagnostics() const noexcept { return m_diagnostics; }
[[nodiscard]] const auto &GetCore() const noexcept { return m_core; }
// The report verifies this unscaled vector. Damping changes the trial
// map, not the vector to which its original linear certificate applies.
[[nodiscard]] const mfem::Vector &GetCorrection() const noexcept { return m_correction; }
[[nodiscard]] const mfem::Vector &GetSelectedCorrection() const noexcept { return m_selectedCorrection; }
template <typename Preparation, typename Metric, typename Evaluate, typename Accepts, typename Better, typename CheckState>
void TryImprove(State &state, double alpha, Preparation &preparation, Metric &metric,
@@ -339,8 +396,7 @@ export namespace mean_field::solver::detail {
const auto communicator = m_core.GetCommunicator();
CoupledRequire(preparation.has_value() && state.normalizedOperator->IsPrepared() &&
CoupledFinite(state.trialNormalizedState) && CoupledFinite(state.trialNormalizedResidual) &&
std::isfinite(metric.merit), communicator,
"Coupled correction requires a successfully prepared original trial.");
ValidMetric(metric), communicator, "Coupled correction requires a successfully prepared original trial.");
m_diagnostics.originalMerit = metric.merit;
if (CoupledAll(accepts(metric), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::original_accepted;
@@ -350,6 +406,27 @@ export namespace mean_field::solver::detail {
m_diagnostics.status = CoupledTrialAttemptStatus::budget_exhausted;
return;
}
if constexpr (std::same_as<Attempt, NearAcceptanceCoupledAttempt>) {
// These are evaluations of saved residual vectors, not new
// physical candidates. No preparation, inverse refresh or
// attempt budget is spent on the prediction gate.
m_core.BuildLinearPrediction(alpha, m_prediction);
const auto baseMetric = evaluate(m_core.GetBaseResidual());
const auto predictionMetric = evaluate(m_prediction);
m_diagnostics.baseMerit = baseMetric.merit;
m_diagnostics.linearPredictionMerit = predictionMetric.merit;
m_diagnostics.predictedDecrease = baseMetric.merit - predictionMetric.merit;
const double decrease = m_diagnostics.predictedDecrease;
const double ratio = decrease > 0.0 ? (metric.merit - baseMetric.merit) / decrease :
std::numeric_limits<double>::quiet_NaN();
m_diagnostics.predictionExcessRatio = ratio;
if (!CoupledAll(ValidMetric(baseMetric) && ValidMetric(predictionMetric) && std::isfinite(decrease) &&
decrease > 0.0 && std::isfinite(ratio) && ratio <= m_attempt.maximumExcessRatio, communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::prediction_deferred;
++m_diagnostics.deferredAttempts;
return;
}
}
m_originalState = state.trialNormalizedState;
m_originalResidual = state.trialNormalizedResidual;
m_originalPhysical = state.candidatePhysicalState;
@@ -368,14 +445,14 @@ export namespace mean_field::solver::detail {
case CoupledTrialCorrectionStatus::norm_limit: m_diagnostics.status = CoupledTrialAttemptStatus::norm_limit; break;
default: m_diagnostics.status = CoupledTrialAttemptStatus::non_finite; break;
}
// No candidate mutation has occurred; its prepared operator and
// freshly refreshed inverse still describe the original trial.
// No candidate mutation has occurred. In particular, a small
// beta cannot rescue an unverified or over-cap full correction.
return;
}
const Metric originalMetric = metric;
bool restorationStarted = false;
auto restore = [&]() {
restorationStarted = true;
bool originalRestorationStarted = false;
auto restoreOriginal = [&]() {
originalRestorationStarted = true;
state.trialNormalizedState = m_originalState;
state.candidatePhysicalState = m_originalPhysical;
auto restored = Prepare(state);
@@ -387,56 +464,155 @@ export namespace mean_field::solver::detail {
metric = evaluate(state.trialNormalizedResidual);
++m_diagnostics.fallbacks;
};
try {
state.trialNormalizedState += m_correction;
const auto formCandidate = [&](double beta) {
state.trialNormalizedState = m_originalState;
state.trialNormalizedState.Add(beta, m_correction);
state.normalizedOperator->DenormalizeState(state.trialNormalizedState, state.candidatePhysicalState);
if (!CoupledAll(CoupledFinite(state.trialNormalizedState) && CoupledFinite(state.candidatePhysicalState), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::non_finite;
restore();
return;
}
if (!CoupledAll(checkState(state.trialNormalizedState), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
restore();
return;
}
auto corrected = Prepare(state);
if (!CoupledAll(corrected.has_value(), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
restore();
return;
}
if (!CoupledAll(CoupledFinite(state.trialNormalizedResidual), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::non_finite;
restore();
return;
}
const auto correctedMetric = evaluate(state.trialNormalizedResidual);
m_diagnostics.correctedMerit = correctedMetric.merit;
if (!CoupledAll(better(correctedMetric, originalMetric), communicator)) {
};
try {
std::optional<Metric> bestMetric;
double bestBeta = 0.0;
int bestStep = -1, lastStep = -1;
double beta = 1.0;
for (int step = 0; step < MaximumSteps(); ++step, beta *= Reduction()) {
lastStep = step;
formCandidate(beta);
m_diagnostics.candidates.push_back({.beta = beta});
++m_diagnostics.candidateCount;
auto &candidate = m_diagnostics.candidates.back();
if (!CoupledAll(CoupledFinite(state.trialNormalizedState) && CoupledFinite(state.candidatePhysicalState), communicator)) {
candidate.status = CoupledTrialCandidateStatus::non_finite;
m_diagnostics.status = CoupledTrialAttemptStatus::non_finite;
continue;
}
if (!CoupledAll(checkState(state.trialNormalizedState), communicator)) {
candidate.status = CoupledTrialCandidateStatus::inadmissible;
m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
continue;
}
const double preparationStart = m_diagnostics.preparationSeconds;
auto corrected = Prepare(state);
candidate.preparationSeconds = m_diagnostics.preparationSeconds - preparationStart;
if (!CoupledAll(corrected.has_value(), communicator)) {
candidate.status = CoupledTrialCandidateStatus::inadmissible;
m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
continue;
}
if (!CoupledAll(CoupledFinite(state.trialNormalizedResidual), communicator)) {
candidate.status = CoupledTrialCandidateStatus::non_finite;
m_diagnostics.status = CoupledTrialAttemptStatus::non_finite;
continue;
}
const auto correctedMetric = evaluate(state.trialNormalizedResidual);
m_diagnostics.correctedMerit = correctedMetric.merit;
if (!CoupledAll(ValidMetric(correctedMetric), communicator)) {
candidate.status = CoupledTrialCandidateStatus::non_finite;
m_diagnostics.status = CoupledTrialAttemptStatus::non_finite;
continue;
}
preparation = std::move(corrected);
candidate.status = CoupledTrialCandidateStatus::evaluated;
candidate.merit = correctedMetric.merit;
const bool improves = CoupledAll(better(correctedMetric, originalMetric), communicator);
candidate.accepted = improves && CoupledAll(accepts(correctedMetric), communicator);
if (improves && (!bestMetric || candidate.accepted || CoupledAll(better(correctedMetric, *bestMetric), communicator))) {
bestMetric = correctedMetric;
bestBeta = beta;
bestStep = step;
}
m_diagnostics.status = CoupledTrialAttemptStatus::not_improved;
restore();
// An improved but unacceptable full correction must not
// suppress the remaining, cheaper physical preparations.
if (candidate.accepted) break;
}
if (!bestMetric) {
restoreOriginal();
return;
}
preparation = std::move(corrected);
metric = correctedMetric;
if (bestStep != lastStep) {
// The current prepared operator belongs to the last beta,
// which can even be inadmissible. Replaying the selected
// state restores its residual, geometry and generation.
formCandidate(bestBeta);
if (!CoupledAll(checkState(state.trialNormalizedState), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
restoreOriginal();
return;
}
auto restored = Prepare(state);
++m_diagnostics.restorations;
if (!CoupledAll(restored.has_value() && CoupledFinite(state.trialNormalizedResidual), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::inadmissible;
restoreOriginal();
return;
}
const auto restoredMetric = evaluate(state.trialNormalizedResidual);
if (!CoupledAll(ValidMetric(restoredMetric), communicator) ||
!CoupledAll(better(restoredMetric, originalMetric), communicator)) {
m_diagnostics.status = CoupledTrialAttemptStatus::not_improved;
restoreOriginal();
return;
}
preparation = std::move(restored);
bestMetric = restoredMetric;
Refresh(state);
}
metric = *bestMetric;
m_selectedCorrection = m_correction;
m_selectedCorrection *= bestBeta;
m_diagnostics.selectedBeta = bestBeta;
m_diagnostics.correctedMerit = metric.merit;
m_diagnostics.status = CoupledTrialAttemptStatus::kept;
++m_diagnostics.kept;
// This is a nonlinear candidate, not an affine gravity completion.
// Its fresh full residual is judged by the caller's normal metric.
// As for an ordinary prepared trial, the caller refreshes the
// inverse only if it commits or needs another auxiliary solve.
// A current selected trial retains the ordinary lifecycle:
// its inverse is refreshed only on commit or another solve.
} catch (...) {
// Numerical rejection is handled above. For collectively
// propagated exceptions, restore the original prepared trial
// before rethrowing. A failed restoration itself propagates to
// the owning Newton context, which restores the accepted base.
if (!restorationStarted) restore();
// Only collectively propagated callback exceptions can be
// rolled back collectively. Failure to restore propagates to
// the owning Newton context's accepted-base restoration.
if (!originalRestorationStarted) restoreOriginal();
throw;
}
}
private:
template <typename Metric> static bool ValidMetric(const Metric &metric) {
bool valid = std::isfinite(metric.merit) && metric.merit >= 0.0;
if constexpr (requires { metric.residualNorm; })
valid = valid && std::isfinite(metric.residualNorm) && metric.residualNorm >= 0.0;
return valid;
}
[[nodiscard]] int MaximumSteps() const noexcept {
if constexpr (std::same_as<Search, BacktrackedCoupledStep>) return m_search.maximumSteps;
else return 1;
}
[[nodiscard]] double Reduction() const noexcept {
if constexpr (std::same_as<Search, BacktrackedCoupledStep>) return m_search.reduction;
else return 1.0;
}
void ValidatePoliciesCollectively() const {
const auto communicator = m_core.GetCommunicator();
bool valid = m_maximumAttempts > 0;
try {
if constexpr (std::same_as<Search, BacktrackedCoupledStep>) m_search.Validate();
if constexpr (std::same_as<Attempt, NearAcceptanceCoupledAttempt>) m_attempt.Validate();
} catch (const std::invalid_argument &) { valid = false; }
CoupledRequire(valid, communicator, "Coupled trial search or attempt controls are invalid on at least one rank.");
double excessRatio = 0.0;
if constexpr (std::same_as<Attempt, NearAcceptanceCoupledAttempt>) excessRatio = m_attempt.maximumExcessRatio;
const std::array<int, 4> integers{m_maximumAttempts, std::same_as<Search, BacktrackedCoupledStep>,
MaximumSteps(), std::same_as<Attempt, NearAcceptanceCoupledAttempt>};
const std::array<double, 2> reals{Reduction(), excessRatio};
std::array<int, 4> minimumIntegers{}, maximumIntegers{};
std::array<double, 2> minimumReals{}, maximumReals{};
if (MPI_Allreduce(integers.data(), minimumIntegers.data(), 4, MPI_INT, MPI_MIN, communicator) != MPI_SUCCESS ||
MPI_Allreduce(integers.data(), maximumIntegers.data(), 4, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS ||
MPI_Allreduce(reals.data(), minimumReals.data(), 2, MPI_DOUBLE, MPI_MIN, communicator) != MPI_SUCCESS ||
MPI_Allreduce(reals.data(), maximumReals.data(), 2, MPI_DOUBLE, MPI_MAX, communicator) != MPI_SUCCESS)
throw std::runtime_error("Coupled trial search and attempt controls could not be compared.");
if (minimumIntegers != maximumIntegers || minimumReals != maximumReals)
throw std::invalid_argument("Coupled trial search and attempt controls must agree across ranks.");
}
void Refresh(State &state) {
const auto start = Clock::now();
state.physicalInverse->Refresh();
@@ -451,7 +627,9 @@ export namespace mean_field::solver::detail {
}
PreparedCoupledTrialCorrection<Operation, Inverse> m_core;
int m_maximumAttempts, m_attemptsThisIteration{0};
mfem::Vector m_originalState, m_originalResidual, m_originalPhysical, m_correction;
[[no_unique_address]] Search m_search;
[[no_unique_address]] Attempt m_attempt;
mfem::Vector m_originalState, m_originalResidual, m_originalPhysical, m_correction, m_selectedCorrection, m_prediction;
CoupledTrialDiagnostics m_diagnostics{};
};
}