DESC icon indicating copy to clipboard operation
DESC copied to clipboard

Gamma_c (Nemov and Velasco et. al) Γ_c

Open unalmis opened this issue 1 year ago • 11 comments

Resolves #522 . Automatically differentiable

unalmis avatar Jun 02 '24 02:06 unalmis

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     -3.30 +/- 4.20     | -1.84e-02 +/- 2.34e-02 |  5.38e-01 +/- 2.0e-02  |  5.57e-01 +/- 1.1e-02  |
 test_equilibrium_init_medres            |     -4.92 +/- 2.24     | -2.18e-01 +/- 9.96e-02 |  4.22e+00 +/- 7.3e-02  |  4.44e+00 +/- 6.8e-02  |
 test_equilibrium_init_highres           |     +2.74 +/- 3.16     | +1.51e-01 +/- 1.74e-01 |  5.66e+00 +/- 1.2e-01  |  5.51e+00 +/- 1.3e-01  |
 test_objective_compile_dshape_current   |     +1.80 +/- 6.13     | +7.10e-02 +/- 2.42e-01 |  4.02e+00 +/- 5.9e-02  |  3.95e+00 +/- 2.3e-01  |
 test_objective_compute_dshape_current   |     +1.06 +/- 2.29     | +5.43e-05 +/- 1.18e-04 |  5.20e-03 +/- 1.0e-04  |  5.15e-03 +/- 6.3e-05  |
 test_objective_jac_dshape_current       |     +2.72 +/- 8.58     | +1.17e-03 +/- 3.68e-03 |  4.41e-02 +/- 3.0e-03  |  4.29e-02 +/- 2.1e-03  |
 test_perturb_2                          |     -1.07 +/- 5.14     | -2.21e-01 +/- 1.06e+00 |  2.04e+01 +/- 7.1e-01  |  2.06e+01 +/- 7.9e-01  |
 test_proximal_freeb_jac                 |     -2.68 +/- 2.73     | -2.04e-01 +/- 2.08e-01 |  7.42e+00 +/- 9.8e-02  |  7.63e+00 +/- 1.8e-01  |
 test_solve_fixed_iter                   |     +1.08 +/- 2.38     | +3.60e-01 +/- 7.91e-01 |  3.36e+01 +/- 5.9e-01  |  3.33e+01 +/- 5.3e-01  |
 test_LinearConstraintProjection_build   |     +5.02 +/- 2.87     | +5.16e-01 +/- 2.95e-01 |  1.08e+01 +/- 1.9e-01  |  1.03e+01 +/- 2.3e-01  |
 test_build_transform_fft_midres         |     -3.38 +/- 3.66     | -2.16e-02 +/- 2.34e-02 |  6.18e-01 +/- 1.6e-02  |  6.40e-01 +/- 1.7e-02  |
 test_build_transform_fft_highres        |     -1.69 +/- 2.63     | -1.70e-02 +/- 2.65e-02 |  9.91e-01 +/- 2.3e-02  |  1.01e+00 +/- 1.3e-02  |
 test_equilibrium_init_lowres            |     -3.63 +/- 4.57     | -1.51e-01 +/- 1.89e-01 |  4.00e+00 +/- 1.7e-01  |  4.15e+00 +/- 8.5e-02  |
 test_objective_compile_atf              |     -1.06 +/- 4.53     | -8.74e-02 +/- 3.74e-01 |  8.17e+00 +/- 2.8e-01  |  8.26e+00 +/- 2.5e-01  |
 test_objective_compute_atf              |     -5.20 +/- 3.34     | -8.76e-04 +/- 5.63e-04 |  1.60e-02 +/- 2.3e-04  |  1.68e-02 +/- 5.1e-04  |
 test_objective_jac_atf                  |     -1.93 +/- 3.29     | -3.81e-02 +/- 6.50e-02 |  1.94e+00 +/- 4.4e-02  |  1.98e+00 +/- 4.8e-02  |
 test_perturb_1                          |     -3.78 +/- 2.76     | -5.84e-01 +/- 4.27e-01 |  1.49e+01 +/- 3.7e-01  |  1.55e+01 +/- 2.2e-01  |
 test_proximal_jac_atf                   |     -0.81 +/- 0.97     | -6.75e-02 +/- 8.06e-02 |  8.25e+00 +/- 6.4e-02  |  8.32e+00 +/- 4.9e-02  |
 test_proximal_freeb_compute             |     +1.40 +/- 0.99     | +2.82e-03 +/- 1.99e-03 |  2.04e-01 +/- 1.2e-03  |  2.01e-01 +/- 1.6e-03  |
 test_solve_fixed_iter_compiled          |     +1.49 +/- 2.54     | +3.26e-01 +/- 5.57e-01 |  2.22e+01 +/- 4.1e-01  |  2.19e+01 +/- 3.8e-01  |

