Ku/bounce alpha

[unalmis]

• Enable tracking multiple fieldlines with Bounce2D
• simplify ability to analyze drift over alpha as requested by a few people.
• other changes requested on previous pull request.

[codecov]

Codecov Report

Attention: Patch coverage is 94.97487% with 10 lines in your changes missing coverage. Please review.

Project coverage is 95.68%. Comparing base (2f057be) to head (9d489df).
Report is 183 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/integrals/bounce_integral.py	92.85%	6 Missing ⚠️
desc/integrals/_bounce_utils.py	92.59%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #1494      +/-   ##
==========================================
+ Coverage   95.62%   95.68%   +0.06%     
==========================================
  Files         101      101              
  Lines       25546    25566      +20     
==========================================
+ Hits        24428    24463      +35     
+ Misses       1118     1103      -15     

Files with missing lines	Coverage Δ
desc/compute/_deprecated.py	100.00% <100.00%> (ø)
desc/compute/_fast_ion.py	100.00% <100.00%> (ø)
desc/compute/_geometry.py	99.07% <ø> (ø)
desc/compute/_neoclassical.py	100.00% <100.00%> (ø)
desc/integrals/_interp_utils.py	100.00% <ø> (ø)
desc/integrals/basis.py	93.12% <100.00%> (+0.07%)	⬆️
desc/integrals/quad_utils.py	100.00% <100.00%> (ø)
desc/objectives/_fast_ion.py	98.27% <100.00%> (+1.72%)	⬆️
desc/objectives/_neoclassical.py	98.24% <100.00%> (+1.75%)	⬆️
desc/integrals/_bounce_utils.py	91.89% <92.59%> (+7.27%)	⬆️
... and 1 more

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -6.33 +/- 4.96     | -3.65e-02 +/- 2.86e-02 |  5.40e-01 +/- 2.2e-02  |  5.77e-01 +/- 1.9e-02  |
 test_equilibrium_init_medres            |     -0.81 +/- 8.62     | -3.49e-02 +/- 3.70e-01 |  4.26e+00 +/- 7.7e-02  |  4.30e+00 +/- 3.6e-01  |
 test_equilibrium_init_highres           |     -0.42 +/- 3.15     | -2.33e-02 +/- 1.77e-01 |  5.60e+00 +/- 5.0e-02  |  5.62e+00 +/- 1.7e-01  |
 test_objective_compile_dshape_current   |     -0.59 +/- 1.37     | -2.36e-02 +/- 5.52e-02 |  4.00e+00 +/- 3.9e-02  |  4.02e+00 +/- 3.9e-02  |
 test_objective_compute_dshape_current   |     -0.96 +/- 2.39     | -4.98e-05 +/- 1.24e-04 |  5.15e-03 +/- 8.2e-05  |  5.20e-03 +/- 9.3e-05  |
 test_objective_jac_dshape_current       |     -3.05 +/- 10.31    | -1.41e-03 +/- 4.77e-03 |  4.48e-02 +/- 3.7e-03  |  4.62e-02 +/- 3.0e-03  |
 test_perturb_2                          |     +1.59 +/- 2.36     | +3.18e-01 +/- 4.73e-01 |  2.04e+01 +/- 3.7e-01  |  2.01e+01 +/- 3.0e-01  |
 test_proximal_freeb_jac                 |     +0.33 +/- 1.04     | +2.45e-02 +/- 7.80e-02 |  7.50e+00 +/- 5.8e-02  |  7.48e+00 +/- 5.2e-02  |
 test_solve_fixed_iter                   |     +1.01 +/- 1.63     | +3.29e-01 +/- 5.33e-01 |  3.30e+01 +/- 2.8e-01  |  3.27e+01 +/- 4.6e-01  |
 test_LinearConstraintProjection_build   |     +2.61 +/- 4.35     | +2.75e-01 +/- 4.58e-01 |  1.08e+01 +/- 3.9e-01  |  1.05e+01 +/- 2.4e-01  |
 test_build_transform_fft_midres         |     -0.77 +/- 2.35     | -4.76e-03 +/- 1.45e-02 |  6.10e-01 +/- 8.9e-03  |  6.15e-01 +/- 1.1e-02  |
 test_build_transform_fft_highres        |     +0.52 +/- 2.52     | +5.12e-03 +/- 2.48e-02 |  9.88e-01 +/- 1.2e-02  |  9.82e-01 +/- 2.2e-02  |
 test_equilibrium_init_lowres            |     -0.28 +/- 2.02     | -1.11e-02 +/- 7.91e-02 |  3.90e+00 +/- 5.7e-02  |  3.91e+00 +/- 5.5e-02  |
 test_objective_compile_atf              |     -1.20 +/- 1.35     | -9.87e-02 +/- 1.12e-01 |  8.15e+00 +/- 7.1e-02  |  8.25e+00 +/- 8.6e-02  |
 test_objective_compute_atf              |     +0.32 +/- 1.75     | +5.06e-05 +/- 2.80e-04 |  1.61e-02 +/- 1.8e-04  |  1.60e-02 +/- 2.1e-04  |
 test_objective_jac_atf                  |     -0.81 +/- 3.09     | -1.60e-02 +/- 6.11e-02 |  1.96e+00 +/- 5.0e-02  |  1.98e+00 +/- 3.5e-02  |
 test_perturb_1                          |     +0.00 +/- 1.65     | +9.03e-05 +/- 2.44e-01 |  1.48e+01 +/- 1.3e-01  |  1.48e+01 +/- 2.1e-01  |
 test_proximal_jac_atf                   |     +0.09 +/- 1.02     | +7.09e-03 +/- 8.48e-02 |  8.30e+00 +/- 3.5e-02  |  8.29e+00 +/- 7.7e-02  |
 test_proximal_freeb_compute             |     +1.32 +/- 0.71     | +2.61e-03 +/- 1.41e-03 |  2.01e-01 +/- 9.0e-04  |  1.99e-01 +/- 1.1e-03  |
 test_solve_fixed_iter_compiled          |     -0.16 +/- 0.55     | -3.21e-02 +/- 1.12e-01 |  2.05e+01 +/- 8.7e-02  |  2.06e+01 +/- 7.1e-02  |

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

