feat(newton): first newton solver implementation

This commit is contained in:
2026-09-08 06:36:39 -04:00
parent 76818f2f82
commit b3c04d507a
98 changed files with 20397 additions and 11040 deletions

View File

@@ -1,6 +1,13 @@
module;
#include <cmath>
#include <expected>
#include <optional>
#include <stdexcept>
#include <utility>
#include <mfem.hpp>
#include <mpi.h>
module mean_field;
@@ -9,6 +16,8 @@ import :operators.prepared_gravity_displacement_force;
namespace {
using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
using Rejection = mean_field::operators::kernels::GravityDisplacementForceRejection;
using Reason = mean_field::operators::kernels::GravityDisplacementForceRejectionReason;
[[nodiscard]] bool relevant_revisions_match(
const mean_field::operators::context::gravity_field::GravityFieldRevisions &left,
@@ -22,6 +31,69 @@ namespace {
return DomainSchema::template attribute_belongs_to<mean_field::utils::domain::Vacuum>(attribute);
}
[[nodiscard]] Rejection mapping_rejection(const mean_field::mapping::MappingStatus status) {
MFEM_VERIFY(
status != mean_field::mapping::MappingStatus::invalid_dimension,
"Prepared gravity force mapping reported an invariant dimension mismatch."
);
return {.reason = Reason::invalid_mapping, .mappingStatus = status};
}
[[nodiscard]] Rejection non_finite_rejection() noexcept {
return {.reason = Reason::non_finite_arithmetic};
}
[[nodiscard]] int encode_rejection(const std::optional<Rejection> &rejection) noexcept {
if (!rejection.has_value()) {
return 0;
}
if (rejection->reason == Reason::non_finite_arithmetic) {
return 256;
}
return static_cast<int>(rejection->mappingStatus) + 1;
}
[[nodiscard]] Rejection decode_rejection(const int encoded) {
if (encoded >= 256) {
return non_finite_rejection();
}
return mapping_rejection(static_cast<mean_field::mapping::MappingStatus>(encoded - 1));
}
[[nodiscard]] std::expected<
void,
Rejection>
synchronize_rejection(
const std::optional<Rejection> &localRejection,
const MPI_Comm communicator
) {
const int localEncoded = encode_rejection(localRejection);
int globalEncoded = 0;
if (MPI_Allreduce(&localEncoded, &globalEncoded, 1, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS) {
throw std::runtime_error("Could not synchronize prepared gravity-force candidate validity.");
}
if (globalEncoded != 0) {
return std::unexpected(decode_rejection(globalEncoded));
}
return {};
}
[[nodiscard]] bool vector_is_finite(const mfem::Vector &vector) noexcept {
for (int index = 0; index < vector.Size(); ++index) {
if (!std::isfinite(vector(index))) {
return false;
}
}
return true;
}
[[noreturn]] void throw_rejection(const Rejection &rejection) {
if (rejection.reason == Reason::non_finite_arithmetic) {
throw std::domain_error("Prepared gravity force produced non-finite arithmetic.");
}
throw std::domain_error("Prepared gravity force encountered an invalid mapped domain.");
}
void true_to_local(
const mfem::ParFiniteElementSpace &finiteElementSpace,
const mfem::Vector &trueVector,
@@ -108,7 +180,10 @@ namespace mean_field::operators {
);
}
void PreparedGravityDisplacementForceOperator::PrepareElementData() {
std::expected<
void,
kernels::GravityDisplacementForceRejection>
PreparedGravityDisplacementForceOperator::TryPrepareElementData() {
m_elements.clear();
m_elements.reserve(m_fem.mesh->GetNE());
@@ -199,9 +274,12 @@ namespace mean_field::operators {
const mapping::MappingStatus status = m_domainMapper.EvaluateVolume(
mappingData, *transformation, integrationPoint, workspace, mappingContext
);
if (status != mapping::MappingStatus::valid) {
return std::unexpected(mapping_rejection(status));
}
MFEM_VERIFY(
status == mapping::MappingStatus::valid && !mappingContext.mapping.compactified,
"Prepared gravity force encountered an invalid stellar mapping."
!mappingContext.mapping.compactified,
"Prepared gravity force encountered compactification on a stellar element."
);
densityElement.CalcShape(integrationPoint, densityShape);
@@ -220,11 +298,38 @@ namespace mean_field::operators {
data.inverseMeshJacobians(quadraturePoint, entry) = inverseMeshJacobian(row, column);
}
}
if (!std::isfinite(data.baseDensityValues(quadraturePoint)) ||
!std::isfinite(data.referenceWeights(quadraturePoint)) ||
!vector_is_finite(baseGravityReferenceValue)) {
return std::unexpected(non_finite_rejection());
}
for (int row = 0; row < dimension; ++row) {
for (int column = 0; column < dimension; ++column) {
const int entry = row * dimension + column;
if (!std::isfinite(data.mappingJacobians(quadraturePoint, entry)) ||
!std::isfinite(data.inverseMeshJacobians(quadraturePoint, entry))) {
return std::unexpected(non_finite_rejection());
}
}
}
}
}
return {};
}
PreparedGravityDisplacementForceReport PreparedGravityDisplacementForceOperator::Prepare() {
auto result = TryPrepare();
if (!result.has_value()) {
throw_rejection(result.error());
}
return std::move(result).value();
}
std::expected<
PreparedGravityDisplacementForceReport,
kernels::GravityDisplacementForceRejection>
PreparedGravityDisplacementForceOperator::TryPrepare() {
MFEM_VERIFY(
m_gravityContext.IsPrepared(), "PreparedGravityDisplacementForceOperator requires the shared "
"gravity linearization context to be prepared first."
@@ -236,19 +341,45 @@ namespace mean_field::operators {
return {};
}
kernels::apply_gravity_displacement_force_residual(
m_isPrepared = false;
/*
* Build the reusable element plan before assembling the residual.
* This pass stops at the first invalid mapped quadrature point, so a
* rejected line-search candidate need not traverse the full stateless
* residual kernel. Synchronize before proceeding so every rank takes
* the same branch.
*/
const auto elementResult = TryPrepareElementData();
const std::optional<Rejection> localElementRejection =
elementResult.has_value() ? std::optional<Rejection>{} : std::optional<Rejection>{elementResult.error()};
auto synchronizedElement = synchronize_rejection(localElementRejection, m_fem.mesh->GetComm());
if (!synchronizedElement.has_value()) {
return std::unexpected(synchronizedElement.error());
}
auto residualResult = kernels::try_apply_gravity_displacement_force_residual(
m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityContext.GetGravityGradientTrue(),
m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
);
if (!residualResult.has_value()) {
return std::unexpected(residualResult.error());
}
m_cachedResidual.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
m_gravityContext.GetDisplacementMap().gather(m_actionTrue, m_cachedResidual);
PrepareElementData();
std::optional<Rejection> localRejection;
if (!vector_is_finite(m_cachedResidual)) {
localRejection = non_finite_rejection();
}
auto synchronized = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
if (!synchronized.has_value()) {
return std::unexpected(synchronized.error());
}
m_preparedRevisions = requestedRevisions;
++m_residualPreparationCount;
m_isPrepared = true;
return {.preparedResidual = true};
return PreparedGravityDisplacementForceReport{.preparedResidual = true};
}
void PreparedGravityDisplacementForceOperator::BuildResidual(mfem::Vector &residual) const {