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                 |    1.38 %    |     4.085e+03      |     4.141e+03      |    56.17     |       43.78        |       38.66        |
  test_proximal_jac_w7x_with_eq_update   |   -0.99 %    |     6.546e+03      |     6.481e+03      |    -65.02    |       162.08       |       160.47       |
  test_proximal_freeb_jac                |   -0.13 %    |     1.338e+04      |     1.336e+04      |    -17.85    |       89.35        |       88.75        |
  test_proximal_freeb_jac_blocked        |    0.08 %    |     7.689e+03      |     7.696e+03      |     6.43     |       78.72        |       77.40        |
  test_proximal_freeb_jac_batched        |   -0.72 %    |     7.727e+03      |     7.671e+03      |    -55.52    |       79.03        |       76.30        |
  test_proximal_jac_ripple               |   -2.29 %    |     3.650e+03      |     3.566e+03      |    -83.51    |       64.28        |       63.68        |
  test_proximal_jac_ripple_bounce1d      |   -0.61 %    |     3.912e+03      |     3.888e+03      |    -23.96    |       78.55        |       77.27        |
  test_eq_solve                          |   -0.14 %    |     2.237e+03      |     2.234e+03      |    -3.07     |       98.36        |       98.78        |

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

[dpanici]

@ddudt In the tutorial deflation notebook, for the QA NAE constrained. the savings on the deflation optimization is almost half. that one is only taken for 15 iterations so the cost is dominated by compilation, which is ~3s long, when the overall optimization once compiled takes only 3s.

Even the QH optimization in that notebook, which is using auglag and thus has many iterations, has a 5-10s savings per optimization when recompilation is avoided, bringing the times down from 25-30s per optimization to 20s per optimization.

While the notebook only runs like 3 outer iterations (i.e. only gets 3 solutions), in practice I've used these to get 20-30 solutions, so the savings would be even higher for real applications.

So while it is problem dependent, it can be a huge timesaver to smartly avoid recompilation using the existing functionality I have. For the user API it is pretty straightforward, so I am fine with how it is. I'll push back on the desire to remove this feature as it seems helpful and effective.