Skip to content

Non-stellarator symmetric configurations in vmec_compute_geometry #575

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
Open

Non-stellarator symmetric configurations in vmec_compute_geometry #575

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
dd4fab3
Adds support for non-stellarator symmetric configurations in vmec_com…
Dec 2, 2025
bdb9150
Added a test for non-stellarator symmetric vmec_compute_geometry() in…
Dec 2, 2025
6185503
Updated comment to better describe the test.
Dec 2, 2025
2d0d7d3
Quick fix to get the CI closer to passing.
akaptano Dec 2, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
132 changes: 86 additions & 46 deletions src/simsopt/mhd/vmec_diagnostics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -944,9 +944,7 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0):
# If given a Vmec object, convert it to vmec_splines:
if isinstance(vs, Vmec):
vs = vmec_splines(vs)
if not vs.stellsym:
raise NotImplementedError("vmec_compute_geometry() does not yet support non-stellarator-symmetric configurations.")


# Make sure s is an array:
try:
ns = len(s)
Expand Down Expand Up @@ -988,6 +986,7 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0):
phi = np.kron(np.ones((ns, ntheta, 1)), phi.reshape(1, 1, nphi))

# Shorthand:
stellsym = vs.stellsym
mnmax = vs.mnmax
xm = vs.xm
xn = vs.xn
Expand All @@ -1009,14 +1008,32 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0):
d_rmnc_d_s = np.zeros((ns, mnmax))
d_zmns_d_s = np.zeros((ns, mnmax))
d_lmns_d_s = np.zeros((ns, mnmax))
# stellsym
rmns = np.zeros((ns, mnmax))
zmnc = np.zeros((ns, mnmax))
lmnc = np.zeros((ns, mnmax))
d_rmns_d_s = np.zeros((ns, mnmax))
d_zmnc_d_s = np.zeros((ns, mnmax))
d_lmnc_d_s = np.zeros((ns, mnmax))


for jmn in range(mnmax):
rmnc[:, jmn] = vs.rmnc[jmn](s)
zmns[:, jmn] = vs.zmns[jmn](s)
lmns[:, jmn] = vs.lmns[jmn](s)
d_rmnc_d_s[:, jmn] = vs.d_rmnc_d_s[jmn](s)
d_zmns_d_s[:, jmn] = vs.d_zmns_d_s[jmn](s)
d_lmns_d_s[:, jmn] = vs.d_lmns_d_s[jmn](s)

if not stellsym:
rmns[:, jmn] = vs.rmns[jmn](s)
zmnc[:, jmn] = vs.zmnc[jmn](s)
lmnc[:, jmn] = vs.lmnc[jmn](s)
d_rmns_d_s[:, jmn] = vs.d_rmns_d_s[jmn](s)
d_zmnc_d_s[:, jmn] = vs.d_zmnc_d_s[jmn](s)
d_lmnc_d_s[:, jmn] = vs.d_lmnc_d_s[jmn](s)



gmnc = np.zeros((ns, mnmax_nyq))
bmnc = np.zeros((ns, mnmax_nyq))
d_bmnc_d_s = np.zeros((ns, mnmax_nyq))
Expand All @@ -1027,6 +1044,17 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0):
bsubvmnc = np.zeros((ns, mnmax_nyq))
d_bsupumnc_d_s = np.zeros((ns, mnmax_nyq))
d_bsupvmnc_d_s = np.zeros((ns, mnmax_nyq))
# stellsym
gmns = np.zeros((ns, mnmax_nyq))
bmns = np.zeros((ns, mnmax_nyq))
d_bmns_d_s = np.zeros((ns, mnmax_nyq))
bsupumns = np.zeros((ns, mnmax_nyq))
bsupvmns = np.zeros((ns, mnmax_nyq))
bsubsmnc = np.zeros((ns, mnmax_nyq))
bsubumns = np.zeros((ns, mnmax_nyq))
bsubvmns = np.zeros((ns, mnmax_nyq))
d_bsupumns_d_s = np.zeros((ns, mnmax_nyq))
d_bsupvmns_d_s = np.zeros((ns, mnmax_nyq))
for jmn in range(mnmax_nyq):
gmnc[:, jmn] = vs.gmnc[jmn](s)
bmnc[:, jmn] = vs.bmnc[jmn](s)
Expand All @@ -1038,7 +1066,19 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0):
bsubvmnc[:, jmn] = vs.bsubvmnc[jmn](s)
d_bsupumnc_d_s[:, jmn] = vs.d_bsupumnc_d_s[jmn](s)
d_bsupvmnc_d_s[:, jmn] = vs.d_bsupvmnc_d_s[jmn](s)

