Fix `objective.compile` in Benchmarks #1483

YigitElma · 2024-12-20T06:51:37Z

Some benchmarks were using obj.compile and jac_scaled_error together but obj.compile only compiles jac_scaled and compute_scaled. This cause some benchmarks to have very different min/max values, like,

With the change it is more consistent,

Misc

Adjusts number of rounds for tests
Adds Action Details step to tests
Resolves Export compiled objectives for common equilibrium resolutions? #1402

github-actions · 2024-12-20T07:45:43Z

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -1.75 +/- 10.22    | -9.57e-03 +/- 5.60e-02 |  5.38e-01 +/- 3.2e-02  |  5.48e-01 +/- 4.6e-02  |
 test_equilibrium_init_medres            |     -2.07 +/- 2.92     | -8.91e-02 +/- 1.26e-01 |  4.21e+00 +/- 4.4e-02  |  4.30e+00 +/- 1.2e-01  |
 test_equilibrium_init_highres           |     -0.01 +/- 5.06     | -3.34e-04 +/- 2.79e-01 |  5.52e+00 +/- 1.3e-01  |  5.52e+00 +/- 2.5e-01  |
 test_objective_compile_dshape_current   |     +3.20 +/- 6.97     | +1.37e-01 +/- 2.99e-01 |  4.42e+00 +/- 1.7e-01  |  4.29e+00 +/- 2.5e-01  |
 test_objective_compute_dshape_current   |     -3.05 +/- 6.92     | -1.73e-04 +/- 3.93e-04 |  5.50e-03 +/- 2.1e-04  |  5.67e-03 +/- 3.3e-04  |
 test_objective_jac_dshape_current       |     -2.59 +/- 7.60     | -1.15e-03 +/- 3.38e-03 |  4.34e-02 +/- 2.2e-03  |  4.45e-02 +/- 2.5e-03  |
 test_perturb_2                          |     -0.17 +/- 2.25     | -3.39e-02 +/- 4.52e-01 |  2.00e+01 +/- 2.2e-01  |  2.01e+01 +/- 4.0e-01  |
 test_proximal_freeb_jac                 |     -2.01 +/- 1.11     | -1.52e-01 +/- 8.43e-02 |  7.42e+00 +/- 3.5e-02  |  7.57e+00 +/- 7.7e-02  |
 test_solve_fixed_iter                   |     +0.07 +/- 2.00     | +2.29e-02 +/- 6.43e-01 |  3.22e+01 +/- 3.2e-01  |  3.22e+01 +/- 5.6e-01  |
 test_LinearConstraintProjection_build   |     +1.01 +/- 2.69     | +1.42e-01 +/- 3.78e-01 |  1.42e+01 +/- 3.6e-01  |  1.41e+01 +/- 1.1e-01  |
 test_build_transform_fft_midres         |     +2.07 +/- 7.80     | +1.32e-02 +/- 4.97e-02 |  6.50e-01 +/- 3.8e-02  |  6.37e-01 +/- 3.3e-02  |
 test_build_transform_fft_highres        |     +1.60 +/- 7.03     | +1.59e-02 +/- 7.00e-02 |  1.01e+00 +/- 6.1e-02  |  9.96e-01 +/- 3.5e-02  |
 test_equilibrium_init_lowres            |     +1.00 +/- 8.94     | +4.02e-02 +/- 3.59e-01 |  4.06e+00 +/- 2.5e-01  |  4.02e+00 +/- 2.6e-01  |
 test_objective_compile_atf              |     +1.41 +/- 6.28     | +1.16e-01 +/- 5.16e-01 |  8.33e+00 +/- 4.9e-01  |  8.22e+00 +/- 1.8e-01  |
 test_objective_compute_atf              |     -2.00 +/- 3.56     | -3.17e-04 +/- 5.65e-04 |  1.56e-02 +/- 1.7e-04  |  1.59e-02 +/- 5.4e-04  |
 test_objective_jac_atf                  |     +1.00 +/- 2.24     | +1.92e-02 +/- 4.32e-02 |  1.95e+00 +/- 3.1e-02  |  1.93e+00 +/- 3.0e-02  |
 test_perturb_1                          |     +0.68 +/- 2.62     | +9.87e-02 +/- 3.79e-01 |  1.46e+01 +/- 3.5e-01  |  1.45e+01 +/- 1.6e-01  |
 test_proximal_jac_atf                   |     +0.81 +/- 1.80     | +6.60e-02 +/- 1.47e-01 |  8.24e+00 +/- 1.4e-01  |  8.18e+00 +/- 4.8e-02  |
 test_proximal_freeb_compute             |     +0.96 +/- 1.09     | +1.91e-03 +/- 2.18e-03 |  2.01e-01 +/- 1.9e-03  |  1.99e-01 +/- 1.1e-03  |
 test_solve_fixed_iter_compiled          |     +0.26 +/- 2.73     | +5.27e-02 +/- 5.50e-01 |  2.03e+01 +/- 3.8e-01  |  2.02e+01 +/- 4.0e-01  |

