`map_coordinates` and boozer angles

[f0uriest]

It seems like map_coordinates doesn't work with boozer coordinates. The use case is computing things on a uniform grid in boozer coordinates. In theory one could write a boozmn file and then load that and interpolate, but it might be worth it to see if we can get the regular coordinate mapping function to work, though may be tricky due to the weird surface averages and fourier transforms needed to compute the boozer angles.

[f0uriest]

For now, here's some code that should work:

import interpax
import desc
import jax.numpy as jnp
import jax
eq = desc.examples.get("precise_QA")
rho=0.5

grid = desc.grid.LinearGrid(rho=rho, M=20, N=20, NFP=eq.NFP)
# we compute the values of theta_B and zeta_B on a uniform grid in desc theta,zeta coordinates
# note that theta_B and zeta_B will NOT be uniformly spaced
data = eq.compute(["theta_B", "zeta_B"], grid=grid)

# construct interpolation of theta_B(theta, zeta) etc.
theta_B = interpax.Interpolator2D(x=grid.nodes[grid.unique_theta_idx, 1],
                                 y=grid.nodes[grid.unique_zeta_idx, 2],
                                 f=grid.meshgrid_reshape(data['theta_B'], "rtz")[0],
                                 period=(2*np.pi, 2*np.pi/eq.NFP)
                                )
zeta_B = interpax.Interpolator2D(x=grid.nodes[grid.unique_theta_idx, 1],
                                 y=grid.nodes[grid.unique_zeta_idx, 2],
                                 f=grid.meshgrid_reshape(data['zeta_B'], "rtz")[0],
                                 period=(2*np.pi, 2*np.pi/eq.NFP)
                                )

# root finding to find theta, zeta coresponding to a particular theta_B, zeta_B
def residual(x, xB):
    t, z = x # current guess
    tB, zB = xB # target boozer angles
    t0 = theta_B(t,z)
    z0 = zeta_B(t,z)
    return jnp.array([(t0-tB), z0-zB])

# here x = initial guess for theta, zeta corresponding to the target xB == [theta_B, zeta_B]
rootfun = jax.vmap(lambda x, xB: desc.backend.root(residual, x0=x, args=(xB,)))

# We want a uniform grid of theta_B, zeta_B, so we can use the grid we already made which is uniformly spaced
# we're just sort of re-labeling it. As a guess we assume that theta ~ theta_B
# t, z are the desc theta,zeta cooresponding to uniformly spaced theta_B,zeta_B.
# note that t,z are generally non-uniform
t, z = rootfun(grid.nodes[:,1:], grid.nodes[:,1:]).T
# because its periodic, sometimes the rootfinder will give values outside of our usual 2pi domain so 
# we mod it back since some stuff may assume everything is in (0,2pi)
t = t%(2*np.pi)
z = z%(2*np.pi/eq.NFP)

# this class allows for custom node placement which we need because t, z are not uniformly 
# spaced in desc coordinates
boozer_grid = desc.grid.Grid(jnp.array([rho*jnp.ones_like(t), t, z]).T, NFP=eq.NFP)

# now compute whatever stuff you want on the new grid
data = eq.compute(["R", "Z", "|B|"], grid=boozer_grid)

[f0uriest]

To get it to work with the existing machinery in map_coordinates would require:

• implementing compute functions for {theta,zeta}_B_{r,t,z} ie derivatives of the boozer coordinates wrt desc coordinates. For the angle derivatives this is pretty easy since we already have nu_{t,z}. The radial derivatives might be trickier
• Also would need to do something to allow you to compute {theta,zeta}_B at a single point. Right now this requires a dense grid to compute w_mn

An alternative option would be to add a special case if inbasis=("rho", "theta_B", "zeta_B") and then basically use the code above. This would likely be the easiest though perhaps more clunky

[YigitElma]

This would be nice to have. I used booz_xform before for this (from Boozer coordinates to lab coordinates), but it was very slow without MPI. This could also help for comparison with other codes for particle tracing, i.e. SIMSOPT traces particles in Boozer coordinates.

[unalmis]

@f0uriest I doubt the code is correct. It has the same oversight I fixed in #1180.

You need to interpolate the stream map: theta - theta_B and zeta - zeta B. Coordinates are not periodic maps. zeta_B(2pi, 2pi/NFP) is not equal to zeta_B(0,0) .

[unalmis]

If your coordinates can parametrize maps that are not constant they must have secular terms.

It is useful to view the root finding code in #1919.