Skip to content

Add Enzyme objective state gradient#62

Open
krystophny wants to merge 34 commits into
developfrom
ad-25-objective-state-gradient
Open

Add Enzyme objective state gradient#62
krystophny wants to merge 34 commits into
developfrom
ad-25-objective-state-gradient

Conversation

@krystophny

Copy link
Copy Markdown
Member

Makes one non-energy objective differentiable with respect to the flat state.

  • Adds EvalBsqVarianceFlatKernel: volume-weighted variance of |B|^2, computed single-pass from the same flat context and with the same argument signature as the energy kernel.
  • Reverse-mode gradient gvec_mhd3d_flat_bsq_variance_gradient through the existing __enzyme_autodiff interface; cached-context wrappers eval_bsq_variance and eval_bsq_variance_gradient.
  • State.objective_value("bsq_variance", q) and State.objective_state_gradient("bsq_variance", q, method="exact"|"fd"); the FD path differentiates the fixed-state objective, never an equilibrium re-solve.
  • Unknown objective names raise MissingQuantityError; non-Enzyme builds raise UnsupportedFeatureError.

Closes #26.

Verification

Test fails before

$ pytest test-CI/unit-pygvec/test_ad_objectives.py -q --require-pygvec
4 passed in 1.41s   # no objective tests, State.objective_value missing

Test passes after

Enzyme-on isolated build (fixed RZ, LLVM/Flang 22.1.6):

$ pytest test-CI/unit-pygvec/test_ad_objectives.py -q --require-pygvec
7 passed in 1.00s

test_objective_state_gradient_matches_fd checks the exact reverse-mode gradient against per-component centered differences at rtol=1e-5. test_objective_value_bsq_variance reconstructs the objective from the exported diagnostics at rtol=1e-10.

Enzyme-off isolated build:

$ pytest test-CI/unit-pygvec/test_ad_objectives.py -q --require-pygvec
4 passed, 3 skipped in 0.93s

Fresh fixed-RZ Enzyme build has exactly one gvec_mhd3d_flat_bsq_variance and one gvec_mhd3d_flat_bsq_variance_gradient symbol; ctest 6/6.

@krystophny krystophny force-pushed the ad-25-objective-state-gradient branch from 909b42c to 51d9773 Compare July 9, 2026 22:48
@krystophny krystophny force-pushed the ad-25-objective-state-gradient branch from 51d9773 to be9d952 Compare July 9, 2026 23:02
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