More efficient field computation

[YigitElma]

API may need changes, but I think the desired behavior cannot be implemented by adding already computed data to the field's compute method.

For field line tracing, the speed-up is quite significant. I will do profiling for other objectives.

Created the PR to have some code to discuss. I will finalize it later.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    1.46 %    |     3.979e+03      |     4.037e+03      |    57.98     |       32.28        |       30.25        |
  test_proximal_jac_w7x_with_eq_update   |   -0.88 %    |     6.654e+03      |     6.595e+03      |    -58.78    |       159.97       |       156.95       |
  test_proximal_freeb_jac                |   -0.02 %    |     1.335e+04      |     1.335e+04      |    -2.16     |       81.16        |       78.93        |
  test_proximal_freeb_jac_blocked        |   -0.35 %    |     7.682e+03      |     7.655e+03      |    -26.82    |       71.72        |       71.11        |
  test_proximal_freeb_jac_batched        |   -0.42 %    |     7.675e+03      |     7.643e+03      |    -31.95    |       71.32        |       70.72        |
  test_proximal_jac_ripple               |    0.42 %    |     3.434e+03      |     3.448e+03      |    14.33     |       55.67        |       55.18        |
  test_proximal_jac_ripple_bounce1d      |   -0.62 %    |     3.787e+03      |     3.764e+03      |    -23.42    |       71.07        |       70.26        |
  test_eq_solve                          |   -0.40 %    |     1.938e+03      |     1.930e+03      |    -7.66     |       60.71        |       59.11        |

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

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -4.23 +/- 3.42     | -3.75e-02 +/- 3.02e-02 |  8.47e-01 +/- 2.2e-02  |  8.84e-01 +/- 2.1e-02  |
 test_equilibrium_init_medres            |     +1.53 +/- 5.26     | +1.06e-01 +/- 3.66e-01 |  7.05e+00 +/- 2.5e-01  |  6.95e+00 +/- 2.6e-01  |
 test_equilibrium_init_highres           |     -2.96 +/- 4.06     | -2.38e-01 +/- 3.25e-01 |  7.78e+00 +/- 2.6e-01  |  8.02e+00 +/- 1.9e-01  |
 test_objective_compile_dshape_current   |     -1.94 +/- 2.26     | -8.56e-02 +/- 9.97e-02 |  4.33e+00 +/- 5.9e-02  |  4.42e+00 +/- 8.0e-02  |
 test_objective_compute_dshape_current   |    +11.65 +/- 11.27    | +8.09e-05 +/- 7.82e-05 |  7.75e-04 +/- 6.4e-05  |  6.94e-04 +/- 4.5e-05  |
 test_objective_jac_dshape_current       |     +6.30 +/- 22.77    | +1.63e-03 +/- 5.88e-03 |  2.75e-02 +/- 4.0e-03  |  2.58e-02 +/- 4.3e-03  |
 test_perturb_2                          |     -1.85 +/- 1.88     | -3.88e-01 +/- 3.94e-01 |  2.05e+01 +/- 3.0e-01  |  2.09e+01 +/- 2.5e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.94 +/- 2.63     | +1.16e-01 +/- 3.25e-01 |  1.25e+01 +/- 2.6e-01  |  1.24e+01 +/- 1.9e-01  |
 test_proximal_freeb_jac                 |     +0.51 +/- 4.47     | +2.51e-02 +/- 2.18e-01 |  4.91e+00 +/- 1.7e-01  |  4.88e+00 +/- 1.4e-01  |
 test_solve_fixed_iter_compiled          |     +5.28 +/- 3.69     | +3.38e-01 +/- 2.36e-01 |  6.74e+00 +/- 2.2e-01  |  6.40e+00 +/- 8.8e-02  |
 test_LinearConstraintProjection_build   |     +0.90 +/- 3.93     | +6.39e-02 +/- 2.80e-01 |  7.18e+00 +/- 2.3e-01  |  7.11e+00 +/- 1.5e-01  |
 test_objective_compute_ripple_bounce1d  |     -2.41 +/- 7.01     | -7.54e-03 +/- 2.19e-02 |  3.05e-01 +/- 1.6e-02  |  3.13e-01 +/- 1.5e-02  |
 test_objective_grad_ripple_bounce1d     |     +0.24 +/- 3.44     | +2.31e-03 +/- 3.30e-02 |  9.60e-01 +/- 2.3e-02  |  9.58e-01 +/- 2.4e-02  |
 test_build_transform_fft_midres         |     -1.12 +/- 3.11     | -9.68e-03 +/- 2.68e-02 |  8.52e-01 +/- 2.2e-02  |  8.62e-01 +/- 1.5e-02  |
 test_build_transform_fft_highres        |     -0.29 +/- 2.50     | -3.19e-03 +/- 2.71e-02 |  1.08e+00 +/- 2.3e-02  |  1.08e+00 +/- 1.5e-02  |
 test_equilibrium_init_lowres            |     -0.12 +/- 3.37     | -7.71e-03 +/- 2.15e-01 |  6.36e+00 +/- 1.7e-01  |  6.37e+00 +/- 1.4e-01  |
 test_objective_compile_atf              |     -0.30 +/- 2.92     | -1.74e-02 +/- 1.71e-01 |  5.85e+00 +/- 1.0e-01  |  5.87e+00 +/- 1.4e-01  |
 test_objective_compute_atf              |     -4.17 +/- 10.54    | -1.10e-04 +/- 2.79e-04 |  2.54e-03 +/- 1.8e-04  |  2.65e-03 +/- 2.2e-04  |
 test_objective_jac_atf                  |     -0.48 +/- 4.87     | -8.46e-03 +/- 8.59e-02 |  1.75e+00 +/- 4.1e-02  |  1.76e+00 +/- 7.5e-02  |
 test_perturb_1                          |     +0.11 +/- 3.38     | +1.62e-02 +/- 5.14e-01 |  1.52e+01 +/- 2.1e-01  |  1.52e+01 +/- 4.7e-01  |
 test_proximal_jac_atf                   |     +0.37 +/- 1.62     | +1.96e-02 +/- 8.61e-02 |  5.33e+00 +/- 5.8e-02  |  5.31e+00 +/- 6.3e-02  |
 test_proximal_freeb_compute             |     +0.81 +/- 4.03     | +1.18e-03 +/- 5.85e-03 |  1.47e-01 +/- 3.4e-03  |  1.45e-01 +/- 4.7e-03  |
 test_solve_fixed_iter                   |     +0.65 +/- 2.37     | +1.59e-01 +/- 5.80e-01 |  2.47e+01 +/- 4.3e-01  |  2.45e+01 +/- 3.9e-01  |
 test_objective_compute_ripple           |     +0.95 +/- 5.18     | +2.05e-03 +/- 1.12e-02 |  2.18e-01 +/- 6.9e-03  |  2.16e-01 +/- 8.8e-03  |
 test_objective_grad_ripple              |     -0.99 +/- 2.61     | -9.91e-03 +/- 2.63e-02 |  9.95e-01 +/- 2.0e-02  |  1.00e+00 +/- 1.6e-02  |

