Make a differentiable version of FourierRZToroidalSurface.constant_offset_surface

[dpanici]

Adds differentiable version of offset surface algorithm. Potentially useful if one desires an objective that concerns a conformal vessel to a changing eq during optimization.

Also fixes what was likely a bug affecting robustness before, which was that I was using arctan instead of a corrected arctan2 when computing the zeta of the offset point for the rootfind. Arctan has a range of [-pi/2, pi/2] while arctan2 will choose the correct quadrant btwn [-pi,pi] (and which can be made to go to [0,2pi] by adding 2pi to the negative part), which more closely adheres to our grid zeta domain which is always positive by convention

notebook showing example use within an objective:
test_differentiable_offset.ipynb

[dpanici]

only is differentiable though if params is passed, otherwise the end will not be differentiatedcorrectlt

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    0.06 %    |     3.938e+03      |     3.941e+03      |     2.51     |       37.96        |       36.25        |
  test_proximal_jac_w7x_with_eq_update   |   -1.31 %    |     6.541e+03      |     6.456e+03      |    -85.43    |       161.22       |       161.93       |
  test_proximal_freeb_jac                |    0.59 %    |     1.320e+04      |     1.328e+04      |    77.27     |       87.31        |       87.73        |
  test_proximal_freeb_jac_blocked        |    1.45 %    |     7.468e+03      |     7.576e+03      |    107.98    |       75.99        |       76.39        |
  test_proximal_freeb_jac_batched        |   -0.86 %    |     7.535e+03      |     7.470e+03      |    -64.88    |       75.89        |       75.90        |
  test_proximal_jac_ripple               |    0.51 %    |     3.576e+03      |     3.594e+03      |    18.38     |       60.37        |       59.86        |
  test_proximal_jac_ripple_bounce1d      |    0.18 %    |     3.713e+03      |     3.720e+03      |     6.71     |       69.29        |       68.89        |
  test_eq_solve                          |   -1.01 %    |     1.985e+03      |     1.965e+03      |    -20.07    |       82.91        |       82.80        |

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

[codecov]

Codecov Report

❌ Patch coverage is 14.77273% with 75 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.52%. Comparing base (a846c9b) to head (36e4beb).

Files with missing lines	Patch %	Lines
desc/objectives/_geometry.py	16.94%	49 Missing ⚠️
desc/geometry/surface.py	10.34%	26 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2016      +/-   ##
==========================================
- Coverage   95.77%   95.52%   -0.26%     
==========================================
  Files         101      101              
  Lines       27796    27881      +85     
==========================================
+ Hits        26622    26632      +10     
- Misses       1174     1249      +75     

Files with missing lines	Coverage Δ
desc/geometry/surface.py	91.70% <10.34%> (-5.84%)	⬇️
desc/objectives/_geometry.py	86.60% <16.94%> (-10.19%)	⬇️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[YigitElma]

Need to add test to make sure this is not compiled again. I had a similar check for field line intergration in test_field_line_integrate_jax_transforms

[YigitElma]

I don't think I follow the code. I know this was added in a previous PR but I wanted to ask to see if this code is correct or not. So, the above rootfind basically tries to find some zeta where if you start from (theta, zeta) and go in the normal direction by offset amount, your zeta is equal to the original grid's zeta. In short, find some points on the original surface with uniform spacing in theta but non-uniform spacing in zeta that give points in offset surface that are uniformly spaced in zeta (and assign a different theta definition that is uniform).

Why do you find the offset zetas here? They should be already equal to grid.nodes[:, 2]?

[dpanici]

They should be equal, assuming the rootfind succeeded. If the rootfind did not succeed they may be different. The rootfind should succeed in most cases tho so I made the change to use the grid.nodes[:,2]

Something weird I found is in some cases, the "fft" transform method is more fragile when doing the fit then the "direct1" method. Not sure why

[YigitElma]

Seems fine to me for now. I will test it later myself