Bounce averaging

[unalmis]

trafficstars

This PR is completed. An alternative PR with more intelligent math (see #1045 ) will likely replace it. We should keep/merge both for comparison though.

• [x] differentiable algorithm to compute bounce points.
• [x] Fixed some bugs with numpy compatibility.
• [x] differentiable algorithm to compute bounce integrals.
• [x] works with any numerical quadrature
• [x] Matched analytic theory.

After this pull request:

• Should be easy to implement #519 (basically computing a weighted of bounce integrals).

[github-actions[bot]]

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

[codecov[bot]]

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 3398 commits behind head on master.

Files with missing lines	Patch %	Lines
desc/integrals/bounce_utils.py	94.41%	10 Missing :warning:
desc/utils.py	76.47%	4 Missing :warning:
desc/equilibrium/equilibrium.py	0.00%	1 Missing :warning:
desc/integrals/basis.py	96.77%	1 Missing :warning:

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%)	:arrow_down:
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%)	:arrow_up:
desc/grid.py	93.07% <100.00%> (+0.17%)	:arrow_up:
desc/integrals/__init__.py	100.00% <100.00%> (ø)
... and 9 more

... and 1 file with indirect coverage changes

[f0uriest]

See also #719 for some other related ideas.

[dpanici]

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

[unalmis]

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

[rahulgaur104]

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]

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 https://github.com/PlasmaControl/DESC/pull/854#discussion_r1502008862

[rahulgaur104]

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.

Done! The test is available in tests/test_equilibrium.py:: test_shifted_circle_geometry

[rahulgaur104]

The test I added probably broke a few things but we can fix them later this week.

[unalmis]

I've added more tests. I would be surprised if there are any bugs in the bounce integrals routine. The bounce points are tested to be computed correctly, after which it is just the quadrature itself. The remaining task is to make sure the test test_bounce_averaged_drifts is working. Would like if I could review this with someone; the test fails before bounce integrals are performed, so it's probably some issue with grids and coordinates.

[rahulgaur104]

Working on it now....

[unalmis]

My jax version allows for excluding keyword arguments in jnp.vectorize. Apparently some don't, so it must be a new feature, and we’ll either need to update the minimum jax version, or remove the excluded arguments with some edits. It's just a few lines so I'll just comment what to change in the PR if we choose the latter option. (My jax version is 0.4.25 and we can see that string arguments are allowed to be excluded as shown in the first line of code here: https://jax.readthedocs.io/en/latest/_modules/jax/_src/numpy/vectorize.html#vectorize).

[rahulgaur104]

There is accidental modding happening somewhere. Here's the result of modding alpha in grad alpha

The curve should be smooth and not jump like that towards the end.

[rahulgaur104]

Since there will only ever be two types of bounce integrals, I think we should have two modes for bi: one corresponding to 1/sqrt(1-lambda B) and the other corresponding to sqrt(1-lambda B) because when I give bi expressions like geometry_terms x 1/sqrt(1-lambda B), I get a runtime warning due to the points where lambda = 1/B even though, in the quadrature, we are never exactly at those points.

With these two modes, we just have to call bi(integrand without sqrt(1-lambda B), mode = 1 or 2) and the sqrt term will be multipled in the bounce average routine.

[rahulgaur104]

I have added fixes to the bounce_average_test. Now, the only thing that we need to fix are:

• The discontinuity of the grid (and the coefficients)
• The sqrt(1-lambda B) terms in the bounce integral operator should be added during the bounce average quadrature. The user just supplies the geometry and maybe the pitch angles.

[unalmis]

Just copying and pasting from a few different comments, hopefully context is still clear:

Since there will only ever be two types of bounce integrals, I think we should have two modes for bi: one corresponding to 1/sqrt(1-lambda B) and the other corresponding to sqrt(1-lambda B) because when I give bi expressions like geometry_terms x 1/sqrt(1-lambda B), I get a runtime warning due to the points where lambda = 1/B even though, in the quadrature, we are never exactly at those points.

With these two modes, we just have to call bi(integrand without sqrt(1-lambda B), mode = 1 or 2) and the sqrt term will be multipled in the bounce average routine.

The sqrt(1-lambda B) terms in the bounce integral operator should be added during the bounce average quadrature. The user just supplies the geometry and maybe the pitch angles.

So when we pass in the function f=$f$ into the bounce integral function, we currently do a bounce integral of the integrand $f / \sqrt{1 - \lambda B}$.

The current API integrates functions which are defined (including singular ones) over the grid. For such functions, which may be singular, you should just evaluate $f$ on a grid with nodes where it is finite and pass that in as input. You can specify the method used to spline $f$ onto the quadrature points. Of course, we can't evaluate $f$ where it's not defined and recover what it evaluates to at the quadrature points.

