Update trust_region_step_exact_qr

[YigitElma]

• Updates the initial guess for $\alpha$ to be 0 instead of 0.001* $\alpha$ _upper. I found out that the solution of $\alpha$ is usually pretty small. This reduces the total amount of QR's in the while loop.
• Removes the last QR outside the while loop of trust_region_step_exact_qr. The loop condition already satisfies the step norm to be around trust radius.

Misc

• Adds execute_on_cpu flag to desc.examples.get

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     +0.98 +/- 3.72     | +6.16e-03 +/- 2.34e-02 |  6.36e-01 +/- 1.9e-02  |  6.29e-01 +/- 1.3e-02  |
 test_build_transform_fft_highres        |     +0.98 +/- 2.65     | +9.67e-03 +/- 2.60e-02 |  9.94e-01 +/- 1.6e-02  |  9.84e-01 +/- 2.1e-02  |
 test_equilibrium_init_lowres            |     -1.54 +/- 3.78     | -6.41e-02 +/- 1.57e-01 |  4.10e+00 +/- 6.1e-02  |  4.16e+00 +/- 1.5e-01  |
 test_objective_compile_atf              |     +0.74 +/- 4.11     | +6.01e-02 +/- 3.31e-01 |  8.13e+00 +/- 2.6e-01  |  8.07e+00 +/- 2.0e-01  |
 test_objective_compute_atf              |     -1.80 +/- 2.28     | -1.97e-04 +/- 2.50e-04 |  1.07e-02 +/- 1.3e-04  |  1.09e-02 +/- 2.1e-04  |
 test_objective_jac_atf                  |     +0.91 +/- 3.45     | +1.75e-02 +/- 6.65e-02 |  1.94e+00 +/- 5.0e-02  |  1.93e+00 +/- 4.4e-02  |
 test_perturb_1                          |     -0.35 +/- 2.76     | -5.20e-02 +/- 4.06e-01 |  1.47e+01 +/- 2.6e-01  |  1.47e+01 +/- 3.2e-01  |
 test_proximal_jac_atf                   |     -0.57 +/- 1.00     | -4.66e-02 +/- 8.29e-02 |  8.20e+00 +/- 7.8e-02  |  8.25e+00 +/- 2.7e-02  |
 test_proximal_freeb_compute             |     +1.03 +/- 1.23     | +2.02e-03 +/- 2.41e-03 |  1.99e-01 +/- 1.6e-03  |  1.97e-01 +/- 1.8e-03  |
 test_solve_fixed_iter_compiled          |     -3.42 +/- 1.95     | -6.02e-01 +/- 3.44e-01 |  1.70e+01 +/- 1.4e-01  |  1.76e+01 +/- 3.1e-01  |
 test_build_transform_fft_lowres         |     -0.66 +/- 6.35     | -3.43e-03 +/- 3.29e-02 |  5.15e-01 +/- 2.1e-02  |  5.19e-01 +/- 2.5e-02  |
 test_equilibrium_init_medres            |     +0.83 +/- 1.00     | +3.37e-02 +/- 4.05e-02 |  4.07e+00 +/- 3.4e-02  |  4.04e+00 +/- 2.2e-02  |
 test_equilibrium_init_highres           |     +0.21 +/- 1.04     | +1.10e-02 +/- 5.55e-02 |  5.35e+00 +/- 3.8e-02  |  5.33e+00 +/- 4.0e-02  |
 test_objective_compile_dshape_current   |     +0.28 +/- 0.93     | +1.07e-02 +/- 3.55e-02 |  3.83e+00 +/- 2.2e-02  |  3.82e+00 +/- 2.8e-02  |
 test_objective_compute_dshape_current   |     +1.05 +/- 1.44     | +3.77e-05 +/- 5.15e-05 |  3.62e-03 +/- 2.9e-05  |  3.58e-03 +/- 4.2e-05  |
 test_objective_jac_dshape_current       |     +4.25 +/- 6.94     | +1.66e-03 +/- 2.71e-03 |  4.07e-02 +/- 1.4e-03  |  3.91e-02 +/- 2.3e-03  |
 test_perturb_2                          |     -0.33 +/- 1.47     | -6.18e-02 +/- 2.78e-01 |  1.89e+01 +/- 2.1e-01  |  1.89e+01 +/- 1.9e-01  |
 test_proximal_freeb_jac                 |     -0.71 +/- 1.04     | -5.24e-02 +/- 7.66e-02 |  7.35e+00 +/- 6.2e-02  |  7.40e+00 +/- 4.5e-02  |
 test_solve_fixed_iter                   |     -1.10 +/- 1.85     | -3.07e-01 +/- 5.17e-01 |  2.76e+01 +/- 4.0e-01  |  2.79e+01 +/- 3.3e-01  |
 test_LinearConstraintProjection_build   |     +0.13 +/- 0.89     | +2.84e-02 +/- 1.99e-01 |  2.25e+01 +/- 1.4e-01  |  2.25e+01 +/- 1.4e-01  |

[f0uriest]

isn't this the same as the cho option but using LU instead of cholesky?

[YigitElma]

isn't this the same as the cho option but using LU instead of cholesky?

Yes, pretty similar. jax.scipy.linlag.solve with symmetric A should call LAPACK ?sysv (docs), and it is something like LU. For finding alpha, I used Secant method.

Cholesky is also really fast (maybe faster than this one), but maybe this is more accurate, that's why I wanted to test it. If I remember correctly the problem with Cholesky was the accuracy.

[f0uriest]

Generally the innaccuracy comes from explitly forming J^T J when J is poorly conditioned, so I don't think you'll see any major improvements with this. Also cholesky should be faster.

[dpanici]

run on multiple threads

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.58%. Comparing base (17d4939) to head (e64f764).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1363   +/-   ##
=======================================
  Coverage   95.58%   95.58%           
=======================================
  Files          96       96           
  Lines       24597    24598    +1     
=======================================
+ Hits        23512    23513    +1     
  Misses       1085     1085           

Files with missing lines	Coverage Δ
desc/examples/__init__.py	100.00% <100.00%> (ø)
desc/optimize/tr_subproblems.py	99.44% <100.00%> (-0.01%)	⬇️

[f0uriest]

do these changes help reduce the number of iterations of the subproblem? If so by how much?

We may want to make the same changes to the other trust region subproblems (svd, cho etc) to keep them consistent.

[f0uriest]

also this could be simplified to alpha = jnp.clip(alpha, alpha_lower, alpha_upper)

[YigitElma]

I don't know the exact effect of clipping alpha rather than setting it to 0.001*alpha_upper. My reasoning for these changes was to make the initial guess of alpha 0, and whenever alpha goes out of the range setting it to 0.001*alpha_upper seemed to be taking alpha away from the actual solution. So, it seemed necessary to clip alpha when I changed the initial guess.

[YigitElma]

The total number iterations for root finding depends on the case but these changes help to reduce it. Previously a root finding was taking 5-10 iterations, now around 3-6.

[f0uriest]

ok in that case can you make the same changes to trust_region_step_exact_{svd,cho}