Barrier-overlap objective with continuous, differentiable-ready variants - #1
Barrier-overlap objective with continuous, differentiable-ready variants#1krystophny wants to merge 112 commits into
Conversation
Objective mode 'barrier' (default) runs two trapped-only SIMPLE fast classification runs at the birth and barrier surfaces and minimizes the mu-resolved barrier-breach fraction from itpplasma/sensopt, which tracks traced late loss across devices where the previous single-surface confined/ideal/jpar score does not. The NEAT-based score remains available via EP_OPT_OBJECTIVE=neat. Reactor mapping follows the ARIES-CS point (a=1.704 m, B=5.865 T from Aminor_p/volavgB); EP_OPT_SCALE_MODEL keeps the small equilibrium and preserves rho/a by scaling alpha energy with (B_scale*RZ_scale)^2 through facE_al, as before. Port verified against the sensopt qi_23_1900 reference row: barrier_overlap 0.20209 vs 0.20078, chaotic trapped fraction at s=0.6 0.22625 vs 0.22678 (ntestpart=3000, newer simple.x).
The 2026-07 conclusion that the fast classifier is slower than direct loss tracing rested on three configuration faults, not on physics. class_plot was .True. in every commit of the old simple_barrier.py. The early exit in SIMPLE's classification.f90 requires .not. class_plot, so no orbit ever left early and every trace ran to trace_time. The later spatial-grid path set class_plot = .False. but also tcut = -1d0, which disabled the classifier outright, since fast_class was absent from SIMPLE's dispatch predicate until itpplasma/SIMPLE#512. The birth-surface run also traced passing particles that the overlap metric then discards, three quarters of the budget, and the surfaces carried far more orbit-seconds than the direct objective they were meant to replace. This objective pins fast_class = .True., class_plot = .False., tcut = -1d0 and notrace_passing = 1 on both surfaces, so resolved orbits exit after nturns tips and no fractal cut runs. Starting points are a deterministic 1024-point (theta, zeta, pitch) product grid per surface rather than random draws, so candidates are compared at identical phase-space points. Bin edges are anchored to the reactor field instead of the common mu range of the two samples. Deriving them from the samples, as before, moved the boundaries with the candidate geometry and made the metric discontinuous in the design variables. Prompt losses carry classifier code 0, not 2, so they are invisible to the overlap. They enter as a third constraint rather than a weighted penalty, which keeps the objective a clean test of the proxy. validate_barrier_optimum.py closes the loop by tracing the winning configuration and the anchor at the direct contract, 100 ms and 256 particles at s = 0.25. Runs use the midpoint symplectic integrator in Boozer coordinates.
A smoke run on a campaign equilibrium left 128 of 1024 starts trapped, about eight per mu bin, because the pitch grid spanned the full unit interval while trapping needs lambda**2 < 1 - B/B_max. The overlap metric discards everything passing, so seven eighths of the grid carried no signal. Bound the grid at |lambda| <= 0.6. The bound is a fixed constant rather than a per-candidate trapping boundary, so every design is still sampled at identical phase-space points.
The optimizer comparison driver submits one job per candidate and selects the worker by METRIC, so barrier-overlap and direct-loss campaigns share the same submission path.
raw_coordinate_contract and write_raw_candidate went through simsopt.mhd.Vmec only to reach the boundary surface and to write a candidate input. That pulls in the compiled VMEC2000 extension and mpi4py, neither of which is installed on the cluster, so the optimizer driver could not start there. Nothing in the driver runs VMEC; equilibria are solved by the standalone xvmec binary on the worker. Read the boundary with SurfaceRZFourier.from_vmec_input instead, and write candidates by substituting the boundary block into the base input, so resolution, profiles and solver controls are inherited verbatim rather than dropped as get_nml alone would.
The ALPES base input ships NS_ARRAY = 3. SIMPLE splines the equilibrium with ns_s = ns_tp = 5, so three flux surfaces overrun the spline construction and abort the process inside spline_vmec_data with heap corruption, several frames below anything the caller can interpret. Check ns against the stencil width before launching SIMPLE and say which knob to raise. The campaign uses the input.ALPES_QUASR_0021326_1_ns16 variant, which matches the NS_ARRAY = 16 the direct campaign solved at.
The driver submits over ssh, so --export=ALL carries only the ssh session's environment, not the pinned toolchain paths and hashes.
The discrete barrier overlap counts integer class codes, so its gradient is zero almost everywhere and a delta on the switching set. Pointing any differentiation tool at it returns 0.0 silently. Three stacked discontinuities cause that: the trapped mask, mu bin membership, and the class code. The first two are mollified here with a logistic in trap_par and a smooth bin kernel. The third cannot be mollified after the fact, because only the code survives; it needs the margins the classifier forms and discards, which SIMPLE now writes to class_scores.dat (itpplasma/SIMPLE#513). Two continuous scores follow. The J_parallel score is a logistic in the tip spread relative to tol_perpinv, so the classifier's threshold is its zero-width limit. The topology score is a logistic in the negated monotonicity margin, falling back to the J_parallel score where the recurrence test decided and no margin exists, since a constant there would flatten the metric over a whole class of orbits. Both are finite-time over nturns bounce periods, so Lyapunov exposure stays a few bounce times rather than the 1e3-1e4 of a slowing-down trace. That is what makes this the differentiable route and not merely the smooth one. The optimizer objective is now selectable between the discrete metric and either smooth variant, and every response records the discrete value whatever is optimized, so the campaigns stay comparable. Runs without class_scores.dat report the smooth block as unavailable rather than failing, so the discrete campaign is unaffected by the SIMPLE build. Tests pin the property that makes the surrogate valid: as the three widths go to zero the smooth overlap converges to the discrete one, for both classifiers, and it moves under a sub-threshold drift change that the discrete metric cannot see at all.
Two faults the first cluster run exposed. The bin edges ran from zero to 1/B with B in Tesla, but SIMPLE carries the field in Gauss, so perp_inv is of order 1e-5 and every trapped particle landed in the first bin. A mu-resolved metric was silently a single-bin one, and the smooth bin width, taken as a fraction of that span, was three orders of magnitude wider than the data. Edges now sit on the Gauss scale and are centred on the reference mu with a band set by the mirror ratio, since trapped particles occupy a narrow band rather than reaching down to zero. The smooth bin width is a fraction of the bin spacing. Unresolved orbits were scored as chaotic, which rewards designs whose orbits merely resolve faster. On the reference equilibrium 57 percent of trapped orbits do not resolve within nturns, so that choice dominated the metric. They are censored instead — dropped from numerator and denominator alike — and the resolved fraction is reported per surface and classifier so the dilution stays visible.
The topology score needs a monotonicity margin, and short traces rarely produce one: on the reference equilibrium no orbit carried it at nturns = 8. The metric is then undefined, which reached the worker as a NaN and failed JSON serialisation, killing the candidate. Report null instead, and name the reason with the resolved fractions when an objective asks for a value that does not exist.
The worker verifies a code manifest and the SIMPLE executable hash on every candidate, but CODE_ROOT and SIMPLE_X pointed at the live checkout and the live build tree. Editing the repository or rebuilding SIMPLE while a campaign ran therefore failed every job in flight: 188 candidates died as worker_failure against a manifest whose files had moved underneath them, and the validation step died on a stale executable hash. stage_campaign.sh copies the executed files and the binary into the campaign root, pins their hashes there, and writes the environment the Slurm job sources. Work can then continue in the repository without touching a run. This is the layout the earlier direct campaigns already used.
The mollified class scores are correct but nearly empty on real runs: on a reactor-scale candidate at nturns = 8, two of 1024 orbits carried a J_parallel spread and none carried a monotonicity margin, because both need the classifier to complete nturns return periods of the tip map within a 20 ms trace. A smooth metric with no data steers nothing. The banana tips carry a quantity that needs no classification: the radial excursion max(s) - min(s), defined for any orbit with two tips. It is the width of the radial band the orbit explores, which is what a phase-space barrier suppresses, so it measures the intended physics directly rather than mollifying a proxy for it. SIMPLE writes it alongside the class margins. The score saturates as one minus an exponential in the excursion over a reference, rather than a logistic, because the excursion has a hard floor at zero and no threshold to sit symmetrically about. The reference is the excursion at which an orbit counts as fully transporting, defaulting to five percent of normalised toroidal flux.
The radial value was computed and dropped: the return dict listed the jpar and topology keys by hand, so smooth-radial raised KeyError in the worker. Build the keys from the values that were computed, and test that each classifier reaches both the value map and the resolved fractions.
There was a problem hiding this comment.
Review verdict: Request changes
Summary: This is a large PR (~87 files, ~110k insertions) that reorients the repository from the old single-stage VMEC+simsopt EP optimizer into a TuRBO/SCBO campaign driver for a direct 100 ms alpha-loss objective (main.py now just calls prime_direct_scbo.main()). The new pipeline — simple_direct.py, evaluate_direct_loss.py, evaluate_vmec_geometry.py, generate_scbo_candidates.py, data_informed_surface.py, smooth_barrier.py/barrier_overlap.py, the Condor/Slurm backend (scbo_backend.py) and the run scripts — is coherent, defensively validated (hash pinning, schema checks, failure-kind propagation, bounded unit coordinates), and shipped with a substantial unit-test suite. The physics changes (Gauss-scale mu bins, censoring unresolved orbits, radial-excursion score) are well documented in commit messages. I could not execute the tests (no numpy/pytest in this environment), but every reviewed module compiles.
Findings:
-
[blocker] prime_direct_scbo.py:41 (via main.py), advance_alpha_scbo.py:10, run_optimizer_comparison.py:149 — the new entry points import
simsopt_dfo.scbo_checkpoint(prime_state,issue_candidates,write_checkpoint,complete_candidates), butsimsopt_dfois declared nowhere: it is absent frompyproject.tomldependencies and[tool.uv.sources], absent fromuv.lock(0 matches), and not vendored in the repo. In a fresh environment the documentedpython3 main.pyand the whole campaign advance loop (advance_alpha_scbo.py) will raiseImportError. Required fix: add thesimsopt-dfo(TuRBO) dependency/source topyproject.tomland regenerateuv.lock, or vendor the package. -
[minor] .DS_Store, ._.DS_Store — macOS Finder metadata files (6.1 KB and 4.1 KB binaries) are committed at the repo root;
.gitignoredoes not cover them. Remove them and add*.DS_Store/._*to.gitignore. -
[minor] runs/plasma_simple_proxy_correlation_v1/ — ~13 MB of campaign output (design.npz, 9k-line start.dat inputs, manifest.json) is committed into the source tree, even though README.md says campaign artifacts live in the companion
alpha-loss-optimization-datarepository. Consider moving this snapshot out of the source repo (or into test data referenced by a test) to keep the PR reviewable.
Verdict: Request changes — the missing simsopt_dfo dependency breaks the documented entry point and the entire campaign driver in any clean environment, so the PR cannot run as described.
There was a problem hiding this comment.
Review verdict: Request changes
Summary: This is a large (~110k insertions) branch that pivots the repository from the old MPI main.py EP/direct optimizer to a new campaign framework for a deterministic 100 ms direct alpha-loss optimization at s=0.25. It replaces the objective with pinned SIMPLE direct loss traced per candidate, adds weighted-PCA surface charts, raw-Fourier comparison, barrier-overlap/continuous-classifier proxies, binary/input hashing, Slurm/Condor orchestration, and a pytest suite. The new Python is generally careful and well-tested (bounds checks, hash pinning, deterministic seeds, schema validation), but there are reproducibility and hygiene gaps around the SCBO engine dependency, the rewritten main.py vs the in-repo AGENTS.md, a stale AGENTS.md Testing section, and committed data artifacts.
Findings:
- [major] pyproject.toml / uv.lock —
simsopt_dfois imported at runtime bymain.py(viaprime_direct_scbo.py:41),advance_alpha_scbo.py:10, andrun_optimizer_comparison.py:149(from simsopt_dfo.scbo_checkpoint import prime_state, issue_candidates, complete_candidates, write_checkpoint), but it is absent frompyproject.tomldependencies,[tool.uv.sources], anduv.lock. A cleanuv sync/pip install -e .thenpython3 main.pyfails withModuleNotFoundError: simsopt_dfoat the first step of the campaign. Required fix: declare the SCBO/TuRBO dependency with a pinned revision (as done forweightedpca). - [major] AGENTS.md (added by this PR) no longer matches the code it documents. It describes
main.pyas the MPI-aware optimizer producingout_s*_NFP*folders/CSV logs and instructspython3 main.py, but this PR replacedmain.pywith a thin wrapper whoseprime_direct_scbo.parser()requires--validation-summary --reference-wout --state --requests --priming(allrequired=True). Runningpython3 main.pyper the docs now exits with an argparse error; the documentedmpirun -n 4/NEAT workflow is gone. The in-repo instructions should describe the actual entrypoints (main.py --validation-summary ...,advance_alpha_scbo.py,run_optimizer_comparison.py,stage_campaign.sh). - [major] Binary/data artifacts and stale inputs committed. The diff adds
.DS_Storeand._.DS_Store(not covered by the new.gitignore),uv.lock(1791 lines),runs/plasma_simple_proxy_correlation_v1/(13 MB of 9216-linestart.dat+design.npzrows), andtest_optimization/(wout.ncbinaries,parvmecinfo.txt). These are campaign artifacts or generated data that belong in the companion data repo README.md already references. Consider removing the junk, adding.DS_Storeto.gitignore, and trimming generated data. - [minor] Inconsistent
--outer-surfacedefault.build_barrier_scbo_response.pydefaults--outer-surfaceto0.6, whilestage_campaign.sh(OUTER_SURFACE=0.7) andrun_barrier_candidate.sh(--outer-surface "${OUTER_SURFACE:-0.7}") use0.7, andbarrier_metricsreportss_outer=0.7. The shell always overrides it, so no runtime bug, but running the builder standalone silently produces a different (0.6) contract. Align the defaults. - [minor]
stage_campaign.shsnapshots only a subset of modules the workers import. Thefiles=(...)list omitsbuild_direct_scbo_response.py,evaluate_direct_loss.py,execution_record.py, andrun_direct_scbo_candidate.sh, whichrun_direct_scbo_candidate.sh(the direct-loss worker) imports/executes viaCODE_ROOT. Benign while only the barrier worker runs, but the direct-loss path isn't pinned by the code-manifest, contradicting the script's own reproducibility comment. Add the direct-loss worker and dependencies to the snapshot list (or document that only the barrier campaign is staged). - [minor] AGENTS.md Testing section is stale. It claims "There is no pytest suite" and lists removed
regression_test.py/test_search_global_minimum.pyas the robustness scripts, even though this PR adds a realtest_*.pysuite andpytest>=9.1.1inpyproject.toml. Update the Testing section to point atpytest.
Verdict: Request changes — the simsopt_dfo dependency is undeclared so the campaign entrypoints (main.py, advance_alpha_scbo.py) cannot run from a clean install, and the in-repo AGENTS.md instructions that ship with this PR now describe a workflow the same PR deleted. Declare the dependency, align the docs, and trim committed artifacts before merge.
Needs SIMPLE#512 and
SIMPLE#513.
What this does
Restores a barrier-overlap objective, fixes the three faults that made the July
campaign conclude the fast classifier was slower than direct loss tracing, and
adds continuous variants that a differentiation tool could act on.
Why the July result was configuration, not physics
class_plot = .True.in every commit of the oldsimple_barrier.pydisabledthe early exit in SIMPLE's
classification.f90, so no orbit ever left early.class_plot = .False.but alsotcut = -1d0,which disabled the classifier outright, because
fast_classwas missing fromSIMPLE's dispatch predicate until #512.
discards, three quarters of the budget.
Measured after fixing all three, same equilibrium, aCluster, 16 threads,
multharm = 5, midpoint integrator in Boozer coordinates: barrier overlap over2x1024 starts takes 188 s against 234 s for the direct 256-particle 100 ms
evaluation.
Continuous metrics
The discrete overlap counts integer class codes, so its gradient is zero almost
everywhere. Three stacked discontinuities: the trapped mask, mu bin membership,
and the class code. The first two are mollified here; the third needs the
margins SIMPLE now exposes.
Measured resolution on the ALPES ns16 reference at
nturns = 8:Both class margins need
nturnsreturn periods of the tip map, which a 20 mstrace almost never reaches, and lowering
nturnsdoes not help. The radialexcursion of the banana tips needs two tips and no threshold, and is the width
of the radial band the orbit explores — what a phase-space barrier suppresses.
It is the variant with data.
Faults found by running it
1/Bwith B in Tesla, but SIMPLE carries thefield in Gauss, so
perp_invis of order 1e-5 and every trapped particlelanded in the first bin. A mu-resolved metric was silently single-bin.
whose orbits merely resolve faster. They are dropped from numerator and
denominator alike, with the resolved fraction reported.
CODE_ROOTandSIMPLE_Xpointed at the livecheckout and build tree, so editing the repository or rebuilding SIMPLE
failed every job in flight — 188 candidates died against a moving manifest.
stage_campaign.shsnapshots the executed files and the binary into thecampaign root, the layout the earlier direct campaigns already used.
simsopt.mhd.Vmeconly for the boundary andto write a candidate, pulling in a compiled extension the cluster lacks.
It now reads with
SurfaceRZFourier.from_vmec_inputand writes bysubstituting the boundary block into the base input, so profiles and solver
controls are inherited rather than dropped.
NS_ARRAY = 3, which abortsSIMPLE inside
spline_vmec_data. Rejected up front with the knob named.Tests
68, all with independent oracles. The load-bearing ones: the smooth overlap
converges to the discrete metric as the three mollifier widths go to zero, for
every classifier; it moves under a sub-threshold drift change the discrete
metric cannot see; censored orbits do not dilute it; and the mu edges are
checked against a measured trapped band so the Gauss-scale error cannot recur.
Status
Two campaigns running on aCluster from staged snapshots, dual annealing and
BOBYQA, 128 evaluations each: one on the discrete metric, one on the smooth
radial metric. Each chains a direct-loss validation of the winner at the 100 ms
256-particle contract.