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

@@ -2,7 +2,13 @@ module;
#include <array>
#include <cmath>
#include <expected>
#include <mfem.hpp>
#include <optional>
#include <stdexcept>
#include <utility>
#include <mpi.h>
module mean_field;
@@ -116,6 +122,109 @@ namespace {
) {
MFEM_VERIFY(prepared.identity == requested.identity || prepared.revision != requested.revision, message);
}
using MassRejection = mean_field::operators::MassNormalizationPreparationRejection;
using MassRejectionReason = mean_field::operators::MassNormalizationPreparationRejectionReason;
[[nodiscard]] int mapping_status_priority(const mean_field::mapping::MappingStatus status) {
using Status = mean_field::mapping::MappingStatus;
switch (status) {
case Status::non_positive_determinant:
return 7;
case Status::non_finite_result:
return 6;
case Status::non_finite_input:
return 5;
case Status::outside_reference_domain:
return 4;
case Status::at_compactified_infinity:
return 3;
case Status::invalid_reference_radius:
return 2;
case Status::valid:
throw std::logic_error("A valid mapping cannot be a mass-normalization candidate rejection.");
case Status::invalid_dimension:
throw std::logic_error("A mapping dimension error cannot be a mass-normalization candidate rejection.");
}
throw std::logic_error("Unknown mapping status in mass-normalization candidate rejection.");
}
[[nodiscard]] mean_field::mapping::MappingStatus mapping_status_from_priority(const int priority) {
using Status = mean_field::mapping::MappingStatus;
switch (priority) {
case 7:
return Status::non_positive_determinant;
case 6:
return Status::non_finite_result;
case 5:
return Status::non_finite_input;
case 4:
return Status::outside_reference_domain;
case 3:
return Status::at_compactified_infinity;
case 2:
return Status::invalid_reference_radius;
default:
throw std::logic_error("Invalid synchronized mapping priority for mass normalization.");
}
}
/*
* Phase priority is explicit and independent of enum representation:
* mapping wins over interpolation, which wins over assembled-mass
* arithmetic. The mapping detail is likewise selected explicitly.
*/
[[nodiscard]] int rejection_priority(const MassRejection &rejection) {
switch (rejection.reason) {
case MassRejectionReason::mapping_failure:
return 300 + mapping_status_priority(rejection.mappingStatus);
case MassRejectionReason::non_finite_density_interpolation:
return 200;
case MassRejectionReason::non_finite_assembled_mass:
return 100;
}
throw std::logic_error("Unknown mass-normalization candidate-rejection reason.");
}
[[nodiscard]] MassRejection rejection_from_priority(const int priority) {
if (priority >= 300) {
return {
.reason = MassRejectionReason::mapping_failure,
.mappingStatus = mapping_status_from_priority(priority - 300)
};
}
if (priority == 200) {
return {.reason = MassRejectionReason::non_finite_density_interpolation};
}
if (priority == 100) {
return {.reason = MassRejectionReason::non_finite_assembled_mass};
}
throw std::logic_error("Invalid synchronized mass-normalization candidate-rejection priority.");
}
void retain_higher_priority_rejection(
std::optional<MassRejection> &current,
const MassRejection candidate
) {
if (!current.has_value() || rejection_priority(candidate) > rejection_priority(*current)) {
current = candidate;
}
}
[[nodiscard]] std::optional<MassRejection> synchronize_rejection(
const std::optional<MassRejection> &local,
const MPI_Comm communicator
) {
const int localPriority = local.has_value() ? rejection_priority(*local) : 0;
int globalPriority = 0;
if (MPI_Allreduce(&localPriority, &globalPriority, 1, MPI_INT, MPI_MAX, communicator) != MPI_SUCCESS) {
throw std::runtime_error("PreparedMassNormalizationOperator could not synchronize candidate validity.");
}
if (globalPriority == 0) {
return std::nullopt;
}
return rejection_from_priority(globalPriority);
}
} // namespace
namespace mean_field::operators {
@@ -154,6 +263,26 @@ namespace mean_field::operators {
PreparedMassNormalizationReport PreparedMassNormalizationOperator::Prepare(
const MassNormalizationStateView &state,
const MassNormalizationDependencies &dependencies
) {
auto result = TryPrepare(state, dependencies);
if (!result.has_value()) {
const MassNormalizationPreparationRejection &rejection = result.error();
switch (rejection.reason) {
case MassNormalizationPreparationRejectionReason::mapping_failure:
throw std::domain_error("PreparedMassNormalizationOperator could not map the candidate geometry.");
case MassNormalizationPreparationRejectionReason::non_finite_density_interpolation:
throw std::domain_error("PreparedMassNormalizationOperator produced a non-finite quadrature density.");
case MassNormalizationPreparationRejectionReason::non_finite_assembled_mass:
throw std::domain_error("PreparedMassNormalizationOperator assembled a non-finite mass residual.");
}
throw std::logic_error("Unknown mass-normalization candidate-rejection reason.");
}
return std::move(result).value();
}
MassNormalizationPreparationResult PreparedMassNormalizationOperator::TryPrepare(
const MassNormalizationStateView &state,
const MassNormalizationDependencies &dependencies
) {
MFEM_VERIFY(
std::isfinite(state.targetMass) && state.targetMass > 0.0,
@@ -195,6 +324,7 @@ namespace mean_field::operators {
m_isPrepared = false;
PreparedMassNormalizationReport report;
std::optional<MassNormalizationPreparationRejection> localRejection;
if (rebuildStaticPlan) {
BuildStaticPlan();
@@ -202,27 +332,37 @@ namespace mean_field::operators {
}
if (refreshGeometry) {
RefreshGeometry(m_gravityContext.GetGeometryContext().GetDisplacementTrue());
localRejection = RefreshGeometry(m_gravityContext.GetGeometryContext().GetDisplacementTrue());
report.refreshedGeometry = true;
}
if (refreshDensity) {
RefreshDensity(m_gravityContext.GetDensityTrue());
if (auto densityRejection = RefreshDensity(m_gravityContext.GetDensityTrue());
densityRejection.has_value()) {
retain_higher_priority_rejection(localRejection, *densityRejection);
}
report.refreshedDensity = true;
}
if (auto globalRejection = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
globalRejection.has_value()) {
return std::unexpected(*globalRejection);
}
if (updateTargetMass) {
m_targetMass = state.targetMass;
report.updatedTargetMass = true;
}
if (refreshGeometry || refreshDensity) {
AssembleResidual();
if (auto rejection = AssembleResidual(); rejection.has_value()) {
return std::unexpected(*rejection);
}
report.assembledResidual = true;
} else if (updateTargetMass) {
m_cachedResidual.SetSize(1);
m_cachedResidual(0) = m_currentMass - m_targetMass;
++m_preparationCount;
if (auto rejection = UpdateResidualForTargetMass(); rejection.has_value()) {
return std::unexpected(*rejection);
}
report.assembledResidual = true;
}
@@ -238,6 +378,13 @@ namespace mean_field::operators {
return Prepare({.targetMass = constraint.targetMass().value()}, dependencies);
}
MassNormalizationPreparationResult PreparedMassNormalizationOperator::TryPrepare(
const models::CompiledFixedMass &constraint,
const MassNormalizationDependencies &dependencies
) {
return TryPrepare({.targetMass = constraint.targetMass().value()}, dependencies);
}
void PreparedMassNormalizationOperator::BuildStaticPlan() {
m_elements.clear();
m_elements.reserve(m_fem.mesh->GetNE());
@@ -287,14 +434,17 @@ namespace mean_field::operators {
}
int globalStellarElementCount = 0;
MPI_Allreduce(
&localStellarElementCount, &globalStellarElementCount, 1, MPI_INT, MPI_SUM, m_fem.mesh->GetComm()
);
if (MPI_Allreduce(
&localStellarElementCount, &globalStellarElementCount, 1, MPI_INT, MPI_SUM, m_fem.mesh->GetComm()
) != MPI_SUCCESS) {
throw std::runtime_error("PreparedMassNormalizationOperator could not count stellar elements.");
}
MFEM_VERIFY(globalStellarElementCount > 0, "PreparedMassNormalizationOperator found no stellar elements.");
}
void PreparedMassNormalizationOperator::RefreshGeometry(const mfem::Vector &displacement) {
std::optional<MassNormalizationPreparationRejection>
PreparedMassNormalizationOperator::RefreshGeometry(const mfem::Vector &displacement) {
MFEM_VERIFY(
displacement.Size() == m_fem.displacementFes->GetTrueVSize(),
"PreparedMassNormalizationOperator received a displacement "
@@ -309,6 +459,7 @@ namespace mean_field::operators {
true_to_local(*m_fem.displacementFes, displacement, displacementLocal);
mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());
std::optional<MassNormalizationPreparationRejection> rejection;
for (ElementPAData &data : m_elements) {
displacementLocal.GetSubVector(data.displacementDofs, data.baseDisplacement);
@@ -346,17 +497,23 @@ namespace mean_field::operators {
);
MFEM_VERIFY(
status == mapping::MappingStatus::valid, "Stateless mapping failed while preparing mass "
"normalization. Element: "
<< data.elementId
<< ", attribute: " << transformation->Attribute
<< ", status: " << static_cast<int>(status)
status != mapping::MappingStatus::invalid_dimension,
"Stateless mapping reported a dimension error while preparing mass normalization."
);
if (status != mapping::MappingStatus::valid) {
retain_higher_priority_rejection(
rejection, {.reason = MassNormalizationPreparationRejectionReason::mapping_failure,
.mappingStatus = status}
);
}
}
}
return rejection;
}
void PreparedMassNormalizationOperator::RefreshDensity(const mfem::Vector &density) {
std::optional<MassNormalizationPreparationRejection>
PreparedMassNormalizationOperator::RefreshDensity(const mfem::Vector &density) {
MFEM_VERIFY(
density.Size() == m_fem.densityFes->GetTrueVSize(),
"PreparedMassNormalizationOperator received a density vector "
@@ -371,6 +528,7 @@ namespace mean_field::operators {
true_to_local(*m_fem.densityFes, density, densityLocal);
mfem::Vector elementDensity;
std::optional<MassNormalizationPreparationRejection> rejection;
for (ElementPAData &data : m_elements) {
densityLocal.GetSubVector(data.densityDofs, elementDensity);
@@ -381,29 +539,62 @@ namespace mean_field::operators {
for (QuadraturePointData &point : data.quadraturePoints) {
point.density = elementDensity * point.densityShape;
MFEM_VERIFY(
std::isfinite(point.density), "PreparedMassNormalizationOperator produced a non-finite "
"quadrature density."
);
if (!std::isfinite(point.density)) {
retain_higher_priority_rejection(
rejection,
{.reason = MassNormalizationPreparationRejectionReason::non_finite_density_interpolation}
);
}
}
}
return rejection;
}
void PreparedMassNormalizationOperator::AssembleResidual() {
std::optional<MassNormalizationPreparationRejection> PreparedMassNormalizationOperator::AssembleResidual() {
double localMass = 0.0;
std::optional<MassNormalizationPreparationRejection> localRejection;
for (const ElementPAData &data : m_elements) {
for (const QuadraturePointData &point : data.quadraturePoints) {
localMass += point.density * point.mappingContext.quadrature.weight;
const double contribution = point.density * point.mappingContext.quadrature.weight;
if (!std::isfinite(contribution) || !std::isfinite(localMass + contribution)) {
localRejection = {.reason = MassNormalizationPreparationRejectionReason::non_finite_assembled_mass};
continue;
}
localMass += contribution;
}
}
m_currentMass = GlobalSum(localMass);
MFEM_VERIFY(std::isfinite(m_currentMass), "PreparedMassNormalizationOperator assembled a non-finite mass.");
if (auto globalRejection = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
globalRejection.has_value()) {
return globalRejection;
}
m_currentMass = GlobalSum(localMass);
if (!std::isfinite(m_currentMass)) {
return MassNormalizationPreparationRejection{
.reason = MassNormalizationPreparationRejectionReason::non_finite_assembled_mass
};
}
return UpdateResidualForTargetMass();
}
std::optional<MassNormalizationPreparationRejection>
PreparedMassNormalizationOperator::UpdateResidualForTargetMass() {
m_cachedResidual.SetSize(1);
m_cachedResidual(0) = m_currentMass - m_targetMass;
std::optional<MassNormalizationPreparationRejection> localRejection;
if (!std::isfinite(m_cachedResidual(0))) {
localRejection = {.reason = MassNormalizationPreparationRejectionReason::non_finite_assembled_mass};
}
if (auto globalRejection = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
globalRejection.has_value()) {
return globalRejection;
}
++m_preparationCount;
return std::nullopt;
}
void PreparedMassNormalizationOperator::BuildResidual(mfem::Vector &residual) const {
@@ -716,7 +907,9 @@ namespace mean_field::operators {
double PreparedMassNormalizationOperator::GlobalSum(const double localValue) const {
double globalValue = 0.0;
MPI_Allreduce(&localValue, &globalValue, 1, MPI_DOUBLE, MPI_SUM, m_fem.mesh->GetComm());
if (MPI_Allreduce(&localValue, &globalValue, 1, MPI_DOUBLE, MPI_SUM, m_fem.mesh->GetComm()) != MPI_SUCCESS) {
throw std::runtime_error("PreparedMassNormalizationOperator could not assemble a distributed scalar.");
}
return globalValue;
}