if not stellsym:
gmns[:, jmn] = vs.gmns[jmn](s)
bmns[:, jmn] = vs.bmns[jmn](s)
d_bmns_d_s[:, jmn] = vs.d_bmns_d_s[jmn](s)
bsupumns[:, jmn] = vs.bsupumns[jmn](s)
bsupvmns[:, jmn] = vs.bsupvmns[jmn](s)
bsubsmnc[:, jmn] = vs.bsubsmnc[jmn](s)
bsubumns[:, jmn] = vs.bsubumns[jmn](s)
bsubvmns[:, jmn] = vs.bsubvmns[jmn](s)
d_bsupumns_d_s[:, jmn] = vs.d_bsupumns_d_s[jmn](s)
d_bsupvmns_d_s[:, jmn] = vs.d_bsupvmns_d_s[jmn](s)


# Now that we know theta_vmec, compute all the geometric quantities
angle = xm[:, None, None, None] * theta_vmec[None, :, :, :] - xn[:, None, None, None] * phi[None, :, :, :]
cosangle = np.cos(angle)
Expand All @@ -1055,30 +1095,30 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0):
n2sinangle = xn[:, None, None, None]**2 * sinangle
# Order of indices in cosangle and sinangle: mn, s, theta, phi
# Order of indices in rmnc, bmnc, etc: s, mn
R = np.einsum('ij,jikl->ikl', rmnc, cosangle)
d_R_d_s = np.einsum('ij,jikl->ikl', d_rmnc_d_s, cosangle)
d_R_d_theta_vmec = np.einsum('ij,jikl->ikl', -rmnc, msinangle)
d_R_d_phi = np.einsum('ij,jikl->ikl', rmnc, nsinangle)
d2_R_d_phi2 = np.einsum('ij,jikl->ikl', -rmnc, n2cosangle)
d2_R_d_theta_vmec2 = np.einsum('ij,jikl->ikl', -rmnc, m2cosangle)
d2_R_d_theta_vmec_d_phi = np.einsum('ij,jikl->ikl', rmnc, mncosangle)
d2_R_d_s_d_theta_vmec = np.einsum('ij,jikl->ikl', -d_rmnc_d_s, msinangle)
d2_R_d_s_d_phi = np.einsum('ij,jikl->ikl', d_rmnc_d_s, nsinangle)

Z = np.einsum('ij,jikl->ikl', zmns, sinangle)
d_Z_d_s = np.einsum('ij,jikl->ikl', d_zmns_d_s, sinangle)
d_Z_d_theta_vmec = np.einsum('ij,jikl->ikl', zmns, mcosangle)
d_Z_d_phi = np.einsum('ij,jikl->ikl', -zmns, ncosangle)
d2_Z_d_theta_vmec2 = np.einsum('ij,jikl->ikl', -zmns, m2sinangle)
d2_Z_d_phi2 = np.einsum('ij,jikl->ikl', -zmns, n2sinangle)
d2_Z_d_theta_vmec_d_phi = np.einsum('ij,jikl->ikl', zmns, mnsinangle)
d2_Z_d_s_d_theta_vmec = np.einsum('ij,jikl->ikl', d_zmns_d_s, mcosangle)
d2_Z_d_s_d_phi = np.einsum('ij,jikl->ikl', -d_zmns_d_s, ncosangle)

