jit compile tangent computation in Proximal

[YigitElma]

• Removes the batched_vectorize call from Proximal, instead conducting the vectorized operation manually
• New function is fully jittable, hence prevents the repeated SVD calls
• Updates the clunky pad/clip/split operations inside Proximal
• Instead of computing JVP of ForceBalance for each tangent, only computes for a maximum of dimc_per_thing[eq_idx] tangents. For example, lets say you have 300 boundary coefficients and 100 coil coefficients inside a free boundary solve, master computes 400 JVPs, this PR computes 300.

Note: Removing the batched_vectorize will help to make #1495 work with parallel force balance constraints too.

The document I shared in #1623 (comment) might help to review the changes in this PR. One change from that document is that in this PR, I removed the prox._feasible_tangents and instead multiplied the related part by eq_feasible_tangents, these are equivalent because D@Z is exactly eq_feasible_tangents.

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -0.16 %    |     4.227e+03      |     4.220e+03      |    -6.81     |       31.40        |       29.14        |
  test_proximal_jac_w7x_with_eq_update   |   -4.51 %    |     6.854e+03      |     6.545e+03      |   -309.19    |       148.84       |       154.65       |
  test_proximal_freeb_jac                |   -0.74 %    |     1.354e+04      |     1.344e+04      |   -100.71    |       76.85        |       80.79        |
  test_proximal_freeb_jac_blocked        |   -1.31 %    |     7.887e+03      |     7.784e+03      |   -103.34    |       64.69        |       72.36        |
  test_proximal_freeb_jac_batched        |   -0.83 %    |     7.859e+03      |     7.794e+03      |    -65.13    |       64.00        |       72.75        |
  test_proximal_jac_ripple               |   -0.94 %    |     3.791e+03      |     3.755e+03      |    -35.49    |       49.96        |       57.38        |
  test_proximal_jac_ripple_bounce1d      |   -1.13 %    |     3.974e+03      |     3.929e+03      |    -44.79    |       63.95        |       72.34        |
  test_eq_solve                          |    1.10 %    |     1.832e+03      |     1.852e+03      |    20.23     |       52.23        |       54.86        |
  test_objective_quadratic_flux_jac      |   -0.02 %    |     1.888e+03      |     1.888e+03      |    -0.38     |       34.34        |       37.15        |

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         |     -3.96 +/- 3.14     | -3.11e-02 +/- 2.47e-02 |  7.54e-01 +/- 2.0e-02  |  7.85e-01 +/- 1.4e-02  |
 test_equilibrium_init_lowres            |     -2.79 +/- 4.14     | -1.68e-01 +/- 2.49e-01 |  5.85e+00 +/- 1.6e-01  |  6.02e+00 +/- 1.9e-01  |
 test_objective_compile_atf              |     -0.88 +/- 2.30     | -4.58e-02 +/- 1.20e-01 |  5.16e+00 +/- 5.8e-02  |  5.21e+00 +/- 1.0e-01  |
 test_objective_compute_atf              |     +0.59 +/- 8.71     | +1.22e-05 +/- 1.81e-04 |  2.09e-03 +/- 1.2e-04  |  2.08e-03 +/- 1.4e-04  |
 test_objective_jac_atf                  |     +0.04 +/- 2.17     | +6.54e-04 +/- 3.24e-02 |  1.49e+00 +/- 2.4e-02  |  1.49e+00 +/- 2.2e-02  |
 test_perturb_1                          |     -1.54 +/- 3.41     | -1.64e-01 +/- 3.63e-01 |  1.05e+01 +/- 2.5e-01  |  1.06e+01 +/- 2.6e-01  |
 test_proximal_jac_atf                   |     +2.27 +/- 2.77     | +1.05e-01 +/- 1.29e-01 |  4.75e+00 +/- 9.8e-02  |  4.65e+00 +/- 8.3e-02  |
 test_proximal_jac_atf_with_eq_update    |     +0.70 +/- 3.03     | +7.31e-02 +/- 3.18e-01 |  1.06e+01 +/- 1.3e-01  |  1.05e+01 +/- 2.9e-01  |
 test_proximal_freeb_compute             |     +2.45 +/- 8.17     | +2.99e-03 +/- 9.95e-03 |  1.25e-01 +/- 8.2e-03  |  1.22e-01 +/- 5.6e-03  |
 test_solve_fixed_iter_compiled          |     +0.76 +/- 2.81     | +4.45e-02 +/- 1.64e-01 |  5.89e+00 +/- 1.5e-01  |  5.84e+00 +/- 7.3e-02  |
 test_LinearConstraintProjection_build   |     -1.40 +/- 2.40     | -8.04e-02 +/- 1.37e-01 |  5.65e+00 +/- 6.5e-02  |  5.73e+00 +/- 1.2e-01  |
 test_objective_compute_ripple           |     +0.95 +/- 5.08     | +1.79e-03 +/- 9.52e-03 |  1.89e-01 +/- 6.1e-03  |  1.87e-01 +/- 7.3e-03  |
 test_objective_grad_ripple              |     +2.61 +/- 3.46     | +2.29e-02 +/- 3.04e-02 |  9.00e-01 +/- 1.9e-02  |  8.77e-01 +/- 2.4e-02  |
 test_objective_quadratic_flux_compute   |     -3.62 +/- 3.36     | -9.67e-04 +/- 8.97e-04 |  2.57e-02 +/- 7.0e-04  |  2.67e-02 +/- 5.5e-04  |
 test_build_transform_fft_midres         |     -1.08 +/- 4.59     | -9.59e-03 +/- 4.10e-02 |  8.82e-01 +/- 3.6e-02  |  8.92e-01 +/- 2.0e-02  |
 test_build_transform_fft_highres        |     -1.28 +/- 3.37     | -1.52e-02 +/- 3.99e-02 |  1.17e+00 +/- 2.5e-02  |  1.18e+00 +/- 3.1e-02  |
 test_equilibrium_init_medres            |     -0.27 +/- 2.17     | -1.88e-02 +/- 1.50e-01 |  6.87e+00 +/- 1.2e-01  |  6.89e+00 +/- 9.3e-02  |
 test_objective_compile_dshape_current   |     +0.97 +/- 2.82     | +4.04e-02 +/- 1.17e-01 |  4.20e+00 +/- 9.8e-02  |  4.16e+00 +/- 6.5e-02  |
 test_objective_compute_dshape_current   |     +5.48 +/- 10.75    | +3.86e-05 +/- 7.57e-05 |  7.42e-04 +/- 5.9e-05  |  7.04e-04 +/- 4.8e-05  |
 test_objective_jac_dshape_current       |     -5.96 +/- 26.95    | -1.54e-03 +/- 6.97e-03 |  2.43e-02 +/- 5.7e-03  |  2.59e-02 +/- 4.0e-03  |
 test_perturb_2                          |     -1.81 +/- 1.93     | -2.86e-01 +/- 3.06e-01 |  1.56e+01 +/- 1.2e-01  |  1.58e+01 +/- 2.8e-01  |
