Inverse Coordinate Transform

[mwiesenberger]

In this document https://arxiv.org/pdf/2502.04374 there is an algorithm desccribed of how to obtain the inverse coordinate transformation $s(R,Z,\varphi), \theta(R,Z,\varphi), \zeta(R,Z,\varphi)$ as Algorithm 4. I wonder if this algorithm is readily implemented in vmecpp or another python package? I was unable to find it, is it not publicly available?
Thank you

[jons-pf]

Hi, I do have a prototype implementation (see below), but not published that yet properly. I might have a look if I can also pull together the tests. Be aware that this inverse coordinate transform needs serious testing, before it can be used in production.

In particular, the radial derivatives need proper regularization, because VMEC++ computes the flux surface geometry only on a discrete set of flux surfaces.

def evaluate_rz(rho_theta, r_cos, r_sin, z_cos, z_sin, modes):
  """Evaluate R(rho, theta) and Z(rho, theta) from Fourier coefficients.

  Args:
    rho_theta: Array of shape (n_points, 2) with (rho, theta) pairs
    r_cos: R^cos coefficients
    r_sin: R^sin coefficients
    z_cos: Z^cos coefficients
    z_sin: Z^sin coefficients
    modes: List of mode dictionaries

  Returns:
    Tuple (R, Z) of arrays with shape (n_points,)
  """

def evaluate_rz_derivatives(rho_theta, r_cos, r_sin, z_cos, z_sin, modes):
  """Evaluate derivatives of R(rho, theta) and Z(rho, theta) from Fourier coefficients.

  Computes the Jacobian matrix components: dR/drho, dR/dtheta, dZ/drho, dZ/dtheta.

  Args:
    rho_theta: Array of shape (n_points, 2) with (rho, theta) pairs
    r_cos: R^cos coefficients
    r_sin: R^sin coefficients
    z_cos: Z^cos coefficients
    z_sin: Z^sin coefficients
    modes: List of mode dictionaries

  Returns:
    Tuple (dR_drho, dR_dtheta, dZ_drho, dZ_dtheta) of arrays with shape (n_points,)
  """

def inverse_transform(R0, Z0, r_cos, r_sin, z_cos, z_sin, modes, eps=1e-12, max_iter=10):
  """Find (rho, theta) for given (R0, Z0) using Newton's method with adaptive
  backtracking.

  Args:
    R0, Z0: Target coordinates
    r_cos: R^cos coefficients
    r_sin: R^sin coefficients
    z_cos: Z^cos coefficients
    z_sin: Z^sin coefficients
    modes: List of mode dictionaries
    eps: Convergence tolerance
    max_iter: Maximum iterations

  Returns:
    Dictionary with keys:
      - rho: final rho coordinate
      - theta: final theta coordinate
      - converged: True if converged
      - iterations: number of iterations taken
      - distance: final distance to target
      - history: list of (rho, theta, R, Z, dist) tuples for each iteration
  """
  # initial guess based on axis and effective radius
  R_axis, Z_axis = evaluate_rz(np.array([[0.0, 0.0]]), r_cos, r_sin, z_cos, z_sin, modes)
  dR_ax = R0 - R_axis[0]
  dZ_ax = Z0 - Z_axis[0]
  d = math.hypot(dR_ax, dZ_ax)

  # estimate effective radius from (m=1) coefficient
  idx_m1 = None
  for i, mode in enumerate(modes):
    if mode["m"] == 1:
      idx_m1 = i
      break

  if idx_m1 is not None:
    a_eff = math.sqrt((abs(r_cos[idx_m1]) or 1e-6) * (abs(z_sin[idx_m1]) or 1e-6)) or 1.0
  else:
    a_eff = 1.0

  rho = d / a_eff
  theta = math.atan2(dZ_ax, dR_ax) % (2 * math.pi)

  # limit initial guess for rho to 1
  rho = min(rho, 1.0)

  beta = 1.0
  prev_dist2 = float("inf")
  tau_prev = float("nan")

  history = []

  for k in range(max_iter):
    Rk, Zk = evaluate_rz(np.array([[rho, theta]]), r_cos, r_sin, z_cos, z_sin, modes)
    Rk, Zk = Rk[0], Zk[0]
    dR = Rk - R0
    dZ = Zk - Z0
    dist2 = dR * dR + dZ * dZ

    history.append((rho, theta, Rk, Zk, math.sqrt(dist2)))

    if dist2 < eps**2:
      return {
        "rho": rho,
        "theta": theta,
        "converged": True,
        "iterations": k,
        "distance": math.sqrt(dist2),
        "history": history,
      }

    dR_dr, dR_dth, dZ_dr, dZ_dth = evaluate_rz_derivatives(
      np.array([[rho, theta]]), r_cos, r_sin, z_cos, z_sin, modes
    )
    dR_dr, dR_dth, dZ_dr, dZ_dth = dR_dr[0], dR_dth[0], dZ_dr[0], dZ_dth[0]

    tau = dR_dr * dZ_dth - dR_dth * dZ_dr
    if k == 0:
      tau_prev = tau

    # back-tracking if worse
    if dist2 > prev_dist2:
      tau = 0.75 * tau_prev + 0.25 * tau
      beta *= 0.5
    else:
      beta = 1.0

    rho_new = rho - beta * (dZ_dth * dR - dR_dth * dZ) / tau
    th_new = theta - beta * (dR_dr * dZ - dZ_dr * dR) / tau

    # handle crossing axis
    if rho_new < 0:
      rho_new = -rho_new
      th_new += math.pi

    # limit to rho <= 1
    rho_new = min(rho_new, 1.0)

    theta = th_new % (2 * math.pi)
    rho = rho_new

    prev_dist2 = dist2
    tau_prev = tau

  return {
    "rho": rho,
    "theta": theta,
    "converged": False,
    "iterations": max_iter,
    "distance": math.sqrt(prev_dist2),
    "history": history,
  }

[mwiesenberger]

Hi, thanks for the reply, in the meantime I managed to quickly throw something together myself.

In case you're interested: I just fit a polynomial to the vmec coefficients in $\rho$, which then defines a derivative. That works very well as long as $\rho \leq 1$ but the suggested extrapolation of the vmec coefficients Eq. 9.105 in the arxiv paper is problematic as it is not continuously differentiable at $\rho=1$. This basically means that any polynomial fit wildly oscillates for $\rho > 1$ because of the discontinuity. In that case it is better to use an extrapolation that matches the derivative at $\rho=1$. Something like
$X_{mn}(\rho) = X_{mn}(1) + \frac{\partial X_{mn}}{\partial \rho}(1) (\rho -1 ) e^{-(\rho -1)}$ and $X_{10}(\rho) = X_{mn}(1) \rho$ makes the extrapolation continuously differentiable and thus smooth. The derivative at $\rho=1$ can be obtained by finite differences.