dpanici commented on 2025-01-02T20:08:58Z
----------------------------------------------------------------

Line #9.    # Seems like these resolutions are more than sufficient.

the first field line here seems to have some smaller wells ignored

and a rho label for the plots would be very useful, or at least for this PR make this comment specify that you are looking at a single alpha for different rho values as that wasn't clear to me

unalmis commented on 2025-01-03T00:30:25Z
----------------------------------------------------------------

Yes typically there is a finer structure near axis, but the integration domain is also much smaller (max B - min B is small)

[dpanici]

just the one comment on the plotting when grid has multiple rho for check_points, either

• change check_points to have a rho/alpha label somewhere in the title
• or just make a comment in that notebook that they are at different rho, same alpha, and make a TODO for changing check_points

[unalmis]

Yes typically there is a finer structure near axis, but the integration domain is also much smaller (max B - min B is small)

---

View entire conversation on ReviewNB

[review-notebook-app]

View / edit / reply to this conversation on ReviewNB

unalmis commented on 2025-01-03T00:32:31Z
----------------------------------------------------------------

This was technically reverse mode optimization, but I added forward mode to the notebook so that it uses less memory on the CI.

[rahulgaur104]

Here's what I am running to test this

#!/usr/bin/env python3
import numpy as np
from desc.examples import get
from desc.grid import LinearGrid
from matplotlib import pyplot as plt
from desc.integrals import Bounce2D

import pdb

eq = get("W7-X")

rho = 0.5
X, Y = 16, 32
theta = Bounce2D.compute_theta(eq, X, Y, rho)
num_transit = 3
num_well = 10 * num_transit
num_quad = 32

Y_B = 32

N_alpha = int(4)
num_pitch = int(32)

alpha_arr = np.linspace(0, 2*np.pi, N_alpha)
grid = LinearGrid(rho=rho, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=False)

data = eq.compute(
            "Gamma_c",
            grid=grid,
            theta=theta,
            Y_B=Y_B,
            num_transit=num_transit,
            num_well=num_well,
            num_quad=num_quad,
            num_pitch=num_pitch,
            alpha =alpha_arr,
            # Can also specify ``pitch_batch_size`` which determines the
            # number of pitch values to compute simultaneously.
            # Reduce this if insufficient memory. If insufficient memory is detected
            # early then the code will exit and return ε = 0 everywhere. If not detected
            # early then typical OOM errors will occur.
            )

pdb.set_trace()

but the shape of gamma_c is still (1, num_pitch). I was expecting gamma_c (inside the compute function) to have a shape (4, 32).

Edit: My bad, gamma_c is the right shape. Approving!

[rahulgaur104]

Codecov short by 0.65%. Merging now.