Add Deflation methods

[dpanici]

Deflation method motivation: find multiple solutions to non-convex optimization problems (which can include certain equilibirum solves)

This PR adds ways to apply deflation methods in stellarator optimization and equilibrium solving through the new DeflationOperator objective

• adds a new generic objective DeflationOperator whose cost is simply M(x;y) = 1/(x-y)^p + sigma (to add as constraints to an optimization like in Tarek 2022 work). This can be used as a standalone metric, or another _Objective can be passed to it to wrap it and return as the cost M(x;y)f(x) where f(x) is that _Objective's compute value, like is done in usual deflation
• Adds tutorial covering these

References:

• Riley 2024 - for the "exp" deflation type
• Farrell 2015 - for the addition of the shift parameter and the general form of deflation used

TODO

• add tests
• update changelog
• update ForceBalanceDeflated to use pytree inputs for params_to_deflate_with
• add option for using single shift, like discussed in Riley 2024
• Add a wrapper objective so that one can multiply the deflation operator with any arbitrary objective, instead of only able to add it as an extra cost when doing stage one/two optimization.
• figure out how to avoid recompilation

Future work for another PR:

• Implement algorithms from Riley 2024
  • Implement deflated line-search Gauss-Newton algorithm like they use
  • Adapt their algorithm for our usual trust-region approach
• add _equilibrium as attribute of DeflationOperator and test using it in proximal-lsq-exact

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    3.78 %    |     3.873e+03      |     4.020e+03      |    146.25    |       39.40        |       36.86        |
  test_proximal_jac_w7x_with_eq_update   |    0.78 %    |     6.484e+03      |     6.535e+03      |    50.72     |       165.44       |       164.23       |
  test_proximal_freeb_jac                |    0.03 %    |     1.320e+04      |     1.320e+04      |     3.70     |       84.43        |       83.44        |
  test_proximal_freeb_jac_blocked        |   -0.17 %    |     7.531e+03      |     7.518e+03      |    -13.11    |       74.11        |       75.43        |
  test_proximal_freeb_jac_batched        |    0.23 %    |     7.492e+03      |     7.509e+03      |    17.27     |       73.00        |       73.61        |
  test_proximal_jac_ripple               |   -2.19 %    |     3.551e+03      |     3.474e+03      |    -77.86    |       64.58        |       65.91        |
  test_proximal_jac_ripple_bounce1d      |   -0.47 %    |     3.558e+03      |     3.541e+03      |    -16.83    |       76.15        |       77.76        |
  test_eq_solve                          |    2.22 %    |     2.001e+03      |     2.045e+03      |    44.43     |       96.20        |       93.77        |

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

[f0uriest]

So as I see it there are two main ways of applying deflation

1. As a multiplicative factor on the objective eg f(x) -> M(x,x*)f(x)
2. As an additional inequality constraint M(x,x*)<r

The new DeflationOperator objective seems to cover the 2nd case, but for the first case ForceBalanceDeflated only works for equilibrium problems. I think it would be better as a sort of "wrapper objective" that can be applied to any objective to multiply the deflation operator.

We could possibly combine the two and make it a single objective like

class DeflationOperator:
    """Multiplicative or constraint type deflation"""

def __init__(self, objective=None, ...):
    self.objective = objective

def compute(self, x):
    if self.objective is not None:
        f = self.objective.compute(x)
    else:
        f = 1
    return M(x,x*)*f

This would cover both cases, either treating the deflation as an extra constraint (with objective=None) or applying multiplicative deflation to an arbitrary objective (eg by passing objective=ForceBalance())

[f0uriest]

also don't forget to add DeflationOperator to the list of objectives in the docs (and double check that we didn't forget any others)

[f0uriest]

i don't think you ever have conditionals on the mask in the compute function? it looks like its just used as the where arg to prod/sum?

[f0uriest]

I think you might want this to be += jnp.invert(self._not_all_things_to_deflate_are_None) that way it only adds 1 if all the deflated things are None

[dpanici]

yes you are right. the logic was getting confusing