feat(preconditioner): major work on preconditioner system

first preconditioner MVP
This commit is contained in:
2026-09-04 07:54:10 -04:00
parent 25510008dd
commit 71423d543f
61 changed files with 15920 additions and 422 deletions

View File

@@ -1,4 +1,5 @@
module;
#include "profile.h"
#include <cmath>
#include <cstdint>
#include <memory>
@@ -19,6 +20,13 @@ namespace {
.reduced_size();
}
bool communicator_has_single_rank(const MPI_Comm communicator) {
int size = 0;
MFEM_VERIFY(MPI_Comm_size(communicator, &size) == MPI_SUCCESS, "Failed to query the MPI communicator size.");
MFEM_VERIFY(size > 0, "The MPI communicator must contain at least one rank.");
return size == 1;
}
void true_to_local(
const mfem::ParFiniteElementSpace &finite_element_space,
const mfem::Vector &true_vector,
@@ -361,7 +369,8 @@ namespace mean_field::operators {
field::Displacement,
DomainSchema>(*f.displacementFes)
),
m_variationWorkspace(domain_mapper.GetDimension()) {
m_variationWorkspace(domain_mapper.GetDimension()),
m_single_rank(communicator_has_single_rank(f.gravityFluxFes->GetComm())) {
MFEM_VERIFY(f.mesh != nullptr, "PreparedMappedHDivMassOperator requires a mesh.");
MFEM_VERIFY(
f.gravityFluxFes != nullptr, "PreparedMappedHDivMassOperator requires the "
@@ -409,6 +418,8 @@ namespace mean_field::operators {
}
void PreparedMappedHDivMassOperator::PrepareVariationData() {
MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::PrepareVariationData", 0);
m_variationElements.clear();
m_variationElements.reserve(m_fem.mesh->GetNE());
@@ -481,6 +492,19 @@ namespace mean_field::operators {
}
void PreparedMappedHDivMassOperator::Prepare(const mfem::Vector &displacement) {
MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::Prepare linearization", 0);
PrepareImpl(displacement, PreparationMode::linearization);
}
void PreparedMappedHDivMassOperator::PreparePrimal(const mfem::Vector &displacement) {
MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedHDivMassOperator::Prepare primal", 0);
PrepareImpl(displacement, PreparationMode::primal);
}
void PreparedMappedHDivMassOperator::PrepareImpl(
const mfem::Vector &displacement,
const PreparationMode mode
) {
MFEM_VERIFY(
displacement.Size() == m_displacement_map.reduced_size(),
"PreparedMappedHDivMassOperator received a displacement vector "
@@ -496,6 +520,9 @@ namespace mean_field::operators {
);
}
m_is_prepared = false;
m_has_variation_data = false;
m_displacement_true.SetSize(m_displacement_map.full_size());
m_displacement_map.scatter(displacement, m_displacement_true);
@@ -541,7 +568,12 @@ namespace mean_field::operators {
m_stellar_mass_form->Assemble();
m_vacuum_mass_form->Assemble();
PrepareVariationData();
if (mode == PreparationMode::linearization) {
PrepareVariationData();
m_has_variation_data = true;
} else {
m_variationElements.clear();
}
m_is_prepared = true;
++m_preparation_count;
@@ -551,6 +583,8 @@ namespace mean_field::operators {
const mfem::Vector &gravity_gradient,
mfem::Vector &action
) const {
MEAN_FIELD_PROFILE_SCOPE("PreparedMappedHDivMassOperator::Mult");
MFEM_VERIFY(
m_is_prepared, "PreparedMappedHDivMassOperator must be prepared "
"before Mult is called."
@@ -564,15 +598,39 @@ namespace mean_field::operators {
"with the wrong size."
);
m_flux_true.SetSize(m_flux_map.full_size());
m_action_true.SetSize(m_flux_map.full_size());
m_domain_action_true.SetSize(m_flux_map.full_size());
m_flux_map.scatter(gravity_gradient, m_flux_true);
m_stellar_mass_form->Mult(m_flux_true, m_action_true);
m_vacuum_mass_form->Mult(m_flux_true, m_domain_action_true);
m_action_true += m_domain_action_true;
action.SetSize(Height());
m_flux_map.gather(m_action_true, action);
const mfem::Vector *gravity_gradient_true = &gravity_gradient;
if (!m_flux_map.is_identity()) [[unlikely]] {
m_flux_true.SetSize(m_flux_map.full_size());
m_flux_map.scatter(gravity_gradient, m_flux_true);
gravity_gradient_true = &m_flux_true;
}
mfem::Vector *action_true = &action;
if (!m_flux_map.is_identity()) [[unlikely]] {
m_action_true.SetSize(m_flux_map.full_size());
action_true = &m_action_true;
}
if (m_single_rank) [[likely]] {
action_true->SetSize(m_flux_map.full_size());
m_domain_action_true.SetSize(m_flux_map.full_size());
m_stellar_mass_form->Mult(*gravity_gradient_true, *action_true);
m_vacuum_mass_form->Mult(*gravity_gradient_true, m_domain_action_true);
*action_true += m_domain_action_true;
} else {
true_to_local(*m_fem.gravityFluxFes, *gravity_gradient_true, m_flux_local);
m_action_local.SetSize(m_fem.gravityFluxFes->GetVSize());
m_domain_action_local.SetSize(m_fem.gravityFluxFes->GetVSize());
m_stellar_mass_form->Mult(m_flux_local, m_action_local);
m_vacuum_mass_form->Mult(m_flux_local, m_domain_action_local);
m_action_local += m_domain_action_local;
local_to_true(*m_fem.gravityFluxFes, m_action_local, *action_true);
}
if (!m_flux_map.is_identity()) [[unlikely]] {
action.SetSize(Height());
m_flux_map.gather(m_action_true, action);
}
}
void PreparedMappedHDivMassOperator::MultDisplacementVariationTrue(
@@ -583,6 +641,11 @@ namespace mean_field::operators {
MFEM_VERIFY(
m_is_prepared, "PreparedMappedHDivMassOperator must be prepared before applying a displacement variation."
);
MFEM_VERIFY(
m_has_variation_data,
"PreparedMappedHDivMassOperator requires linearization preparation before applying a displacement "
"variation."
);
MFEM_VERIFY(
gravityGradientTrue.Size() == m_fem.gravityFluxFes->GetTrueVSize(),
"The full gravity-gradient vector has the wrong size."
@@ -707,6 +770,10 @@ namespace mean_field::operators {
return m_is_prepared;
}
bool PreparedMappedHDivMassOperator::HasVariationData() const noexcept {
return m_has_variation_data;
}
std::uint64_t PreparedMappedHDivMassOperator::GetPreparationCount() const noexcept {
return m_preparation_count;
}