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

CHANGELOG.md

@@ -94,6 +94,8 @@ New Features
 - Parallelized ideal ballooning stability and Newcomb ballooning metrics and [other improvements](https://github.com/PlasmaControl/DESC/pull/1763).
 - Adds ``FourierXYCoil`` to compatible coils for ``CoilSetArclengthVariance`` objective.
 - Separated ``gamma_c`` calculation from ``Gamma_c``. User can also plot ``gamma_c`` using the ``plot_gammac`` function.
+- Adds ability to compute rotational transform by integrating magnetic field object (for example, from coils) by passing `iota=True` and a `desc.geometry.core.Curve` object as `axis=curve`to the `desc.magnetic_fields.field_line_integrate` function.
+
 
 Bug Fixes
 

desc/magnetic_fields/_core.py

@@ -30,7 +30,7 @@
 from desc.compute.utils import get_params, get_transforms
 from desc.derivatives import Derivative
 from desc.equilibrium import EquilibriaFamily, Equilibrium
-from desc.grid import LinearGrid, _Grid
+from desc.grid import Grid, LinearGrid, _Grid
 from desc.integrals import compute_B_plasma
 from desc.io import IOAble
 from desc.optimizable import Optimizable, OptimizableCollection, optimizable_parameter
@@ -2598,6 +2598,8 @@
     bounds_Z=(-np.inf, np.inf),
     chunk_size=None,
     bs_chunk_size=None,
+    iota=False,
+    axis=None,
     options=None,
     return_aux=False,
 ):
@@ -2641,6 +2643,14 @@
     bs_chunk_size: int or None
         Chunk size to use when evaluating Biot-Savart for the magnetic field. If None,
         evaluates all the source grid points at once. Defaults to None.
+    iota: bool
+        Whether or not to also find the rotational transform of each field line. If
+        True, requires axis to be passed in. iota is computed assuming the poloidal
+        angle is increasing counter-clockwise about the axis as the field line
+        moves along increasing toroidal angle phi.
+    axis: Curve
+        The magnetic axis relative to which the rotational transform is computed. Only
+        required if iota=True.
     options : dict, optional
         Additional arguments to pass to the diffrax diffeqsolve.
     return_aux : bool, optional
@@ -2650,15 +2660,18 @@
 
     Returns
     -------
-    r, z : ndarray
+    r, z : ndarray of shape [n_phis, n_initial_points]
         arrays of r and z coordinates of the field line, corresponding to the
         input phis
+    iotas : 1-D ndarray of shape [n_initial_points]
+        if iota=True, the rotational transform computed for each field line traced.
     """
     if options is None:
         options = {}
 
     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")
 
     min_step_size = jnp.where(
         phis[-1] > phis[0], min_step_size, -jnp.abs(min_step_size)
@@ -2693,6 +2706,8 @@
         adjoint=adjoint,
         chunk_size=chunk_size,
         bs_chunk_size=bs_chunk_size,
+        iota=iota,
+        axis=axis,
         options=options,
         return_aux=return_aux,
     )
@@ -2715,6 +2730,8 @@
     options,
     chunk_size=None,
     bs_chunk_size=None,
+    iota=False,
+    axis=None,
     return_aux=False,
 ):
     """JIT/AD friendly field line integrator.
@@ -2744,35 +2761,90 @@
     x0 = jnp.array([r0, phis[0] * jnp.ones_like(r0), z0]).T
     # scale to make toroidal field (bp) positive
     scale = jnp.sign(field.compute_magnetic_field(x0)[0, 1])
+    if iota:
+        data_ax = axis.compute(["R", "Z"], grid=LinearGrid(zeta=phis[0]))
+        dR = x0[:, 0] - data_ax["R"].squeeze()
+        dZ = x0[:, 2] - data_ax["Z"].squeeze()
+        theta0 = np.arctan2(dZ, dR)
+        x0 = np.hstack([x0, theta0[:, None]])
 
     # suppress warnings till its fixed upstream:
     # https://github.com/patrick-kidger/diffrax/issues/445
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", message="unhashable type")
-        out = vmap_chunked(
-            _intfun_wrapper, in_axes=(0,) + 14 * (None,), chunk_size=chunk_size
-        )(
-            x0,
-            field,
-            params,
-            source_grid,
-            phis,
-            scale,
-            solver,
-            max_steps,
-            min_step_size,
-            saveat,
-            stepsize_controller,
-            event,
-            adjoint,
-            bs_chunk_size,
-            options,
-        )
+        if iota is False:
+            out = vmap_chunked(
+                _intfun_wrapper, in_axes=(0,) + 14 * (None,), chunk_size=chunk_size
+            )(
+                x0,
+                field,
+                params,
+                source_grid,
+                phis,
+                scale,
+                solver,
+                max_steps,
+                min_step_size,
+                saveat,
+                stepsize_controller,
+                event,
+                adjoint,
+                bs_chunk_size,
+                options,
+            )
+        else:
+            out = vmap_chunked(
+                _intfun_wrapper_iota, in_axes=(0,) + 15 * (None,), chunk_size=chunk_size
+            )(
+                x0,
+                field,
+                params,
+                source_grid,
+                phis,
+                scale,
+                solver,
+                max_steps,
+                min_step_size,
+                saveat,
+                stepsize_controller,
+                event,
+                adjoint,
+                bs_chunk_size,
+                options,
+                axis,
+            )
 
     x = out.ys
     x = jnp.where(jnp.isinf(x), jnp.nan, x)
     r = x[:, :, 0].squeeze().T.reshape((phis.size, *rshape))
     z = x[:, :, 2].squeeze().T.reshape((phis.size, *rshape))
