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                 |   -5.16 %    |     4.211e+03      |     3.994e+03      |   -217.42    |       34.14        |       32.75        |
  test_proximal_jac_w7x_with_eq_update   |   -0.37 %    |     6.564e+03      |     6.540e+03      |    -24.17    |       153.06       |       159.71       |
  test_proximal_freeb_jac                |    0.17 %    |     1.337e+04      |     1.339e+04      |    23.33     |       82.43        |       83.14        |
  test_proximal_freeb_jac_blocked        |    0.38 %    |     7.724e+03      |     7.754e+03      |    29.23     |       71.99        |       73.21        |
  test_proximal_freeb_jac_batched        |    0.14 %    |     7.674e+03      |     7.685e+03      |    10.71     |       72.20        |       72.10        |
  test_proximal_jac_ripple               |   -1.49 %    |     3.658e+03      |     3.603e+03      |    -54.36    |       58.38        |       58.38        |
  test_proximal_jac_ripple_bounce1d      |   -0.32 %    |     3.813e+03      |     3.801e+03      |    -12.31    |       73.75        |       72.96        |
  test_eq_solve                          |    2.95 %    |     2.119e+03      |     2.181e+03      |    62.43     |       92.05        |       92.46        |

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())

[YigitElma]

@ddudt @dpanici This was ready before with 2 approvals, merge conflict caused the stale review