Skip to content

Commit

Permalink
2X faster vmec_fieldlines + sign flip
Browse files Browse the repository at this point in the history 
 Date:      Sun Jul 3 21:25:39 2022 -0400
 Committer: Rahul Gaur <rg6256@stellar-intel.princeton.edu>
	modified:   ../../src/simsopt/mhd/vmec_diagnostics.py
	modified:   test_vmec_diagnostics.py
  • Loading branch information
Rahul Gaur committed Jul 15, 2022
1 parent 472cc44 commit bf6ea86
Show file tree
Hide file tree
Showing 2 changed files with 159 additions and 12 deletions.
165 changes: 156 additions & 9 deletions src/simsopt/mhd/vmec_diagnostics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -932,6 +932,14 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F
d_rmnc_d_s = np.zeros((ns, mnmax))
d_zmns_d_s = np.zeros((ns, mnmax))
d_lmns_d_s = np.zeros((ns, mnmax))

######## CAREFUL!!###########################################################
# Everything here and in vmec_splines is designed for up-down symmetric eqlbia
# When we start optimizing equilibria with lasym = "True"
# we should edit this as well as vmec_splines
lmnc = np.zeros((ns, mnmax))
lasym = 0

for jmn in range(mnmax):
rmnc[:, jmn] = vs.rmnc[jmn](s)
zmns[:, jmn] = vs.zmns[jmn](s)
Expand Down Expand Up @@ -987,16 +995,153 @@ def residual(theta_v, phi0, theta_p_target, jradius):
"""
return theta_p_target - (theta_v + np.sum(lmns[jradius, :] * np.sin(xm * theta_v - xn * phi0)))

# Solve for theta_vmec corresponding to theta_pest:
# A helper function that returns 0 when theta_vmec_trys are close to the
# proper toroidal angle in magnetic coordinates
# vmec_fieldlines new residual routine (written by Alan Goodman)
def fzero_residuals_function(theta_vmec_trys, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym):

angle = xm * theta_vmec_trys - xn * phis
#pdb.set_trace()
fzero_residual = t
print("WAT")
cosangle = np.cos(angle)
fzero_residual += np.sum(lmnc * np.cos(angle), axis=1)


return fzero_residual

# A helper function that numerically finds values that minimize fzero_residuals_function
# This is the non-linear root finder that turns out to be ~8X faster than scipy root_scalar
# written by Alan Goodman
def get_roots(a0, b0, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym):
converged = False
nconv = 0
a = 1 * a0
b = 1 * b0
roots = np.zeros(a.shape)
toteval = len(a.flatten())
itmax_bracket = 10
itmax_root = 100
tol = 1e-10

fa = fzero_residuals_function(a, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym)
fb = fzero_residuals_function(b, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym)
for it in range(itmax_bracket):
eps = np.finfo(a[0, 0]).eps
inds = (((fa > 0.0) * (fb > 0.0)) + ((fa < 0.0) * (fb < 0.0)))
if np.sum(inds) != 0:
a[inds] = a[inds] - 0.3 * np.ones(a[inds].shape)
b[inds] = b[inds] + 0.3 * np.ones(a[inds].shape)
fa = fzero_residuals_function(a, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym)
fb = fzero_residuals_function(b, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym)
else:
break

c =1 * b
fc = 1 * fb
d = 0 * fb
e = 0 * fb
convds = False * np.ones(b.shape)

for it in range(itmax_root):
inds = (((fb > 0.0) * (fc > 0.0)) + ((fb < 0.0) * (fc < 0.0)))
c[inds] = a[inds]
fc[inds] = fa[inds]
d[inds] = b[inds] - a[inds]
e[inds] = d[inds]

inds = (np.abs(fc) < np.abs(fb))
a[inds] = b[inds]
b[inds] = c[inds]
c[inds] = a[inds]
fa[inds] = fb[inds]
fb[inds] = fc[inds]
fc[inds] = fa[inds]

tol1 = 2.0 * eps * np.abs(b) + 0.5 * tol
Xm = 0.5 * (c - b)
indsC = ((np.abs(Xm) <= tol1) + (fb == 0.0))

roots[indsC] = b[indsC]
convds[indsC] = True
#nconv = len(inds[0])
nconv = np.sum(indsC)
if nconv == toteval:
converged = True
return roots, converged, nconv

p = 0 * a
s = 0 * p
q = 0 * p
r = 0 * p

inds1 = (((np.abs(e) >= tol1) * (np.abs(fa) > np.abs(fb))))
s[inds1] = (fb[inds1])/(fa[inds1])

inds2 = (a == c)
inds = inds2*inds1
s[inds] = (fb[inds])/(fa[inds])
p[inds] = 2.0 * Xm[inds] * s[inds]
q[inds] = 1.0 - s[inds]

inds = np.invert(inds2)*inds1
q[inds] = fa[inds] / fc[inds]
r[inds] = fb[inds] / fc[inds]
p[inds] = s[inds] * (2.0 * Xm[inds] * q[inds] * (q[inds] - r[inds]) - (b[inds] - a[inds]) * (r[inds] - 1.0))
q[inds] = (q[inds] - 1.0) * (r[inds] - 1.0) * (s[inds] - 1.0)

inds = (p > 0.0)*inds1
q[inds] = -q[inds]
p[inds1] = np.abs(p[inds1])

inds2 = (2.0 * p < np.minimum(3.0 * Xm * q - np.abs(tol1 * q), np.abs(e * q)))
inds = inds2*inds1
e[inds] = d[inds]
d[inds] = (p[inds])/(q[inds])

inds = np.invert(inds2)*inds1
d[inds] = Xm[inds]
e[inds] = d[inds]

inds = np.invert(inds1)
d[inds] = Xm[inds]
e[inds] = d[inds]

a = 1 * b
fa = 1 * fb

inds = np.where((np.abs(d) > tol1))
b[inds] += d[inds]

inds = np.where((np.abs(d) <= tol1))
b[inds] += np.copysign(tol1[inds], Xm[inds])
fb = fzero_residuals_function(b, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym)
return roots, converged, nconv

#theta_vmec = np.zeros((ns, nalpha, nl))
#for js in range(ns):
# for jalpha in range(nalpha):
# for jl in range(nl):
# theta_guess = theta_pest[js, jalpha, jl]
# solution = root_scalar(residual,args=(phi[js, jalpha, jl], theta_pest[js, jalpha, jl], js),
# bracket=(theta_guess - 1.0, theta_guess + 1.0))
# theta_vmec[js, jalpha, jl] = solution.root
## Solve for theta_vmec corresponding to theta_pest (new routine)
## Does the same calculation as the commented code above but faster
theta_vmec = np.zeros((ns, nalpha, nl))
for js in range(ns):
for jalpha in range(nalpha):
for jl in range(nl):
theta_guess = theta_pest[js, jalpha, jl]
solution = root_scalar(residual,
args=(phi[js, jalpha, jl], theta_pest[js, jalpha, jl], js),
bracket=(theta_guess - 1.0, theta_guess + 1.0))
theta_vmec[js, jalpha, jl] = solution.root
theta_guess = np.reshape(theta_pest[js, jalpha, :], (-1, 1))
theta_vmec_mins = theta_guess - 0.3*np.ones(theta_guess.shape)
theta_vmec_maxs = theta_guess + 0.3*np.ones(theta_guess.shape)
theta_test, converged, nconv = get_roots(theta_vmec_mins, theta_vmec_maxs, theta_guess, np.reshape(phi[js, jalpha, :], (-1,1)), xm, xn, lmns[js], lmnc[js], 'False')
theta_vmec[js, jalpha] = np.reshape(theta_test, (-1,))

if converged == False:
print("* Error! Conversion from theta_pest to theta_vmec failed to converge")
print("===========================================================================")
print("")
err = True

# 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, :, :, :]
Expand Down Expand Up @@ -1169,12 +1314,14 @@ def residual(theta_v, phi0, theta_p_target, jradius):

gds22 = grad_psi_dot_grad_psi * shat[:, None, None] * shat[:, None, None] / (L_reference * L_reference * B_reference * B_reference * s[:, None, None])

gbdrift = 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * B_cross_grad_B_dot_grad_alpha \
# temporary fix. Please see issue #238 and the discussion therein
gbdrift = -1 * 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * B_cross_grad_B_dot_grad_alpha \
/ (modB * modB * modB) * toroidal_flux_sign

gbdrift0 = B_cross_grad_B_dot_grad_psi * 2 * shat[:, None, None] / (modB * modB * modB * sqrt_s[:, None, None]) * toroidal_flux_sign

cvdrift = gbdrift + 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * mu_0 * d_pressure_d_s[:, None, None] \
# temporary fix. Please see issue #238 and the discussion therein
cvdrift = -1 * gbdrift + -1 * 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * mu_0 * d_pressure_d_s[:, None, None] \
* toroidal_flux_sign / (edge_toroidal_flux_over_2pi * modB * modB)

cvdrift0 = gbdrift0
Expand Down
6 changes: 3 additions & 3 deletions tests/mhd/test_vmec_diagnostics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -545,7 +545,7 @@ def test_consistency(self):
np.testing.assert_allclose(fl.B_cross_grad_B_dot_grad_alpha, fl.B_cross_grad_B_dot_grad_alpha_alternate, atol=0.02)

# Check 2 ways of computing cvdrift:
cvdrift_alt = 2 * fl.B_reference * fl.L_reference * fl.L_reference \
cvdrift_alt = -1 * 2 * fl.B_reference * fl.L_reference * fl.L_reference \
* np.sqrt(fl.s)[:, None, None] * fl.B_cross_kappa_dot_grad_alpha \
/ (fl.modB * fl.modB) * fl.toroidal_flux_sign
np.testing.assert_allclose(fl.cvdrift, cvdrift_alt)
Expand Down Expand Up @@ -601,9 +601,9 @@ def test_stella_regression(self):
np.testing.assert_allclose(fl.gds2[0, j, :], np.fromstring(lines[16 + j], sep=' '), rtol=2e-4)
np.testing.assert_allclose(fl.gds21[0, j, :], np.fromstring(lines[20 + j], sep=' '), atol=2e-4)
np.testing.assert_allclose(fl.gds22[0, j, :], np.fromstring(lines[24 + j], sep=' '), atol=2e-4)
np.testing.assert_allclose(fl.gbdrift[0, j, :], np.fromstring(lines[28 + j], sep=' '), atol=2e-4)
np.testing.assert_allclose(-1 * fl.gbdrift[0, j, :], np.fromstring(lines[28 + j], sep=' '), atol=2e-4)
np.testing.assert_allclose(fl.gbdrift0[0, j, :], np.fromstring(lines[32 + j], sep=' '), atol=1e-4)
np.testing.assert_allclose(fl.cvdrift[0, j, :], np.fromstring(lines[36 + j], sep=' '), atol=2e-4)
np.testing.assert_allclose(-1 * fl.cvdrift[0, j, :], np.fromstring(lines[36 + j], sep=' '), atol=2e-4)
np.testing.assert_allclose(fl.cvdrift0[0, j, :], np.fromstring(lines[40 + j], sep=' '), atol=1e-4)

def test_axisymm(self):
Expand Down

0 comments on commit bf6ea86

Please sign in to comment.