Yes, before bounce_averaging, for a given lambda there will be bmag values where lambda*bmag > 1(where sqrt(1-lambda B) = nan) which is why I am saying that maybe the sqrt(1-lambda B) and 1/sqrt(1-lambda B) should be multiplied during the quadrature when you have already accounted for the bounce points and quadrature etc.

• f(theta) * sqrt(1-lambda B)
• f(theta) * 1/sqrt(1-lambda B)
• f(theta)

I see that you want to integrate a set of locally defined functions, e.g. $[ f= g \sqrt{1 - \lambda B} \mid \lambda ]$ each of which are only defined over the particular well where $\lambda B < 1$. The proper way to do this is to decompose $f$ into functions which are defined over the grid, interpolate those well-defined values onto the quadrature points, then do the quadrature.

The cleanest, and also general way to do this is to allow the input to the bounce integral function f an unpackable argument *f and add an optional callable argument F. Then *f represents a set of functions which will be interpolated from the grid to the quadrature points. F will be a callable that defines the composition of those functions. We can simply interpolate *f and then evaluate the desired function via calling F(*f). For example, suppose $F = B / g$ where $g$ and $B$ are well-defined over the grid but their composition is only locally defined. Then the user creates the function defining the composition

def F(*f):
    g, B = f
    return B / g

which can be passed into bounce integral function as

bi = bounce_integral(...)
result = bi(g, B, F=F)

This will also work for the special case that $g$ is a function of the pitch angles, so long as F encodes that dependence. e.g. baking in the pitch angles to the function via

def F(*f):
    g, B = f
    return B / g * pitch

where the user can obtain nice pitch angles via the function call pitch_of_extrema. This will extend the current functionality to the quadrature of any locally defined functions, rather than those which are defined over the entire grid. It will also make it easier to merge with existing desc.compute api stuff.

Separately it was mentioned,

I am only checking after bi. After bi everything is nan.

In general there will be some nan. This is because it is not possible to write a differntiable algorithm to encode into F the ability to adaptively choose the pitch angle based on which well the inputs were interpolated onto. Or even one that works with numpy since this required jagged arrays.

The issue with the test right now is that