+test_proximal_jac_atf_chunked           |    -67.92 +/- 0.68     | -1.14e+01 +/- 1.14e-01 |  5.38e+00 +/- 3.7e-02  |  1.68e+01 +/- 1.1e-01  |
 test_proximal_freeb_jac                 |     +0.72 +/- 2.85     | +3.44e-02 +/- 1.37e-01 |  4.83e+00 +/- 1.2e-01  |  4.80e+00 +/- 7.1e-02  |
 test_solve_fixed_iter                   |     -0.64 +/- 2.43     | -1.56e-01 +/- 5.90e-01 |  2.42e+01 +/- 4.5e-01  |  2.43e+01 +/- 3.8e-01  |
 test_objective_compute_ripple_bounce1d  |     -1.63 +/- 3.96     | -5.01e-03 +/- 1.22e-02 |  3.02e-01 +/- 9.6e-03  |  3.07e-01 +/- 7.5e-03  |
 test_objective_grad_ripple_bounce1d     |     +1.17 +/- 2.32     | +1.10e-02 +/- 2.20e-02 |  9.58e-01 +/- 2.1e-02  |  9.47e-01 +/- 6.4e-03  |
 test_objective_quadratic_flux_jac       |     +0.81 +/- 1.23     | +1.80e-02 +/- 2.73e-02 |  2.25e+00 +/- 1.6e-02  |  2.23e+00 +/- 2.2e-02  |

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.35%. Comparing base (7582fb3) to head (87b886a).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2239      +/-   ##
==========================================
- Coverage   94.35%   94.35%   -0.01%     
==========================================
  Files         101      101              
  Lines       29036    29034       -2     
