Add exponential spectral scaling for optimization, an alternative to Fourier continuation

[byoungj]

Add an exponential decaying x_scale as a preconditioner for stellarator optimizations.

[github-actions[bot]]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -0.77 %    |     3.917e+03      |     3.887e+03      |    -29.96    |       39.55        |       36.15        |
  test_proximal_jac_w7x_with_eq_update   |   -1.46 %    |     6.541e+03      |     6.446e+03      |    -95.36    |       160.47       |       161.17       |
  test_proximal_freeb_jac                |    0.31 %    |     1.316e+04      |     1.320e+04      |    40.77     |       85.13        |       82.32        |
  test_proximal_freeb_jac_blocked        |    0.66 %    |     7.482e+03      |     7.532e+03      |    49.54     |       72.39        |       72.18        |
  test_proximal_freeb_jac_batched        |   -1.32 %    |     7.544e+03      |     7.445e+03      |    -99.30    |       71.89        |       72.18        |
  test_proximal_jac_ripple               |   -0.49 %    |     3.491e+03      |     3.474e+03      |    -16.96    |       64.18        |       63.88        |
  test_proximal_jac_ripple_bounce1d      |    0.72 %    |     3.532e+03      |     3.557e+03      |    25.47     |       76.64        |       74.84        |
  test_eq_solve                          |    2.57 %    |     1.975e+03      |     2.026e+03      |    50.81     |       93.26        |       92.85        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[codecov[bot]]

Codecov Report

:x: Patch coverage is 94.73684% with 5 lines in your changes missing coverage. Please review. :white_check_mark: Project coverage is 95.75%. Comparing base (a34746c) to head (973d557). :warning: Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/optimize/aug_lagrangian.py	0.00%	2 Missing :warning:
desc/optimize/aug_lagrangian_ls.py	0.00%	2 Missing :warning:
desc/optimizable.py	80.00%	1 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1736      +/-   ##
==========================================
+ Coverage   95.73%   95.75%   +0.01%     
==========================================
  Files         102      102              
  Lines       28256    28340      +84     
==========================================
+ Hits        27052    27136      +84     
  Misses       1204     1204              

Files with missing lines	Coverage Δ
desc/equilibrium/equilibrium.py	96.08% <100.00%> (+0.06%)	:arrow_up:
desc/geometry/curve.py	96.44% <100.00%> (+0.10%)	:arrow_up:
desc/geometry/surface.py	97.60% <100.00%> (+0.05%)	:arrow_up:
desc/magnetic_fields/_current_potential.py	99.60% <100.00%> (+<0.01%)	:arrow_up:
desc/optimize/optimizer.py	96.20% <100.00%> (+1.47%)	:arrow_up:
desc/utils.py	93.61% <100.00%> (+0.09%)	:arrow_up:
desc/optimizable.py	96.90% <80.00%> (-0.95%)	:arrow_down:
desc/optimize/aug_lagrangian.py	96.93% <0.00%> (-0.38%)	:arrow_down:
desc/optimize/aug_lagrangian_ls.py	95.81% <0.00%> (-0.41%)	:arrow_down:

... and 2 files with indirect coverage changes

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ddudt]

There are two different x_scale arguments that are related to scaling the optimization variables: "x_scale" and "linear_constraint_options": "x_scale". You are using the former, but you might also want to look at the latter. See PR #1697 for details.

[f0uriest]

One thought that might make future versions of this easier, is if we allow the x_scale the user passes in to be a pytree with the same structure as things.params_dict that we then parse into the correct vector to hand off to the optimizer. eg instead of

x_scale = np.array([1,2,3,4,5,6])

we do

x_scale = {"R_lmn": [1,2,3], "Z_lmn": [4,5,6]}

[dpanici]

Another thought that i realized: right now if ess is used, then np.ones is set for x_scale of any other optimizable DOF in the problem, this means that if one is say doing single stage where both boundary and the coil geometry can vary, the coil DOFs will not be scaled at all. It would be nice to either have some combination of like ess-auto or to be able to specify per-DOF which scaling is preferred (maybe at that point, expose ess scaling function to the user and then let them pass it in manually, perhaps as a pytree like Rory suggested)

[dpanici]

@byoungj any update on this? this would be really nice to have in master, I've found that it actually helps free boundary calculations converge more quickly and more robustly, so I'd like to push this even if just for that main purpose

edit: adding notebook showing the improvement for the vacuum stellarator free bdry we have in our notebooks: free_boundary_equilibrium.ipynb

[byoungj]

Thanks for checking in @dpanici! I will get back to this today/tomorrow so we can push to main :)

[dpanici]

Thanks for checking in @dpanici! I will get back to this today/tomorrow so we can push to main :)