codecov · 2024-12-20T07:46:53Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.64%. Comparing base (f8c66c7) to head (cb34cdd).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1483   +/-   ##
=======================================
  Coverage   95.63%   95.64%           
=======================================
  Files         101      101           
  Lines       25542    25542           
=======================================
+ Hits        24428    24430    +2     
+ Misses       1114     1112    -2

Files with missing lines	Coverage Δ
desc/objectives/objective_funs.py	`94.74% <100.00%> (ø)`

... and 2 files with indirect coverage changes

review-notebook-app · 2024-12-20T07:54:03Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

sinaatalay

This PR looks good to me, although others should review it as I am looking at it with only my Python and GitHub knowledge, without DESC knowledge.

Workflows (benchmark.yaml, notebook_tests.yaml, and regression_tests.yaml works the same way as its previous version, except a new step is added, Action Details, which moves all the debugging-related commands to a separate step. It makes sense.
There is a change in desc.objectives.objective_funcs.ObjectiveFunction.compile method, which uses compute_scaled_error method instead of compute_scaled in lsq and all modes. This hasn't been explained in the PR or commit messages. Maybe @YigitElma should explain it, but I am sure it's okay.
The documentation is updated and seems okay.
Tests haven't been changed algorithmically (except changing rounds in benchmarks) but have been cleaned. It looks better, I don't see any errors.

docs/performance_tips.rst

f0uriest · 2024-12-21T19:50:46Z

tests/benchmarks/benchmark_cpu_small.py

@@ -131,67 +131,53 @@ def build():
        N = 25
        _ = Equilibrium(L=L, M=M, N=N)

-    benchmark.pedantic(build, setup=setup, iterations=1, rounds=50)
+    benchmark.pedantic(build, setup=setup, iterations=1, rounds=10)


reducing the number of rounds will make the statistics more noisy, and may lead to more false positives

I agree. For example, in the last benchmark, this PR shows speed improvement for perturb but it shouldn't. We can decide on the exact number of rounds, but my intention is to balance the time spent on tests. Previously, these equilibrium initialization tests took more time than fixed_iter_solve and perturb tests. Given that benchmark workflow started to take around 50mins, I wanted to reduce them from 50, which is a bit overkill.

f0uriest · 2024-12-21T19:52:48Z

tests/benchmarks/benchmark_cpu_small.py



 @pytest.mark.slow
 @pytest.mark.benchmark
 def test_proximal_freeb_compute(benchmark):
    """Benchmark computing free boundary objective with proximal constraint."""
-    jax.clear_caches()


why remove this?

This doesn't cause too much difference,(I can re add it) but technically only run is benchmarked, so clearing the cache here doesn't have much purpose.

f0uriest · 2024-12-21T19:53:53Z

tests/benchmarks/benchmark_cpu_small.py


-    def setup():
+    def run():


moving everything to the run function is changing what its actually profiling, this now includes a bunch of other stuff besides building the linear constraints, is that what we want?

The extra stuff is just building individual constraints and objectives, right? Previously, this was effectively just benchmarking factorize_linear_constraints. Do we usually pass built constraints to the LinearConstraintProjection? Then I can revert it.

I think we usually pass built constraints to it, yes

DESC/desc/optimize/optimizer.py

Line 218 in f8c66c7

if linear_constraint is not None and not linear_constraint.built:

tests/benchmarks/benchmark_cpu_small.py

dpanici · 2025-01-03T21:38:32Z

docs/performance_tips.rst

+
+Caching the Compiled (Jitted) Code
+----------------------------------
+Although the compiled code is fast, it still takes time to compile. If you are running the same optimization, or some similar optimization, multiple times, you can save time by caching the compiled code. This automatically happens for a single session (for example, until you restart your kernel in Jupyter Notebook) but once you start using another session, the code will need to be recompiled. Fortunately, there is a way to bypass this. First create a cache directory (i.e. ``jax-caches``), and put the following code at the beginning of your script:


you have this locally in your backend.py right? so you would not usually need this in every script you run if it is there?

Yeah, if you put it to backend.py, you shouldn't need it every time. But it can be messy for users, so I didn't mention that.

fix obj.compile and adjust benchmarks

e5fb89a

YigitElma added easy Short and simple to code or review skip_changelog No need to update changelog on this PR labels Dec 20, 2024

YigitElma added 2 commits December 20, 2024 02:06

fix typo

9e76b3f

add cache code option to docs

3cf2283

YigitElma added 2 commits December 20, 2024 02:52

fix some issues

2b87c78

fix nae notebook issue

a6179ac

YigitElma added 4 commits December 20, 2024 03:14

fix the url link

2ab9f0d

fix docs

8ff1c59

fix docs

8737f7f

still trying to fix docs

601da8a

YigitElma requested review from a team, rahulgaur104, f0uriest, ddudt, dpanici, kianorr, sinaatalay and unalmis and removed request for a team December 20, 2024 20:34

sinaatalay previously approved these changes Dec 20, 2024

View reviewed changes

f0uriest reviewed Dec 21, 2024

View reviewed changes

docs/performance_tips.rst Outdated Show resolved Hide resolved

f0uriest reviewed Dec 21, 2024

View reviewed changes

tests/benchmarks/benchmark_cpu_small.py Show resolved Hide resolved

address Rory's comments

a4ab21e

YigitElma dismissed sinaatalay’s stale review via a4ab21e December 21, 2024 20:37

YigitElma requested review from f0uriest and sinaatalay December 22, 2024 18:44

YigitElma added 3 commits December 29, 2024 13:16

Merge branch 'master' into yge/benchmarks

33df48c

Merge branch 'master' into yge/benchmarks

50e141b

Merge branch 'master' into yge/benchmarks

65d0b8a

dpanici reviewed Jan 3, 2025

View reviewed changes

Merge branch 'master' into yge/benchmarks

cb34cdd

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix `objective.compile` in Benchmarks #1483

Fix `objective.compile` in Benchmarks #1483

YigitElma commented Dec 20, 2024 •

edited

Loading

github-actions bot commented Dec 20, 2024 •

edited

Loading

codecov bot commented Dec 20, 2024 •

edited

Loading

review-notebook-app bot commented Dec 20, 2024

sinaatalay left a comment •

edited

Loading

f0uriest Dec 21, 2024

YigitElma Dec 21, 2024

f0uriest Dec 21, 2024

YigitElma Dec 21, 2024

f0uriest Dec 21, 2024

YigitElma Dec 21, 2024

dpanici Jan 3, 2025

dpanici Jan 3, 2025

dpanici Jan 3, 2025

YigitElma Jan 3, 2025

Fix objective.compile in Benchmarks #1483

Are you sure you want to change the base?

Fix objective.compile in Benchmarks #1483

Conversation

YigitElma commented Dec 20, 2024 • edited Loading

Misc

github-actions bot commented Dec 20, 2024 • edited Loading

codecov bot commented Dec 20, 2024 • edited Loading

Codecov Report

review-notebook-app bot commented Dec 20, 2024

sinaatalay left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Fix `objective.compile` in Benchmarks #1483

Fix `objective.compile` in Benchmarks #1483

YigitElma commented Dec 20, 2024 •

edited

Loading

github-actions bot commented Dec 20, 2024 •

edited

Loading

codecov bot commented Dec 20, 2024 •

edited

Loading

sinaatalay left a comment •

edited

Loading