==========================================
- Hits        27396    27394       -2     
  Misses       1640     1640              

Files with missing lines	Coverage Δ
desc/optimize/_constraint_wrappers.py	96.89% <100.00%> (-0.02%)	⬇️

🚀 New features to boost your workflow:

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

[YigitElma]

Previously, I was trying to be smart by removing the transpose from the Fxh and handle that at the return operation (transpose of u and v etc) but that cause slowdown on CPU. Maybe SVD lowering of JAX on CPU is dependent on tall/wide?!

[dpanici]

I don't think I understand the difference bwtn these two paths

[YigitElma]

• dim_opt = number of free optimization variables (Rb+Zb+profiles+coils etc after linear constraints)
• prox.dim_x = total number of optimization variables (Rb+Zb+profiles+coils etc)
• dim_c = total number of eq optimization variables (Rb+Zb+profiles)
• dim_xeq = size of full eq state vector (Rlmn+Zlmn+Llmn+Rb+Zb etc)

_proximal_eq_tangents calls the JVP of the constraint with eq_feasible_tangents and tangents for the optimization variables. vs is a dim_opt by prox.dim_x sized array, when we take vs[eq_idx], it becomes dim_opt by dim_c. So, the indexing operation doesn't change the number of tangents but makes the tangent vectors shorter.

We know that JVP wrt non-equilibrium optimization variables are 0, so if dim_opt > dim_c, we are going to have some trivially 0 JVPs. We can either take JVP wrt to vs[eq_idx] @ dxdc.T which is dim_opt by dim_xeq or wrt dxdc.T which is dim_c by dim_xeq. These branches compute the lower one with a static check.

[YigitElma]

I originilly had a comment like this

    # dxdc maps the eq DoFs (Rb_lmn, Zb_lmn, profiles) to the full eq state vector.
    # Everything in _proximal_eq_tangents is linear in the directions it gets, so we
    # can either take the JVPs for the columns of dxdc and apply v afterwards with a
    # matmul (first branch), or apply v first and take the JVPs only for the
    # directions in v (second branch). The JVPs dominate the cost, so we use the one
    # with fewer directions. The first branch never needs more than
    # dimc_per_thing[eq_idx] of them, since that is all the eq DoFs there are, so it
    # only pays off when v has more directions than that. Usually v holds the feasible
    # tangents of the outer LinearConstraintProjection and we optimize only some of
    # the boundary modes, so the second branch is the usual one and the first only
    # kicks in with a lot of coil, surface etc DoFs, important for single stage.

But it was very long, I can keep it if this is not clear.

[ddudt]

Do these changes mean that it is faster to use a relatively small jac_chunk_size with ForceBalance as a nonlinear constraint in single-stage optimizations? Or is it still best to use the default (no chunking) for that objective if it can fit in memory?

[YigitElma]

Do these changes mean that it is faster to use a relatively small jac_chunk_size with ForceBalance as a nonlinear constraint in single-stage optimizations? Or is it still best to use the default (no chunking) for that objective if it can fit in memory?

If it can fit in the memory, no chunking should still be faster. These changes basically remove the overhead of unjitted batched_vectorize call, i.e. multiple SVD's. Now, the ProximalProjection constraint chunk size act like any other objective's chunk size.