Add surf_batch_size to Boozer transform to control memory usage

[dpanici]

• Add surf_batch_size as a kwarg to the Boozer transform compute variables to allow control of memory usage when using with large number of surfaces
• Added as a plumb-thru kwarg to
  • QuasisymmetryBoozer and Omnigenity objectives
  • make_boozmn_output utility function

Especially helpful when making high res boozmn.nc using DESC

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    0.73 %    |     4.054e+03      |     4.083e+03      |    29.46     |       31.64        |       29.63        |
  test_proximal_jac_w7x_with_eq_update   |   -1.44 %    |     6.623e+03      |     6.528e+03      |    -95.15    |       154.67       |       152.75       |
  test_proximal_freeb_jac                |    0.15 %    |     1.336e+04      |     1.338e+04      |    20.04     |       81.42        |       79.84        |
  test_proximal_freeb_jac_blocked        |    0.69 %    |     7.674e+03      |     7.727e+03      |    52.86     |       71.00        |       69.81        |
  test_proximal_freeb_jac_batched        |   -0.39 %    |     7.677e+03      |     7.647e+03      |    -29.62    |       69.49        |       70.11        |
  test_proximal_jac_ripple               |   -1.63 %    |     3.574e+03      |     3.515e+03      |    -58.40    |       55.06        |       54.44        |
  test_proximal_jac_ripple_bounce1d      |   -3.43 %    |     3.759e+03      |     3.630e+03      |   -129.09    |       68.41        |       68.42        |
  test_eq_solve                          |   -1.22 %    |     1.947e+03      |     1.923e+03      |    -23.84    |       58.08        |       57.98        |

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.32%. Comparing base (5bf4543) to head (328ba34).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2267      +/-   ##
==========================================
- Coverage   94.32%   94.32%   -0.01%     
==========================================
  Files         101      101              
  Lines       28866    28870       +4     
==========================================
+ Hits        27229    27231       +2     
- Misses       1637     1639       +2     

Files with missing lines	Coverage Δ
desc/compute/_omnigenity.py	100.00% <100.00%> (+0.39%)	⬆️
desc/objectives/_omnigenity.py	97.04% <100.00%> (+0.02%)	⬆️
desc/vmec_utils.py	90.78% <100.00%> (ø)

... 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.

[ddudt]

Looks fine, but has it been tested? Like running a test with multiple surfaces, and using this to only compute a few surfaces at a time.

[ddudt]

Really it defaults to chunk_size=None and that defaults to no chunking, right?

[YigitElma]

Yes, there is no explicit default, but kwargs.get("surf_batch_size") returns None and it computes everything without chunking.