Add ability to compute iota from field line integration given a magnetic axis

[dpanici]

• Adds iota computation from external fields as an optional output from field_line_integrate

Resolve #491

Math

Adds an extra variable to the state, which is the polar angle $\theta$ of the field line w.r.t. some input magnetic axis $(R_0(s), Z_0(s)$ represented by a DESC Curve object defined as

$$ \begin{equation} \theta = tan^{-1} \left(\frac{Z-Z_0}{R - R_0}\right) =: tan^{-1} \left(\frac{\Delta Z}{\Delta R}\right) \end{equation} $$

We want to essentially keep a running sum of how theta has changed across the field line's trajectory w.r.t., so we choose to follow how FOCUS does it (see #491 ) and find the differential $d\theta/ds$ ($s$ being arclength along field line), then just integrate that along with our field line trajectories (which we need to integrate anyways to compute theta and thus iota).

The differential becomes

$$ \begin{align} d\theta/ds =& \frac{1}{1+(\Delta Z)^2 / (\Delta R)^2} d/ds (\frac{\Delta Z}{\Delta R})\\ &= \frac{(\Delta R)^2}{(\Delta R)^2+(\Delta Z)^2 } \left( d \Delta Z/ds \frac{1}{\Delta R} - \Delta Z d\Delta R / ds \frac{1}{(\Delta R)^2} \right)\\ &= \frac{1}{(\Delta R)^2+(\Delta Z)^2 } \left(\Delta R \frac{d\Delta Z}{ds} - \Delta Z \frac{d\Delta R}{ds}\right) \end{align} $$

Where $d\Delta X/ds = \frac{dX}{ds} - \frac{dX_0}{ds}$ ($X \in {R,Z}$), and $dR/ds$ and $dZ/ds$ are just the RHS of the magnetic field line trajectory ODE that we already integrate. To account for the fact that this expression came from a polar angle definition which increases in the opposite direction of DESC's poloidal angle, we must add a negative sign. Then, to account for when $B_{\phi}<0$, we add $sign(B_{\phi})$ (this is needed I believe because at its most basic, iota is defined as $d\chi/d\psi$, so if the geometry all stays the same (i.e. same dR_ds etc) and so does the poloidal field, but the toroidal flux changes sign, so must iota, my tests for the various scenarios seem to confirm this):

$$ \begin{align} d\theta/ds &= \frac{-sign(B_{\phi})}{(\Delta R)^2+(\Delta Z)^2 }\left(\Delta R \frac{d\Delta Z}{ds} - \Delta Z \frac{d\Delta R}{ds}\right) \end{align} $$

At the end, iota is very simply found as $\theta_f/\phi_f$, the final poloidal and toroidal angles, where no further sums or anything are needed as the integration itself took care of the $\sum_N d\theta / \sum_N d\phi$

[dpanici]

Not completely sure if this description actually describes how we are defining iota. If the field line is pointing opposite the direction of increasing phi, then the toroidal flux is negative, and if the B_poloidal in DESC coords is positive, this results in a negative iota since d\chi/d\psi < 0 (chi increases as psi decreases).

[dpanici]

Had to change this to always be 1D, otherwise the grabbing of the last theta value later will throw an index error when r0 and z0 are passed in as only 0-D arrays

[unalmis]

forget map and just do jnp.atleast_1d(r0, z0, phis)

[f0uriest]

I think this should be a seperate thing, after rshape = r0.shape. That way we still know the original input shape so that we return things in that shape, and adding an extra dimension is only done internally

[f0uriest]

ie the whole block would be

r0, z0, phis = map(jnp.asarray, (r0, z0, phis))
assert r0.shape == z0.shape, "r0 and z0 must have the same shape"
errorif(iota and axis is None, ValueError, "Must pass in an axis if iota=True")
rshape = r0.shape
r0 = jnp.atleast_1d(r0).flatten()
z0 = jnp.atleast_1d(z0).flatten()

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    4.22 %    |     3.901e+03      |     4.065e+03      |    164.62    |       35.25        |       31.13        |
  test_proximal_jac_w7x_with_eq_update   |    0.16 %    |     6.750e+03      |     6.761e+03      |    10.83     |       160.12       |       159.29       |
  test_proximal_freeb_jac                |   -0.19 %    |     1.322e+04      |     1.319e+04      |    -24.47    |       77.82        |       77.31        |
  test_proximal_freeb_jac_blocked        |   -0.32 %    |     7.607e+03      |     7.582e+03      |    -24.62    |       68.79        |       67.68        |
  test_proximal_freeb_jac_batched        |   -0.61 %    |     7.611e+03      |     7.564e+03      |    -46.39    |       69.10        |       67.98        |
  test_proximal_jac_ripple               |   -0.89 %    |     7.664e+03      |     7.596e+03      |    -68.21    |       70.11        |       68.49        |
  test_proximal_jac_ripple_spline        |   -2.37 %    |     3.522e+03      |     3.439e+03      |    -83.48    |       70.49        |       70.19        |
  test_eq_solve                          |    1.45 %    |     2.002e+03      |     2.031e+03      |    28.95     |       125.45       |       124.66       |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.

[github-actions]

|             benchmark_name             |         dt(%)          |         dt(s)          |        t_new(s)        |        t_old(s)        | 
| -------------------------------------- | ---------------------- | ---------------------- | ---------------------- | ---------------------- |
 test_build_transform_fft_lowres         |     +0.50 +/- 4.02     | +2.75e-03 +/- 2.22e-02 |  5.56e-01 +/- 2.0e-02  |  5.53e-01 +/- 9.3e-03  |
 test_equilibrium_init_medres            |     -0.41 +/- 1.27     | -2.00e-02 +/- 6.19e-02 |  4.84e+00 +/- 5.2e-02  |  4.86e+00 +/- 3.4e-02  |
 test_equilibrium_init_highres           |     +0.11 +/- 1.37     | +5.89e-03 +/- 7.49e-02 |  5.46e+00 +/- 5.1e-02  |  5.46e+00 +/- 5.5e-02  |
 test_objective_compile_dshape_current   |     -2.04 +/- 1.74     | -6.94e-02 +/- 5.90e-02 |  3.33e+00 +/- 5.0e-02  |  3.40e+00 +/- 3.2e-02  |
 test_objective_compute_dshape_current   |     -1.81 +/- 5.28     | -1.42e-05 +/- 4.16e-05 |  7.72e-04 +/- 2.9e-05  |  7.86e-04 +/- 3.0e-05  |
 test_objective_jac_dshape_current       |     +0.13 +/- 20.78    | +4.27e-05 +/- 6.75e-03 |  3.25e-02 +/- 4.2e-03  |  3.25e-02 +/- 5.3e-03  |
 test_perturb_2                          |     +0.30 +/- 2.52     | +5.03e-02 +/- 4.21e-01 |  1.68e+01 +/- 2.6e-01  |  1.67e+01 +/- 3.3e-01  |
 test_proximal_jac_atf_with_eq_update    |     +0.39 +/- 1.32     | +5.27e-02 +/- 1.80e-01 |  1.37e+01 +/- 1.4e-01  |  1.37e+01 +/- 1.1e-01  |
 test_proximal_freeb_jac                 |     -0.46 +/- 1.11     | -2.33e-02 +/- 5.56e-02 |  4.99e+00 +/- 4.9e-02  |  5.01e+00 +/- 2.6e-02  |
 test_solve_fixed_iter_compiled          |     +0.29 +/- 2.11     | +4.83e-02 +/- 3.55e-01 |  1.69e+01 +/- 1.6e-01  |  1.68e+01 +/- 3.2e-01  |
 test_LinearConstraintProjection_build   |     +0.15 +/- 3.19     | +1.21e-02 +/- 2.63e-01 |  8.26e+00 +/- 2.2e-01  |  8.25e+00 +/- 1.4e-01  |
 test_objective_compute_ripple_spline    |     +0.58 +/- 2.61     | +1.67e-03 +/- 7.55e-03 |  2.91e-01 +/- 5.1e-03  |  2.89e-01 +/- 5.5e-03  |
 test_objective_grad_ripple_spline       |     +1.28 +/- 2.06     | +1.42e-02 +/- 2.28e-02 |  1.12e+00 +/- 1.3e-02  |  1.11e+00 +/- 1.9e-02  |
 test_build_transform_fft_midres         |     -0.70 +/- 3.46     | -4.98e-03 +/- 2.46e-02 |  7.07e-01 +/- 2.1e-02  |  7.12e-01 +/- 1.2e-02  |
 test_build_transform_fft_highres        |     +0.29 +/- 1.74     | +2.76e-03 +/- 1.65e-02 |  9.49e-01 +/- 1.5e-02  |  9.46e-01 +/- 7.4e-03  |
 test_equilibrium_init_lowres            |     -0.92 +/- 1.47     | -4.46e-02 +/- 7.13e-02 |  4.80e+00 +/- 3.2e-02  |  4.84e+00 +/- 6.4e-02  |
 test_objective_compile_atf              |     -0.88 +/- 3.63     | -5.56e-02 +/- 2.29e-01 |  6.26e+00 +/- 1.7e-01  |  6.31e+00 +/- 1.5e-01  |
 test_objective_compute_atf              |     +3.85 +/- 9.13     | +7.77e-05 +/- 1.84e-04 |  2.10e-03 +/- 1.5e-04  |  2.02e-03 +/- 1.0e-04  |
 test_objective_jac_atf                  |     +0.81 +/- 1.98     | +1.39e-02 +/- 3.38e-02 |  1.72e+00 +/- 1.1e-02  |  1.71e+00 +/- 3.2e-02  |
 test_perturb_1                          |     +2.35 +/- 2.50     | +3.31e-01 +/- 3.53e-01 |  1.44e+01 +/- 2.3e-01  |  1.41e+01 +/- 2.7e-01  |
 test_proximal_jac_atf                   |     +1.14 +/- 1.24     | +6.24e-02 +/- 6.80e-02 |  5.55e+00 +/- 4.5e-02  |  5.49e+00 +/- 5.1e-02  |
 test_proximal_freeb_compute             |     +1.54 +/- 2.67     | +2.49e-03 +/- 4.33e-03 |  1.64e-01 +/- 3.5e-03  |  1.62e-01 +/- 2.6e-03  |
 test_solve_fixed_iter                   |     +1.70 +/- 1.72     | +4.99e-01 +/- 5.05e-01 |  2.98e+01 +/- 3.4e-01  |  2.93e+01 +/- 3.8e-01  |
 test_objective_compute_ripple           |     +0.40 +/- 1.06     | +1.03e-02 +/- 2.76e-02 |  2.61e+00 +/- 2.5e-02  |  2.60e+00 +/- 1.1e-02  |
 test_objective_grad_ripple              |     +0.33 +/- 1.35     | +1.53e-02 +/- 6.27e-02 |  4.68e+00 +/- 4.7e-02  |  4.66e+00 +/- 4.2e-02  |

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[YigitElma]

Tbh, I would rather have this as a separate function like "iota_from_field_lines"

[f0uriest]

per discussion last week this should probably be moved to a global scope to avoid recompiling each time. Same for the original odefun

[f0uriest]

should add a note that we'll need to revisit this after #891

[f0uriest]

warning about unequal NFP?

[f0uriest]

also the new plot of iota doesn't look that great, do we expect better agreement? Do we need to integrate the fieldlines for more transits?

Maybe its just the axis limits make it look worse than it is?

[f0uriest]

OK I think I'm a bit confused about the math here. Suppose we're in PEST coordinates $(\vartheta, \phi)$. Then $\iota = \frac{d\vartheta}{d\phi}$ along a field line. But you're using polar $\theta = \vartheta - \lambda(\theta, \phi)$ for some unknown stream function $\lambda$ (which generally isnt the same as our usual one because our usual theta isn't polar).

