Switch compute from recursive to looped by f0uriest · Pull Request #1557 · PlasmaControl/DESC

f0uriest · 2025-02-02T03:08:35Z

This PR builds upon #2156

_compute now uses a for loop instead of recursion
- We still need to call _get_deps to consider given data
Cleans up data dependency getters and removes the redundant functions

Resolves #1336

github-actions · 2025-02-02T03:39:08Z

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -1.70 +/- 4.64     | -1.48e-02 +/- 4.04e-02 |  8.57e-01 +/- 3.8e-02  |  8.72e-01 +/- 1.4e-02  |
 test_equilibrium_init_lowres            |     -1.06 +/- 4.68     | -7.03e-02 +/- 3.10e-01 |  6.54e+00 +/- 8.4e-02  |  6.61e+00 +/- 3.0e-01  |
 test_objective_compile_atf              |     -4.49 +/- 4.93     | -2.89e-01 +/- 3.17e-01 |  6.14e+00 +/- 2.3e-01  |  6.43e+00 +/- 2.2e-01  |
 test_objective_compute_atf              |     +1.83 +/- 11.14    | +3.98e-05 +/- 2.43e-04 |  2.22e-03 +/- 1.4e-04  |  2.18e-03 +/- 2.0e-04  |
 test_objective_jac_atf                  |     -0.06 +/- 4.64     | -9.38e-04 +/- 7.44e-02 |  1.60e+00 +/- 6.0e-02  |  1.60e+00 +/- 4.4e-02  |
 test_perturb_1                          |     -1.12 +/- 2.16     | -1.36e-01 +/- 2.63e-01 |  1.20e+01 +/- 1.3e-01  |  1.22e+01 +/- 2.3e-01  |
 test_proximal_jac_atf                   |     -0.62 +/- 2.27     | -3.28e-02 +/- 1.20e-01 |  5.25e+00 +/- 5.6e-02  |  5.28e+00 +/- 1.1e-01  |
 test_proximal_freeb_compute             |     -0.14 +/- 3.47     | -2.33e-04 +/- 5.90e-03 |  1.70e-01 +/- 4.1e-03  |  1.70e-01 +/- 4.3e-03  |
 test_solve_fixed_iter                   |     +0.36 +/- 2.78     | +8.95e-02 +/- 6.96e-01 |  2.51e+01 +/- 5.7e-01  |  2.50e+01 +/- 4.0e-01  |
 test_LinearConstraintProjection_build   |     -4.58 +/- 6.72     | -3.50e-01 +/- 5.13e-01 |  7.28e+00 +/- 2.4e-01  |  7.64e+00 +/- 4.6e-01  |
 test_objective_compute_ripple           |     -4.25 +/- 6.52     | -1.08e-02 +/- 1.65e-02 |  2.42e-01 +/- 1.0e-02  |  2.53e-01 +/- 1.3e-02  |
 test_objective_grad_ripple              |     +0.67 +/- 2.74     | +5.93e-03 +/- 2.42e-02 |  8.91e-01 +/- 1.9e-02  |  8.85e-01 +/- 1.5e-02  |
 test_objective_quadratic_flux_compute   |     -2.13 +/- 13.50    | -1.14e-03 +/- 7.22e-03 |  5.24e-02 +/- 4.4e-03  |  5.35e-02 +/- 5.8e-03  |
 test_build_transform_fft_midres         |     +0.13 +/- 3.84     | +1.15e-03 +/- 3.49e-02 |  9.10e-01 +/- 2.1e-02  |  9.09e-01 +/- 2.8e-02  |
 test_build_transform_fft_highres        |     +0.12 +/- 2.64     | +1.42e-03 +/- 3.18e-02 |  1.20e+00 +/- 1.9e-02  |  1.20e+00 +/- 2.5e-02  |
 test_equilibrium_init_medres            |     -0.14 +/- 3.67     | -9.61e-03 +/- 2.60e-01 |  7.08e+00 +/- 1.7e-01  |  7.08e+00 +/- 2.0e-01  |
 test_objective_compile_dshape_current   |     -1.91 +/- 2.22     | -8.29e-02 +/- 9.62e-02 |  4.25e+00 +/- 7.7e-02  |  4.33e+00 +/- 5.8e-02  |
 test_objective_compute_dshape_current   |     +2.96 +/- 15.04    | +2.02e-05 +/- 1.03e-04 |  7.02e-04 +/- 9.7e-05  |  6.82e-04 +/- 3.2e-05  |
 test_objective_jac_dshape_current       |     +0.77 +/- 24.76    | +1.86e-04 +/- 5.93e-03 |  2.41e-02 +/- 4.9e-03  |  2.40e-02 +/- 3.3e-03  |
 test_perturb_2                          |     +0.69 +/- 1.12     | +1.10e-01 +/- 1.79e-01 |  1.61e+01 +/- 1.4e-01  |  1.59e+01 +/- 1.1e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.98 +/- 1.97     | +1.18e-01 +/- 2.38e-01 |  1.22e+01 +/- 1.9e-01  |  1.21e+01 +/- 1.5e-01  |
 test_proximal_freeb_jac                 |     +2.01 +/- 5.36     | +9.62e-02 +/- 2.57e-01 |  4.89e+00 +/- 2.0e-01  |  4.80e+00 +/- 1.6e-01  |
 test_solve_fixed_iter_compiled          |     -0.52 +/- 3.84     | -3.11e-02 +/- 2.32e-01 |  6.01e+00 +/- 1.4e-01  |  6.04e+00 +/- 1.8e-01  |
 test_objective_compute_ripple_bounce1d  |     -1.14 +/- 3.42     | -3.46e-03 +/- 1.04e-02 |  3.00e-01 +/- 3.8e-03  |  3.03e-01 +/- 9.7e-03  |
 test_objective_grad_ripple_bounce1d     |     +2.37 +/- 2.72     | +2.22e-02 +/- 2.55e-02 |  9.61e-01 +/- 2.2e-02  |  9.39e-01 +/- 1.3e-02  |
 test_objective_quadratic_flux_jac       |     +0.36 +/- 0.47     | +3.05e-02 +/- 3.96e-02 |  8.40e+00 +/- 2.5e-02  |  8.37e+00 +/- 3.1e-02  |

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

