Add option to choose Proximal inversion method #2031

YigitElma · 2025-12-05T19:52:27Z

Adds new options to use during the computation of projected tangents. The operation is used to solve:

$$dx =\left(\frac{\partial F}{\partial x}\right)^{-1}\frac{\partial F}{\partial c}$$

svd : just computes SVD of $\frac{\partial F}{\partial x}$
svd-reg : to reproduce the old behavior, adds the smallest eigenvalue to all eigenvalues as regularization
qr : uses QR to solve the system (default)

Resolves #2000

github-actions · 2025-12-05T20:17:34Z

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    4.32 %    |     3.890e+03      |     4.058e+03      |    168.16    |       38.89        |       35.94        |
  test_proximal_jac_w7x_with_eq_update   |   -0.08 %    |     6.480e+03      |     6.475e+03      |    -4.94     |       148.11       |       161.57       |
  test_proximal_freeb_jac                |   -0.40 %    |     1.320e+04      |     1.315e+04      |    -52.20    |       83.11        |       83.32        |
  test_proximal_freeb_jac_blocked        |    0.31 %    |     7.474e+03      |     7.498e+03      |    23.32     |       71.89        |       72.39        |
  test_proximal_freeb_jac_batched        |   -0.59 %    |     7.449e+03      |     7.406e+03      |    -43.65    |       71.78        |       72.41        |
  test_proximal_jac_ripple               |   -2.01 %    |     3.515e+03      |     3.445e+03      |    -70.55    |       61.67        |       64.49        |
  test_proximal_jac_ripple_bounce1d      |   -5.34 %    |     3.660e+03      |     3.465e+03      |   -195.51    |       72.51        |       75.12        |
  test_eq_solve                          |   -3.40 %    |     2.050e+03      |     1.980e+03      |    -69.62    |       94.08        |       93.58        |

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

…gle step

YigitElma · 2025-12-05T21:05:38Z

tests/test_examples.py

    eq.solve(verbose=3)

-    eq1 = run_qh_step(0, eq)
+    eq1 = run_qh_step(2, eq)


I should check the notebooks to see if anything is changed with this

@dpanici also reported this kind of effect in #2000 but I think here doing continuation with the current settings overcommits the lower order modes to a local minimum and cause issues in later steps. Since without continuation it converges immediately, I don't think this is a problem.

If this test doesn't work with the new option that suggests that there may be something wrong. We should figure out what the actual issue is rather than changing the test.

(note that this is such low resolution that it would already converge without the fourier continuation, but we want to test that part as well)

codecov · 2025-12-05T22:56:38Z

Codecov Report

❌ Patch coverage is 78.26087% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 95.75%. Comparing base (3753ff0) to head (0eb8369).

Files with missing lines	Patch %	Lines
desc/optimize/_constraint_wrappers.py	77.27%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2031      +/-   ##
==========================================
- Coverage   95.75%   95.75%   -0.01%     
==========================================
  Files         102      102              
  Lines       28344    28356      +12     
==========================================
+ Hits        27142    27152      +10     
- Misses       1202     1204       +2

Files with missing lines	Coverage Δ
desc/optimize/optimizer.py	`96.21% <100.00%> (+0.01%)`	⬆️
desc/optimize/_constraint_wrappers.py	`96.14% <77.27%> (-1.22%)`	⬇️

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

desc/optimize/_constraint_wrappers.py

desc/optimize/optimizer.py

github-actions · 2025-12-06T22:13:04Z

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_midres         |     +0.77 +/- 4.79     | +6.28e-03 +/- 3.90e-02 |  8.21e-01 +/- 2.7e-02  |  8.15e-01 +/- 2.8e-02  |
 test_build_transform_fft_highres        |     -0.20 +/- 2.04     | -2.11e-03 +/- 2.21e-02 |  1.08e+00 +/- 1.4e-02  |  1.08e+00 +/- 1.7e-02  |
 test_equilibrium_init_lowres            |     -1.89 +/- 2.98     | -1.09e-01 +/- 1.72e-01 |  5.66e+00 +/- 1.1e-01  |  5.77e+00 +/- 1.3e-01  |
+test_objective_compile_atf              |     -4.79 +/- 1.11     | -3.81e-01 +/- 8.84e-02 |  7.57e+00 +/- 6.3e-02  |  7.96e+00 +/- 6.2e-02  |
 test_objective_compute_atf              |     -7.79 +/- 10.91    | -1.75e-04 +/- 2.45e-04 |  2.07e-03 +/- 1.9e-04  |  2.25e-03 +/- 1.6e-04  |
 test_objective_jac_atf                  |     +0.09 +/- 3.15     | +1.55e-03 +/- 5.45e-02 |  1.73e+00 +/- 2.9e-02  |  1.73e+00 +/- 4.6e-02  |
 test_perturb_1                          |     +2.57 +/- 1.64     | +3.76e-01 +/- 2.39e-01 |  1.50e+01 +/- 6.0e-02  |  1.46e+01 +/- 2.3e-01  |