You're computing $\frac{d\theta}{d\phi} = \frac{d\vartheta}{d\phi} + \frac{d\lambda}{d\phi} + \frac{d\lambda}{d\vartheta} \frac{d\vartheta}{d\phi} = \iota + \frac{d\lambda}{d\phi} + \iota \frac{d\lambda}{d\vartheta}$

So it seems you're getting some extra stuff on top of iota. I can see that when averaged over a surface, those other bits cancel out, but that may require a loooot more transits than we usually do?

[dpanici]

OK I think I'm a bit confused about the math here. Suppose we're in PEST coordinates ( ϑ , ϕ ) . Then ι = d ϑ d ϕ along a field line. But you're using polar θ = ϑ − λ ( θ , ϕ ) for some unknown stream function λ (which generally isnt the same as our usual one because our usual theta isn't polar).

You're computing d θ d ϕ = d ϑ d ϕ + d λ d ϕ + d λ d ϑ d ϑ d ϕ = ι + d λ d ϕ + ι d λ d ϑ

So it seems you're getting some extra stuff on top of iota. I can see that when averaged over a surface, those other bits cancel out, but that may require a loooot more transits than we usually do?

That's right, it is essentially taking the basic definition of iota. Definitely we need enough transits to accurately compute it, though I've not done any studies of convergence