YigitElma · 2026-04-18T16:07:40Z

The logic should take into account that a given data may make certain sub dependencies obselete

…from before

… always have these keys in, some clean up

YigitElma · 2026-04-19T05:05:49Z

        Type of object to compute for, eg Equilibrium, Curve, etc.
    names : str or array-like of str
        Name(s) of the quantity(s) to compute.
-    deps : set[str]


Previously this function was recursive, so this deps was needed but not anymore

YigitElma · 2026-04-19T05:07:24Z

-
-    def _get_derivs_1_key(key):
-        if has_axis:
-            if "full_with_axis_dependencies" in data_index[p][key]:


We used to call this function in _build_data_index, but now data_index is already built when we call this function

github-actions · 2026-04-19T05:25:10Z

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |   -2.00 %    |     4.045e+03      |     3.964e+03      |    -81.07    |       33.83        |       24.22        |
  test_proximal_jac_w7x_with_eq_update   |   -1.13 %    |     6.612e+03      |     6.537e+03      |    -74.67    |       126.32       |       123.80       |
  test_proximal_freeb_jac                |    0.60 %    |     1.332e+04      |     1.340e+04      |    80.38     |       74.82        |       72.22        |
  test_proximal_freeb_jac_blocked        |    0.79 %    |     7.651e+03      |     7.711e+03      |    60.36     |       69.65        |       67.28        |
  test_proximal_freeb_jac_batched        |    1.85 %    |     7.619e+03      |     7.760e+03      |    140.88    |       67.89        |       65.19        |
  test_proximal_jac_ripple               |   -1.26 %    |     3.588e+03      |     3.543e+03      |    -45.04    |       44.45        |       41.54        |
  test_proximal_jac_ripple_bounce1d      |    1.04 %    |     3.664e+03      |     3.702e+03      |    38.16     |       55.24        |       49.39        |
  test_eq_solve                          |    5.41 %    |     2.004e+03      |     2.113e+03      |    108.36    |       44.45        |       42.15        |
  test_objective_quadratic_flux_jac      |   -0.09 %    |     2.503e+03      |     2.501e+03      |    -2.29     |       42.60        |       42.48        |

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

