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

@@ -35,6 +35,128 @@ namespace {
}
}
void local_to_true(
const mfem::ParFiniteElementSpace &finite_element_space,
const mfem::Vector &local_vector,
mfem::Vector &true_vector
) {
true_vector.SetSize(finite_element_space.GetTrueVSize());
true_vector = 0.0;
const mfem::Operator *prolongation = finite_element_space.GetProlongationMatrix();
if (prolongation != nullptr) {
prolongation->MultTranspose(local_vector, true_vector);
} else {
true_vector = local_vector;
}
}
mean_field::quadrature::MappingKind get_mapping_kind(
const mean_field::mapping::DomainMapper &domain_mapper,
const mfem::ElementTransformation &transformation
) {
return domain_mapper.IsCompactifiedElement(transformation) ? mean_field::quadrature::MappingKind::kelvin
: mean_field::quadrature::MappingKind::general;
}
const mfem::IntegrationRule &get_hdiv_mass_rule(
const mean_field::fem::FEM &f,
const mean_field::mapping::DomainMapper &domain_mapper,
const mfem::FiniteElement &element,
const mfem::ElementTransformation &transformation
) {
using GravityField = mean_field::field::Field<mean_field::field::Gravity>;
const mean_field::quadrature::Query query =
GravityField::make_query<mean_field::field::Gravity::Form::HDivMass>(
mean_field::quadrature::QuadratureRole::discretization, transformation.OrderW(), {},
mean_field::utils::DOMAINS::ALL, get_mapping_kind(domain_mapper, transformation)
);
const auto resolution = f.quadratureFactory->get(query, transformation.GetGeometryType());
MFEM_VERIFY(
resolution.integration_rule != nullptr,
"The quadrature policy did not return an H(div) mass integration rule."
);
return *resolution.integration_rule;
}
int frozen_mapping_width(const int dimension) {
return 3 * dimension + 4 * dimension * dimension + 3;
}
void freeze_mapping_context(
const mean_field::mapping::VolumeMappingContext &context,
const int quadrature_point,
mfem::DenseMatrix &data
) {
const int dimension = context.mapping.reference_position.Size();
const int displacement_jacobian_start = 3 * dimension;
const int mapping_jacobian_start = displacement_jacobian_start + dimension * dimension;
const int inverse_mapping_start = mapping_jacobian_start + dimension * dimension;
const int inverse_element_start = inverse_mapping_start + dimension * dimension;
const int scalar_start = inverse_element_start + dimension * dimension;
for (int component = 0; component < dimension; ++component) {
data(quadrature_point, component) = context.mapping.reference_position(component);
data(quadrature_point, dimension + component) = context.mapping.displaced_position(component);
data(quadrature_point, 2 * dimension + component) = context.mapping.physical_position(component);
}
for (int row = 0; row < dimension; ++row) {
for (int column = 0; column < dimension; ++column) {
const int entry = row * dimension + column;
data(quadrature_point, displacement_jacobian_start + entry) =
context.mapping.displacement_jacobian(row, column);
data(quadrature_point, mapping_jacobian_start + entry) = context.mapping.mapping_jacobian(row, column);
data(quadrature_point, inverse_mapping_start + entry) =
context.mapping.inverse_mapping_jacobian(row, column);
data(quadrature_point, inverse_element_start + entry) = context.quadrature.J_inv(row, column);
}
}
data(quadrature_point, scalar_start) = context.mapping.mapping_determinant;
data(quadrature_point, scalar_start + 1) = context.quadrature.weight;
data(quadrature_point, scalar_start + 2) = context.mapping.compactified ? 1.0 : 0.0;
}
void thaw_mapping_context(
const mfem::DenseMatrix &data,
const int quadrature_point,
const int dimension,
mean_field::mapping::VolumeMappingContext &context
) {
const int displacement_jacobian_start = 3 * dimension;
const int mapping_jacobian_start = displacement_jacobian_start + dimension * dimension;
const int inverse_mapping_start = mapping_jacobian_start + dimension * dimension;
const int inverse_element_start = inverse_mapping_start + dimension * dimension;
const int scalar_start = inverse_element_start + dimension * dimension;
context.mapping.reference_position.SetSize(dimension);
context.mapping.displaced_position.SetSize(dimension);
context.mapping.physical_position.SetSize(dimension);
context.mapping.displacement_jacobian.SetSize(dimension, dimension);
context.mapping.mapping_jacobian.SetSize(dimension, dimension);
context.mapping.inverse_mapping_jacobian.SetSize(dimension, dimension);
context.quadrature.J_inv.SetSize(dimension, dimension);
for (int component = 0; component < dimension; ++component) {
context.mapping.reference_position(component) = data(quadrature_point, component);
context.mapping.displaced_position(component) = data(quadrature_point, dimension + component);
context.mapping.physical_position(component) = data(quadrature_point, 2 * dimension + component);
}
for (int row = 0; row < dimension; ++row) {
for (int column = 0; column < dimension; ++column) {
const int entry = row * dimension + column;
context.mapping.displacement_jacobian(row, column) =
data(quadrature_point, displacement_jacobian_start + entry);
context.mapping.mapping_jacobian(row, column) = data(quadrature_point, mapping_jacobian_start + entry);
context.mapping.inverse_mapping_jacobian(row, column) =
data(quadrature_point, inverse_mapping_start + entry);
context.quadrature.J_inv(row, column) = data(quadrature_point, inverse_element_start + entry);
}
}
context.mapping.mapping_determinant = data(quadrature_point, scalar_start);
context.mapping.compactified = data(quadrature_point, scalar_start + 2) != 0.0;
context.quadrature.detJ = context.mapping.mapping_determinant;
context.quadrature.weight = data(quadrature_point, scalar_start + 1);
}
int find_representative_element(
const mean_field::fem::FEM &f,
const mfem::Array<int> &marker
@@ -238,7 +360,8 @@ namespace mean_field::operators {
field::make_field_dof_map<
field::Displacement,
DomainSchema>(*f.displacementFes)
) {
),
m_variationWorkspace(domain_mapper.GetDimension()) {
MFEM_VERIFY(f.mesh != nullptr, "PreparedMappedHDivMassOperator requires a mesh.");
MFEM_VERIFY(
f.gravityFluxFes != nullptr, "PreparedMappedHDivMassOperator requires the "
@@ -285,6 +408,78 @@ namespace mean_field::operators {
validate_uniform_domain_discretization(f, m_vacuum_marker, vacuum_element_id);
}
void PreparedMappedHDivMassOperator::PrepareVariationData() {
m_variationElements.clear();
m_variationElements.reserve(m_fem.mesh->GetNE());
mfem::Vector displacementLocal;
true_to_local(*m_fem.displacementFes, m_displacement_true, displacementLocal);
mfem::Vector elementDisplacement;
mfem::Vector elementCompactification;
mapping::VolumeMappingContext mappingContext;
for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
m_variationElements.emplace_back();
ElementVariationData &data = m_variationElements.back();
data.elementId = elementId;
data.gravityGradientDofTransformation =
m_fem.gravityFluxFes->GetElementVDofs(elementId, data.gravityGradientDofs);
data.displacementDofTransformation =
m_fem.displacementFes->GetElementVDofs(elementId, data.displacementDofs);
mfem::DofTransformation *compactificationDofTransformation =
m_fem.compactificationFes->GetElementDofs(elementId, data.compactificationDofs);
displacementLocal.GetSubVector(data.displacementDofs, elementDisplacement);
m_fem.compactificationCoordinate->GetSubVector(data.compactificationDofs, elementCompactification);
if (data.displacementDofTransformation != nullptr) {
data.displacementDofTransformation->InvTransformPrimal(elementDisplacement);
}
if (compactificationDofTransformation != nullptr) {
compactificationDofTransformation->InvTransformPrimal(elementCompactification);
}
data.baseDisplacement = elementDisplacement;
data.compactification = elementCompactification;
const mfem::FiniteElement &gravityGradientElement = *m_fem.gravityFluxFes->GetFE(elementId);
const mfem::FiniteElement &displacementElement = *m_fem.displacementFes->GetFE(elementId);
const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(elementId);
mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(elementId);
MFEM_VERIFY(
transformation != nullptr, "Prepared H(div) variation data received a null element transformation."
);
data.integrationRule = &get_hdiv_mass_rule(m_fem, m_domain_mapper, gravityGradientElement, *transformation);
data.frozenMappingData.SetSize(
data.integrationRule->GetNPoints(), frozen_mapping_width(m_domain_mapper.GetDimension())
);
const mapping::ElementDisplacementData displacementData =
mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
const mapping::ElementCompactificationData compactificationData(
compactificationElement, data.compactification
);
const mapping::ElementMappingData mappingData{
.displacement = displacementData, .compactification = compactificationData
};
for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
const mapping::MappingStatus status = m_domain_mapper.EvaluateVolume(
mappingData, *transformation, integrationPoint, m_variationWorkspace, mappingContext
);
MFEM_VERIFY(
status == mapping::MappingStatus::valid,
"Prepared H(div) variation data encountered an invalid mapping. Element: "
<< elementId << ", quadrature point: " << quadraturePoint
<< ", status: " << static_cast<int>(status)
);
freeze_mapping_context(mappingContext, quadraturePoint, data.frozenMappingData);
}
}
}
void PreparedMappedHDivMassOperator::Prepare(const mfem::Vector &displacement) {
MFEM_VERIFY(
displacement.Size() == m_displacement_map.reduced_size(),
@@ -346,6 +541,8 @@ namespace mean_field::operators {
m_stellar_mass_form->Assemble();
m_vacuum_mass_form->Assemble();
PrepareVariationData();
m_is_prepared = true;
++m_preparation_count;
}
@@ -378,6 +575,110 @@ namespace mean_field::operators {
m_flux_map.gather(m_action_true, action);
}
void PreparedMappedHDivMassOperator::MultDisplacementVariationTrue(
const mfem::Vector &gravityGradientTrue,
const mfem::Vector &displacementVariationTrue,
mfem::Vector &actionVariationTrue
) const {
MFEM_VERIFY(
m_is_prepared, "PreparedMappedHDivMassOperator must be prepared before applying a displacement variation."
);
MFEM_VERIFY(
gravityGradientTrue.Size() == m_fem.gravityFluxFes->GetTrueVSize(),
"The full gravity-gradient vector has the wrong size."
);
MFEM_VERIFY(
displacementVariationTrue.Size() == m_fem.displacementFes->GetTrueVSize(),
"The full displacement variation has the wrong size."
);
true_to_local(*m_fem.gravityFluxFes, gravityGradientTrue, m_gravityGradientLocal);
true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacementVariationLocal);
m_localVariationAction.SetSize(m_fem.gravityFluxFes->GetVSize());
m_localVariationAction = 0.0;
const int dimension = m_domain_mapper.GetDimension();
for (const ElementVariationData &data : m_variationElements) {
MFEM_VERIFY(
data.integrationRule != nullptr &&
data.frozenMappingData.Height() == data.integrationRule->GetNPoints() &&
data.frozenMappingData.Width() == frozen_mapping_width(dimension),
"Prepared H(div) variation data is incomplete."
);
m_gravityGradientLocal.GetSubVector(data.gravityGradientDofs, m_elementGravityGradient);
m_displacementVariationLocal.GetSubVector(data.displacementDofs, m_elementDisplacementVariation);
if (data.gravityGradientDofTransformation != nullptr) {
data.gravityGradientDofTransformation->InvTransformPrimal(m_elementGravityGradient);
}
if (data.displacementDofTransformation != nullptr) {
data.displacementDofTransformation->InvTransformPrimal(m_elementDisplacementVariation);
}
const mfem::FiniteElement &gravityGradientElement = *m_fem.gravityFluxFes->GetFE(data.elementId);
const mfem::FiniteElement &displacementElement = *m_fem.displacementFes->GetFE(data.elementId);
const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);
mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);
MFEM_VERIFY(
transformation != nullptr,
"Prepared H(div) displacement variation received a null element transformation."
);
const mapping::ElementDisplacementData baseDisplacementData =
mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
const mapping::ElementDisplacementData directionData =
mapping::ElementDisplacementDataFromElementVDofs(displacementElement, m_elementDisplacementVariation);
const mapping::ElementCompactificationData compactificationData(
compactificationElement, data.compactification
);
const mapping::ElementMappingData mappingData{
.displacement = baseDisplacementData, .compactification = compactificationData
};
m_elementVariationAction.SetSize(gravityGradientElement.GetDof());
m_elementVariationAction = 0.0;
m_gravityGradientValue.SetSize(dimension);
m_massTensorVariationAction.SetSize(dimension);
m_gravityGradientShape.SetSize(gravityGradientElement.GetDof(), dimension);
m_massTensorVariation.SetSize(dimension, dimension);
for (int quadraturePoint = 0; quadraturePoint < data.integrationRule->GetNPoints(); ++quadraturePoint) {
const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
thaw_mapping_context(data.frozenMappingData, quadraturePoint, dimension, m_baseMappingContext);
const mapping::MappingStatus status = m_domain_mapper.EvaluateVolumeVariation(
mappingData, directionData, *transformation, integrationPoint, m_baseMappingContext,
m_variationWorkspace, m_mappingVariation
);
MFEM_VERIFY(
status == mapping::MappingStatus::valid,
"Prepared H(div) displacement variation encountered an invalid mapping variation. Element: "
<< data.elementId << ", quadrature point: " << quadraturePoint
<< ", status: " << static_cast<int>(status)
);
mapping::ComputeHDivMassTensorVariation(
m_baseMappingContext.mapping, m_mappingVariation.mapping, m_massTensorVariation
);
transformation->SetIntPoint(&integrationPoint);
gravityGradientElement.CalcVShape(*transformation, m_gravityGradientShape);
m_gravityGradientShape.MultTranspose(m_elementGravityGradient, m_gravityGradientValue);
m_massTensorVariation.Mult(m_gravityGradientValue, m_massTensorVariationAction);
const double referenceWeight = integrationPoint.weight * transformation->Weight();
m_gravityGradientShape.AddMult(m_massTensorVariationAction, m_elementVariationAction, referenceWeight);
}
if (data.gravityGradientDofTransformation != nullptr) {
data.gravityGradientDofTransformation->TransformDual(m_elementVariationAction);
}
m_localVariationAction.AddElementVector(data.gravityGradientDofs, m_elementVariationAction);
}
local_to_true(*m_fem.gravityFluxFes, m_localVariationAction, actionVariationTrue);
}
void PreparedMappedHDivMassOperator::AssembleDiagonal(mfem::Vector &diagonal) const {
mfem::Vector true_diagonal;
AssembleTrueDiagonal(true_diagonal);