+    theta = x[:, :, 3].squeeze().T.reshape((phis.size, *rshape))
+    if iota:
+        # compute the iota = 1/N sum(dtheta_n / dphi_n)
+        # using WBA: iota = 1/normalization * sum(w(n/N) dtheta_n / dphi_n)
+        # where w(x) = exp(-[x(1-x)]^-1) and normalization is the
+        # normalization for that weight fxn
+        N = phis.size
+
+        def weight(n, N):
+            x = n / N
+            return jnp.exp(-(x ** (-1)) * (1 - x) ** (-1))
+
+        dts = jnp.diff(theta, axis=0)  # [nphi-1 x nfieldlines]
+        dphis = jnp.diff(phis, axis=0)
+        normalization = np.sum([weight(n, N) for n in range(1, N)])
+        # n+1 here is bc the sum is really from n=0 to N-1 inclusive,
+        #  but n=0 contributes nothing so we ignore it
+        weighted_sum = np.sum(
+            jnp.asarray([weight(n + 1, N) * dts[n] / dphis[n] for n in range(N - 1)]),
+            axis=0,
+        )
+        iota = weighted_sum / normalization
+
+        if return_aux:
+            return r, z, iota, (out.stats, out.result)
+        else:
+            return r, z, iota
 
     if return_aux:
         return r, z, (out.stats, out.result)
@@ -2830,6 +2902,91 @@
     ).squeeze()
 
 
