Add a lot of optimizers including optax support to sgd

[YigitElma]

• Unifies all the SGD type optimizers, developers only need to implement the update rule
• x_scale is now used with SGD methods too
• Adds wrappers for optax optimizers, and they can be called by optax-name
• Any custom optax optimizer can be used via

        import optax
        from desc.optimize import Optimizer
        from desc.examples import get

        eq = get("DSHAPE")

        # Optimizer
        opt = optax.chain(
            optax.sgd(learning_rate=1.0),
            optax.scale_by_zoom_linesearch(max_linesearch_steps=15),
        )
        optimizer = Optimizer("optax-custom")
        eq.solve(optimizer=optimizer, options={"optax-options": {"update_rule": opt}})

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    3.00 %    |     3.895e+03      |     4.012e+03      |    117.03    |       39.74        |       36.33        |
  test_proximal_jac_w7x_with_eq_update   |    0.76 %    |     6.585e+03      |     6.634e+03      |    49.79     |       162.70       |       163.71       |
  test_proximal_freeb_jac                |   -0.37 %    |     1.323e+04      |     1.318e+04      |    -48.81    |       85.70        |       84.40        |
  test_proximal_freeb_jac_blocked        |   -0.48 %    |     7.537e+03      |     7.502e+03      |    -35.93    |       75.27        |       74.88        |
  test_proximal_freeb_jac_batched        |    0.32 %    |     7.484e+03      |     7.508e+03      |    24.32     |       73.76        |       75.69        |
  test_proximal_jac_ripple               |    1.47 %    |     3.468e+03      |     3.519e+03      |    51.07     |       65.97        |       68.48        |
  test_proximal_jac_ripple_bounce1d      |    2.40 %    |     3.571e+03      |     3.656e+03      |    85.59     |       77.62        |       80.36        |
  test_eq_solve                          |    4.16 %    |     2.044e+03      |     2.129e+03      |    84.94     |       94.81        |       96.74        |

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

[codecov]

Codecov Report

❌ Patch coverage is 98.07692% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 94.51%. Comparing base (5bb92a4) to head (8278279).
⚠️ Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/optimize/stochastic.py	98.03%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #2041   +/-   ##
=======================================
  Coverage   94.51%   94.51%           
=======================================
  Files         102      102           
  Lines       28640    28675   +35     
=======================================
+ Hits        27068    27103   +35     
  Misses       1572     1572           

Files with missing lines	Coverage Δ
desc/optimize/_desc_wrappers.py	91.17% <100.00%> (+0.13%)	⬆️
desc/optimize/stochastic.py	98.01% <98.03%> (+1.00%)	⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[dpanici]

just the small docstring fix, should be explitict that x_scale='"auto"` does no scaling here

[f0uriest]

I'd double check that the x_scale logic is correct

Also, did you look at whether we could just wrap stuff from optax?

From the examples eg https://optax.readthedocs.io/en/latest/api/optimizers.html#optax.adam

it looks like the user could just pass in an optax.solver and then we can just do

opt_state = solver.init(x0)
...
g = grad(x)*x_scale
updates, opt_state = solver.update(g, opt_state, x)
x = optax.apply_updates(x, x_scale*updates)

or something similar. That would give users access to a much wider array of first order optimizers, and save us having to do it all ourselves

[f0uriest]

sgd is technically public (https://desc-docs.readthedocs.io/en/stable/_api/optimize/desc.optimize.sgd.html#desc.optimize.sgd) so if we want to change the name we should keep an alias to the old one with a deprecation warning. That said, I'm not sure we really need to change the name. "SGD" is already used pretty generically in the ML community for a bunch of first order stochastic methods like ADAM, ADAGRAD, RMSPROP, etc

[YigitElma]

Yeah, SGD is in fact the general name. I can revert to the old name, and just emphasize that "sgd" option is with nesterov momentum. I was trying to make a distinction I guess

[f0uriest]

this seems pretty large (steps would be 10% of x), have you checked how robust this is?

[YigitElma]

I was trying to solve eq with these and even though none of them converged 10% was better for a variety of equilibrium. I haven't checked other optimization problems. Reverted the change and added a backguard against 0 and NaNs.

[f0uriest]

personally I prefer unicode for stuff like this. TeX looks nice in the rendered html docs, but is much harder to read as code.

[YigitElma]

I mostly agree and don't have a strong stance either way. My general preference is to use LaTeX for complex equations or public-facing objectives that users will first encounter in the documentation (like guiding center equations, optimization algorithms). For internal development notes or specific compute functions that aren't usually viewed on the web, I’m fine with Unicode since it keeps the source code more readable.

[unalmis]

Also, did you look at whether we could just wrap stuff from optax?

or something similar. That would give users access to a much wider array of first order optimizers, and save us having to do it all ourselves

Also optimistix has trust region methods with easy to use linear solvers.., e.g. normal conjugate gradient etc.