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

feat(preconditioner): major work on preconditioner system

Browse Source 
first preconditioner MVP
This commit is contained in:
Emily Boudreaux
2026-09-04 07:54:10 -04:00
parent 25510008dd
commit 71423d543f
61 changed files with 15920 additions and 422 deletions
Download Patch File Download Diff File Expand all files Collapse all files

180
libmeanfield/impl/operators/contexts/gravity_field_context.cpp
View File

@@ -9,7 +9,132 @@ import :operators.context.gravity_field;
namespace {
using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
[[nodiscard]] std::unique_ptr<mfem::ParMixedBilinearForm> make_divergence_operator(const mean_field::fem::FEM &f) {
void true_to_local(
const mfem::ParFiniteElementSpace &finite_element_space,
const mfem::Vector &true_vector,
mfem::Vector &local_vector
) {
MFEM_VERIFY(
true_vector.Size() == finite_element_space.GetTrueVSize(),
"True-DOF operator received an input vector with the wrong size."
);
local_vector.SetSize(finite_element_space.GetVSize());
const mfem::Operator *prolongation = finite_element_space.GetProlongationMatrix();
if (prolongation != nullptr) {
prolongation->Mult(true_vector, local_vector);
} else {
local_vector = true_vector;
}
}
void local_to_true(
const mfem::ParFiniteElementSpace &finite_element_space,
const mfem::Vector &local_vector,
mfem::Vector &true_vector
) {
MFEM_VERIFY(
local_vector.Size() == finite_element_space.GetVSize(),
"True-DOF operator produced a local vector with the wrong size."
);
true_vector.SetSize(finite_element_space.GetTrueVSize());
const mfem::Operator *prolongation = finite_element_space.GetProlongationMatrix();
if (prolongation != nullptr) {
prolongation->MultTranspose(local_vector, true_vector);
} else {
true_vector = local_vector;
}
}
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;
}
class TrueDofParMixedBilinearFormOperator final : public mfem::Operator {
public:
TrueDofParMixedBilinearFormOperator(
const mfem::ParFiniteElementSpace &trial_space,
const mfem::ParFiniteElementSpace &test_space,
std::unique_ptr<mfem::ParMixedBilinearForm> local_form
)
: Operator(
test_space.GetTrueVSize(),
trial_space.GetTrueVSize()
),
m_trial_space(trial_space),
m_test_space(test_space),
m_local_form(std::move(local_form)),
m_single_rank(communicator_has_single_rank(trial_space.GetComm())) {
int communicators_compare = MPI_UNEQUAL;
MFEM_VERIFY(
MPI_Comm_compare(trial_space.GetComm(), test_space.GetComm(), &communicators_compare) == MPI_SUCCESS,
"Failed to compare mixed-operator MPI communicators."
);
MFEM_VERIFY(
communicators_compare == MPI_IDENT || communicators_compare == MPI_CONGRUENT,
"True-DOF mixed operator requires congruent trial and test communicators."
);
MFEM_VERIFY(m_local_form != nullptr, "True-DOF mixed operator requires a local bilinear form.");
MFEM_VERIFY(
m_local_form->Width() == m_trial_space.GetVSize(),
"True-DOF mixed operator received an incompatible trial space."
);
MFEM_VERIFY(
m_local_form->Height() == m_test_space.GetVSize(),
"True-DOF mixed operator received an incompatible test space."
);
}
void Mult(
const mfem::Vector &input,
mfem::Vector &output
) const override {
MFEM_VERIFY(input.Size() == Width(), "True-DOF mixed operator received an input with the wrong size.");
if (m_single_rank) [[likely]] {
output.SetSize(Height());
m_local_form->Mult(input, output);
return;
}
true_to_local(m_trial_space, input, m_trial_local);
m_test_local.SetSize(m_test_space.GetVSize());
m_local_form->Mult(m_trial_local, m_test_local);
local_to_true(m_test_space, m_test_local, output);
}
void MultTranspose(
const mfem::Vector &input,
mfem::Vector &output
) const override {
MFEM_VERIFY(input.Size() == Height(), "True-DOF mixed transpose received an input with the wrong size.");
if (m_single_rank) [[likely]] {
output.SetSize(Width());
m_local_form->MultTranspose(input, output);
return;
}
true_to_local(m_test_space, input, m_test_local);
m_trial_local.SetSize(m_trial_space.GetVSize());
m_local_form->MultTranspose(m_test_local, m_trial_local);
local_to_true(m_trial_space, m_trial_local, output);
}
private:
const mfem::ParFiniteElementSpace &m_trial_space;
const mfem::ParFiniteElementSpace &m_test_space;
std::unique_ptr<mfem::ParMixedBilinearForm> m_local_form;
mutable mfem::Vector m_trial_local;
mutable mfem::Vector m_test_local;
bool m_single_rank;
};
[[nodiscard]] std::unique_ptr<mfem::Operator> make_divergence_operator(const mean_field::fem::FEM &f) {
auto divergence =
std::make_unique<mfem::ParMixedBilinearForm>(f.gravityFluxFes.get(), f.gravityPotentialFes.get());
@@ -29,7 +154,9 @@ namespace {
divergence->AddDomainIntegrator(integrator.release());
divergence->Assemble();
return divergence;
return std::make_unique<TrueDofParMixedBilinearFormOperator>(
*f.gravityFluxFes, *f.gravityPotentialFes, std::move(divergence)
);
}
void validate_displacement(
@@ -167,6 +294,25 @@ namespace mean_field::operators::context::gravity_field {
const mfem::Vector &displacement,
const DiscretizationRevision discretization_revision,
const DisplacementRevision displacement_revision
) {
return PrepareImpl(
displacement, discretization_revision, displacement_revision, PreparationMode::linearization
);
}
GravityFieldGeometryPreparation GravityFieldGeometryContext::PreparePrimal(
const mfem::Vector &displacement,
const DiscretizationRevision discretization_revision,
const DisplacementRevision displacement_revision
) {
return PrepareImpl(displacement, discretization_revision, displacement_revision, PreparationMode::primal);
}
GravityFieldGeometryPreparation GravityFieldGeometryContext::PrepareImpl(
const mfem::Vector &displacement,
const DiscretizationRevision discretization_revision,
const DisplacementRevision displacement_revision,
const PreparationMode mode
) {
validate_displacement(m_displacement_map, displacement);
@@ -185,21 +331,38 @@ namespace mean_field::operators::context::gravity_field {
const bool discretization_changed = !m_is_prepared || discretization_revision != m_discretization_revision;
const bool displacement_changed = !m_is_prepared || displacement_revision != m_displacement_revision;
const bool requires_variation = mode == PreparationMode::linearization;
const bool variation_upgrade = requires_variation && !m_variation_state_prepared;
GravityFieldGeometryPreparation preparation;
if (!discretization_changed && !displacement_changed) {
if (!discretization_changed && !displacement_changed && !variation_upgrade) {
return preparation;
}
const auto prepare_mass = [&](PreparedMappedHDivMassOperator &mass_operator) {
if (requires_variation) {
mass_operator.Prepare(displacement);
} else {
mass_operator.PreparePrimal(displacement);
}
};
const auto prepare_source = [&](PreparedMappedGravitySourceOperator &source_operator) {
if (requires_variation) {
source_operator.Prepare(displacement);
} else {
source_operator.PreparePrimal(displacement);
}
};
if (discretization_changed) {
auto mass_operator = std::make_unique<PreparedMappedHDivMassOperator>(m_fem, m_domain_mapper);
auto source_operator = std::make_unique<PreparedMappedGravitySourceOperator>(m_fem, m_domain_mapper);
auto divergence_operator = make_divergence_operator(m_fem);
auto transpose_divergence_operator = std::make_unique<mfem::TransposeOperator>(divergence_operator.get());
mass_operator->Prepare(displacement);
source_operator->Prepare(displacement);
prepare_mass(*mass_operator);
prepare_source(*source_operator);
m_mass_operator = std::move(mass_operator);
m_source_operator = std::move(source_operator);
@@ -220,8 +383,8 @@ namespace mean_field::operators::context::gravity_field {
"operator."
);
m_mass_operator->Prepare(displacement);
m_source_operator->Prepare(displacement);
prepare_mass(*m_mass_operator);
prepare_source(*m_source_operator);
preparation.rebuilt_mass_operator = true;
preparation.rebuilt_source_operator = true;
@@ -232,8 +395,9 @@ namespace mean_field::operators::context::gravity_field {
m_discretization_revision = discretization_revision;
m_displacement_revision = displacement_revision;
m_is_prepared = true;
m_variation_state_prepared = requires_variation;
preparation.refreshed_variation_state = true;
preparation.refreshed_variation_state = requires_variation;
return preparation;
}

38
libmeanfield/impl/operators/gravity_field.cpp
View File

@@ -346,20 +346,32 @@ namespace mean_field::operators {
make_residual_view(action, m_residual_offsets, gravity_poisson_residual_block);
const field::FieldDofMap &flux_map = geometry_context.GetMassOperator().GetFluxMap();
const field::FieldDofMap &potential_map = geometry_context.GetSourceOperator().GetPotentialMap();
mfem::Vector potential_true(potential_map.full_size());
mfem::Vector transpose_divergence_action_true(flux_map.full_size());
mfem::Vector transpose_divergence_action(flux_map.reduced_size());
mfem::Vector gradient_true(flux_map.full_size());
mfem::Vector divergence_action_true(potential_map.full_size());
geometry_context.GetMassOperator().Mult(gravity_gradient, gravity_gradient_action);
potential_map.scatter(gravity_potential, potential_true);
geometry_context.GetTransposeDivergenceOperator().Mult(potential_true, transpose_divergence_action_true);
flux_map.gather(transpose_divergence_action_true, transpose_divergence_action);
gravity_gradient_action += transpose_divergence_action;
flux_map.scatter(gravity_gradient, gradient_true);
geometry_context.GetDivergenceOperator().Mult(gradient_true, divergence_action_true);
potential_map.gather(divergence_action_true, gravity_poisson_action);
if (flux_map.is_identity() && potential_map.is_identity()) [[likely]] {
m_transpose_divergence_action_true.SetSize(flux_map.full_size());
geometry_context.GetTransposeDivergenceOperator().Mult(
gravity_potential, m_transpose_divergence_action_true
);
gravity_gradient_action += m_transpose_divergence_action_true;
geometry_context.GetDivergenceOperator().Mult(gravity_gradient, gravity_poisson_action);
return;
}
m_potential_true.SetSize(potential_map.full_size());
m_transpose_divergence_action_true.SetSize(flux_map.full_size());
m_transpose_divergence_action.SetSize(flux_map.reduced_size());
m_gradient_true.SetSize(flux_map.full_size());
m_divergence_action_true.SetSize(potential_map.full_size());
potential_map.scatter(gravity_potential, m_potential_true);
geometry_context.GetTransposeDivergenceOperator().Mult(m_potential_true, m_transpose_divergence_action_true);
flux_map.gather(m_transpose_divergence_action_true, m_transpose_divergence_action);
gravity_gradient_action += m_transpose_divergence_action;
flux_map.scatter(gravity_gradient, m_gradient_true);
geometry_context.GetDivergenceOperator().Mult(m_gradient_true, m_divergence_action_true);
potential_map.gather(m_divergence_action_true, gravity_poisson_action);
}
void GravityFieldOperator::ApplyDensitySource(
@@ -529,7 +541,7 @@ namespace mean_field::operators {
++displacement_revision.value;
}
m_gravity_field_geometry_context.Prepare(displacement, discretization_revision, displacement_revision);
m_gravity_field_geometry_context.PreparePrimal(displacement, discretization_revision, displacement_revision);
m_displacement = displacement;
}

3
libmeanfield/impl/operators/gravity_field_jacobian.cpp
View File

@@ -267,6 +267,9 @@ namespace mean_field::operators {
gravity_poisson_action -= source_action;
gravity_poisson_action -= source_variation_action;
gravity_gradient_action.SyncAliasMemory(action);
gravity_poisson_action.SyncAliasMemory(action);
}
const context::gravity_field::GravityFieldLinearizationContext &

29
libmeanfield/impl/operators/prepared_barotropic_closure.cpp
View File

@@ -469,6 +469,35 @@ namespace mean_field::operators {
m_densityMap.gather(m_fullResidual, residual);
}
void PreparedBarotropicClosureOperator::AssembleDensityJacobianDiagonal(mfem::Vector &diagonal) const {
VerifyPrepared();
mfem::Vector localDiagonal(m_fem.densityFes->GetVSize());
localDiagonal = 0.0;
mfem::Vector elementDiagonal;
for (const ElementPAData &data : m_elements) {
MFEM_VERIFY(
data.densityDofTransformation == nullptr,
"Density mass-diagonal assembly currently requires scalar L2 element DOFs without a DOF transform."
);
elementDiagonal.SetSize(data.densityDofs.Size());
elementDiagonal = 0.0;
for (int trialDof = 0; trialDof < data.densityDofs.Size(); ++trialDof) {
for (int quadraturePoint = 0; quadraturePoint < data.quadratureWeights.Size(); ++quadraturePoint) {
const double basis = data.densityBasis(quadraturePoint, trialDof);
elementDiagonal(trialDof) += data.quadratureWeights(quadraturePoint) * basis * basis;
}
}
localDiagonal.AddElementVector(data.densityDofs, elementDiagonal);
}
mfem::Vector trueDiagonal;
local_to_true(*m_fem.densityFes, localDiagonal, trueDiagonal);
diagonal.SetSize(m_densityMap.reduced_size());
m_densityMap.gather(trueDiagonal, diagonal);
}
void PreparedBarotropicClosureOperator::Mult(
const mfem::Vector &densityVariation,
const mfem::Vector &enthalpyVariation,

170
libmeanfield/impl/operators/prepared_gravity_source.cpp
View File

@@ -1,4 +1,5 @@
module;
#include "profile.h"
#include <cmath>
#include <cstdint>
#include <memory>
@@ -304,6 +305,19 @@ namespace mean_field::operators {
}
void PreparedMappedGravitySourceOperator::Prepare(const mfem::Vector &displacement) {
MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedGravitySourceOperator::Prepare linearization", 0);
PrepareImpl(displacement, PreparationMode::linearization);
}
void PreparedMappedGravitySourceOperator::PreparePrimal(const mfem::Vector &displacement) {
MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedGravitySourceOperator::Prepare primal", 0);
PrepareImpl(displacement, PreparationMode::primal);
}
void PreparedMappedGravitySourceOperator::PrepareImpl(
const mfem::Vector &displacement,
const PreparationMode mode
) {
MFEM_VERIFY(
displacement.Size() == m_displacement_map.reduced_size(),
"PreparedMappedGravitySourceOperator received a displacement "
@@ -318,7 +332,8 @@ namespace mean_field::operators {
);
}
m_is_prepared = false;
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);
m_elements.clear();
@@ -343,8 +358,10 @@ namespace mean_field::operators {
data.potential_dof_transformation =
m_fem.gravityPotentialFes->GetElementDofs(element_id, data.potential_dofs);
data.displacement_dof_transformation =
m_fem.displacementFes->GetElementVDofs(element_id, data.displacement_dofs);
if (mode == PreparationMode::linearization) {
data.displacement_dof_transformation =
m_fem.displacementFes->GetElementVDofs(element_id, data.displacement_dofs);
}
const mfem::FiniteElement &density_element = *m_fem.densityFes->GetFE(element_id);
@@ -367,7 +384,9 @@ namespace mean_field::operators {
data.potential_basis.SetSize(quadrature_point_count, potential_dof_count);
const int dimension = m_fem.mesh->Dimension();
data.inverse_element_jacobians.SetSize(quadrature_point_count, dimension * dimension);
if (mode == PreparationMode::linearization) {
data.inverse_element_jacobians.SetSize(quadrature_point_count, dimension * dimension);
}
data.quadrature_data.SetSize(quadrature_point_count);
@@ -395,11 +414,13 @@ namespace mean_field::operators {
const double coefficient_value = source_coefficient.Eval(transformation, integration_point);
const mfem::DenseMatrix &inverse_element_jacobian = source_coefficient.GetInverseElementJacobian();
for (int row = 0; row < dimension; ++row) {
for (int column = 0; column < dimension; ++column) {
data.inverse_element_jacobians(quadrature_point, row * dimension + column) =
inverse_element_jacobian(row, column);
if (mode == PreparationMode::linearization) {
const mfem::DenseMatrix &inverse_element_jacobian = source_coefficient.GetInverseElementJacobian();
for (int row = 0; row < dimension; ++row) {
for (int column = 0; column < dimension; ++column) {
data.inverse_element_jacobians(quadrature_point, row * dimension + column) =
inverse_element_jacobian(row, column);
}
}
}
@@ -420,13 +441,16 @@ namespace mean_field::operators {
MFEM_VERIFY(!m_elements.empty(), "PreparedMappedGravitySourceOperator found no stellar elements.");
m_is_prepared = true;
m_is_prepared = true;
m_has_variation_data = mode == PreparationMode::linearization;
++m_preparation_count;
}
void PreparedMappedGravitySourceOperator::Mult(
const mfem::Vector &density,
mfem::Vector &action
) const {
MEAN_FIELD_PROFILE_SCOPE("PreparedMappedGravitySourceOperator::Mult");
MFEM_VERIFY(
m_is_prepared, "PreparedMappedGravitySourceOperator must be prepared before "
"Mult is called."
@@ -440,49 +464,47 @@ namespace mean_field::operators {
m_density_true.SetSize(m_density_map.full_size());
m_density_map.scatter(density, m_density_true);
mfem::Vector density_local;
true_to_local(*m_fem.densityFes, m_density_true, m_density_local);
true_to_local(*m_fem.densityFes, m_density_true, density_local);
mfem::Vector local_action(m_fem.gravityPotentialFes->GetVSize());
local_action = 0.0;
mfem::Vector element_density;
mfem::Vector quadrature_density;
mfem::Vector element_action;
m_local_action.SetSize(m_fem.gravityPotentialFes->GetVSize());
m_local_action = 0.0;
for (const ElementPAData &data : m_elements) {
density_local.GetSubVector(data.density_dofs, element_density);
m_density_local.GetSubVector(data.density_dofs, m_element_input);
if (data.density_dof_transformation != nullptr) {
data.density_dof_transformation->InvTransformPrimal(element_density);
data.density_dof_transformation->InvTransformPrimal(m_element_input);
}
quadrature_density.SetSize(data.quadrature_data.Size());
m_quadrature_action.SetSize(data.quadrature_data.Size());
// B_density * x_e
data.density_basis.Mult(element_density, quadrature_density);
data.density_basis.Mult(m_element_input, m_quadrature_action);
// D * B_density * x_e
for (int q = 0; q < quadrature_density.Size(); ++q) {
quadrature_density(q) *= data.quadrature_data(q);
for (int q = 0; q < m_quadrature_action.Size(); ++q) {
m_quadrature_action(q) *= data.quadrature_data(q);
}
element_action.SetSize(data.potential_dofs.Size());
m_element_action.SetSize(data.potential_dofs.Size());
// B_potential^T * D * B_density * x_e
data.potential_basis.MultTranspose(quadrature_density, element_action);
data.potential_basis.MultTranspose(m_quadrature_action, m_element_action);
if (data.potential_dof_transformation != nullptr) {
data.potential_dof_transformation->TransformDual(element_action);
data.potential_dof_transformation->TransformDual(m_element_action);
}
local_action.AddElementVector(data.potential_dofs, element_action);
m_local_action.AddElementVector(data.potential_dofs, m_element_action);
}
local_to_true(*m_fem.gravityPotentialFes, local_action, m_action_true);
action.SetSize(Height());
m_potential_map.gather(m_action_true, action);
if (m_potential_map.is_identity()) {
local_to_true(*m_fem.gravityPotentialFes, m_local_action, action);
} else {
local_to_true(*m_fem.gravityPotentialFes, m_local_action, m_action_true);
action.SetSize(Height());
m_potential_map.gather(m_action_true, action);
}
}
void PreparedMappedGravitySourceOperator::MultDisplacementVariationTrue(
@@ -494,6 +516,11 @@ namespace mean_field::operators {
m_is_prepared,
"PreparedMappedGravitySourceOperator must be prepared before applying a displacement variation."
);
MFEM_VERIFY(
m_has_variation_data,
"PreparedMappedGravitySourceOperator requires linearization preparation before applying a displacement "
"variation."
);
MFEM_VERIFY(
densityTrue.Size() == m_fem.densityFes->GetTrueVSize(), "The full density vector has the wrong size."
);
@@ -591,47 +618,44 @@ namespace mean_field::operators {
"with the wrong size."
);
m_potential_true.SetSize(m_potential_map.full_size());
m_potential_map.scatter(potential, m_potential_true);
mfem::Vector potential_local;
true_to_local(*m_fem.gravityPotentialFes, m_potential_true, potential_local);
mfem::Vector local_action(m_fem.densityFes->GetVSize());
local_action = 0.0;
mfem::Vector element_potential;
mfem::Vector quadrature_potential;
mfem::Vector element_action;
for (const ElementPAData &data : m_elements) {
potential_local.GetSubVector(data.potential_dofs, element_potential);
if (data.potential_dof_transformation != nullptr) {
data.potential_dof_transformation->InvTransformPrimal(element_potential);
}
quadrature_potential.SetSize(data.quadrature_data.Size());
data.potential_basis.Mult(element_potential, quadrature_potential);
for (int q = 0; q < quadrature_potential.Size(); ++q) {
quadrature_potential(q) *= data.quadrature_data(q);
}
element_action.SetSize(data.density_dofs.Size());
data.density_basis.MultTranspose(quadrature_potential, element_action);
if (data.density_dof_transformation != nullptr) {
data.density_dof_transformation->TransformDual(element_action);
}
local_action.AddElementVector(data.density_dofs, element_action);
if (m_potential_map.is_identity()) {
true_to_local(*m_fem.gravityPotentialFes, potential, m_potential_local);
} else {
m_potential_true.SetSize(m_potential_map.full_size());
m_potential_map.scatter(potential, m_potential_true);
true_to_local(*m_fem.gravityPotentialFes, m_potential_true, m_potential_local);
}
local_to_true(*m_fem.densityFes, local_action, m_action_true);
m_local_action.SetSize(m_fem.densityFes->GetVSize());
m_local_action = 0.0;
for (const ElementPAData &data : m_elements) {
m_potential_local.GetSubVector(data.potential_dofs, m_element_input);
if (data.potential_dof_transformation != nullptr) {
data.potential_dof_transformation->InvTransformPrimal(m_element_input);
}
m_quadrature_action.SetSize(data.quadrature_data.Size());
data.potential_basis.Mult(m_element_input, m_quadrature_action);
for (int q = 0; q < m_quadrature_action.Size(); ++q) {
m_quadrature_action(q) *= data.quadrature_data(q);
}
m_element_action.SetSize(data.density_dofs.Size());
data.density_basis.MultTranspose(m_quadrature_action, m_element_action);
if (data.density_dof_transformation != nullptr) {
data.density_dof_transformation->TransformDual(m_element_action);
}
m_local_action.AddElementVector(data.density_dofs, m_element_action);
}
local_to_true(*m_fem.densityFes, m_local_action, m_action_true);
action.SetSize(Width());
m_density_map.gather(m_action_true, action);
}
@@ -639,6 +663,10 @@ namespace mean_field::operators {
return m_is_prepared;
}
bool PreparedMappedGravitySourceOperator::HasVariationData() const noexcept {
return m_has_variation_data;
}
std::uint64_t PreparedMappedGravitySourceOperator::GetPreparationCount() const noexcept {
return m_preparation_count;
}

89
libmeanfield/impl/operators/prepared_hdiv_mass.cpp
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;
}

30
libmeanfield/impl/operators/prepared_hydrostatic_equilibrium.cpp
View File

@@ -787,6 +787,36 @@ namespace mean_field::operators {
++m_algebraicJacobianStatistics.enthalpyApplications;
}
void PreparedHydrostaticEquilibriumOperator::AssembleEnthalpyJacobianDiagonal(mfem::Vector &diagonal) const {
VerifyPrepared();
mfem::Vector localDiagonal(m_fem.enthalpyFes->GetVSize());
localDiagonal = 0.0;
mfem::Vector elementDiagonal;
for (const ElementPAData &data : m_elements) {
MFEM_VERIFY(
data.enthalpyDofTransformation == nullptr,
"Enthalpy mass-diagonal assembly currently requires scalar H1 element DOFs without a DOF transform."
);
MFEM_VERIFY(
data.enthalpyJacobian.Height() == data.enthalpyDofs.Size() &&
data.enthalpyJacobian.Width() == data.enthalpyDofs.Size(),
"The prepared enthalpy Jacobian block is not square on an element."
);
elementDiagonal.SetSize(data.enthalpyDofs.Size());
for (int dof = 0; dof < data.enthalpyDofs.Size(); ++dof) {
elementDiagonal(dof) = data.enthalpyJacobian(dof, dof);
}
localDiagonal.AddElementVector(data.enthalpyDofs, elementDiagonal);
}
mfem::Vector trueDiagonal;
local_to_true(*m_fem.enthalpyFes, localDiagonal, trueDiagonal);
diagonal.SetSize(m_context.GetEnthalpyMap().reduced_size());
m_context.GetEnthalpyMap().gather(trueDiagonal, diagonal);
}
void PreparedHydrostaticEquilibriumOperator::ApplyGravityPotentialJacobianAction(
const mfem::Vector &gravityPotentialVariation,
mfem::Vector &action