Controlling radial discretization

[EdoAlvarezR]

I'm trying to generate somewhat of a uniform grid (which I hope to use as the initial condition for an non-equilibrium unsteady MHD sim), but the flux surfaces that come out of the VMEC++ solution are not uniformly spaced. Is there an input that controls the radial spacing of the flux surfaces or is that completely controlled by the internal computation of VMEC?

Here is the code that generated the solution shown above:

import os, shutil, copy, re
from pathlib import Path

import time
import warnings

import json, math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

import netCDF4
import vmecpp

import sys
sys.path.append('/home/edoalvar/Dropbox/MHDResearch/OtherCodes/vmecpp/examples')
from visualize_magnetic_field import calculate_magnetic_field

import pyvista as pv



# ----------- RUN VMEC++ ---------------------------------------------------

# Construct a VmecInput object, e.g. from a classic Fortran input file or VMEC++'s json format
input_file = "data/w7x.json"
vmec_input = vmecpp.VmecInput.from_file(input_file)

vmec_input.ns_array = np.array([30])

# Run VMEC++:
vmec_output = vmecpp.run(vmec_input)

# ----------- CONSTRUCT MAGNETIC FIELD OVER PLASMA VOLUME -----------------

# Resolution of plasma volume
num_s = vmec_output.wout.ns
num_theta = 61
num_phi = 361

# Grid in theta and phi along the flux surface
grid_j = np.arange(0, num_s)
grid_theta = np.linspace(0.0, 2 * np.pi, num_theta, endpoint=True)
grid_phi = np.linspace(0.0, 2 * np.pi, num_phi, endpoint=True)

# Compute Cartesian coordinates of flux surface geometry and magnetic field
x, y, z = (np.zeros([num_s, num_theta, num_phi]) for i in range(3))
bx, by, bz = (np.zeros([num_s, num_theta, num_phi]) for i in range(3))
jx, jy, jz = (np.zeros([num_s, num_theta, num_phi]) for i in range(3))

for idx_j, j in enumerate(grid_j):
    for idx_theta, theta in enumerate(grid_theta):
        for idx_phi, phi in enumerate(grid_phi):
            
            # Compute
            X, B, J = calculate_magnetic_field(vmec_output, j, theta, phi)
    
            # Unpack coordinates
            x[idx_j, idx_theta, idx_phi] = X[0]
            y[idx_j, idx_theta, idx_phi] = X[1]
            z[idx_j, idx_theta, idx_phi] = X[2]
    
            bx[idx_j, idx_theta, idx_phi] = B[0]
            by[idx_j, idx_theta, idx_phi] = B[1]
            bz[idx_j, idx_theta, idx_phi] = B[2]
    
            jx[idx_j, idx_theta, idx_phi] = J[0]
            jy[idx_j, idx_theta, idx_phi] = J[1]
            jz[idx_j, idx_theta, idx_phi] = J[2]

# ----------- VISUALIZE WITH pyVista ---------------------------------------

# Create structured grid
grid = pv.StructuredGrid(x, y, z)

# Combine vector components
B = np.c_[   bx.ravel(order='F'),
             by.ravel(order='F'),
             bz.ravel(order='F')]

J = np.c_[   jx.ravel(order='F'),
             jy.ravel(order='F'),
             jz.ravel(order='F')]

# Attach vector field
grid["B"] = B
grid["J"] = J

# Output VTK file for visualization in ParaView
grid.save("temps/vmec_volume002.vts")

In advance, thanks for your help! 🙌

-Ed Alvarez

[jons-pf]

Well, the flux surface are equally spaced in the radial direction, but in terms of enclosed toroidal magnetic flux and not in a radius-like coordinate ;-)
There is an option though in VMEC++ to change this. It is called aphi in the inputs and refers to a polynomial, which controls the mapping of the normalized radial coordinate s to normalized toroidal flux phi/phiedge (screenshots from a talk by John Schmitt (Auburn Uni) on this):


and in The Numerics of VMEC++ in Sec. (5.4).

[EdoAlvarezR]

That worked well. I got a more uniform grid using vmec_input.aphi = np.array([0.05, 0.09, 0.05]) :

Thanks @jons-pf!