experiment: replay captured candidates and run focused frozen probes

This commit is contained in:
2026-09-17 08:18:09 -04:00
parent c10725bece
commit 9e1b65dbde
3 changed files with 321 additions and 13 deletions

View File

@@ -18,6 +18,8 @@ import mean_field;
#include "plateau_enthalpy_step_diagnostics.hpp"
#include "plateau_physical_diagnostics.hpp"
#include "plateau_preconditioner_diagnostics.hpp"
#include "plateau_pressure_obstruction_diagnostics.hpp"
#include "plateau_shape_curvature_diagnostics.hpp"
using namespace mean_field;
namespace {
@@ -26,9 +28,10 @@ struct Options {
double n = 1., angularMomentum = 0., atol = 0., rtol = 1e-8, linearTolerance = .03;
int maxNewton = 12, maxLinear = 200, restart = 40, probeRestart = 40;
bool probes = false, noCentral = false;
bool pressureProbe = false, shapeCurvature = false;
std::vector<int> captureIterations;
std::string mesh = "sandbox.smesh";
std::filesystem::path output, probeState;
std::filesystem::path output, probeState, probeDirection, replayCandidate;
};
Options Parse(int argc, char **argv) {
Options o;
@@ -63,6 +66,14 @@ Options Parse(int argc, char **argv) {
o.rtol = std::stod(v());
else if (a == "--probe-state")
o.probeState = v();
else if (a == "--probe-direction")
o.probeDirection = v();
else if (a == "--replay-candidate")
o.replayCandidate = v();
else if (a == "--pressure-probe")
o.pressureProbe = true;
else if (a == "--shape-curvature")
o.shapeCurvature = true;
else if (a == "--probes")
o.probes = true;
else if (a == "--no-central")
@@ -87,6 +98,10 @@ Options Parse(int argc, char **argv) {
throw std::invalid_argument("Invalid options");
if (o.probeRestart != 40 && o.probeRestart != 160)
throw std::invalid_argument("--probe-restart must be 40 or 160");
if ((o.pressureProbe || o.shapeCurvature) && (o.probeState.empty() || o.probeDirection.empty()))
throw std::invalid_argument("Focused probes require --probe-state and --probe-direction");
if ((!o.probeDirection.empty() || !o.replayCandidate.empty()) && o.probeState.empty())
throw std::invalid_argument("A supplied direction or candidate requires --probe-state");
if (std::filesystem::exists(o.output))
throw std::runtime_error("Output directory already exists");
std::filesystem::create_directories(o.output);
@@ -125,6 +140,101 @@ void SaveVector(const std::filesystem::path &p, const mfem::Vector &v) {
s << v.Size() << '\n';
v.Print(s, 1);
}
mfem::Vector LoadVector(const std::filesystem::path &path, int expectedSize) {
std::ifstream input(path);
int count = -1;
input >> count;
if (!input || count != expectedSize)
throw std::runtime_error("Snapshot size mismatch: " + path.string());
mfem::Vector result(count);
for (int i = 0; i < count; ++i) {
input >> result(i);
if (!input || !std::isfinite(result(i)))
throw std::runtime_error("Invalid snapshot: " + path.string());
}
std::string extra;
if (input >> extra)
throw std::runtime_error("Unexpected trailing snapshot data: " + path.string());
return result;
}
template <class State>
void FocusedProbes(State &s, const fem::FEM &f, const Options &o, const mfem::Vector &physicalBase) {
mfem::Vector base;
s.normalizedOperator->NormalizeState(physicalBase, base);
auto restore = [&]() {
s.trialNormalizedState = base;
if (!s.PrepareTrial())
throw std::runtime_error("Focused probe base is inadmissible");
};
restore();
const mfem::Vector residual(s.trialNormalizedResidual);
SaveVector(o.output / "frozen_physical_base.txt", physicalBase);
SaveVector(o.output / "frozen_normalized_base.txt", base);
SaveVector(o.output / "frozen_normalized_residual.txt", residual);
auto layout = File(o.output / "layout.csv");
layout << "kind,block,offset,size,scale\n";
for (const auto &b : s.Problem().GetManifest().residualBlocks())
layout << "residual," << b.stableId << ',' << b.offset << ',' << b.size << ',' << b.scale << '\n';
for (const auto &b : s.Problem().GetManifest().valueBlocks())
layout << "value," << b.stableId << ',' << b.offset << ',' << b.size << ',' << b.scale << '\n';
auto checks = File(o.output / "frozen_input_checks.csv");
checks << "metric,value\nbase_residual_norm," << residual.Norml2() << '\n';
const auto capturedBase = o.probeState.parent_path() / "normalized_base.txt";
if (std::filesystem::is_regular_file(capturedBase)) {
mfem::Vector difference = LoadVector(capturedBase, base.Size());
difference -= base;
checks << "captured_normalized_base_difference," << difference.Norml2() << '\n';
if (difference.Norml2() > 64. * std::numeric_limits<double>::epsilon() * std::max(1., base.Norml2()))
throw std::runtime_error("Frozen normalization disagrees with the captured base");
}
const auto capturedResidual = o.probeState.parent_path() / "normalized_residual.txt";
if (std::filesystem::is_regular_file(capturedResidual)) {
mfem::Vector difference = LoadVector(capturedResidual, residual.Size());
difference -= residual;
checks << "captured_residual_difference," << difference.Norml2() << '\n';
}
if (!o.replayCandidate.empty()) {
s.trialNormalizedState = LoadVector(o.replayCandidate, base.Size());
const auto prepared = s.PrepareTrial();
auto replay = File(o.output / "replayed_candidate.csv");
replay << "admissible,reason,stage,eos_code,minimum_jacobian,residual_norm\n";
if (prepared)
replay << "1,-1,-1,-1," << prepared->minimumJacobianDeterminant << ','
<< s.trialNormalizedResidual.Norml2() << '\n';
else
replay << "0," << int(prepared.error().reason) << ',' << int(prepared.error().stage) << ','
<< int(prepared.error().thermodynamicErrorCode) << ','
<< prepared.error().minimumJacobianDeterminant << ",nan\n";
restore();
}
if (!o.probeDirection.empty()) {
const auto direction = LoadVector(o.probeDirection, base.Size());
mfem::Vector action(base.Size());
restore();
s.normalizedOperator->Mult(direction, action);
SaveVector(o.output / "probe_normalized_correction.txt", direction);
checks << "direction_norm," << direction.Norml2() << "\naction_norm," << action.Norml2() << '\n';
action += residual;
checks << "linear_defect_norm," << action.Norml2() << "\nlinear_defect_relative,"
<< action.Norml2() / std::max(residual.Norml2(), 1e-300) << '\n';
checks.flush();
s.normalizedCorrection = direction;
auto bounds = File(o.output / "frozen_enthalpy_bound.csv");
bounds << "metric,value\n";
for (const auto &[key, value] : experiment::plateau_diagnostics::MeasureEnthalpyStep(s, f, o.n, &physicalBase))
bounds << key << ',' << value << '\n';
bounds.flush();
if (o.pressureProbe)
experiment::plateau_diagnostics::ProbePressureObstruction(s, f, o.n, base, direction, o.output);
if (o.shapeCurvature)
experiment::plateau_diagnostics::ProbeShapeCurvature(s, f, o.n, base, direction, o.output);
}
restore();
mfem::Vector difference(s.trialNormalizedResidual);
difference -= residual;
checks << "final_restoration_residual_difference," << difference.Norml2() << '\n';
}
template <class State>
void FrozenProbes(State &s, const fem::FEM &f, const Options &o, const mfem::Vector *physicalBase = nullptr) {
@@ -240,7 +350,7 @@ template <class Context> int RunContext(Context &context, const Options &o, cons
metadata << "n=" << o.n << "\nK=" << scales.K << "\nrho_c=" << scales.rho << "\nJ=" << o.angularMomentum
<< "\nM=1\nR_reference=1\nxi1=" << scales.xi << "\nmu1=" << scales.mu << "\ncentral_phase=" << !o.noCentral
<< "\nrestart=" << o.restart << "\nprobe_restart=" << o.probeRestart
<< "\nfrozen_probe_backend=direct_mfem_fgmres_cold"
<< "\nfrozen_probe_backend=" << (o.probeDirection.empty() ? "direct_mfem_fgmres_cold" : "supplied_direction_no_linear_solve")
<< "\nlinear_tol=" << o.linearTolerance << "\nmax_linear=" << o.maxLinear
<< "\nmax_newton=" << o.maxNewton << "\natol=" << o.atol << "\nrtol=" << o.rtol
<< "\ncompiler=" << __VERSION__ << "\ncompiled=" << __DATE__ << ' ' << __TIME__ << '\n';
@@ -249,24 +359,23 @@ template <class Context> int RunContext(Context &context, const Options &o, cons
if (o.maxNewton != 0)
throw std::invalid_argument("--probe-state requires --max-newton 0");
metadata << "probe_state=" << o.probeState << '\n';
if (!o.probeDirection.empty())
metadata << "probe_direction=" << o.probeDirection << '\n';
if (!o.replayCandidate.empty())
metadata << "replay_candidate=" << o.replayCandidate << '\n';
metadata << "pressure_probe=" << o.pressureProbe << "\nshape_curvature=" << o.shapeCurvature << '\n';
metadata.flush();
solver::detail::StellarEquilibriumContextDiagnostics::WithState(context, [&](auto &s, const fem::FEM &f) {
std::ifstream input(o.probeState);
int count = 0;
input >> count;
if (count != s.AcceptedPhysicalState().Size())
throw std::runtime_error("Snapshot size mismatch");
mfem::Vector snapshot(count);
for (int i = 0; i < count; ++i) {
input >> snapshot(i);
if (!input || !std::isfinite(snapshot(i)))
throw std::runtime_error("Invalid snapshot");
}
const auto snapshot = LoadVector(o.probeState, s.AcceptedPhysicalState().Size());
s.normalizedOperator->NormalizeState(snapshot, s.trialNormalizedState);
auto prep = s.PrepareTrial();
if (!prep)
throw std::runtime_error("Snapshot is inadmissible");
std::cout << "Snapshot F=" << s.trialNormalizedResidual.Norml2() << std::endl;
if (!o.probeDirection.empty() || !o.replayCandidate.empty()) {
FocusedProbes(s, f, o, snapshot);
return;
}
experiment::plateau_diagnostics::InspectPreconditioner(s, o.output);
FrozenProbes(s, f, o, &snapshot);
});