• the stuff before bounce integral fails (which i don't know why - something with coordinate conversions and modding I agree).
• The input to the bounce integral function raises an error because numpy doesn't take sqrt of negatives. If this error was supressed and somehow fed the argument into the bounce integral function and saw nan output, then the code then splined a function which was nan at a knot so continuity implies it is nan everywhere, which is correctly reproduced.

[rahulgaur104]

When I run pytest -k test_bounce_averaged_drifts, I get the following error

tests/test_bounce_integral.py:43: in <module>                                                     
    @np.vectorize(signature="(m)->()")                                                          
E   TypeError: vectorize.__init__() missing 1 required positional argument: 'pyfunc' 

Is the signature correct?

When you commend these lines you can run the bounce integral test until it throws an error at bi.

[unalmis]

When I run pytest -k test_bounce_averaged_drifts, I get the following error

tests/test_bounce_integral.py:43: in <module>                                                     
    @np.vectorize(signature="(m)->()")                                                          
E   TypeError: vectorize.__init__() missing 1 required positional argument: 'pyfunc' 

However, when I copy the test into a separate, external python file, it don't get the same error

Another versioning thing. I'll update it to be more compatible with @partial(np.vectorize, signature="(m)->()")

[unalmis]

There is accidental modding happening somewhere. Here's the result of modding alpha in grad alpha

The curve should be smooth and not jump like that towards the end.

My hunch is the issue is in https://github.com/PlasmaControl/DESC/pull/854/files#r1545524478. For context this was a method I added before Rory's recent PR that lets you transform from field line coords to DESC coords with a good initial guess.

There I use compute_theta_coords to go from straight field line coords to DESC coords. However to obtain the straight field line coords in the first place, I transform from Clebsh field line coords to straight field line coords via using the definition of iota and alpha, but I do nothing to handle the periodicity. So when I call compute_theta_coords I'm feeding it in non-periodic zeta, whereas I think compute_theta_coords assumes the input zeta is the same as the periodic DESC zeta. Thoughts?

[rahulgaur104]

I checked the zeta that goes into compute_theta_coords and it is continuous. However, the theta_PEST theta comes out of compute_theta_coords is discontinuous. However, that shouldn't be a problem because the same thing happens with test_shifted_circle_geometry as well.

[rahulgaur104]

override_grid=False fixes the discontinuity.

Only thing left to do is case wise quadrature for sqrt(1-lambda B) and 1/sqrt(1-lambda B) singularities.

I think the best way to implement this to add the "weight" would be multiply the geometry terms with (sqrt(1-lambda B))^n where n is an integer specified by the user during the quadrature. This will account for all the different weights.

[rahulgaur104]

The test fails on my copy of desc for a different reason

 ImportError while importing test module '/home/rgaur/DESC/tests/test_grid.py'.
E   ImportError: cannot import name 'meshgrid_expand' from 'desc.grid' (/home/rgaur/DESC/desc/grid.py)

I am assuming this is due to version incompatibility or something about meshgrid_expand in desc.grid.

[unalmis]

The test fails on my copy of desc for a different reason

 ImportError while importing test module '/home/rgaur/DESC/tests/test_grid.py'.
E   ImportError: cannot import name 'meshgrid_expand' from 'desc.grid' (/home/rgaur/DESC/desc/grid.py)

I am assuming this is due to version incompatibility or something about meshgrid_expand in desc.grid.

I think I deleted that function from the branch. Oh I forgot to remove it from the test.

[unalmis]

Since I have fixed the floating point issues wedneday, I am confident there are no longer bugs in bounce_integral.py and the other tests corroborate this.

For the bounce average test I've fixed some issues with it, and isolated what I believe is the last issue. The mismatch between analytic and numerical is now reduced to a sign error which is fine since the sign of the bounce integral is arbitrary (the bounce average has a sign of course)

Now here is the remaining issue: When we correctly use psi_boundary at the lcfs boundary for normalizing B_ref we reproduce that analytic result in agreement with the paper. When we incorrectly use psi for normalization we get that numerical matches analytic up to a sign error. So I think final issue is just an inconsistent use of psi and psi boundary or in computing the drifts in the test or normalization elsewhere in the DESC compute funs for these quantities.

These plots are generated automatically by running pytest -k test_bounce_averaged_drifts tests/test_bounce_integral.py

[rahulgaur104]

|B| does not look like it's normalized because if it were 1/lambda should go from <1 to >1. I don't know how the normalization can affect the output. I'll probably revert back the normalization.

[unalmis]

|B| does not look like it's normalized because if it were 1/lambda should go from <1 to >1. I don't know how the normalization can affect the output. I'll probably revert back the normalization.

If it were normalized by it's actual maximum value, e.g. if we defined B_ref = np.max(data["|B|"]) then 1/lambda would go from min(|B|) to 1. Currently the approximation B_ref = 2 * np.abs(psi_boundary) / L_ref**2 is used in the test, so it is only approximately normalized to 1. This is why I referred to using psi instead of psi boundary in B_ref as the incorrect normalization, even though in theory the normalization shouldn't matter.

Technically, normalization could have an affect on the chosen coefficients of the spline coefficients and hence affect the numerical results, but this should be a high order effect. We can ignore that for now since the normalization also affects the analytic result. Isn't the issue that the analytic formula for dp_drho is using unnormalized |B|?

[rahulgaur104]

Ok, looking at the test now...

So, my objective is to revert back to the old normalization and keep splitting the problem into smaller and smaller bit until I find the issue. I'll report back my observations in a few hours.

Clearly, the integrals work without any geometric terms in front and the analytical and numerical geometric terms match and don't affect the numerical result. This mean that there is something going on with the analytical solution.

Also, the splined |B| should only be used for finding the nodes and determining their values. I am assuming the splined |B| and |B| from the equilibrium match wall at the non-quadrature points.

Btw, did you check the incomplete integrals against a correct source, like mathematica?

[unalmis]

Btw, did you check the incomplete integrals against a correct source, like mathematica?

If you mean what's computed in test_integral_0 and test_integral_1, yes. They are tested against scipy's adaptive quadrature, and since scipy's eliptic integral functions failed, I felt obligated to check it's adaptive quadrature against mathematica. They agree.

Also, the splined |B| should only be used for finding the nodes and determining their values. I am assuming the splined |B| and |B| from the equilibrium match wall at the non-quadrature points.

|B| is first splined to compute the bounce points which determine the quadrature points. Then |B| is again interpolated from the original |B| to the quadrature points using an equivalent interpolation method.

So, my objective is to revert back to the old normalization and keep splitting the problem into smaller and smaller bit until I find the issue. I'll report back my observations in a few hours.

Thanks! I'll write a method to plot the numerically computed integrand as a function of pitch and zeta for whatever is given to the bounce integrate method. It will help us see if the integrand is actually well-behaved, which could diagnose bugs.

[unalmis]

Another reason the normalization could affect the result is that the compute functions for gbdrift and cvdrift in desc.compute._metric don't use normalized B. They just use the actual B field.

[rahulgaur104]

So, I have made a lot of changes to the analytical terms but I think there may be more mistakes which I will check later.

Note that the bounce integral can have dtheta/b dot grad theta or d zeta / b dot grad zeta or dl (where l is the length along a file line), but you can't mix them.