github-actions[bot] avatar Jun 02 '24 03:06 github-actions[bot]

Two places where we can improve GammaC:

  1. When we bounce average the radial and binormal drifts, we can eliminate the denominator integral(tau) since we only care about bounce-avg(normal)/bounce-avg(binormal).
  2. The normal and binormal averages are the radial and binormal derivatives of the quantity J = int v_|| dl/tau, the second adiabatic invariant and it doesn't have a non-integrable singularity so it's more robust.

I'll add the analytical model from Hegna to the tests probably in the bounce branch.

rahulgaur104 avatar Jul 24 '24 19:07 rahulgaur104

Two places where we can improve GammaC:

1. When we bounce average the radial and binormal drifts, we can eliminate the denominator integral(tau) since we only care about bounce-avg(normal)/bounce-avg(binormal).

2. The normal and binormal averages  are the radial and binormal derivatives of the quantity J = int v_|| dl/tau, the second adiabatic invariant and it doesn't have a non-integrable singularity so it's more robust.

I'll add the analytical model from Hegna to the tests probably in the bounce branch.

Documenting for others after discussion

  1. was already done
  2. It's not possible to approximate a locally defined quantity such as J with a finite number of spectral coefficients

These were good ideas though thanks!

unalmis avatar Jul 25 '24 20:07 unalmis

Use the built in utitilities now. Don't use this method. @ddudt Here is that plotting script I mentioned. Useful to check how the amount of knots per toroidal transit you use to reconstruct |B| along a field line effects the number of ripples you can detect.

On my machine, matplotlib will generate an interactive plot that I can zoom in/out of when I run this function from a .py file. I'd recommend that if you're plotting since Jupyter notebook will take longer to render a high resolution plot.

Can run with pytest -k test_plot_fieldline or python filename.py.

@pytest.mark.unit  # remove this line if pytest not installed.
def test_plot_fieldline(
    # Change these to whatever you want
    eq=get("W7-X"),
    rho=np.array([0.9]),
    alpha=np.array([0]),
    zeta_min=-np.pi,
    zeta_max=np.pi,
    knots_per_toroidal_transit=75,
    num_pitch=100,
    num_points_plot=50000,
    **kwargs,
):
    """Utility to plot

    Parameters
    ----------
    eq : Equilibrium
        Equilibrium to compute and plot on.
    rho : ndarray
        Flux surfaces to compute on.
    alpha : ndarray
        Field line labels to compute on for each flux surface.
        For irrational surface, you can just pick a single value and make zeta_max
        large.
    zeta_min, zeta_max : float
        Specifies the width over which to follow a field line.
        Following the field line for one toroidal transit would correspond
        to |zeta_max - zeta_min| = 2pi. For axisymmetric devices, a field line
        average over one toroidal transit is equivalent to a flux surface average.
        For non-axisymmetric devices, a field line average converges to a flux
        surface average in the infinite limit. Both statements assume irrational
        surface.
    knots_per_toroidal_transit : int
        A good reference density is 100 knots per toroidal transit.
        This function is useful to see how low this parameter can be set while
        still maintaining a good function approximation to |B| along the field line.
    num_pitch : int
        Number of pitch values uniformly spaced from 1/max B to 1/min B to detect
        bounce points. This is proportional to mu / E.
    num_points_plot : int
        Number of points to plot. This is just the resolution for matplotlib to plot |B|
        along field line. Purely graphical, does not affect convergence.
    kwargs
        Keyword arguments to pass to plotting function.
        https://github.com/PlasmaControl/DESC/blob/
        2c54f9d2d61ca10cb000b46a3d081687f9f8cbe2/desc/compute/bounce_integral.py#L304
        Some options are to `include_knots=True` to plot vertical lines at points where
        the spline of |B| (the blue plot) was constrained to interpolate the true
        values of |B| as computed by DESC.
        ``alpha_knot=0.3`` - This just controls the transparency of the vertical lines.
        ``alpha_pitch=0.3`` - Controls transparancey of the horizontal lines that denote
        a particles 1/pitch or E/mu.

    """

    # Create a grid in (rho, theta, zeta) from input (rho, alpha, zeta) coordinates.
    zeta = np.linspace(
        start=zeta_min,
        stop=zeta_max,
        num=knots_per_toroidal_transit * int((zeta_max - zeta_min) / (2 * np.pi)),
    )
    grid = eq.get_rtz_grid(
        rho,
        alpha,
        zeta,
        coordinates="raz",
        period=(np.inf, 2 * np.pi, np.inf),
    )
    data = eq.compute(["|B|", "|B|_z|r,a", "min_tz |B|", "max_tz |B|"], grid=grid)

    bounce_integrate, spline = bounce_integral(
        B_sup_z=np.ones_like(data["|B|"]),  # don't need correct values to plot |B|.
        B=data["|B|"],
        B_z_ra=data["|B|_z|r,a"],
        knots=zeta,
    )
    pitch = get_pitch(
        grid.compress(data["min_tz |B|"]), grid.compress(data["max_tz |B|"]), num_pitch
    )
    _, _ = bounce_points(
        pitch,
        **spline,
        check=True,
        plot=True,
        num=num_points_plot,
        include_knots=kwargs.pop("include_knots", False),
        **kwargs,
    )


