Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bounce averaging #854

Merged
merged 355 commits into from
Sep 3, 2024
Merged

Bounce averaging #854

merged 355 commits into from
Sep 3, 2024

Conversation

unalmis
Copy link
Collaborator

@unalmis unalmis commented Feb 4, 2024

This PR adds functionality to compute bounce averages in DESC

  • Differentiable algorithm to compute bounce points and integrals.
  • Works with any numerical quadrature
  • Fixed bugs with numpy compatibility.

Related

@unalmis unalmis linked an issue Feb 4, 2024 that may be closed by this pull request
Copy link
Contributor

github-actions bot commented Feb 4, 2024

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -0.22 +/- 6.07     | -1.19e-03 +/- 3.21e-02 |  5.27e-01 +/- 2.7e-02  |  5.28e-01 +/- 1.7e-02  |
 test_equilibrium_init_medres            |     +2.15 +/- 5.65     | +8.92e-02 +/- 2.34e-01 |  4.23e+00 +/- 2.1e-01  |  4.14e+00 +/- 1.0e-01  |
 test_equilibrium_init_highres           |     +4.24 +/- 2.78     | +2.31e-01 +/- 1.52e-01 |  5.69e+00 +/- 1.1e-01  |  5.46e+00 +/- 1.1e-01  |
 test_objective_compile_dshape_current   |     +0.64 +/- 1.83     | +2.43e-02 +/- 6.96e-02 |  3.83e+00 +/- 6.2e-02  |  3.81e+00 +/- 3.2e-02  |
 test_objective_compute_dshape_current   |     +1.17 +/- 1.91     | +4.02e-05 +/- 6.57e-05 |  3.47e-03 +/- 5.5e-05  |  3.43e-03 +/- 3.5e-05  |
 test_objective_jac_dshape_current       |    +10.18 +/- 8.18     | +3.91e-03 +/- 3.14e-03 |  4.23e-02 +/- 2.4e-03  |  3.84e-02 +/- 2.0e-03  |
 test_perturb_2                          |     -1.77 +/- 2.43     | -3.13e-01 +/- 4.30e-01 |  1.74e+01 +/- 2.9e-01  |  1.77e+01 +/- 3.2e-01  |
 test_proximal_freeb_jac                 |     +0.44 +/- 1.20     | +3.30e-02 +/- 8.96e-02 |  7.49e+00 +/- 6.2e-02  |  7.45e+00 +/- 6.5e-02  |
 test_solve_fixed_iter                   |     +0.84 +/- 60.70    | +4.20e-02 +/- 3.03e+00 |  5.04e+00 +/- 2.1e+00  |  5.00e+00 +/- 2.2e+00  |

Copy link

codecov bot commented Feb 4, 2024

Codecov Report

Attention: Patch coverage is 96.49123% with 16 lines in your changes missing coverage. Please review.

Project coverage is 95.42%. Comparing base (60de6fc) to head (917ad1c).
Report is 356 commits behind head on master.

Files with missing lines Patch % Lines
desc/integrals/bounce_utils.py 94.41% 10 Missing ⚠️
desc/utils.py 76.47% 4 Missing ⚠️
desc/equilibrium/equilibrium.py 0.00% 1 Missing ⚠️
desc/integrals/basis.py 96.77% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##           master     #854      +/-   ##
==========================================
+ Coverage   95.35%   95.42%   +0.07%     
==========================================
  Files          90       95       +5     
  Lines       22755    23186     +431     
==========================================
+ Hits        21698    22126     +428     
- Misses       1057     1060       +3     
Files with missing lines Coverage Δ
desc/backend.py 90.12% <100.00%> (-0.13%) ⬇️
desc/compute/_bootstrap.py 100.00% <100.00%> (ø)
desc/compute/_equil.py 100.00% <100.00%> (ø)
desc/compute/_field.py 100.00% <100.00%> (ø)
desc/compute/_metric.py 100.00% <100.00%> (ø)
desc/compute/_profiles.py 99.19% <100.00%> (ø)
desc/compute/_stability.py 100.00% <100.00%> (ø)
desc/equilibrium/coords.py 88.75% <100.00%> (+4.81%) ⬆️
desc/grid.py 93.07% <100.00%> (+0.17%) ⬆️
desc/integrals/__init__.py 100.00% <100.00%> (ø)
... and 9 more

... and 1 file with indirect coverage changes

@f0uriest
Copy link
Member

f0uriest commented Feb 5, 2024

See also #719 for some other related ideas.

Copy link
Collaborator

@rahulgaur104 rahulgaur104 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a few suggestions. I'll add the analytical test soon. It's somewhere in my old laptop.

desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
@dpanici
Copy link
Collaborator

dpanici commented Feb 14, 2024

  • Use splines for both rootfind and integration
  • Keep exact method you have here for comparison

@unalmis unalmis marked this pull request as ready for review February 18, 2024 12:03
tests/test_compute_utils.py Outdated Show resolved Hide resolved
@unalmis unalmis marked this pull request as draft February 18, 2024 21:43
Copy link
Collaborator

@dpanici dpanici left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Spline num/den separately, evaluate integrals with quadratures

@unalmis unalmis marked this pull request as ready for review February 25, 2024 09:03
@unalmis
Copy link
Collaborator Author

unalmis commented Feb 25, 2024