[unalmis]

OK I think I'm a bit confused about the math here. Suppose we're in PEST coordinates ( ϑ , ϕ ) . Then ι = d ϑ d ϕ along a field line. But you're using polar θ = ϑ − λ ( θ , ϕ ) for some unknown stream function λ (which generally isnt the same as our usual one because our usual theta isn't polar).

You're computing d θ d ϕ = d ϑ d ϕ + d λ d ϕ + d λ d ϑ d ϑ d ϕ = ι + d λ d ϕ + ι d λ d ϑ

So it seems you're getting some extra stuff on top of iota. I can see that when averaged over a surface, those other bits cancel out, but that may require a loooot more transits than we usually do?

That's right, it is essentially taking the basic definition of iota. Definitely we need enough transits to accurately compute it, though I've not done any studies of convergence

Do a convergence study by calling Bounce2D.plot_theta. You will see theta(alpha,zeta). can check how long until mean slope matches what you expect. Edit: ok see discussion in slack

[unalmis]

I couldn’t use this to avoid approximating theta in field line coords by the way because you still need to evaluate this quantity along field lines: d θ d ϕ = d ϑ d ϕ + d λ d ϕ + d λ d ϑ d ϑ d ϕ = ι + d λ d ϕ + ι d λ d ϑ. in your case you know the field lines but not iota