lambd = np.einsum('ij,jikl->ikl', lmns, sinangle)
d_lambda_d_s = np.einsum('ij,jikl->ikl', d_lmns_d_s, sinangle)
d_lambda_d_theta_vmec = np.einsum('ij,jikl->ikl', lmns, mcosangle)
d_lambda_d_phi = np.einsum('ij,jikl->ikl', -lmns, ncosangle)
R = np.einsum('ij,jikl->ikl', rmnc, cosangle) + np.einsum('ij,jikl->ikl', rmns, sinangle)
d_R_d_s = np.einsum('ij,jikl->ikl', d_rmnc_d_s, cosangle) + np.einsum('ij,jikl->ikl', d_rmns_d_s, sinangle)
d_R_d_theta_vmec = np.einsum('ij,jikl->ikl', -rmnc, msinangle) + np.einsum('ij,jikl->ikl', rmns, mcosangle)
d_R_d_phi = np.einsum('ij,jikl->ikl', rmnc, nsinangle) + np.einsum('ij,jikl->ikl', -rmns, ncosangle)
d2_R_d_phi2 = np.einsum('ij,jikl->ikl', -rmnc, n2cosangle) + np.einsum('ij,jikl->ikl', -rmns, n2sinangle)
d2_R_d_theta_vmec2 = np.einsum('ij,jikl->ikl', -rmnc, m2cosangle) + np.einsum('ij,jikl->ikl', -rmns, m2sinangle)
d2_R_d_theta_vmec_d_phi = np.einsum('ij,jikl->ikl', rmnc, mncosangle) + np.einsum('ij,jikl->ikl', rmns, mnsinangle)
d2_R_d_s_d_theta_vmec = np.einsum('ij,jikl->ikl', -d_rmnc_d_s, msinangle) + np.einsum('ij,jikl->ikl', d_rmns_d_s, mcosangle)
d2_R_d_s_d_phi = np.einsum('ij,jikl->ikl', d_rmnc_d_s, nsinangle) + np.einsum('ij,jikl->ikl', -d_rmns_d_s, ncosangle)

Z = np.einsum('ij,jikl->ikl', zmns, sinangle) + np.einsum('ij,jikl->ikl', zmnc, cosangle)
d_Z_d_s = np.einsum('ij,jikl->ikl', d_zmns_d_s, sinangle) + np.einsum('ij,jikl->ikl', d_zmnc_d_s, cosangle)
d_Z_d_theta_vmec = np.einsum('ij,jikl->ikl', zmns, mcosangle) + np.einsum('ij,jikl->ikl', -zmnc, msinangle)
d_Z_d_phi = np.einsum('ij,jikl->ikl', -zmns, ncosangle) + np.einsum('ij,jikl->ikl', zmnc, nsinangle)
d2_Z_d_phi2 = np.einsum('ij,jikl->ikl', -zmns, n2sinangle) + np.einsum('ij,jikl->ikl', -zmnc, n2cosangle)
d2_Z_d_theta_vmec2 = np.einsum('ij,jikl->ikl', -zmns, m2sinangle) + np.einsum('ij,jikl->ikl', -zmnc, m2cosangle)
d2_Z_d_theta_vmec_d_phi = np.einsum('ij,jikl->ikl', zmns, mnsinangle) + np.einsum('ij,jikl->ikl', zmnc, mncosangle)
d2_Z_d_s_d_theta_vmec = np.einsum('ij,jikl->ikl', d_zmns_d_s, mcosangle) + np.einsum('ij,jikl->ikl', -d_zmnc_d_s, msinangle)
d2_Z_d_s_d_phi = np.einsum('ij,jikl->ikl', -d_zmns_d_s, ncosangle) + np.einsum('ij,jikl->ikl', d_zmnc_d_s, nsinangle)

lambd = np.einsum('ij,jikl->ikl', lmns, sinangle) + np.einsum('ij,jikl->ikl', lmnc, cosangle)
d_lambda_d_s = np.einsum('ij,jikl->ikl', d_lmns_d_s, sinangle) + np.einsum('ij,jikl->ikl', d_lmnc_d_s, cosangle)
d_lambda_d_theta_vmec = np.einsum('ij,jikl->ikl', lmns, mcosangle) + + np.einsum('ij,jikl->ikl', -lmnc, msinangle)
d_lambda_d_phi = np.einsum('ij,jikl->ikl', -lmns, ncosangle) + np.einsum('ij,jikl->ikl', lmnc, nsinangle)
theta_pest = theta_vmec + lambd

# Now handle the Nyquist quantities:
Expand All @@ -1090,24 +1130,24 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0):
msinangle = xm_nyq[:, None, None, None] * sinangle
nsinangle = xn_nyq[:, None, None, None] * sinangle

