hiddenSymmetries / simsopt

Simons Stellarator Optimizer Code
https://simsopt.readthedocs.io
MIT License
83 stars 43 forks source link

vectorized BoozerLS residual evaluation #409

Closed andrewgiuliani closed 2 months ago

andrewgiuliani commented 2 months ago

This PR introduces two capabilities:

  1. calculating the Boozer residual for the BoozerLS algorithm using vectorized operations in cpp. This is especially important for the Hessian evaluation, since one can quickly run out of memory using the original python implementation.
  2. the boozer residual is the integrated squared PDE residual: \int residual(varphi, theta)^2 dvarphi dtheta. This PR allows the user to weight the residual with 1/modB, preventing the residual from scaling with the coil currents, i.e., \int [residual(varphi, theta)/modB(varphi, theta)]^2 dvarphi dtheta.

Computing the Hessian for surfaces with mpol, ntor = N, I obtain the following timings in seconds:

N, non-vectorized, vectorized
1  0.12            0.06
2  0.82            0.10
3  2.95            0.29
4  7.59            0.67
5  17.56           1.25
6  killed          2.33
7  killed          4.02
8  killed          6.51
9  killed          9.98
codecov[bot] commented 2 months ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 91.28%. Comparing base (d830719) to head (f4c9931). Report is 45 commits behind head on master.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## master #409 +/- ## ========================================== - Coverage 91.55% 91.28% -0.28% ========================================== Files 74 74 Lines 12911 13010 +99 ========================================== + Hits 11821 11876 +55 - Misses 1090 1134 +44 ``` | [Flag](https://app.codecov.io/gh/hiddenSymmetries/simsopt/pull/409/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=hiddenSymmetries) | Coverage Δ | | |---|---|---| | [unittests](https://app.codecov.io/gh/hiddenSymmetries/simsopt/pull/409/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=hiddenSymmetries) | `91.28% <100.00%> (-0.28%)` | :arrow_down: | Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=hiddenSymmetries#carryforward-flags-in-the-pull-request-comment) to find out more.

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

andrewgiuliani commented 2 months ago

If the new vectorized routine is much better, did you consider removing the old routine?

The algorithms in BoozerSurface.minimize_boozer_penalty_constraints_ls rely on the old routine and I preferred to leave it alone, rather than to do a larger refactor of the BoozerSurface class. Specifically, the old routines can return the residual evaluated at all collocation points (a vector) or the sum of all residuals squared (a scalar). The new routines only return the scalarized residual, while BoozerSurface.minimize_boozer_penalty_constraints_ls require the vector residual.

Similarly, it sounds like the new weight_inv_modB option improves robustness, so you could consider making it the only option, so there are fewer cases to maintain.

The BoozerExact algorithm assumes that the residuals are unweighted, so I think we should maintain both options.