+test_proximal_jac_atf                   |    -37.05 +/- 1.16     | -1.99e+00 +/- 6.22e-02 |  3.38e+00 +/- 5.4e-02  |  5.36e+00 +/- 3.0e-02  |
 test_proximal_freeb_compute             |     +1.28 +/- 3.60     | +2.03e-03 +/- 5.68e-03 |  1.60e-01 +/- 4.6e-03  |  1.58e-01 +/- 3.3e-03  |
 test_solve_fixed_iter                   |     +2.16 +/- 3.27     | +5.97e-01 +/- 9.03e-01 |  2.82e+01 +/- 6.3e-01  |  2.76e+01 +/- 6.5e-01  |
 test_objective_compute_ripple           |     +2.36 +/- 3.02     | +4.81e-03 +/- 6.16e-03 |  2.09e-01 +/- 3.3e-03  |  2.04e-01 +/- 5.2e-03  |
 test_objective_grad_ripple              |     -4.18 +/- 3.22     | -3.71e-02 +/- 2.86e-02 |  8.52e-01 +/- 1.9e-02  |  8.89e-01 +/- 2.1e-02  |
 test_build_transform_fft_lowres         |     -0.17 +/- 2.91     | -1.27e-03 +/- 2.15e-02 |  7.37e-01 +/- 2.0e-02  |  7.38e-01 +/- 6.5e-03  |
 test_equilibrium_init_medres            |     -0.08 +/- 0.95     | -4.85e-03 +/- 5.75e-02 |  6.03e+00 +/- 4.7e-02  |  6.03e+00 +/- 3.3e-02  |
 test_equilibrium_init_highres           |     -1.35 +/- 0.88     | -9.21e-02 +/- 6.04e-02 |  6.75e+00 +/- 4.2e-02  |  6.84e+00 +/- 4.4e-02  |
 test_objective_compile_dshape_current   |     -0.59 +/- 6.18     | -2.42e-02 +/- 2.53e-01 |  4.06e+00 +/- 2.2e-01  |  4.09e+00 +/- 1.2e-01  |
 test_objective_compute_dshape_current   |     +4.04 +/- 13.68    | +2.90e-05 +/- 9.79e-05 |  7.45e-04 +/- 9.1e-05  |  7.16e-04 +/- 3.6e-05  |
 test_objective_jac_dshape_current       |     -2.07 +/- 15.93    | -5.43e-04 +/- 4.18e-03 |  2.57e-02 +/- 3.5e-03  |  2.63e-02 +/- 2.3e-03  |
 test_perturb_2                          |     +0.05 +/- 1.49     | +9.82e-03 +/- 2.69e-01 |  1.81e+01 +/- 2.3e-01  |  1.81e+01 +/- 1.3e-01  |
+test_proximal_jac_atf_with_eq_update    |     -8.81 +/- 0.85     | -1.17e+00 +/- 1.13e-01 |  1.21e+01 +/- 9.5e-02  |  1.33e+01 +/- 6.1e-02  |
 test_proximal_freeb_jac                 |     -0.47 +/- 2.49     | -2.30e-02 +/- 1.21e-01 |  4.86e+00 +/- 5.6e-02  |  4.88e+00 +/- 1.1e-01  |
 test_solve_fixed_iter_compiled          |     +0.90 +/- 1.70     | +7.06e-02 +/- 1.34e-01 |  7.95e+00 +/- 5.8e-02  |  7.88e+00 +/- 1.2e-01  |
 test_LinearConstraintProjection_build   |     -0.65 +/- 2.03     | -5.47e-02 +/- 1.70e-01 |  8.34e+00 +/- 9.2e-02  |  8.39e+00 +/- 1.4e-01  |
 test_objective_compute_ripple_bounce1d  |     -2.31 +/- 2.73     | -6.88e-03 +/- 8.14e-03 |  2.91e-01 +/- 6.7e-03  |  2.98e-01 +/- 4.6e-03  |
 test_objective_grad_ripple_bounce1d     |     -2.48 +/- 2.43     | -2.37e-02 +/- 2.33e-02 |  9.34e-01 +/- 1.9e-02  |  9.57e-01 +/- 1.3e-02  |

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

YigitElma · 2025-12-08T05:42:25Z

If people are ok with the changes, I will update the notebooks. Given the single step works just fine, maybe we can make the advanced optimization notebook a single step?