I would like to use a Hermite spline for |B|, but it looks like interpax doesn't support that. Should I add that?

desc/compute/utils.py Outdated Show resolved Hide resolved
desc/backend.py Outdated Show resolved Hide resolved
Copy link
Collaborator

@rahulgaur104 rahulgaur104 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have an idea for a simple tokamak test that we can use to test this PR. I'll add the details of that test and more comments soon.

desc/backend.py Outdated Show resolved Hide resolved
desc/backend.py Show resolved Hide resolved
desc/compute/_field.py Outdated Show resolved Hide resolved
tests/test_compute_utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
@unalmis unalmis mentioned this pull request Feb 26, 2024
@rahulgaur104
Copy link
Collaborator

rahulgaur104 commented Feb 26, 2024

The equations for the elliptic integrals K and E are obtained in the limit of a large-aspect ratio, low-beta tokamak. Therefore, if we create a large-aspect ratio tokamak, our bounce integrals should match the elliptic integrals while calling all the functions that you have defined. So full code coverage + simple fast tests.

Here's what to use

surf = FourierRZToroidalSurface(
R_lmn=[1.0, 0.1],
Z_lmn=[0.0, -0.1],
modes_R=np.array([[0, 0], [1, 0]]),
modes_Z=np.array([[0, 0], [-1, 0]]),
sym=True,
NFP=eq.NFP,
)
pressure = PowerSeriesProfile([1e2, 0, -1e2])
eq = Equilibrium(L=6, M=6, N=6, surface=surf, Psi=1.0, pressure=pressure)

Solve the equilibrium and try to adjust the beta to be around 1%. Then, we can pick a surface, say, rho = 0.5, get the |B| and pick a few values of "pitch" on that surface and calculate the bounce average of an equilibrium quantity, let's take g_zz or g_rr.

This should be close to the elliptic integrals (maybe differs by a constant factor). If you are busy, I can add this test on Tuesday or Wednesday.

@unalmis
Copy link
Collaborator Author

unalmis commented Feb 26, 2024

This should be close to the elliptic integrals (maybe differs by a constant factor). If you are busy, I can add this test on Tuesday or Wednesday.

If you want to, I would appreciate that. I'll be focusing on some coursework until Wednesday, and then I will fix this #854 (comment)

desc/compute/_field.py Outdated Show resolved Hide resolved
desc/compute/_field.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/compute/utils.py Outdated Show resolved Hide resolved
desc/integrals/bounce_integral.py Outdated Show resolved Hide resolved
desc/integrals/bounce_utils.py Outdated Show resolved Hide resolved
tests/test_integrals.py Show resolved Hide resolved
…ent commit

Ensure batch=False bounce integration is done in test to catch future regressions
@unalmis unalmis requested a review from f0uriest August 30, 2024 05:18
@unalmis unalmis requested a review from f0uriest August 30, 2024 18:25
f0uriest
f0uriest previously approved these changes Aug 30, 2024
@rahulgaur104
Copy link
Collaborator

Reviewing now....

desc/backend.py Show resolved Hide resolved
desc/grid.py Outdated Show resolved Hide resolved
desc/integrals/basis.py Show resolved Hide resolved
desc/integrals/bounce_integral.py Outdated Show resolved Hide resolved
desc/grid.py Show resolved Hide resolved
desc/utils.py Show resolved Hide resolved
tests/test_integrals.py Show resolved Hide resolved
tests/test_integrals.py Show resolved Hide resolved
tests/test_integrals.py Show resolved Hide resolved
tests/test_integrals.py Show resolved Hide resolved
After the recent refactoring to the `Bounce1D` class that resulted from
#1214, the API is a little too strict for computations like effective
ripple etc. where we vectorize the computation over over some dimensions
and loop over others to save memory.

This PR changes the expected shape of the pitch angle input to
`Bounce1D` in #854 from `(P, M, L)` to `(M, L, P)`. With this change,
the two leading axes of all inputs to the methods in that class is `(M,
L)`.

These changes are tested and already included in downstream branches. I
am making new PR instead of directly committing to the `bounce` branch
for people who have already reviewed the `bounce` PR.

 This is better because
1. Easier usage for end users. (Previously, you'd have to manually add
trailing axes to pitch angle array).
2. Makes it much simpler to use with JAX's new batched map.
3. Previously we would loop over the pitch angles to save memory.
However, this means some computation is repeated because interpax would
interpolate multiple times. By looping over the field lines instead and
doing the interpolation for all the pitch angles at once, both
`_bounce_quadrature` and `interp_to_argmin` are faster. (I'm seeing 30%
faster speed just from computing effective ripple (no optimization), but
I don't plan to do any benchmarking to see whether that is from recent
changes like #1154 or #1043 , or others).
Copy link
Collaborator

@rahulgaur104 rahulgaur104 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Look good to me.

@unalmis
Copy link
Collaborator Author

unalmis commented Sep 3, 2024

@f0uriest Besides adding comments, the changes from your last review are in #1242 , which was done so that computing effective ripple can be made simpler: https://github.com/PlasmaControl/DESC/pull/1003/files#r1741405180

@f0uriest
Copy link
Member

f0uriest commented Sep 3, 2024

Update description before merging

@unalmis unalmis merged commit d2e9a2c into master Sep 3, 2024
24 checks passed
@unalmis unalmis deleted the bounce branch September 3, 2024 04:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
4 participants