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

@@ -0,0 +1,209 @@
module;
#include <algorithm>
#include <cmath>
#include <limits>
#include <stdexcept>
#include <mfem.hpp>
#include <mpi.h>
module mean_field;
import :field.mfem;
import :seed.stellar_equilibrium_projection;
import :utils.domain;
import :utils.misc;
namespace {
using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
void validate_profile(const mean_field::seed::RadialProfile &profile) {
const int sampleCount = profile.radius.Size();
if (sampleCount < 2 || profile.density.Size() != sampleCount ||
profile.specificEnthalpy.Size() != sampleCount) {
throw std::invalid_argument("A radial seed projection requires equally sized profiles with two samples.");
}
if (!std::isfinite(profile.stellarRadius.value()) || profile.stellarRadius.value() <= 0.0 ||
!std::isfinite(profile.centralDensity.value()) || profile.centralDensity.value() <= 0.0 ||
!std::isfinite(profile.centralSpecificEnthalpy.value()) || profile.centralSpecificEnthalpy.value() <= 0.0) {
throw std::invalid_argument("A radial seed projection requires finite, positive physical scales.");
}
for (int index = 0; index < sampleCount; ++index) {
if (!std::isfinite(profile.radius(index)) || !std::isfinite(profile.density(index)) ||
!std::isfinite(profile.specificEnthalpy(index)) || profile.density(index) < 0.0 ||
profile.specificEnthalpy(index) < 0.0) {
throw std::invalid_argument("A radial seed projection received a non-finite or negative profile.");
}
if (index > 0 && profile.radius(index) <= profile.radius(index - 1)) {
throw std::invalid_argument("A radial seed projection requires strictly increasing radii.");
}
}
const int surfaceIndex = sampleCount - 1;
const double radialScale = std::max(profile.stellarRadius.value(), 1.0);
if (std::abs(profile.radius(0)) > 64.0 * std::numeric_limits<double>::epsilon() * radialScale ||
std::abs(profile.radius(surfaceIndex) - profile.stellarRadius.value()) >
64.0 * std::numeric_limits<double>::epsilon() * radialScale ||
profile.density(0) != profile.centralDensity.value() ||
profile.specificEnthalpy(0) != profile.centralSpecificEnthalpy.value() ||
profile.density(surfaceIndex) != 0.0 || profile.specificEnthalpy(surfaceIndex) != 0.0) {
throw std::invalid_argument("A radial seed projection received inconsistent center or surface metadata.");
}
}
[[nodiscard]] double interpolate_profile(
const mfem::Vector &radius,
const mfem::Vector &values,
const double requestedRadius
) {
if (requestedRadius <= radius(0)) {
return values(0);
}
const int finalIndex = radius.Size() - 1;
if (requestedRadius >= radius(finalIndex)) {
return values(finalIndex);
}
int lowerIndex = 0;
int upperIndex = finalIndex;
while (upperIndex - lowerIndex > 1) {
const int middleIndex = lowerIndex + (upperIndex - lowerIndex) / 2;
if (radius(middleIndex) <= requestedRadius) {
lowerIndex = middleIndex;
} else {
upperIndex = middleIndex;
}
}
const double fraction = (requestedRadius - radius(lowerIndex)) / (radius(upperIndex) - radius(lowerIndex));
return (1.0 - fraction) * values(lowerIndex) + fraction * values(upperIndex);
}
struct SurfaceRadiusRange final {
double minimum;
double maximum;
};
[[nodiscard]] SurfaceRadiusRange measure_surface_radius(const mean_field::fem::FEM &finiteElementModel) {
if (finiteElementModel.surfaceDeformationFes == nullptr) {
throw std::invalid_argument("Radial seed projection requires the surface-deformation space.");
}
mfem::ParFiniteElementSpace &surfaceSpace = *finiteElementModel.surfaceDeformationFes;
const mean_field::field::ScalarBoundaryDofMap surfaceMap =
mean_field::field::make_stellar_surface_scalar_dof_map<DomainSchema>(surfaceSpace);
mfem::Vector radiusSquared(surfaceMap.local_size());
radiusSquared = 0.0;
mfem::ParGridFunction coordinateField(&surfaceSpace);
for (int component = 0; component < surfaceSpace.GetMesh()->SpaceDimension(); ++component) {
mfem::FunctionCoefficient coordinateCoefficient([component](const mfem::Vector &position) {
return position(component);
});
coordinateField.ProjectCoefficient(coordinateCoefficient);
mfem::Vector coordinateTrue;
coordinateField.GetTrueDofs(coordinateTrue);
const mfem::Vector surfaceCoordinate = surfaceMap.gather(coordinateTrue);
for (int index = 0; index < radiusSquared.Size(); ++index) {
radiusSquared(index) += surfaceCoordinate(index) * surfaceCoordinate(index);
}
}
double localMinimum = std::numeric_limits<double>::infinity();
double localMaximum = 0.0;
for (int index = 0; index < radiusSquared.Size(); ++index) {
const double radius = std::sqrt(radiusSquared(index));
localMinimum = std::min(localMinimum, radius);
localMaximum = std::max(localMaximum, radius);
}
double globalMinimum = 0.0;
double globalMaximum = 0.0;
MPI_Allreduce(&localMinimum, &globalMinimum, 1, MPI_DOUBLE, MPI_MIN, surfaceSpace.GetComm());
MPI_Allreduce(&localMaximum, &globalMaximum, 1, MPI_DOUBLE, MPI_MAX, surfaceSpace.GetComm());
if (!std::isfinite(globalMinimum) || !std::isfinite(globalMaximum) || globalMinimum <= 0.0 ||
globalMaximum < globalMinimum) {
throw std::runtime_error("The stellar surface has no finite, positive radial extent.");
}
return {.minimum = globalMinimum, .maximum = globalMaximum};
}
} // namespace
namespace mean_field::seed::detail {
ProjectedRadialFields projectRadialFields(
const equilibrium::StellarDiscretization &discretization,
const RadialProfile &profile,
const dimensions::MassValue targetMass,
const dimensions::PressureValue targetSurfacePressure,
const StellarEquilibriumProjectionOptions &options
) {
validate_profile(profile);
if (!std::isfinite(options.surfaceRadiusRelativeTolerance) || options.surfaceRadiusRelativeTolerance < 0.0) {
throw std::invalid_argument("The surface-radius projection tolerance must be finite and nonnegative.");
}
if (targetSurfacePressure.value() != 0.0) {
throw std::invalid_argument("A Lane-Emden radial seed requires a zero-pressure isobaric surface.");
}
fem::FEM &finiteElementModel = discretization.finiteElementModel();
const SurfaceRadiusRange surfaceRadius = measure_surface_radius(finiteElementModel);
const double targetRadius = profile.stellarRadius.value();
const double comparisonScale = std::max({targetRadius, surfaceRadius.maximum, 1.0e-300});
const double relativeMismatch =
std::max(std::abs(surfaceRadius.minimum - targetRadius), std::abs(surfaceRadius.maximum - targetRadius)) /
comparisonScale;
if (relativeMismatch > options.surfaceRadiusRelativeTolerance) {
throw std::invalid_argument(
"The radial seed surface does not coincide with the spherical reference discretization."
);
}
if (finiteElementModel.densityFes == nullptr || finiteElementModel.enthalpyFes == nullptr ||
finiteElementModel.displacementFes == nullptr || finiteElementModel.gravityFluxFes == nullptr ||
finiteElementModel.gravityPotentialFes == nullptr) {
throw std::invalid_argument("Radial seed projection requires the complete equilibrium discretization.");
}
mfem::FunctionCoefficient densityCoefficient([&profile](const mfem::Vector &position) {
return interpolate_profile(profile.radius, profile.density, position.Norml2());
});
mfem::FunctionCoefficient enthalpyCoefficient([&profile](const mfem::Vector &position) {
return interpolate_profile(profile.radius, profile.specificEnthalpy, position.Norml2());
});
mfem::ParGridFunction densityField(finiteElementModel.densityFes.get());
mfem::ParGridFunction enthalpyField(finiteElementModel.enthalpyFes.get());
mfem::ParGridFunction displacementField(finiteElementModel.displacementFes.get());
densityField = 0.0;
enthalpyField = 0.0;
displacementField = 0.0;
densityField.ProjectCoefficient(densityCoefficient);
enthalpyField.ProjectCoefficient(enthalpyCoefficient);
const physics::GravitySolution gravitySolution =
physics::solve_gravity_field(finiteElementModel, options.gravity, densityField, displacementField);
const field::FieldDofGridFunctionAdapter densityAdapter =
field::make_field_dof_grid_function_adapter<field::Density, DomainSchema>(*finiteElementModel.densityFes);
const field::FieldDofGridFunctionAdapter enthalpyAdapter =
field::make_field_dof_grid_function_adapter<field::Enthalpy, DomainSchema>(*finiteElementModel.enthalpyFes);
const field::FieldDofGridFunctionAdapter gravityFluxAdapter =
field::make_field_dof_grid_function_adapter<field::Gravity, DomainSchema>(
*finiteElementModel.gravityFluxFes
);
const field::FieldDofGridFunctionAdapter gravityPotentialAdapter =
field::make_field_dof_grid_function_adapter<field::Gravity, DomainSchema>(
*finiteElementModel.gravityPotentialFes
);
return {
.density = densityAdapter.gather(densityField),
.gravityGradient = gravityFluxAdapter.gather(gravitySolution.gradPhi),
.gravityPotential = gravityPotentialAdapter.gather(gravitySolution.phi),
.specificEnthalpy = enthalpyAdapter.gather(enthalpyField),
.bernoulliConstant = -utils::G * targetMass.value() / targetRadius
};
}
} // namespace mean_field::seed::detail