I'm also happy to help with any of the dev work!

[YigitElma]

I would also be happy to help!

[byoungj]

For this PR, the plan is

• add logic to accept both Equilibrium objects and FourierRZToroidalSurface objects for ESS.
• add warning/error to prevent ESS usage in eq.solve()
• improve documentation clarity on ESS parameters and defaults
• add basic tests for ESS.

For future PRs:

• support ESS scaling for other optimizable DOFs (coils, etc.) beyond boundary.
• allow x_scale to be a pytree matching things.params_dict structure for per-DOF control making it more user-accessible/usable.
• expand the ESS to be used within eq.solve()
• add more comprehensive test coverage including multi things optimization.

I can't really think of a great way to create test. So, I will just add a simple solovev case that verify the optimization runs without broadcast/shape errors.

[dpanici]

I think a decent regression test could be justusing ess for a stell optimization like precise QA where it seems to make it work better/work without fourier continuation. But I also worry that would be pretty fragile of a test...

Or better yet, checking that the x_scale of the optimization matches the expected scaling given alpha etc (might have to have optimizer return x_scale in results but that is fine to add I think)

[dpanici]

Also would be good to check if ess works with lsq-auglag and if not, to add a more graceful error than whatever shape mismatch that would probably crop up if it does not work with auglag

[dpanici]

https://github.com/PlasmaControl/DESC/blob/master/CONTRIBUTING.rst#python-styleguide try installing pre-commit to prevent the formatting comments!

[dpanici]

@f0uriest Add to Optimizable a (private) method to compute ess scaling per object (takes in alpha), we can then use this in optimizer.py when "ess" is used

[f0uriest]

I played around with ESS for regular fixed boundary solve and it seems to work quite well, surfaces look identical by eye and force error is more or less the same. These are 3 highly shaped equilibria where we usually needed to use our continuation/perturbation method to get good solutions.

(this was using the default alpha=1.2 and norm=max(|l|, |m|, |n|))

HSX

W7-X

NCSX

[dpanici]

Something I noticed: right now using x_scale="ess" and doing an augmented lagrangian optimization with inequality constraints, you will get an error due to shape mismatch bc the ess x_scale used is just for x, but when inequality constraints exist we have additional slack variables due to the constraints. The usual "auto" jac scale takes care of this bc it uses lagjac to compute the jacobian, but we will either have to add ones for the slack variables or add some heuristic for them (or use jac scaling for just them?)

specifically doing something like adding to here

to this

    if jac_scale:
        scale, scale_inv = compute_jac_scale(J)
    else:
        x_scale_with_z = jnp.concatenate([x_scale, jnp.ones(z0.size - x0.size)])
        x_scale = jnp.broadcast_to(x_scale_with_z, z.shape)
        scale, scale_inv = x_scale, 1 / x_scale

(but this does not work well if x_scale is given as a single float...)

maybe this is better:

    if jac_scale:
        scale, scale_inv = compute_jac_scale(J)
    else:
        x_scale = jnp.concatenate([x_scale, jnp.ones(z0.size - x0.size)]) if x_scale.size>1 else x_scale
        x_scale = jnp.broadcast_to(x_scale_with_z, z.shape)
        scale, scale_inv = x_scale, 1 / x_scale

[dpanici]

Something I noticed: right now using x_scale="ess" and doing an augmented lagrangian optimization with inequality constraints, you will get an error due to shape mismatch bc the ess x_scale used is just for x, but when inequality constraints exist we have additional slack variables due to the constraints. The usual "auto" jac scale takes care of this bc it uses lagjac to compute the jacobian, but we will either have to add ones for the slack variables or add some heuristic for them (or use jac scaling for just them?)

specifically doing something like adding to here

to this

    if jac_scale:
        scale, scale_inv = compute_jac_scale(J)
    else:
        x_scale_with_z = jnp.concatenate([x_scale, jnp.ones(z0.size - x0.size)])
        x_scale = jnp.broadcast_to(x_scale_with_z, z.shape)
        scale, scale_inv = x_scale, 1 / x_scale

(but this does not work well if x_scale is given as a single float...)

maybe this is better:

    if jac_scale:
        scale, scale_inv = compute_jac_scale(J)
    else:
        x_scale = jnp.concatenate([x_scale, jnp.ones(z0.size - x0.size)]) if x_scale.size>1 else x_scale
        x_scale = jnp.broadcast_to(x_scale_with_z, z.shape)
        scale, scale_inv = x_scale, 1 / x_scale

I made this change locally and it does seem help lsq-auglag out a lot I think! now that both boundary and interior modes are scaled with ess

[dpanici]

Testing for Omnigenous field? - future work

[dpanici]

Any future todos I missed please place in this issue: #1939