codecov · 2026-04-30T04:58:37Z

Codecov Report

❌ Patch coverage is 98.11321% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 94.35%. Comparing base (979839e) to head (5286a45).

Files with missing lines	Patch %	Lines
desc/compute/utils.py	98.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1557      +/-   ##
==========================================
+ Coverage   94.30%   94.35%   +0.05%     
==========================================
  Files         101      101              
  Lines       28934    28897      -37     
==========================================
- Hits        27286    27267      -19     
+ Misses       1648     1630      -18

Files with missing lines	Coverage Δ
desc/compute/__init__.py	`100.00% <100.00%> (ø)`
desc/compute/data_index.py	`95.94% <100.00%> (+0.05%)`	⬆️
desc/compute/utils.py	`96.59% <98.00%> (+7.32%)`	⬆️

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

YigitElma · 2026-04-30T05:40:06Z

I am not sure if this PR is necessary, given that there is no performance difference. I guess it is cleaner than before, anyway it is ready for review.

![PR Age](https://img.shields.io/github/issues/detail/age/PlasmaControl/DESC/2157?label=Ready%20for%20merge%20since) - Adds FFT grids and field aligned grids to `test_compute_everything`. Should be merged before stuff that changes grids #2053 or stuff that changes how the dependencies are ordered/selected for computation #1557 to ensure they do not cause bugs in the computations. - Deletes a bunch of files that are now not necessary and removes bounce objective docstring boilerplate - uses `nondiff_argnums` instead of `nondiff_argnames` so that the code works on old jax versions - Adds comments to explain per bounce_comments PR - Resolves #2162 . --------- Co-authored-by: Rory Conlin <mail@fouriest.net> Co-authored-by: Dario Panici <37969854+dpanici@users.noreply.github.com>

ddudt

need to fix the failing test, otherwise this looks good

YigitElma · 2026-05-14T22:00:50Z

need to fix the failing test

Btw, I cannot recreate the failing test locally.

YigitElma · 2026-05-28T17:27:29Z

Try fixing the iota profile instead of the bounds
Maybe also remove the bounds for the force balance constraint too

YigitElma · 2026-05-29T13:29:25Z

@f0uriest let me know if you approve, since you opened the PR, I can approve for you

YigitElma · 2026-07-01T09:38:44Z

 @pytest.mark.optimize
 def test_constrained_AL_lsq():
    """Tests that the least squares augmented Lagrangian optimizer does something."""
+    # Note: This test is known to be sensitive to small numerical differences


Related comment about this test
#2247 (comment)

YigitElma · 2026-07-13T23:17:10Z

@dpanici can you check this once you have some time?

Switch compute from recursive to looped

12a8e53

dpanici added 2 commits March 26, 2025 15:27

Merge branch 'master' into rc/compute_tiers

0732219

Merge branch 'master' into rc/compute_tiers

75a72cf

dpanici added the run_benchmarks Run timing benchmarks on this PR against current master branch label Apr 9, 2026

unalmis mentioned this pull request Apr 10, 2026

Add fft grid and raz grid to test against master #2157

Merged

YigitElma reviewed Apr 11, 2026

View reviewed changes

Comment thread desc/compute/utils.py Outdated

YigitElma mentioned this pull request Apr 16, 2026

Improve import time #2156

Merged

YigitElma and others added 4 commits April 17, 2026 15:22

Merge branch 'master' into rc/compute_tiers

e0fff1d

use topological order computed in build_data_index instead of tiers

6a66f85

fix _compute, remove set_tiers

a75ec0c

move topo_order to data_index

e5916d2

YigitElma and others added 5 commits April 19, 2026 00:01

remove recursion in _get_deps, remove redundant deps, fix data issue …

ab47753

…from before

remove _get_deps_1_key and other redundant functions, data_index will…

1a959bc

… always have these keys in, some clean up

minor comments

abff730

skip extra get_data_deps calls, remove redundant functions

377a30e

Merge branch 'master' into rc/compute_tiers

662b157

YigitElma reviewed Apr 19, 2026

View reviewed changes

Comment thread desc/compute/__init__.py

YigitElma added 2 commits April 19, 2026 12:18

Merge branch 'master' into rc/compute_tiers

293bcba

Merge branch 'master' into rc/compute_tiers

1a49974

YigitElma marked this pull request as ready for review April 30, 2026 05:40

YigitElma added the skip_changelog No need to update changelog on this PR label Apr 30, 2026

YigitElma requested review from ddudt and dpanici April 30, 2026 05:41

Merge branch 'master' into rc/compute_tiers

54efa1b

Merge branch 'master' into rc/compute_tiers

5b4ee74

ddudt reviewed May 7, 2026

View reviewed changes

Comment thread desc/compute/__init__.py

try more interesting objective

be2728b

YigitElma and others added 2 commits May 29, 2026 01:35

try to fix the test

2151b83

Merge branch 'master' into rc/compute_tiers

e3473a6

YigitElma requested review from YigitElma, ddudt, dpanici and rahulgaur104 and removed request for dpanici May 29, 2026 13:28

YigitElma and others added 2 commits June 10, 2026 22:54

Merge branch 'master' into rc/compute_tiers

329032f

Merge branch 'master' into rc/compute_tiers

0fb61d9

ddudt approved these changes Jun 25, 2026

View reviewed changes

Merge branch 'master' into rc/compute_tiers

e104579

YigitElma reviewed Jul 1, 2026

View reviewed changes

YigitElma added the easy Short and simple to code or review label Jul 9, 2026

Merge branch 'master' into rc/compute_tiers

3a6e8c7

YigitElma added 3 commits July 17, 2026 19:02

Merge branch 'master' into rc/compute_tiers

40b1b21

Merge branch 'master' into rc/compute_tiers

df67447

Merge branch 'master' into rc/compute_tiers

5286a45

dpanici reviewed Jul 30, 2026

View reviewed changes

Comment thread desc/compute/utils.py

dpanici approved these changes Jul 30, 2026

View reviewed changes

YigitElma merged commit d4ddab0 into master Jul 30, 2026
28 checks passed

YigitElma deleted the rc/compute_tiers branch July 30, 2026 18:19

Uh oh!

Conversation

f0uriest commented Feb 2, 2025 • edited by YigitElma Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions Bot commented Feb 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

YigitElma commented Apr 18, 2026

Uh oh!

YigitElma Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

YigitElma Apr 19, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

github-actions Bot commented Apr 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Memory benchmark result

Uh oh!

codecov Bot commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

YigitElma commented Apr 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ddudt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

YigitElma commented May 14, 2026

Uh oh!

YigitElma commented May 28, 2026

Uh oh!

YigitElma commented May 29, 2026

Uh oh!

YigitElma Jul 1, 2026

Choose a reason for hiding this comment

Uh oh!

YigitElma commented Jul 13, 2026

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

f0uriest commented Feb 2, 2025 •

edited by YigitElma

Loading

github-actions Bot commented Feb 2, 2025 •

edited

Loading

github-actions Bot commented Apr 19, 2026 •

edited

Loading

codecov Bot commented Apr 30, 2026 •

edited

Loading

YigitElma commented Apr 30, 2026 •

edited

Loading