+def _intfun_wrapper_iota(
+    x,
+    field,
+    params,
+    source_grid,
+    phis,
+    scale,
+    solver,
+    max_steps,
+    min_step_size,
+    saveat,
+    stepsize_controller,
+    event,
+    adjoint,
+    bs_chunk_size,
+    options,
+    axis,
+):
+    """Wrapper for field line integration + iota."""
+    return diffeqsolve(
+        terms=ODETerm(_odefun_iota),
+        solver=solver,
+        y0=x,
+        t0=phis[0],
+        t1=phis[-1],
+        saveat=saveat,
+        max_steps=max_steps,
+        dt0=min_step_size,
+        stepsize_controller=stepsize_controller,
+        args=[field, params, scale, source_grid, bs_chunk_size, axis],
+        event=event,
+        adjoint=adjoint,
+        **options,
+    )
+
+
+@jit
+def _odefun_iota(s, rpzt, args):
+    field, params, scale, source_grid, bs_chunk_size, axis = args
+    r = rpzt[0]
+    phi = rpzt[1]
+    z = rpzt[2]
+    trans = get_transforms(
+        ["x", "x_s"],
+        axis,
+        grid=Grid(
+            jnp.array([jnp.zeros_like(phi), jnp.zeros_like(phi), phi]).squeeze(),
+            jitable=True,
+        ),
+        jitable=True,
+    )
+    data_axis = axis.compute(["x", "x_s"], transforms=trans, params=axis.params_dict)
+    br, bp, bz = (
+        scale
+        * field.compute_magnetic_field(
+            rpzt[:-1],
+            params,
+            basis="rpz",
+            source_grid=source_grid,
+            chunk_size=bs_chunk_size,
+        ).squeeze()
+    )
+
+    dR_ds = r * br / bp * jnp.sign(bp)
+    dphi_ds = jnp.sign(bp)
+    dZ_ds = r * bz / bp * jnp.sign(bp)
+    # delta_X is the rel. position w.r.t. the magnetic axis
+    delta_R = r - data_axis["x"][:, 0].squeeze()
+    delta_Z = z - data_axis["x"][:, 2].squeeze()
+    d_delta_R_ds = dR_ds - data_axis["x_s"][:, 0].squeeze()
+    d_delta_Z_ds = dZ_ds - data_axis["x_s"][:, 2].squeeze()
+    # one negative sign bc this is based off taking a deriv of
+    # arctan(delta_Z/delta_R), and the arctan angle definition
+    # is opposite what we define as increasing poloidal angle
+    # (i.e. CCW field line rotation is decreasing arctan angle,
+    # but increasing DESC poloidal angle)
+    # sign(bp) to account for when s and Bphi are opposing
+    dtheta_ds = (
+        -jnp.sign(bp)
+        * (delta_R * d_delta_Z_ds - delta_Z * d_delta_R_ds)
+        / (delta_R**2 + delta_Z**2)
+    )
+    return jnp.array([dR_ds, dphi_ds, dZ_ds, dtheta_ds]).squeeze()
+
+
 class OmnigenousField(Optimizable, IOAble):
     """A magnetic field with perfect omnigenity (but is not necessarily analytic).
 

desc/plotting.py

@@ -2121,6 +2121,7 @@
         Axes being plotted to.
     plot_data : dict
         Dictionary of the data plotted, only returned if ``return_data=True``
+        Will include iota if ``iota=True`` was passed in.
 
     Examples
     --------
@@ -2181,8 +2182,9 @@
     phis = (phi + np.arange(0, ntransit)[:, None] * 2 * np.pi / NFP).flatten()
 
     R0, Z0 = np.atleast_1d(R0, Z0)
+    compute_iota = fli_kwargs.get("iota", False)
 
-    fieldR, fieldZ, (_, result) = field_line_integrate(
+    out = field_line_integrate(
         r0=R0,
         z0=Z0,
         phis=phis,
@@ -2191,6 +2193,15 @@
         return_aux=True,
         **fli_kwargs,
     )
+    fieldR = out[0]
+    fieldZ = out[1]
+
+    if compute_iota:
+        iota = out[2]
+        result = out[3][1]
+    else:
+        result = out[2][1]
+
     # result._value is 0 if integration completed successfully
     # for a field line, and >0 if it failed or hit bounds.
     if any(result._value > 0):
@@ -2214,6 +2225,8 @@
         "R": rs,
         "Z": zs,
     }
+    if compute_iota:
+        data["iota"] = iota
 
     rows = np.floor(np.sqrt(nplanes)).astype(int)
     cols = np.ceil(nplanes / rows).astype(int)

tests/test_examples.py

@@ -33,6 +33,7 @@
     SplineMagneticField,
     ToroidalMagneticField,
     VerticalMagneticField,
+    field_line_integrate,
     solve_regularized_surface_current,
 )
 from desc.objectives import (
@@ -1560,9 +1561,11 @@ def test_regcoil_modular_check_B(regcoil_modular_coils):
     surface_current_field = initial_surface_current_field.copy()
 
     np.testing.assert_array_less(chi_B, 1e-6)
-    coords = eq.compute(["R", "phi", "Z", "B"], grid=LinearGrid(M=20, N=20, NFP=eq.NFP))
-    B = coords["B"]
-    coords = np.vstack([coords["R"], coords["phi"], coords["Z"]]).T
+    data_eq = eq.compute(
+        ["R", "phi", "Z", "B", "iota"], grid=LinearGrid(M=20, N=20, NFP=eq.NFP)
+    )
+    B = data_eq["B"]
+    coords = np.vstack([data_eq["R"], data_eq["phi"], data_eq["Z"]]).T
     B_from_surf = surface_current_field.compute_magnetic_field(
         coords,
         source_grid=LinearGrid(M=60, N=60, NFP=surface_current_field.NFP),
@@ -1571,6 +1574,26 @@ def test_regcoil_modular_check_B(regcoil_modular_coils):
     )
     np.testing.assert_allclose(B, B_from_surf, rtol=5e-2, atol=5e-4)
 
+    # test iota against the equilibrium iota
+    data_eq = eq.compute(
+        ["R", "phi", "Z", "B", "iota"], grid=LinearGrid(rho=0.6, NFP=eq.NFP)
+    )
+    iota_eq = data_eq["iota"][0]
+    r0 = data_eq["R"][0]
+    z0 = data_eq["Z"][0]
+    phi0 = data_eq["phi"][0]
+    phis = np.linspace(phi0, 2 * np.pi / eq.NFP * 25, 2)
+    _, _, iota_field = field_line_integrate(
+        r0,
+        z0,
+        phis,
+        field=surface_current_field,
+        iota=True,
+        axis=eq.axis,
+        chunk_size=20,
+    )
+    np.testing.assert_allclose(iota_field, iota_eq, rtol=1e-2)
+
 
 @pytest.mark.regression
 @pytest.mark.solve