if __name__ == "__main__":
    test_plot_fieldline()

unalmis avatar Aug 01 '24 00:08 unalmis

@ddudt I have added Nemov's Gamma_c in 921f4d5

unalmis avatar Aug 05 '24 21:08 unalmis

From looking through STELLOPTs version of Gamma_C here I see a few possible issues with their implementation that might lead to differences:

  • Ntransits and Nzetapertransit hard coded to 400 each, and npitch hard coded to 80, so will need to recompile if we want to scan resolution (from a comment it looks like ROSE uses nzetapertransit=290?)
  • They don't seem to do any root finding to get bounce points, or even interpolation. They take the bounce points as the first and last zeta points in the fixed zeta grid where B < 1/lambda. I think this means they probably miss a lot of the contribution to Gamma_C near the tops of wells
  • The integration (over both lambda and zeta) seems to only be a first order Reimann sum so may need super high resolution.

In contrast, the WISTELL implementation here seems a lot better

  • They spline qtys along a field line using cubic splines
  • bounce points are found through rootfinding on these splines
  • bounce integrals are done using adaptive tanh-sinh quadrature over splines
  • integral over lambda also seems to be done with adaptive gauss-kronrod, by splining gamma_c(lambda)
  • seems like they also can calculate eps_eff

Might be worth reaching out to Aaron and Ben to talk about benchmarking vs WISTELL

f0uriest avatar Sep 28 '24 17:09 f0uriest

From looking through STELLOPTs version of Gamma_C here

Thanks! I know code is not always a leisure read. I agree it seems problematic. The integrals are strong functions of the bounce points; so it's necessary to obtain them accurately. A specialized quadrature is also recommended. Prior to writing the DESC bounce integrals, I tested this (because if we could just run along the field line using a reiman sum, less expensive algorithms could be made to compute this). Assuming the bounce points are known with accuracy, a non-adaptive trapezoidal scheme to evaluate these bounce integrals required orders of magnitude more quadrature points. Might become a race with the floating point plateau.

Note that there are two STELLOPT implementations; the second one is better. They mention use of tanh-sinh, but it's been a while since I looked at it, and I don't remember where it used now.

unalmis avatar Sep 28 '24 20:09 unalmis