[codecov]

Codecov Report

❌ Patch coverage is 97.56098% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 95.83%. Comparing base (e0dcf9e) to head (7f2e61e).

Files with missing lines	Patch %	Lines
desc/plotting.py	83.33%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #1855   +/-   ##
=======================================
  Coverage   95.83%   95.83%           
=======================================
  Files         100      100           
  Lines       27511    27548   +37     
=======================================
+ Hits        26365    26401   +36     
- Misses       1146     1147    +1     

Files with missing lines	Coverage Δ
desc/magnetic_fields/_core.py	96.57% <100.00%> (+0.12%)	⬆️
desc/plotting.py	96.05% <83.33%> (-0.06%)	⬇️

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dpanici]

perhaps could be improved by WBA along the orbit: https://arxiv.org/pdf/2403.19003

[dpanici]

OK I think I'm a bit confused about the math here. Suppose we're in PEST coordinates ( ϑ , ϕ ) . Then ι = d ϑ d ϕ along a field line. But you're using polar θ = ϑ − λ ( θ , ϕ ) for some unknown stream function λ (which generally isnt the same as our usual one because our usual theta isn't polar).
You're computing d θ d ϕ = d ϑ d ϕ + d λ d ϕ + d λ d ϑ d ϑ d ϕ = ι + d λ d ϕ + ι d λ d ϑ
So it seems you're getting some extra stuff on top of iota. I can see that when averaged over a surface, those other bits cancel out, but that may require a loooot more transits than we usually do?

That's right, it is essentially taking the basic definition of iota. Definitely we need enough transits to accurately compute it, though I've not done any studies of convergence

at least for the tutorial notebook, seems that by ntransit=50 most points are converged but yes other methods may work better than this and converge faster