sqrt_g_vmec = np.einsum('ij,jikl->ikl', gmnc, cosangle)
modB = np.einsum('ij,jikl->ikl', bmnc, cosangle)
d_B_d_s = np.einsum('ij,jikl->ikl', d_bmnc_d_s, cosangle)
d_B_d_theta_vmec = np.einsum('ij,jikl->ikl', -bmnc, msinangle)
d_B_d_phi = np.einsum('ij,jikl->ikl', bmnc, nsinangle)

B_sup_theta_vmec = np.einsum('ij,jikl->ikl', bsupumnc, cosangle)
B_sup_phi = np.einsum('ij,jikl->ikl', bsupvmnc, cosangle)
B_sub_s = np.einsum('ij,jikl->ikl', bsubsmns, sinangle)
B_sub_theta_vmec = np.einsum('ij,jikl->ikl', bsubumnc, cosangle)
B_sub_phi = np.einsum('ij,jikl->ikl', bsubvmnc, cosangle)
sqrt_g_vmec = np.einsum('ij,jikl->ikl', gmnc, cosangle) + np.einsum('ij,jikl->ikl', gmns, sinangle)
modB = np.einsum('ij,jikl->ikl', bmnc, cosangle) + np.einsum('ij,jikl->ikl', bmns, sinangle)
d_B_d_s = np.einsum('ij,jikl->ikl', d_bmnc_d_s, cosangle) + np.einsum('ij,jikl->ikl', d_bmns_d_s, sinangle)
d_B_d_theta_vmec = np.einsum('ij,jikl->ikl', -bmnc, msinangle) + np.einsum('ij,jikl->ikl', bmns, mcosangle)
d_B_d_phi = np.einsum('ij,jikl->ikl', bmnc, nsinangle) + np.einsum('ij,jikl->ikl', -bmns, ncosangle)

B_sup_theta_vmec = np.einsum('ij,jikl->ikl', bsupumnc, cosangle) + np.einsum('ij,jikl->ikl', bsupumns, sinangle)
B_sup_phi = np.einsum('ij,jikl->ikl', bsupvmnc, cosangle) + np.einsum('ij,jikl->ikl', bsupvmns, sinangle)
B_sub_s = np.einsum('ij,jikl->ikl', bsubsmns, sinangle) + np.einsum('ij,jikl->ikl', bsubsmnc, cosangle)
B_sub_theta_vmec = np.einsum('ij,jikl->ikl', bsubumnc, cosangle) + np.einsum('ij,jikl->ikl', bsubumns, sinangle)
B_sub_phi = np.einsum('ij,jikl->ikl', bsubvmnc, cosangle) + np.einsum('ij,jikl->ikl', bsubvmns, sinangle)
B_sup_theta_pest = iota[:, None, None] * B_sup_phi
d_B_sup_phi_d_theta_vmec = np.einsum('ij,jikl->ikl', -bsupvmnc, msinangle)
d_B_sup_phi_d_phi = np.einsum('ij,jikl->ikl', bsupvmnc, nsinangle)
d_B_sup_theta_vmec_d_theta_vmec = np.einsum('ij,jikl->ikl', -bsupumnc, msinangle)
d_B_sup_theta_vmec_d_phi = np.einsum('ij,jikl->ikl', bsupumnc, nsinangle)
d_B_sup_theta_vmec_d_s = np.einsum('ij,jikl->ikl', d_bsupumnc_d_s, cosangle)
d_B_sup_phi_d_s = np.einsum('ij,jikl->ikl', d_bsupvmnc_d_s, cosangle)
d_B_sup_phi_d_theta_vmec = np.einsum('ij,jikl->ikl', -bsupvmnc, msinangle) + np.einsum('ij,jikl->ikl', bsupvmns, mcosangle)
d_B_sup_phi_d_phi = np.einsum('ij,jikl->ikl', bsupvmnc, nsinangle) + np.einsum('ij,jikl->ikl', -bsupvmns, ncosangle)
d_B_sup_theta_vmec_d_theta_vmec = np.einsum('ij,jikl->ikl', -bsupumnc, msinangle) + np.einsum('ij,jikl->ikl', bsupumns, mcosangle)
d_B_sup_theta_vmec_d_phi = np.einsum('ij,jikl->ikl', bsupumnc, nsinangle) + np.einsum('ij,jikl->ikl', -bsupumns, ncosangle)
d_B_sup_theta_vmec_d_s = np.einsum('ij,jikl->ikl', d_bsupumnc_d_s, cosangle) + np.einsum('ij,jikl->ikl', d_bsupumns_d_s, sinangle)
d_B_sup_phi_d_s = np.einsum('ij,jikl->ikl', d_bsupvmnc_d_s, cosangle) + np.einsum('ij,jikl->ikl', d_bsupvmns_d_s, sinangle)

