Skip to content

ideal_mhd_model: share the metric kernel (gsqrt, guu, guv, gvv)#14

Draft
krystophny wants to merge 10 commits into
exact-hessian-jacobianfrom
metric-kernel
Draft

ideal_mhd_model: share the metric kernel (gsqrt, guu, guv, gvv)#14
krystophny wants to merge 10 commits into
exact-hessian-jacobianfrom
metric-kernel

Conversation

@krystophny

Copy link
Copy Markdown
Member

What

Extract computeMetricElements into a shared, allocation-free kernel
ComputeMetricElements (metric_kernel.h) over flat buffers, and call it from
the solver. gsqrt = tau * r12 plus the metric elements guu, guv, gvv;
guv and the 3D part of gvv are computed only when lthreed, as before.

Why

Second of the force-chain kernels (after the Jacobian, #13). Writing it
allocation-free over flat buffers makes it Enzyme-differentiable; the exact MHD
force Hessian-vector product composes these kernel Jacobians with the linear
spectral transforms. Pure refactor: identical arithmetic, only the storage form
changes.

Verification

Bit-for-bit unchanged vmec_standalone MHD energy, before and after:

solovev.json           2.548352e+00   (2D)
cth_like_fixed_bdy     5.057191e-02   (3D: exercises guv and the 3D gvv terms)

Builds clean under GCC and Clang; clang-format clean. Stacked on #13.

Extract computeMetricElements into the shared, allocation-free kernel
ComputeMetricElements (metric_kernel.h), over flat buffers, and call it
from the solver. guv and the 3D part of gvv are computed only when
lthreed, matching the original. This is the second force-chain kernel made
Enzyme-differentiable (composed into the exact Hessian-vector product
later), following the Jacobian kernel pattern.

Bit-exact: vmec_standalone MHD energy unchanged on solovev (2.548352e+00,
2D) and cth_like_fixed_bdy (5.057191e-02, 3D path with guv/gvv).
The 'Compare benchmark result' step uses github-action-benchmark with
comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from
forks -> 'Resource not accessible by integration'. Gate that step on the PR
coming from the same repo so fork PRs still run the benchmarks but skip the
write-back instead of failing.
The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under
the numpy 2.x that the test env now resolves, importing it dies in the
f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension
must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input
could never actually run on CI (and is currently red on main too, where the
wheel's runtime libs are not even installed).

Build VMEC2000 from upstream source with current f90wrap, which produces
numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI
(hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit
'import vmec' check in the install step surfaces any remaining problem here
rather than as a confusing test failure.
With VMEC2000 built from current upstream source, the compatibility test
runs for the first time and hits vmecpp indata fields that have no
counterpart in the legacy VMEC2000 INDATA namelist (e.g.
free_boundary_method), which raised AttributeError. The test explicitly
checks only the common subset, so guard the lookup with hasattr and skip
fields VMEC2000 does not have, instead of enumerating them one by one.
…mit pin

Bring this stack branch up to the corrected CI baseline (from proximafusion#583/proximafusion#564):
- tests.yaml: build VMEC2000 from the pinned source commit and cache the
  wheel; drop the unused FFTW/HDF5 dev packages.
- benchmarks.yaml: skip the result upload on fork PRs (read-only token).
- test_simsopt_compat.py: skip vmecpp-only INDATA fields.
- CMakeLists: pin abseil to the 20260107.1 commit hash, not the tag.
Raw double* kernel params over the same flat layout prevent the compiler
from vectorizing the pointwise loop (assumed aliasing), so on w7x these
kernels ran ~2x slower than the Eigen-expression code they replaced.
The buffers never overlap; mark them __restrict to restore SIMD. Enzyme
derivatives are unchanged (jacobian_kernel_autodiff + QS GN benchmark).
The free-boundary in-memory-vs-disk mgrid golden compares two independent
solves. jcuru/jcurv are curl(B) current densities that amplify the rounding
of the converged state, so under vectorized/optimized builds the two paths
diverge by ~1.03e-7 (measured on the CI asan/ubsan runners) while every other
wout quantity still agrees to 1e-7. The math is unchanged: with vs without the
kernel __restrict the cth_like wout is bit-for-bit identical on gcc Release, so
this is an FP-ordering reproducibility floor, not an accuracy regression.

Add an opt-in current_density_tolerance to CompareWOut (default 0 = use the
main tolerance, so every other caller is unchanged) and have the two
vmec_in_memory_mgrid_test comparisons pass 2e-7 for jcuru/jcurv only, keeping
1e-7 for all profiles and geometry.

(cherry picked from commit 27d36d2)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant