Misc

[YigitElma]

Improve the LCP build method when dealing with large arrays, the changes removes intermediate large arrays, replace redundant matmuls by just column scaling and in-place operations. It also removes some of the big arrays completely, like versions of _Z to cut down the base memory.

Includes changes from #2287

Same spirit as #2281, but this one also improves the optimization initialization by roughly half on GPU. It should help a lot for high-res equilibrium solves.

Just a rough eq.solve memory benchmark:

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -3.30 %    |     4.228e+03      |     4.088e+03      |   -139.53    |       33.63        |       27.51        |
  test_proximal_jac_w7x_with_eq_update   |   -1.97 %    |     6.847e+03      |     6.712e+03      |   -134.58    |       131.45       |       130.58       |
  test_proximal_freeb_jac                |   -0.38 %    |     1.353e+04      |     1.348e+04      |    -51.05    |       84.06        |       82.77        |
  test_proximal_freeb_jac_blocked        |   -0.32 %    |     7.863e+03      |     7.838e+03      |    -25.40    |       76.41        |       73.56        |
  test_proximal_freeb_jac_batched        |   -0.95 %    |     7.879e+03      |     7.804e+03      |    -74.82    |       74.06        |       72.86        |
  test_proximal_jac_ripple               |    0.00 %    |     3.711e+03      |     3.711e+03      |     0.02     |       49.61        |       48.74        |
  test_proximal_jac_ripple_bounce1d      |   -2.40 %    |     4.018e+03      |     3.922e+03      |    -96.24    |       61.52        |       62.57        |
  test_eq_solve                          |   -3.76 %    |     1.829e+03      |     1.760e+03      |    -68.82    |       46.50        |       47.01        |
  test_objective_quadratic_flux_jac      |    0.28 %    |     2.710e+03      |     2.717e+03      |     7.71     |       51.41        |       51.59        |

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         |     -0.68 +/- 4.10     | -5.81e-03 +/- 3.48e-02 |  8.43e-01 +/- 3.1e-02  |  8.49e-01 +/- 1.6e-02  |
 test_equilibrium_init_lowres            |     -3.13 +/- 2.33     | -2.05e-01 +/- 1.52e-01 |  6.33e+00 +/- 1.1e-01  |  6.54e+00 +/- 1.1e-01  |
 test_objective_compile_atf              |     -1.02 +/- 4.14     | -6.31e-02 +/- 2.56e-01 |  6.11e+00 +/- 1.6e-01  |  6.17e+00 +/- 2.0e-01  |
 test_objective_compute_atf              |     -5.01 +/- 15.95    | -1.08e-04 +/- 3.44e-04 |  2.05e-03 +/- 8.0e-05  |  2.16e-03 +/- 3.3e-04  |
 test_objective_jac_atf                  |     -0.93 +/- 3.79     | -1.53e-02 +/- 6.21e-02 |  1.62e+00 +/- 4.1e-02  |  1.64e+00 +/- 4.7e-02  |
 test_perturb_1                          |     -2.87 +/- 4.19     | -3.45e-01 +/- 5.04e-01 |  1.17e+01 +/- 2.4e-01  |  1.20e+01 +/- 4.5e-01  |
 test_proximal_jac_atf                   |     +2.40 +/- 3.48     | +1.27e-01 +/- 1.83e-01 |  5.40e+00 +/- 1.7e-01  |  5.28e+00 +/- 6.0e-02  |
 test_proximal_freeb_compute             |     -1.34 +/- 2.56     | -2.26e-03 +/- 4.32e-03 |  1.67e-01 +/- 3.2e-03  |  1.69e-01 +/- 2.9e-03  |
 test_solve_fixed_iter                   |     +1.56 +/- 2.54     | +3.75e-01 +/- 6.12e-01 |  2.44e+01 +/- 4.6e-01  |  2.41e+01 +/- 4.0e-01  |
 test_LinearConstraintProjection_build   |     -9.63 +/- 3.48     | -6.74e-01 +/- 2.44e-01 |  6.32e+00 +/- 1.9e-01  |  7.00e+00 +/- 1.6e-01  |
 test_objective_compute_ripple           |     -1.21 +/- 3.70     | -2.89e-03 +/- 8.81e-03 |  2.35e-01 +/- 6.1e-03  |  2.38e-01 +/- 6.3e-03  |
 test_objective_grad_ripple              |     -0.05 +/- 2.52     | -4.86e-04 +/- 2.27e-02 |  8.97e-01 +/- 1.3e-02  |  8.98e-01 +/- 1.9e-02  |
 test_objective_quadratic_flux_compute   |     +1.05 +/- 16.24    | +5.36e-04 +/- 8.27e-03 |  5.14e-02 +/- 5.1e-03  |  5.09e-02 +/- 6.5e-03  |
 test_build_transform_fft_midres         |     +0.48 +/- 3.05     | +3.18e-03 +/- 2.02e-02 |  6.67e-01 +/- 1.3e-02  |  6.64e-01 +/- 1.5e-02  |
 test_build_transform_fft_highres        |     +0.18 +/- 2.32     | +1.57e-03 +/- 2.07e-02 |  8.96e-01 +/- 1.5e-02  |  8.94e-01 +/- 1.4e-02  |
 test_equilibrium_init_medres            |     -0.12 +/- 3.85     | -6.54e-03 +/- 2.08e-01 |  5.38e+00 +/- 1.4e-01  |  5.39e+00 +/- 1.5e-01  |
 test_objective_compile_dshape_current   |     -1.15 +/- 2.47     | -3.48e-02 +/- 7.45e-02 |  2.99e+00 +/- 6.9e-02  |  3.02e+00 +/- 2.8e-02  |
 test_objective_compute_dshape_current   |     -5.38 +/- 14.69    | -2.59e-05 +/- 7.06e-05 |  4.55e-04 +/- 2.3e-05  |  4.81e-04 +/- 6.7e-05  |
 test_objective_jac_dshape_current       |    -17.86 +/- 26.99    | -4.13e-03 +/- 6.25e-03 |  1.90e-02 +/- 3.2e-03  |  2.32e-02 +/- 5.4e-03  |
 test_perturb_2                          |     -0.64 +/- 1.63     | -7.40e-02 +/- 1.89e-01 |  1.15e+01 +/- 1.5e-01  |  1.16e+01 +/- 1.2e-01  |
 test_proximal_jac_atf_with_eq_update    |     -3.61 +/- 1.33     | -3.53e-01 +/- 1.30e-01 |  9.43e+00 +/- 1.0e-01  |  9.78e+00 +/- 8.1e-02  |
 test_proximal_freeb_jac                 |     +0.41 +/- 2.35     | +1.58e-02 +/- 8.99e-02 |  3.83e+00 +/- 7.9e-02  |  3.82e+00 +/- 4.3e-02  |
 test_solve_fixed_iter_compiled          |     -0.16 +/- 2.46     | -7.77e-03 +/- 1.22e-01 |  4.93e+00 +/- 9.3e-02  |  4.94e+00 +/- 7.8e-02  |
 test_objective_compute_ripple_bounce1d  |     -0.03 +/- 4.10     | -6.73e-05 +/- 9.16e-03 |  2.23e-01 +/- 8.1e-03  |  2.23e-01 +/- 4.3e-03  |
 test_objective_grad_ripple_bounce1d     |     +0.50 +/- 3.68     | +3.88e-03 +/- 2.86e-02 |  7.80e-01 +/- 2.3e-02  |  7.76e-01 +/- 1.7e-02  |
 test_objective_quadratic_flux_jac       |     +0.06 +/- 0.57     | +4.58e-03 +/- 4.41e-02 |  7.78e+00 +/- 1.7e-02  |  7.77e+00 +/- 4.1e-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.38%. Comparing base (de77cff) to head (b0ddec3).

Additional details and impacted files

@@                           Coverage Diff                           @@
##           yge/fix-compute-everything-jax-finufft    #2293   +/-   ##
=======================================================================
  Coverage                                   94.37%   94.38%           
=======================================================================
  Files                                         101      101           
  Lines                                       28972    28975    +3     
=======================================================================
+ Hits                                        27342    27347    +5     
+ Misses                                       1630     1628    -2     

Files with missing lines	Coverage Δ
desc/objectives/utils.py	100.00% <100.00%> (ø)
desc/optimize/_constraint_wrappers.py	96.89% <100.00%> (-0.02%)	⬇️
desc/optimize/optimizer.py	96.23% <100.00%> (ø)
desc/optimize/utils.py	95.63% <100.00%> (+0.10%)	⬆️
desc/perturbations.py	89.21% <100.00%> (+0.05%)	⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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