Necessity of choice of regularization in proximal jvp step

[dpanici]

In ProximalProjection we use a regularization in this line

DESC/desc/optimize/_constraint_wrappers.py

Line 1314 in 548a83b

sf += sf[-1] # add a tiny bit of regularization

This specifically is the part that is doing the (dF/dx^{-1}) @ dF/dc in the proximal projection (i.e. in the term on LHS of eq 2.9 in dudt QS paper )

while I understand that we are typically oversampling force error, and so this matrix may not be very well-conditioned, it seems like in practice having a blanket regularization of the smallest singular value may not always be the best choice. If we take our vacuum stellarator free boundary tutorial and run the free boundary problem without using any continuation stepping, we get two very different results with vs without this regularization: (left with the regularization, right without)

of course, with Fourier continuation we get good results in both cases, but this is showing that perhaps there are cases where the regularization is not helping us.

Maybe a different approach could be to not include this regularization and instead rely on the cutoff to remove the problematically small singular values, or adding a threshold based approach based on the condition number like sf += jnp.where(sf[-1]/sf[0] < 1e-8, sf[-1], 1e-8) or something?

notebooks with the runs (also has runs with the ess scaling as that seems to help free boundary too)

free_boundary_equilibrium_vac_no_regularization.ipynb
free_boundary_equilibrium_vac.ipynb

[dpanici]

@dpanici check singular values for different cases and check condition number

[dpanici]

Checked these for basically every optimization or free bdry example in our notebooks. The condition number of Fxh was on average no more than 1e6, and the cutoff threshold for the singular values was consistently many orders (3-8) of magnitude below the lowest singular value (before adding the regularization, which does not change that statement).

The worst condition number was the vacuum free bdry stellarator first step, which was initially 1e9 but after the first step decreased.

So it seems that at least from these examples, the small singular value truncation does not do anything for regularization of this term's inverse. I will also check these examples without using the regularization and see how the optimization is affected

[YigitElma]

The condition number of Fxh was on average no more than 1e6

One thing I realized before is that the jacobian of the force balance objective (Fxh) has higher condition number as the force error gets higher. My understanding is if you are closer to equilibrium, the derivatives gets closer to 0 in each direction hence singular values become closer, hence lower condition number. With the recent changes to make the proximal projection always have a good force balance (i.e. solve before optimization, and reduced tolerances) we might have eliminated the extremely high condition number problem here.

[dpanici]

I still think it would be interesting to see the condition number of the actual proximally projected jacobian (e.g. the one used in the optimizer) with and without this regularization

[YigitElma]

I will take a closer look at this after APS, along with adding a QR option too (see #1569)