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

@@ -22,26 +22,34 @@ namespace {
);
}
[[nodiscard]] mean_field::operators::StellarEquilibriumLayout make_layout(
using StellarRootForm = mean_field::utils::blocks::surface_deformed_stellar_equilibrium_form;
[[nodiscard]] std::array<
int,
StellarRootForm::value_block_count>
make_value_sizes(
const mean_field::field::FieldDofMap &densityMap,
const int surfaceDeformationParameterCount,
const mean_field::field::FieldDofMap &gravityFluxMap,
const mean_field::field::FieldDofMap &gravityPotentialMap,
const mean_field::field::FieldDofMap &enthalpyMap
) {
using Form = mean_field::utils::blocks::surface_deformed_stellar_equilibrium_form;
return {densityMap.reduced_size(), surfaceDeformationParameterCount, gravityFluxMap.reduced_size(),
gravityPotentialMap.reduced_size(), enthalpyMap.reduced_size(), 1};
}
const std::array<int, Form::value_block_count> valueSizes{
densityMap.reduced_size(), surfaceDeformationParameterCount, gravityFluxMap.reduced_size(),
gravityPotentialMap.reduced_size(), enthalpyMap.reduced_size(), 1
};
const std::array<int, Form::residual_block_count> residualSizes{
gravityFluxMap.reduced_size(), gravityPotentialMap.reduced_size(), densityMap.reduced_size(),
surfaceDeformationParameterCount, enthalpyMap.reduced_size(), 1
};
return {valueSizes, residualSizes};
[[nodiscard]] std::array<
int,
StellarRootForm::residual_block_count>
make_residual_sizes(
const mean_field::field::FieldDofMap &densityMap,
const int surfaceDeformationParameterCount,
const mean_field::field::FieldDofMap &gravityFluxMap,
const mean_field::field::FieldDofMap &gravityPotentialMap,
const mean_field::field::FieldDofMap &enthalpyMap
) {
return {gravityFluxMap.reduced_size(), gravityPotentialMap.reduced_size(), densityMap.reduced_size(),
surfaceDeformationParameterCount, enthalpyMap.reduced_size(), 1};
}
[[nodiscard]] mfem::Array<int> make_gravity_state_offsets(
@@ -70,47 +78,6 @@ namespace {
return offsets;
}
template <int index>
[[nodiscard]] mfem::Vector make_value_view(
const mfem::Vector &vector,
const mean_field::operators::StellarEquilibriumLayout &layout,
const mean_field::utils::blocks::value_block<index> block
) {
MFEM_VERIFY(
vector.Size() == layout.value_offsets().Last(),
"The coupled vector does not match the stellar-equilibrium value layout."
);
return mfem::Vector(const_cast<mfem::real_t *>(vector.GetData()) + layout.offset(block), layout.size(block));
}
template <int index>
[[nodiscard]] mfem::Vector make_residual_view(
mfem::Vector &vector,
const mean_field::operators::StellarEquilibriumLayout &layout,
const mean_field::utils::blocks::residual_block<index> block
) {
MFEM_VERIFY(
vector.Size() == layout.residual_offsets().Last(),
"The coupled vector does not match the stellar-equilibrium residual layout."
);
return mfem::Vector(vector.GetData() + layout.offset(block), layout.size(block));
}
template <int index>
void assign_residual_block(
mfem::Vector &coupledResidual,
const mean_field::operators::StellarEquilibriumLayout &layout,
const mean_field::utils::blocks::residual_block<index> block,
const mfem::Vector &blockResidual,
const char *message
) {
MFEM_VERIFY(layout.size(block) == blockResidual.Size(), message);
mfem::Vector destination = make_residual_view(coupledResidual, layout, block);
destination = blockResidual;
}
void assign_gravity_block(
mfem::Vector &gravityState,
const mfem::Array<int> &offsets,
@@ -264,7 +231,8 @@ namespace mean_field::operators {
field::FieldDofMap enthalpyMap;
field::FieldBoundaryDofMap pressureSurfaceRows;
StellarEquilibriumLayout layout;
std::array<int, StellarRootForm::value_block_count> valueSizes;
std::array<int, StellarRootForm::residual_block_count> residualSizes;
mfem::Array<int> gravityStateOffsets;
mfem::Array<int> gravityResidualOffsets;
@@ -307,7 +275,14 @@ namespace mean_field::operators {
enthalpyMap
)
),
layout(make_layout(
valueSizes(make_value_sizes(
densityMap,
domainDeformation.parameterCount(),
gravityFluxMap,
gravityPotentialMap,
enthalpyMap
)),
residualSizes(make_residual_sizes(
densityMap,
domainDeformation.parameterCount(),
gravityFluxMap,
@@ -339,7 +314,7 @@ namespace mean_field::operators {
fem::FEM &f,
const mapping::DomainMapper &domainMapper,
const eos::Polytrope &equationOfState,
const double targetMass,
models::CompiledFixedMass fixedMassConstraint,
const PressureSurfaceConstraintView surfaceConstraint,
deformation::PreparedDomainDeformationRuntime domainDeformation
)
@@ -347,7 +322,7 @@ namespace mean_field::operators {
f,
domainMapper,
equationOfState,
targetMass,
std::move(fixedMassConstraint),
surfaceConstraint,
MakeConstructionData(
f,
@@ -360,15 +335,31 @@ namespace mean_field::operators {
fem::FEM &f,
const mapping::DomainMapper &domainMapper,
const eos::Polytrope &equationOfState,
const double targetMass,
models::CompiledFixedMass fixedMassConstraint,
const PressureSurfaceConstraintView surfaceConstraint,
ConstructionData constructionData
)
: mfem::Operator(
constructionData.layout.residual_offsets().Last(),
constructionData.layout.value_offsets().Last()
StellarEquilibriumLayout(
constructionData.valueSizes,
constructionData.residualSizes
)
.residual_offsets()
.Last(),
StellarEquilibriumLayout(
constructionData.valueSizes,
constructionData.residualSizes
)
.value_offsets()
.Last()
),
m_rootManifest(
constructionData.valueSizes,
constructionData.residualSizes,
fixedMassConstraint.targetMass().value(),
surfaceConstraint.descriptor().targetPressure,
constructionData.pressureSurfaceRows.size()
),
m_layout(constructionData.layout),
m_gravityStateOffsets(constructionData.gravityStateOffsets),
m_gravityContext(
f,
@@ -413,14 +404,11 @@ namespace mean_field::operators {
surfaceConstraint
),
m_domainDeformation(std::move(constructionData.domainDeformation)),
m_targetMass(targetMass) {
MFEM_VERIFY(
std::isfinite(m_targetMass) && m_targetMass > 0.0,
"PreparedStellarEquilibriumOperator requires a finite, positive target mass."
);
m_fixedMassConstraint(std::move(fixedMassConstraint)) {
MFEM_VERIFY(
Width() == m_layout.value_offsets().Last() && Height() == m_layout.residual_offsets().Last(),
Width() == m_rootManifest.layout().value_offsets().Last() &&
Height() == m_rootManifest.layout().residual_offsets().Last(),
"PreparedStellarEquilibriumOperator has inconsistent block dimensions."
);
@@ -502,27 +490,18 @@ namespace mean_field::operators {
);
}
m_isPrepared = false;
m_isPrepared = false;
using Form = utils::blocks::surface_deformed_stellar_equilibrium_form;
constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);
constexpr auto surfaceDeformationValue =
utils::blocks::get_value_block<Form>(utils::blocks::surface_deformation_field.parameters_term);
constexpr auto gravityGradientValue =
utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);
constexpr auto gravityPotentialValue =
utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.poisson_term);
constexpr auto enthalpyValue =
utils::blocks::get_value_block<Form>(utils::blocks::enthalpy_field.specific_term);
constexpr auto bernoulliValue =
utils::blocks::get_value_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);
const auto rootState = m_rootManifest.stateView(state);
const mfem::Vector reducedDensity = make_value_view(state, m_layout, densityValue);
const mfem::Vector surfaceDeformationParameters = make_value_view(state, m_layout, surfaceDeformationValue);
const mfem::Vector gravityGradient = make_value_view(state, m_layout, gravityGradientValue);
const mfem::Vector gravityPotential = make_value_view(state, m_layout, gravityPotentialValue);
const mfem::Vector reducedEnthalpy = make_value_view(state, m_layout, enthalpyValue);
const mfem::Vector bernoulli = make_value_view(state, m_layout, bernoulliValue);
const mfem::Vector reducedDensity = rootState.block(utils::blocks::density_field.mass_term);
const mfem::Vector surfaceDeformationParameters =
rootState.block(utils::blocks::surface_deformation_field.parameters_term);
const mfem::Vector gravityGradient = rootState.block(utils::blocks::gravity_field.gradient_term);
const mfem::Vector gravityPotential = rootState.block(utils::blocks::gravity_field.poisson_term);
const mfem::Vector reducedEnthalpy = rootState.block(utils::blocks::enthalpy_field.specific_term);
const mfem::Vector bernoulli =
rootState.block(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);
const bool generatedGeometryChanged =
!wasPrepared || dependencies.discretization != m_preparedDependencies.discretization ||
@@ -569,7 +548,7 @@ namespace mean_field::operators {
);
report.massNormalization = m_massNormalizationOperator.Prepare(
{.targetMass = m_targetMass}, make_mass_dependencies(dependencies, m_generatedDisplacementDependency)
m_fixedMassConstraint, make_mass_dependencies(dependencies, m_generatedDisplacementDependency)
);
report.surfaceConstraint = m_surfaceConstraintOperator.Prepare(
@@ -588,21 +567,6 @@ namespace mean_field::operators {
}
void PreparedStellarEquilibriumOperator::AssembleResidual() {
using Form = utils::blocks::surface_deformed_stellar_equilibrium_form;
constexpr auto gravityGradientResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.gradient_term);
constexpr auto gravityPotentialResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.poisson_term);
constexpr auto densityResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::density_field.mass_term);
constexpr auto surfaceShapeResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::surface_deformation_field.shape_equilibrium_term);
constexpr auto enthalpyResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term);
constexpr auto massResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);
mfem::Vector gravity;
mfem::Vector closure;
mfem::Vector surfaceShape;
@@ -621,41 +585,29 @@ namespace mean_field::operators {
m_massNormalizationOperator.BuildResidual(mass);
m_cachedResidual.SetSize(Height());
m_cachedResidual = 0.0;
m_cachedResidual = 0.0;
const auto residualView = m_rootManifest.residualView(m_cachedResidual);
MFEM_VERIFY(
gravity.Size() == m_layout.size(gravityGradientResidual) + m_layout.size(gravityPotentialResidual),
gravity.Size() == residualView.block(utils::blocks::gravity_field.gradient_term).Size() +
residualView.block(utils::blocks::gravity_field.poisson_term).Size(),
"The gravity residual has the wrong size."
);
mfem::Vector gravityGradient(gravity.GetData(), m_layout.size(gravityGradientResidual));
mfem::Vector gravityGradient(
gravity.GetData(), residualView.block(utils::blocks::gravity_field.gradient_term).Size()
);
mfem::Vector gravityPotential(
gravity.GetData() + m_layout.size(gravityGradientResidual), m_layout.size(gravityPotentialResidual)
gravity.GetData() + gravityGradient.Size(),
residualView.block(utils::blocks::gravity_field.poisson_term).Size()
);
assign_residual_block(
m_cachedResidual, m_layout, gravityGradientResidual, gravityGradient,
"The gravity-gradient residual has the wrong size."
);
assign_residual_block(
m_cachedResidual, m_layout, gravityPotentialResidual, gravityPotential,
"The gravity-potential residual has the wrong size."
);
assign_residual_block(
m_cachedResidual, m_layout, densityResidual, closure, "The closure residual has the wrong size."
);
assign_residual_block(
m_cachedResidual, m_layout, surfaceShapeResidual, surfaceShape,
"The surface-shape residual has the wrong size."
);
assign_residual_block(
m_cachedResidual, m_layout, enthalpyResidual, hydrostatic, "The hydrostatic residual has the wrong size."
);
assign_residual_block(
m_cachedResidual, m_layout, massResidual, mass, "The mass-normalization residual has the wrong size."
);
residualView.assign(utils::blocks::gravity_field.gradient_term, gravityGradient);
residualView.assign(utils::blocks::gravity_field.poisson_term, gravityPotential);
residualView.assign(utils::blocks::density_field.mass_term, closure);
residualView.assign(utils::blocks::surface_deformation_field.shape_equilibrium_term, surfaceShape);
residualView.assign(utils::blocks::enthalpy_field.specific_term, hydrostatic);
residualView.assign(utils::blocks::fixed_total_mass_constraint.mass_normalization_term, mass);
++m_statistics.residualAssemblies;
}
@@ -679,39 +631,16 @@ namespace mean_field::operators {
direction, "PreparedStellarEquilibriumOperator received a non-finite Jacobian direction."
);
using Form = utils::blocks::surface_deformed_stellar_equilibrium_form;
const auto rootDirection = m_rootManifest.directionView(direction);
constexpr auto densityValue = utils::blocks::get_value_block<Form>(utils::blocks::density_field.mass_term);
constexpr auto surfaceDeformationValue =
utils::blocks::get_value_block<Form>(utils::blocks::surface_deformation_field.parameters_term);
constexpr auto gravityGradientValue =
utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.gradient_term);
constexpr auto gravityPotentialValue =
utils::blocks::get_value_block<Form>(utils::blocks::gravity_field.poisson_term);
constexpr auto enthalpyValue =
utils::blocks::get_value_block<Form>(utils::blocks::enthalpy_field.specific_term);
constexpr auto bernoulliValue =
utils::blocks::get_value_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);
constexpr auto gravityGradientResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.gradient_term);
constexpr auto gravityPotentialResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::gravity_field.poisson_term);
constexpr auto densityResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::density_field.mass_term);
constexpr auto surfaceShapeResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::surface_deformation_field.shape_equilibrium_term);
constexpr auto enthalpyResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::enthalpy_field.specific_term);
constexpr auto massResidual =
utils::blocks::get_residual_block<Form>(utils::blocks::barotropic_constant_field.mass_normalization_term);
const mfem::Vector reducedDensityDirection = make_value_view(direction, m_layout, densityValue);
const mfem::Vector surfaceDeformationDirection = make_value_view(direction, m_layout, surfaceDeformationValue);
const mfem::Vector gravityGradientDirection = make_value_view(direction, m_layout, gravityGradientValue);
const mfem::Vector gravityPotentialDirection = make_value_view(direction, m_layout, gravityPotentialValue);
const mfem::Vector reducedEnthalpyDirection = make_value_view(direction, m_layout, enthalpyValue);
const mfem::Vector bernoulliDirection = make_value_view(direction, m_layout, bernoulliValue);
const mfem::Vector reducedDensityDirection = rootDirection.block(utils::blocks::density_field.mass_term);
const mfem::Vector surfaceDeformationDirection =
rootDirection.block(utils::blocks::surface_deformation_field.parameters_term);
const mfem::Vector gravityGradientDirection = rootDirection.block(utils::blocks::gravity_field.gradient_term);
const mfem::Vector gravityPotentialDirection = rootDirection.block(utils::blocks::gravity_field.poisson_term);
const mfem::Vector reducedEnthalpyDirection = rootDirection.block(utils::blocks::enthalpy_field.specific_term);
const mfem::Vector bernoulliDirection =
rootDirection.block(utils::blocks::fixed_total_mass_constraint.mass_normalization_term);
m_domainDeformation.applyJacobian(
m_surfaceDeformationParameters, surfaceDeformationDirection, m_volumeDisplacementDirection
@@ -757,41 +686,29 @@ namespace mean_field::operators {
);
action.SetSize(Height());
action = 0.0;
action = 0.0;
const auto actionView = m_rootManifest.residualView(action);
MFEM_VERIFY(
gravityAction.Size() == m_layout.size(gravityGradientResidual) + m_layout.size(gravityPotentialResidual),
gravityAction.Size() == actionView.block(utils::blocks::gravity_field.gradient_term).Size() +
actionView.block(utils::blocks::gravity_field.poisson_term).Size(),
"The gravity Jacobian action has the wrong size."
);
mfem::Vector gravityGradientAction(gravityAction.GetData(), m_layout.size(gravityGradientResidual));
mfem::Vector gravityGradientAction(
gravityAction.GetData(), actionView.block(utils::blocks::gravity_field.gradient_term).Size()
);
mfem::Vector gravityPotentialAction(
gravityAction.GetData() + m_layout.size(gravityGradientResidual), m_layout.size(gravityPotentialResidual)
gravityAction.GetData() + gravityGradientAction.Size(),
actionView.block(utils::blocks::gravity_field.poisson_term).Size()
);
assign_residual_block(
action, m_layout, gravityGradientResidual, gravityGradientAction,
"The gravity-gradient Jacobian action has the wrong size."
);
assign_residual_block(
action, m_layout, gravityPotentialResidual, gravityPotentialAction,
"The gravity-potential Jacobian action has the wrong size."
);
assign_residual_block(
action, m_layout, densityResidual, closureAction, "The closure Jacobian action has the wrong size."
);
assign_residual_block(
action, m_layout, surfaceShapeResidual, m_surfaceShapeAction,
"The surface-shape Jacobian action has the wrong size."
);
assign_residual_block(
action, m_layout, enthalpyResidual, hydrostaticAction, "The hydrostatic Jacobian action has the wrong size."
);
assign_residual_block(
action, m_layout, massResidual, massAction, "The mass-normalization Jacobian action has the wrong size."
);
actionView.assign(utils::blocks::gravity_field.gradient_term, gravityGradientAction);
actionView.assign(utils::blocks::gravity_field.poisson_term, gravityPotentialAction);
actionView.assign(utils::blocks::density_field.mass_term, closureAction);
actionView.assign(utils::blocks::surface_deformation_field.shape_equilibrium_term, m_surfaceShapeAction);
actionView.assign(utils::blocks::enthalpy_field.specific_term, hydrostaticAction);
actionView.assign(utils::blocks::fixed_total_mass_constraint.mass_normalization_term, massAction);
++m_statistics.jacobianApplications;
}
@@ -803,11 +720,30 @@ namespace mean_field::operators {
}
double PreparedStellarEquilibriumOperator::GetTargetMass() const noexcept {
return m_targetMass;
return m_fixedMassConstraint.targetMass().value();
}
const StellarEquilibriumLayout &PreparedStellarEquilibriumOperator::GetLayout() const noexcept {
return m_layout;
return m_rootManifest.layout();
}
const StellarEquilibriumRootManifest &PreparedStellarEquilibriumOperator::GetRootManifest() const noexcept {
return m_rootManifest;
}
RootStateView<utils::blocks::surface_deformed_stellar_equilibrium_form>
PreparedStellarEquilibriumOperator::GetRootStateView(const mfem::Vector &state) const {
return m_rootManifest.stateView(state);
}
ResidualView<utils::blocks::surface_deformed_stellar_equilibrium_form>
PreparedStellarEquilibriumOperator::GetResidualView(mfem::Vector &residual) const {
return m_rootManifest.residualView(residual);
}
RootConstraintReport PreparedStellarEquilibriumOperator::GetFixedMassReport() const {
VerifyPrepared();
return m_rootManifest.fixedMassReport(m_massNormalizationOperator.GetCurrentMass());
}
const StellarEquilibriumDependencies &PreparedStellarEquilibriumOperator::GetDependencies() const {