perf(jacobian-action): major updates to jacobian action application by removing redudant quadrature work. ~5x increase in speed

This commit is contained in:
2026-09-02 17:01:50 -04:00
parent 85500fef3b
commit 25510008dd
74 changed files with 8967 additions and 814 deletions

View File

@@ -42,6 +42,25 @@ namespace {
}
}
void local_to_true(
const mfem::ParFiniteElementSpace &finiteElementSpace,
const mfem::Vector &localVector,
mfem::Vector &trueVector
) {
MFEM_VERIFY(localVector.Size() == finiteElementSpace.GetVSize(), "Local vector has the wrong size.");
trueVector.SetSize(finiteElementSpace.GetTrueVSize());
trueVector = 0.0;
const mfem::Operator *prolongation = finiteElementSpace.GetProlongationMatrix();
if (prolongation != nullptr) {
prolongation->MultTranspose(localVector, trueVector);
} else {
trueVector = localVector;
}
}
const mfem::IntegrationRule &get_mass_normalization_rule(
const mean_field::fem::FEM &f,
const mfem::FiniteElement &densityElement,
@@ -212,6 +231,13 @@ namespace mean_field::operators {
return report;
}
PreparedMassNormalizationReport PreparedMassNormalizationOperator::Prepare(
const models::CompiledFixedMass &constraint,
const MassNormalizationDependencies &dependencies
) {
return Prepare({.targetMass = constraint.targetMass().value()}, dependencies);
}
void PreparedMassNormalizationOperator::BuildStaticPlan() {
m_elements.clear();
m_elements.reserve(m_fem.mesh->GetNE());
@@ -432,6 +458,7 @@ namespace mean_field::operators {
true_to_local(*m_fem.displacementFes, displacementVariation, displacementVariationLocal);
mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());
mapping::VolumeMappingVariation variation;
mfem::Vector elementDisplacementVariation;
double localAction = 0.0;
@@ -464,8 +491,6 @@ namespace mean_field::operators {
mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);
for (const QuadraturePointData &point : data.quadraturePoints) {
mapping::VolumeMappingVariation variation;
const mapping::MappingStatus status = m_domainMapper.EvaluateVolumeVariation(
mappingData, directionData, *transformation, point.integrationPoint, point.mappingContext,
workspace, variation
@@ -558,6 +583,137 @@ namespace mean_field::operators {
++m_actionStatistics.completeApplications;
}
void PreparedMassNormalizationOperator::ApplyJacobian(
const FixedMassJacobianInput &input,
mfem::Vector &action
) const {
ApplyCompleteJacobianAction(input.densityVariation, input.displacementVariation, action);
}
void PreparedMassNormalizationOperator::AssembleDensityTransposeAction(
const double residualDual,
mfem::Vector &densityDual
) const {
mfem::Vector localDual(m_fem.densityFes->GetVSize());
localDual = 0.0;
mfem::Vector elementDual;
for (const ElementPAData &data : m_elements) {
elementDual.SetSize(data.densityDofs.Size());
elementDual = 0.0;
for (const QuadraturePointData &point : data.quadraturePoints) {
elementDual.Add(residualDual * point.mappingContext.quadrature.weight, point.densityShape);
}
if (data.densityDofTransformation != nullptr) {
data.densityDofTransformation->TransformDual(elementDual);
}
localDual.AddElementVector(data.densityDofs, elementDual);
}
mfem::Vector trueDual;
local_to_true(*m_fem.densityFes, localDual, trueDual);
densityDual.SetSize(m_gravityContext.GetDensityMap().reduced_size());
m_gravityContext.GetDensityMap().gather(trueDual, densityDual);
}
void PreparedMassNormalizationOperator::AssembleDisplacementTransposeAction(
const double residualDual,
mfem::Vector &displacementDual
) const {
mfem::Vector localDual(m_fem.displacementFes->GetVSize());
localDual = 0.0;
mapping::DomainMapper::Workspace workspace(m_fem.mesh->Dimension());
mapping::VolumeMappingVariation variation;
mfem::Vector elementDirection;
mfem::Vector elementDual;
for (const ElementPAData &data : m_elements) {
const mfem::FiniteElement &displacementElement = *m_fem.displacementFes->GetFE(data.elementId);
const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);
const mapping::ElementDisplacementData baseDisplacementData =
mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
const mapping::ElementCompactificationData compactificationData(
compactificationElement, data.compactification
);
const mapping::ElementMappingData mappingData{
.displacement = baseDisplacementData, .compactification = compactificationData
};
mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);
elementDirection.SetSize(data.displacementDofs.Size());
elementDual.SetSize(data.displacementDofs.Size());
elementDual = 0.0;
for (int elementDof = 0; elementDof < elementDirection.Size(); ++elementDof) {
elementDirection = 0.0;
elementDirection(elementDof) = 1.0;
const mapping::ElementDisplacementData directionData =
mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementDirection);
double elementDofAction = 0.0;
for (const QuadraturePointData &point : data.quadraturePoints) {
const mapping::MappingStatus status = m_domainMapper.EvaluateVolumeVariation(
mappingData, directionData, *transformation, point.integrationPoint, point.mappingContext,
workspace, variation
);
MFEM_VERIFY(
status == mapping::MappingStatus::valid,
"Stateless mapping variation failed in the mass-normalization transpose action. Element: "
<< data.elementId << ", status: " << static_cast<int>(status)
);
elementDofAction += point.density * variation.weight_variation;
}
elementDual(elementDof) = residualDual * elementDofAction;
}
if (data.displacementDofTransformation != nullptr) {
data.displacementDofTransformation->TransformDual(elementDual);
}
localDual.AddElementVector(data.displacementDofs, elementDual);
}
mfem::Vector trueDual;
local_to_true(*m_fem.displacementFes, localDual, trueDual);
displacementDual.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
m_gravityContext.GetDisplacementMap().gather(trueDual, displacementDual);
}
void PreparedMassNormalizationOperator::ApplyCompleteJacobianTransposeAction(
const double residualDual,
mfem::Vector &densityDual,
mfem::Vector &displacementDual
) const {
VerifyPrepared();
MFEM_VERIFY(std::isfinite(residualDual), "Mass-normalization transpose action received a non-finite dual.");
AssembleDensityTransposeAction(residualDual, densityDual);
AssembleDisplacementTransposeAction(residualDual, displacementDual);
++m_actionStatistics.transposeApplications;
}
void PreparedMassNormalizationOperator::ApplyJacobianTranspose(
const mfem::Vector &residualDual,
FixedMassJacobianTransposeOutput output
) const {
MFEM_VERIFY(residualDual.Size() == 1, "Fixed-mass transpose action requires one residual dual value.");
ApplyCompleteJacobianTransposeAction(residualDual(0), output.densityDual, output.displacementDual);
}
double PreparedMassNormalizationOperator::GlobalSum(const double localValue) const {
double globalValue = 0.0;
MPI_Allreduce(&localValue, &globalValue, 1, MPI_DOUBLE, MPI_SUM, m_fem.mesh->GetComm());
@@ -720,6 +876,44 @@ namespace mean_field::operators {
action(m_layout.offset(massResidual)) = massAction(0);
}
void PreparedMassNormalizationJacobianOperator::MultTranspose(
const mfem::Vector &residualDual,
mfem::Vector &stateDual
) const {
MFEM_VERIFY(
m_preparedOperator.IsPrepared(),
"Prepared mass-normalization MFEM adapter requires a prepared row operator."
);
MFEM_VERIFY(
residualDual.Size() == Height(),
"Prepared mass-normalization MFEM adapter received a residual dual with the wrong size."
);
using Form = utils::blocks::barotropic_equilibrium_form;
constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);
constexpr auto displacementValue =
utils::blocks::get_value_block<Form>(utils::blocks::displacement_field.geometry_term);
constexpr auto massResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);
mfem::Vector densityDual;
mfem::Vector displacementDual;
m_preparedOperator.ApplyCompleteJacobianTransposeAction(
residualDual(m_layout.offset(massResidual)), densityDual, displacementDual
);
stateDual.SetSize(Width());
stateDual = 0.0;
mfem::Vector densityBlock(stateDual.GetData() + m_layout.offset(densityValue), m_layout.size(densityValue));
densityBlock = densityDual;
mfem::Vector displacementBlock(
stateDual.GetData() + m_layout.offset(displacementValue), m_layout.size(displacementValue)
);
displacementBlock = displacementDual;
}
const MassNormalizationLayout &PreparedMassNormalizationJacobianOperator::GetLayout() const noexcept {
return m_layout;
}