sqrt_g_vmec_alt = R * (d_Z_d_s * d_R_d_theta_vmec - d_R_d_s * d_Z_d_theta_vmec)

Expand Down
33 changes: 33 additions & 0 deletions tests/mhd/test_vmec_diagnostics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -626,6 +626,39 @@ def test_L_grad_B_axisymmetry(self):
div_B = data.grad_B__XX + data.grad_B__YY + data.grad_B__ZZ
np.testing.assert_allclose(div_B, 0, atol=3e-11)

@unittest.skipIf(vmec is None, "vmec python package is not found")
def test_basic_non_stellsym(self):
"""
Perform a sqrt(g) calculation comparison and a few regression tests
for a non-stellarator symmetric configuration.
"""
vmec = Vmec(os.path.join(TEST_DIR, "input.basic_non_stellsym"))
vmec.run()
s = 1.0
ntheta = 5
nphi = 4
theta = np.linspace(0, 2 * np.pi, ntheta, endpoint=True)
phi = np.linspace(0, 2 * np.pi / vmec.wout.nfp, nphi, endpoint=True)

#
R_precalc = np.array([[[6.85, 4.571410161513775, 3.8785898384862243, 6.85], [4.5, 3.721410161513776, 3.0285898384862238, 4.5], [5.65, 4.005384757729337, 5.044615242270662, 5.65], [8.0, 4.855384757729337, 5.894615242270662, 8.0], [6.85, 4.571410161513775, 3.8785898384862247, 6.85]]])
d2_Z_d_theta_vmec_d_phi_precalc = np.array([[[-0.5, -0.18301270189221946, 0.6830127018922194, -0.4999999999999999], [0.49999999999999994, -0.6830127018922193, 0.18301270189221908, 0.5000000000000001], [0.5000000000000001, 0.1830127018922194, -0.6830127018922194, 0.49999999999999994], [-0.4999999999999999, 0.6830127018922193, -0.18301270189221913, -0.5], [-0.5000000000000001, -0.1830127018922196, 0.6830127018922194, -0.5]]])


d_B_d_theta_vmec_precalc = np.array([[[3.797382850490039, 8.62225895484851, 14.159181806821593, 3.7973828504900395], [2.8919186555758207, -3.998872017181843, -3.458246653557826, 2.89191865557582], [-4.854840544126677, -7.566177748008503, -11.874600310783727, -4.854840544126677], [-1.5912894561501743, 2.815018903778379, 0.6849479261949663, -1.5912894561501754], [3.797382850490038, 8.622258954848505, 14.159181806821596, 3.7973828504900387]]])

data = vmec_compute_geometry(vmec, s, theta, phi)

# More arrays to check against can be printed by uncommenting this:
Copy link
Contributor

@smiet smiet Dec 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Given this comment, I gather that the precalc arrays are computed from the current state of the code. This makes them regression tests. Have you verified that these precalculated arrays are correct against physical properties of the input.basic_non_stellsym?

Could some test be added that is clearly non-stellarator-symmetric? Not sure of the interpretation of these arrays, but you could test whether the evaluation of quantities at +phi are not equal to that at minus phi.

# print('np.array(' + str(data.d_B_d_theta_vmec.tolist()) + ')')
# and copy+pasting the output as 'precalc' array above

np.testing.assert_allclose(data.R, R_precalc, rtol=1e-5)
np.testing.assert_allclose(data.d2_Z_d_theta_vmec_d_phi, d2_Z_d_theta_vmec_d_phi_precalc, rtol=1e-5)
np.testing.assert_allclose(data.d_B_d_theta_vmec, d_B_d_theta_vmec_precalc, rtol=1e-5)
# Non-regression test, comparing the two ways to calculate sqrt(g)
np.testing.assert_allclose(data.sqrt_g_vmec_alt,data.sqrt_g_vmec, rtol=1e-4, atol=1e-4)


class VmecFieldlinesTests(unittest.TestCase):
def test_fieldline_grids(self):
Expand Down
Loading