review-notebook-app · 2025-12-08T06:08:58Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

f0uriest · 2025-12-08T21:32:40Z

desc/optimize/optimizer.py

    if wrapper is not None and wrapper.lower() in ["prox", "proximal"]:
        perturb_options = options.pop("perturb_options", {})
        solve_options = options.pop("solve_options", {})
+        inv_method = options.pop("prox_inv_method", "qr")


I'd like to see more testing on a wide class of problems before we change the default.

agreed, I will try it on some trickier free boundary problems I have

f0uriest · 2025-12-08T21:34:51Z

tests/test_examples.py

    eq.solve(verbose=3)

-    eq1 = run_qh_step(0, eq)
+    eq1 = run_qh_step(2, eq)


If this test doesn't work with the new option that suggests that there may be something wrong. We should figure out what the actual issue is rather than changing the test.

(note that this is such low resolution that it would already converge without the fourier continuation, but we want to test that part as well)

f0uriest · 2025-12-08T21:36:48Z

CHANGELOG.md


 - Adds new option for ``loss_function``, ``"sum"``, which changes an objective to target the sum of the values computed.
 - Adds ``"scipy-l-bfgs-b"`` optimizer option as a wrapper to scipy's ``"l-bfgs-b"`` method.
+- Adds ``prox_inv_method`` option to be used in equilibrium constrained optimization problems. ``prox_inv_method`` can be ``qr``, ``svd`` or ``svd-reg``. ``svd-reg`` is the legacy behaviour which had an unnecessary regularization term. The new default is ``qr`` which is faster than ``svd`` without loss of significant accuracy. Note that this may change the results of the optimization scripts. For number of example cases, it is been seen that the new default facilitates the optimization without continuation.


these example cases are super low res, so I wouldn't take that as representative of the true behavior

dpanici · 2025-12-08T22:37:13Z

desc/optimize/_constraint_wrappers.py

+        uf, sf, vtf = jnp.linalg.svd(Fxh, full_matrices=False)
+        sfi = jnp.where(sf < cutoff * sf[0], 0, 1 / sf)
+        return vtf.T @ (sfi * (uf.T @ Fc))
+    elif inv_method == "svd-reg":  # this option will be deleted


why the comment on deleting the option?

I was planning to remove it, but I will keep it at least for foreseable future

YigitElma added 2 commits December 4, 2025 18:36

remove the regularization

7d7537e

add option to choose internal proximal inversion routine

5314810

YigitElma self-assigned this Dec 5, 2025

YigitElma and others added 2 commits December 5, 2025 16:02

update the test, continuation causes bad approx. just optimize in sin…

64c7a7f

…gle step

Merge branch 'master' into yge/proximal-qr

71686cf

YigitElma commented Dec 5, 2025

View reviewed changes

YigitElma commented Dec 6, 2025

View reviewed changes

desc/optimize/_constraint_wrappers.py Outdated Show resolved Hide resolved

YigitElma added 2 commits December 6, 2025 12:36

Merge branch 'master' into yge/proximal-qr

52a08a4

Update desc/optimize/_constraint_wrappers.py

6d39253

YigitElma commented Dec 6, 2025

View reviewed changes

desc/optimize/optimizer.py Outdated Show resolved Hide resolved

Update desc/optimize/optimizer.py

1dde9c2

YigitElma added the run_benchmarks Run timing benchmarks on this PR against current master branch label Dec 6, 2025

update changelog, add test

7fc4383

YigitElma requested review from ddudt, dpanici, f0uriest and rahulgaur104 December 8, 2025 05:39

YigitElma marked this pull request as ready for review December 8, 2025 05:40

update the notebook to use single step

7dc09f2

f0uriest requested changes Dec 8, 2025

View reviewed changes

dpanici reviewed Dec 8, 2025

View reviewed changes

YigitElma marked this pull request as draft December 9, 2025 02:29

YigitElma added 2 commits December 8, 2025 18:29

Merge branch 'master' into yge/proximal-qr

71c10ee

Merge branch 'master' into yge/proximal-qr

0eb8369

Add option to choose Proximal inversion method #2031

Are you sure you want to change the base?

Add option to choose Proximal inversion method #2031

Conversation

YigitElma commented Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Memory benchmark result

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov bot commented Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Dec 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

YigitElma commented Dec 8, 2025

Uh oh!

review-notebook-app bot commented Dec 8, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

YigitElma commented Dec 5, 2025 •

edited

Loading

github-actions bot commented Dec 5, 2025 •

edited

Loading

codecov bot commented Dec 5, 2025 •

edited

Loading

github-actions bot commented Dec 6, 2025 •

edited

Loading