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,248 @@
module;
#include <algorithm>
#include <cmath>
#include <numbers>
#include <optional>
#include <stdexcept>
#include <vector>
#include <mfem.hpp>
module mean_field;
import :seed.lane_emden;
import :utils.misc;
namespace {
struct LaneEmdenPoint final {
double coordinate{0.0};
double value{0.0};
double derivative{0.0};
};
struct LaneEmdenDerivative final {
double value{0.0};
double derivative{0.0};
};
[[nodiscard]] LaneEmdenDerivative evaluate_lane_emden_rhs(
const double coordinate,
const double value,
const double derivative,
const double polytropicIndex
) {
const double nonnegativeValue = std::max(value, 0.0);
return {
.value = derivative,
.derivative = -2.0 * derivative / coordinate - std::pow(nonnegativeValue, polytropicIndex)
};
}
[[nodiscard]] LaneEmdenPoint take_lane_emden_step(
const LaneEmdenPoint &point,
const double step,
const double polytropicIndex
) {
const LaneEmdenDerivative first =
evaluate_lane_emden_rhs(point.coordinate, point.value, point.derivative, polytropicIndex);
const LaneEmdenDerivative second = evaluate_lane_emden_rhs(
point.coordinate + 0.5 * step, point.value + 0.5 * step * first.value,
point.derivative + 0.5 * step * first.derivative, polytropicIndex
);
const LaneEmdenDerivative third = evaluate_lane_emden_rhs(
point.coordinate + 0.5 * step, point.value + 0.5 * step * second.value,
point.derivative + 0.5 * step * second.derivative, polytropicIndex
);
const LaneEmdenDerivative fourth = evaluate_lane_emden_rhs(
point.coordinate + step, point.value + step * third.value, point.derivative + step * third.derivative,
polytropicIndex
);
return {
.coordinate = point.coordinate + step,
.value = point.value + step / 6.0 * (first.value + 2.0 * second.value + 2.0 * third.value + fourth.value),
.derivative =
point.derivative +
step / 6.0 * (first.derivative + 2.0 * second.derivative + 2.0 * third.derivative + fourth.derivative)
};
}
[[nodiscard]] std::vector<LaneEmdenPoint> solve_lane_emden(
const double polytropicIndex,
const double coordinateLimit,
const double integrationStep
) {
if (!std::isfinite(polytropicIndex) || polytropicIndex < 0.0) {
throw std::invalid_argument("Lane-Emden integration requires a finite, nonnegative polytropic index.");
}
if (!std::isfinite(coordinateLimit) || coordinateLimit <= 0.0) {
throw std::invalid_argument("The Lane-Emden coordinate limit must be finite and positive.");
}
if (!std::isfinite(integrationStep) || integrationStep <= 0.0) {
throw std::invalid_argument("The Lane-Emden integration step must be finite and positive.");
}
constexpr int maximumStepCount = 2'000'000;
if (std::ceil(coordinateLimit / integrationStep) > static_cast<double>(maximumStepCount)) {
throw std::invalid_argument("The requested Lane-Emden interval exceeds the integration step limit.");
}
const double initialCoordinate = std::min(1.0e-6, coordinateLimit);
const double coordinateSquared = initialCoordinate * initialCoordinate;
const double coordinateCubed = coordinateSquared * initialCoordinate;
const double coordinateFourth = coordinateSquared * coordinateSquared;
LaneEmdenPoint point{
.coordinate = initialCoordinate,
.value = 1.0 - coordinateSquared / 6.0 + polytropicIndex * coordinateFourth / 120.0,
.derivative = -initialCoordinate / 3.0 + polytropicIndex * coordinateCubed / 30.0
};
std::vector<LaneEmdenPoint> solution;
solution.reserve(8192);
solution.push_back({.coordinate = 0.0, .value = 1.0, .derivative = 0.0});
solution.push_back(point);
for (int stepIndex = 0; stepIndex < maximumStepCount && point.coordinate < coordinateLimit; ++stepIndex) {
const double step = std::min(integrationStep, coordinateLimit - point.coordinate);
LaneEmdenPoint nextPoint = take_lane_emden_step(point, step, polytropicIndex);
if (!std::isfinite(nextPoint.value)) {
throw std::runtime_error(
"The Lane-Emden integration produced a non-finite solution before reaching its termination."
);
}
if (nextPoint.value <= 0.0) {
const double rootFraction = point.value / (point.value - nextPoint.value);
solution.push_back(
{.coordinate = point.coordinate + rootFraction * (nextPoint.coordinate - point.coordinate),
.value = 0.0,
.derivative = point.derivative + rootFraction * (nextPoint.derivative - point.derivative)}
);
return solution;
}
solution.push_back(nextPoint);
point = nextPoint;
}
if (point.coordinate < coordinateLimit) {
throw std::runtime_error("The Lane-Emden integration exceeded its step limit.");
}
return solution;
}
[[nodiscard]] double interpolate_lane_emden_value(
const std::vector<LaneEmdenPoint> &solution,
const double coordinate,
std::size_t &lowerIndex
) {
while (lowerIndex + 1 < solution.size() && solution[lowerIndex + 1].coordinate < coordinate) {
++lowerIndex;
}
if (lowerIndex + 1 >= solution.size()) {
return 0.0;
}
const LaneEmdenPoint &lower = solution[lowerIndex];
const LaneEmdenPoint &upper = solution[lowerIndex + 1];
const double interval = upper.coordinate - lower.coordinate;
if (interval <= 0.0) {
throw std::runtime_error("The Lane-Emden interpolation grid is not strictly increasing.");
}
const double fraction = (coordinate - lower.coordinate) / interval;
return std::clamp(lower.value + fraction * (upper.value - lower.value), 0.0, 1.0);
}
} // namespace
namespace mean_field::seed {
DimensionlessLaneEmdenSolution integrateLaneEmden(
const double polytropicIndex,
const double coordinateLimit,
const double integrationStep
) {
const std::vector<LaneEmdenPoint> points = solve_lane_emden(polytropicIndex, coordinateLimit, integrationStep);
DimensionlessLaneEmdenSolution solution{
.coordinate = mfem::Vector(static_cast<int>(points.size())),
.theta = mfem::Vector(static_cast<int>(points.size())),
.thetaDerivative = mfem::Vector(static_cast<int>(points.size())),
.firstZeroCoordinate = std::nullopt
};
for (int index = 0; index < static_cast<int>(points.size()); ++index) {
solution.coordinate(index) = points[static_cast<std::size_t>(index)].coordinate;
solution.theta(index) = points[static_cast<std::size_t>(index)].value;
solution.thetaDerivative(index) = points[static_cast<std::size_t>(index)].derivative;
}
if (points.back().value == 0.0) {
solution.firstZeroCoordinate = points.back().coordinate;
}
return solution;
}
RadialProfile generateLaneEmdenProfile(
const eos::Polytrope &equationOfState,
const dimensions::DensityValue centralDensity,
const int radialSampleCount
) {
if (!std::isfinite(centralDensity.value()) || centralDensity.value() <= 0.0) {
throw std::invalid_argument("A Lane-Emden seed central density must be finite and positive.");
}
if (radialSampleCount < 2) {
throw std::invalid_argument("A Lane-Emden seed requires at least two radial samples.");
}
const double polytropicIndex = equationOfState.polytropic_index();
if (!std::isfinite(polytropicIndex) || polytropicIndex < 1.0 || polytropicIndex >= 5.0) {
throw std::invalid_argument("Lane-Emden seeds require a finite-radius polytrope with 1 <= n < 5.");
}
constexpr double seedCoordinateLimit = 2'000.0;
constexpr double integrationStep = 1.0e-3;
const std::vector solution = solve_lane_emden(polytropicIndex, seedCoordinateLimit, integrationStep);
if (solution.back().value != 0.0) {
throw std::runtime_error("The Lane-Emden integration did not reach its first zero within the step limit.");
}
const double surfaceCoordinate = solution.back().coordinate;
const dimensions::SpecificEnthalpyValue centralEnthalpy =
eos::evaluate<dimensions::quantity::SpecificEnthalpy>(equationOfState, centralDensity);
const double radialScaleSquared = centralEnthalpy.value() / (4.0 * std::numbers::pi_v<double> *
mean_field::utils::G * centralDensity.value());
if (!std::isfinite(radialScaleSquared) || radialScaleSquared <= 0.0) {
throw std::runtime_error("The polytropic Lane-Emden radial scale is not finite and positive.");
}
const double radialScale = std::sqrt(radialScaleSquared);
RadialProfile profile{
.radius = mfem::Vector(radialSampleCount),
.density = mfem::Vector(radialSampleCount),
.specificEnthalpy = mfem::Vector(radialSampleCount),
.stellarRadius = dimensions::LengthValue{radialScale * surfaceCoordinate},
.centralDensity = centralDensity,
.centralSpecificEnthalpy = centralEnthalpy
};
std::size_t interpolationIndex = 0;
for (int sampleIndex = 0; sampleIndex < radialSampleCount; ++sampleIndex) {
const double fraction = static_cast<double>(sampleIndex) / static_cast<double>(radialSampleCount - 1);
const double dimensionlessRadius = fraction * surfaceCoordinate;
const double laneEmdenValue =
interpolate_lane_emden_value(solution, dimensionlessRadius, interpolationIndex);
const dimensions::DensityValue density{centralDensity.value() * std::pow(laneEmdenValue, polytropicIndex)};
profile.radius(sampleIndex) = radialScale * dimensionlessRadius;
profile.density(sampleIndex) = density.value();
profile.specificEnthalpy(sampleIndex) =
eos::evaluate<dimensions::quantity::SpecificEnthalpy>(equationOfState, density).value();
}
profile.radius(0) = 0.0;
profile.density(0) = centralDensity.value();
profile.specificEnthalpy(0) = centralEnthalpy.value();
const int surfaceIndex = radialSampleCount - 1;
profile.radius(surfaceIndex) = profile.stellarRadius.value();
profile.density(surfaceIndex) = 0.0;
profile.specificEnthalpy(surfaceIndex) = 0.0;
return profile;
}
} // namespace mean_field::seed

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