Github CI performance can be noisy. When evaluating the benchmarks, developers should take this into account.

[codecov]

Codecov Report

❌ Patch coverage is 87.62887% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.29%. Comparing base (24aa7b9) to head (ccc14a1).

Files with missing lines	Patch %	Lines
desc/coils.py	90.14%	7 Missing ⚠️
desc/magnetic_fields/_core.py	86.95%	3 Missing ⚠️
desc/particles.py	33.33%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2262      +/-   ##
==========================================
- Coverage   94.32%   94.29%   -0.04%     
==========================================
  Files         101      101              
  Lines       28866    28949      +83     
==========================================
+ Hits        27228    27297      +69     
- Misses       1638     1652      +14     

Files with missing lines	Coverage Δ
desc/particles.py	92.35% <33.33%> (-0.55%)	⬇️
desc/magnetic_fields/_core.py	95.77% <86.95%> (-0.28%)	⬇️
desc/coils.py	97.30% <90.14%> (-0.67%)	⬇️

... and 3 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]

Related to #1827

[ddudt]

We often have a MixedCoilSet where we want to pass a list of different grids for each sub-coil type. So we need to remove this check or change it to allow for lists of grids instead of a single grid.

[YigitElma]

I think this one is fine, but I need to update _as_precomputed_source method of CoilSet and MixedCoilSet to accept a list of grids and pass them properly

[ddudt]

Is there ever a reason why you wouldn't want to use this? I think we can safely make this change without giving the user an option to revert back to the old slow method.

[YigitElma]

I haven't benchmarked the case where we differentiate over this function (the docstring is an optimistic hope). The trick to make the biot savart a single operation rather than multiple biot-savart for individual coils didn't make objectives like QuadraticFlux any faster, but it increased the memory consumption a lot, so I wanted to have an option to opt out of this.