Yeah I looked at that as well but it seems the same? There are more comments explaining some of the coordinate transformations but the bounce points are still rounded to the nearest grid point (there is some stuff about root finding but it's inside a "if 0 > 1" block), and the following stuff doesn't seem to use it. Also looks like it's still just doing first order integration over both zeta and pitch, there's a lone comment about tanh sinh but I didn't see anything that actually looked like it was being used

f0uriest avatar Sep 28 '24 22:09 f0uriest

the WISTELL implementation

I'll look over this more thoroughly later, but from a quick glance, some points to consider

integral over lambda also seems to be done with adaptive gauss-kronrod, by splining gamma_c(lambda)

I see that a cubic spline is fit to the integrand, i.e. onto $f$ in $\int f(\lambda) d\lambda$, and then an adaptive quadrature is used to integrate this spline. Simpson's rule integrates a cubic spline exactly, so I don't understand that strategy.

Likewise the integrand $f(\lambda)$ is not smooth. Each derivative of $f(\lambda)$ is less smooth than its primitive, and for $\Gamma_c$, we have $f(\lambda) \to \infty$ if it intersects a local maximum of |B| (unless there is some magical cancellation between the diverging integrals among the different ripples), so any usual function approximation on $f$ will underestimate the integral.

bounce integrals are done using adaptive tanh-sinh quadrature over splines

FYI, some reasons I chose not to use an adaptive quadrature for the bounce integrals

  1. A fixed quadrature makes the computation homogenous, and this will improve performance on GPUs/numpy/jax etc.
  2. The goal is computing a double integral (1. over pitch $\lambda$ and 2. some bounce integrals). For $\Gamma_c$, in theory, my hunch is an adaptive quadrature on the bounce integrals will over-fit on the diverging integrals and require more pitch values to capture the true double integral
  3. it's more expensive

unalmis avatar Sep 28 '24 22:09 unalmis

Unrelated, but since we're discussing quadrature, I'm not a fan of regular tanh-sinh because it packs nodes too densely near bounce points. This could make floating point errors in the bounce points punishing. Nemov's effective ripple paper mentions how there's additional phenomena to capture by considering particles trapped in wells with ripples, so intuitively that means when we do these bounce integrals we should try to capture more than just the features near the bounce points (hence reduce clustering of quadrature near bounce points as much as possible once the singularity is resolved).

unalmis avatar Sep 28 '24 22:09 unalmis

In previous commits, I had an kwarg option to perform the integral over $\lambda$ with any adaptive quadrature from quadax (without splining the integrand, so the functions behavior is truly captured). Do we want this option back?

unalmis avatar Sep 28 '24 22:09 unalmis

Fix #1288

gretahibbard avatar Dec 03 '24 00:12 gretahibbard

@gretahibbard Thanks for this. When you get a chance could you please

  • Document in #1412 as well what the issue was.
  • If it's a matter of not storing things in constants, I'd prefer just doing that. Right now it looks like objective.build api was changed. Is that necessary?
  • In the last commit, Gamma_d was added. I'd prefer to keep that in #1229.

unalmis avatar Dec 03 '24 04:12 unalmis

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

Codecov Report

Attention: Patch coverage is 98.70968% with 2 lines in your changes missing coverage. Please review.

Project coverage is 95.60%. Comparing base (9ef0b38) to head (ce98913). Report is 1369 commits behind head on master.

Files with missing lines Patch % Lines
desc/compute/_neoclassical.py 98.11% 1 Missing :warning:
desc/objectives/_neoclassical.py 98.00% 1 Missing :warning:
Additional details and impacted files
@@            Coverage Diff             @@
##           master    #1042      +/-   ##
==========================================
+ Coverage   95.57%   95.60%   +0.02%     
==========================================
  Files          98       98              
  Lines       25156    25300     +144     
==========================================
+ Hits        24044    24187     +143     
- Misses       1112     1113       +1     
Files with missing lines Coverage Δ
desc/compute/_basis_vectors.py 100.00% <100.00%> (ø)
desc/compute/_core.py 100.00% <100.00%> (ø)
desc/compute/_field.py 100.00% <100.00%> (ø)
desc/equilibrium/coords.py 88.25% <ø> (ø)
desc/objectives/__init__.py 100.00% <100.00%> (ø)
desc/compute/_neoclassical.py 99.08% <98.11%> (-0.92%) :arrow_down:
desc/objectives/_neoclassical.py 97.82% <98.00%> (+4.49%) :arrow_up:

... and 3 files with indirect coverage changes

:rocket: New features to boost your workflow:
  • :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

codecov[bot] avatar Dec 09 '24 20:12 codecov[bot]