From bf6ea86d69727ab66f017b04536b964b81efeaf6 Mon Sep 17 00:00:00 2001 From: Rahul Gaur Date: Sun, 3 Jul 2022 21:25:39 -0400 Subject: [PATCH 1/5] 2X faster vmec_fieldlines + sign flip Date: Sun Jul 3 21:25:39 2022 -0400 Committer: Rahul Gaur modified: ../../src/simsopt/mhd/vmec_diagnostics.py modified: test_vmec_diagnostics.py --- src/simsopt/mhd/vmec_diagnostics.py | 165 ++++++++++++++++++++++++++-- tests/mhd/test_vmec_diagnostics.py | 6 +- 2 files changed, 159 insertions(+), 12 deletions(-) diff --git a/src/simsopt/mhd/vmec_diagnostics.py b/src/simsopt/mhd/vmec_diagnostics.py index d0671b548..64f42cc30 100644 --- a/src/simsopt/mhd/vmec_diagnostics.py +++ b/src/simsopt/mhd/vmec_diagnostics.py @@ -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) @@ -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, :, :, :] @@ -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 diff --git a/tests/mhd/test_vmec_diagnostics.py b/tests/mhd/test_vmec_diagnostics.py index 127e0378b..4f6dc937e 100755 --- a/tests/mhd/test_vmec_diagnostics.py +++ b/tests/mhd/test_vmec_diagnostics.py @@ -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) @@ -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): From f6e9c4f563ad7fe86aca02a9a8c871d400eb0683 Mon Sep 17 00:00:00 2001 From: rahulgaur104 Date: Sun, 21 Aug 2022 20:55:33 -0400 Subject: [PATCH 2/5] added_vectorized_newton --- src/simsopt/mhd/vmec_diagnostics.py | 153 +++------------------------- 1 file changed, 15 insertions(+), 138 deletions(-) diff --git a/src/simsopt/mhd/vmec_diagnostics.py b/src/simsopt/mhd/vmec_diagnostics.py index 7e53c539c..5177af22d 100644 --- a/src/simsopt/mhd/vmec_diagnostics.py +++ b/src/simsopt/mhd/vmec_diagnostics.py @@ -11,7 +11,7 @@ import numpy as np from scipy.interpolate import interp1d, InterpolatedUnivariateSpline -from scipy.optimize import root_scalar +from scipy.optimize import newton from .vmec import Vmec from .._core.util import Struct @@ -995,145 +995,23 @@ 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))) - # 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 - fzero_residual = theta_pest_targets - (theta_vmec_trys + np.reshape(np.sum(lmns * np.sin(angle), axis=1), (-1, 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 + + def residual(theta_v, phi0, theta_p_target, xm, xn, lmns): + """ + This function is used for computing the value of theta_vmec that + gives a desired theta_pest. + """ - 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 = 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]) + return theta_p_target - (theta_v + np.sum(lmns[:, None] * np.sin(xm[:, None] * theta_v - xn[:, None] * phi0), axis=0)) - 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): - 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 + theta_guess = theta_pest[js, jalpha, :] + solution = newton(residual, x0 = theta_guess - 0.5, args=(phi[js, jalpha, :], theta_pest[js, jalpha, :], xm, xn, lmns[js]), x1 = theta_guess + 0.5) + #pdb.set_trace() + theta_vmec[js, jalpha, :] = solution # 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, :, :, :] @@ -1307,14 +1185,13 @@ def get_roots(a0, b0, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym): gds22 = grad_psi_dot_grad_psi * shat[:, None, None] * shat[:, None, None] / (L_reference * L_reference * B_reference * B_reference * s[:, None, None]) # 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 + 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 # temporary fix. Please see issue #238 and the discussion therein - cvdrift = 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) - + cvdrift = gbdrift - 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 # Package results into a structure to return: From c5b09ba50b18b6119f8b18bf3b6f410d9649d298 Mon Sep 17 00:00:00 2001 From: rahulgaur104 Date: Sun, 21 Aug 2022 20:55:33 -0400 Subject: [PATCH 3/5] faster-root-solver --- src/simsopt/mhd/vmec_diagnostics.py | 158 ++-------------------------- 1 file changed, 9 insertions(+), 149 deletions(-) diff --git a/src/simsopt/mhd/vmec_diagnostics.py b/src/simsopt/mhd/vmec_diagnostics.py index 7e53c539c..4cdc8ccb0 100644 --- a/src/simsopt/mhd/vmec_diagnostics.py +++ b/src/simsopt/mhd/vmec_diagnostics.py @@ -11,7 +11,7 @@ import numpy as np from scipy.interpolate import interp1d, InterpolatedUnivariateSpline -from scipy.optimize import root_scalar +from scipy.optimize import newton from .vmec import Vmec from .._core.util import Struct @@ -980,160 +980,21 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F theta_pest[js, :, :] = theta1d[None, :] phi[js, :, :] = phi_center + (theta1d[None, :] - alpha[:, None]) / iota[js] + def residual(theta_v, phi0, theta_p_target, jradius): """ This function is used for computing the value of theta_vmec that gives a desired theta_pest. """ + return theta_p_target - (theta_v + np.sum(lmns[js, :, None] * np.sin(xm[:, None] * theta_v - xn[:, None] * phi0), axis=0)) - """ - theta_p = theta_v - for jmn in range(len(xn)): - angle = xm[jmn] * theta_v - xn[jmn] * phi0 - theta_p += lmns[jradius, jmn] * np.sin(angle) - return theta_p_target - theta_p - """ - return theta_p_target - (theta_v + np.sum(lmns[jradius, :] * np.sin(xm * theta_v - xn * phi0))) - - # 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 - fzero_residual = theta_pest_targets - (theta_vmec_trys + np.reshape(np.sum(lmns * np.sin(angle), axis=1), (-1, 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 = 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): - 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 + theta_guess = theta_pest[js, jalpha, :] + solution = newton(residual, x0=theta_guess - 1.0, args=(phi[js, jalpha, :], theta_pest[js, jalpha, :], js), x1=theta_guess + 1.0) + theta_vmec[js, jalpha, :] = solution # 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, :, :, :] @@ -1307,14 +1168,13 @@ def get_roots(a0, b0, theta_pest_targets, phis, xm, xn, lmns, lmnc, lasym): gds22 = grad_psi_dot_grad_psi * shat[:, None, None] * shat[:, None, None] / (L_reference * L_reference * B_reference * B_reference * s[:, None, None]) # 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 + 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 # temporary fix. Please see issue #238 and the discussion therein - cvdrift = 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) - + cvdrift = gbdrift - 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 # Package results into a structure to return: From e0059f05c3a7bf8e9c8d42fd3e2f1ebc7a001d60 Mon Sep 17 00:00:00 2001 From: rahulgaur104 Date: Mon, 22 Aug 2022 06:26:27 -0400 Subject: [PATCH 4/5] faster-root-solver --- tests/mhd/test_vmec_diagnostics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mhd/test_vmec_diagnostics.py b/tests/mhd/test_vmec_diagnostics.py index eaca6f30c..e18fcd120 100755 --- a/tests/mhd/test_vmec_diagnostics.py +++ b/tests/mhd/test_vmec_diagnostics.py @@ -604,7 +604,7 @@ def test_stella_regression(self): np.testing.assert_allclose(fl.gds22[0, j, :], np.fromstring(lines[24 + 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(-1 * 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): From eddac271ffd251a2a1a43ffe7767686c680a0de6 Mon Sep 17 00:00:00 2001 From: Matt Landreman Date: Sat, 27 Aug 2022 17:00:49 -0400 Subject: [PATCH 5/5] vmec_fieldlines: reduced number of iterations for root finding by changing x0 --- src/simsopt/mhd/vmec_diagnostics.py | 123 +++++++++++++++------------- 1 file changed, 64 insertions(+), 59 deletions(-) diff --git a/src/simsopt/mhd/vmec_diagnostics.py b/src/simsopt/mhd/vmec_diagnostics.py index e46302b37..5d7296a06 100644 --- a/src/simsopt/mhd/vmec_diagnostics.py +++ b/src/simsopt/mhd/vmec_diagnostics.py @@ -884,7 +884,7 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F # If given a Vmec object, convert it to vmec_splines: if isinstance(vs, Vmec): vs = vmec_splines(vs) - + # Make sure s is an array: try: ns = len(s) @@ -892,7 +892,7 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F s = [s] s = np.array(s) ns = len(s) - + # Make sure alpha is an array try: nalpha = len(alpha) @@ -900,7 +900,7 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F alpha = [alpha] alpha = np.array(alpha) nalpha = len(alpha) - + if (theta1d is not None) and (phi1d is not None): raise ValueError('You cannot specify both theta and phi') if (theta1d is None) and (phi1d is None): @@ -909,7 +909,7 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F nl = len(phi1d) else: nl = len(theta1d) - + # Shorthand: mnmax = vs.mnmax xm = vs.xm @@ -917,7 +917,7 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F mnmax_nyq = vs.mnmax_nyq xm_nyq = vs.xm_nyq xn_nyq = vs.xn_nyq - + # Now that we have an s grid, evaluate everything on that grid: d_pressure_d_s = vs.d_pressure_d_s(s) iota = vs.iota(s) @@ -925,21 +925,21 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F # shat = (r/q)(dq/dr) where r = a sqrt(s) # = - (r/iota) (d iota / d r) = -2 (s/iota) (d iota / d s) shat = (-2 * s / iota) * d_iota_d_s - + rmnc = np.zeros((ns, mnmax)) zmns = np.zeros((ns, mnmax)) lmns = np.zeros((ns, mnmax)) 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 = False - + for jmn in range(mnmax): rmnc[:, jmn] = vs.rmnc[jmn](s) zmns[:, jmn] = vs.zmns[jmn](s) @@ -947,7 +947,7 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F 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) - + gmnc = np.zeros((ns, mnmax_nyq)) bmnc = np.zeros((ns, mnmax_nyq)) d_bmnc_d_s = np.zeros((ns, mnmax_nyq)) @@ -965,10 +965,10 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F bsubsmns[:, jmn] = vs.bsubsmns[jmn](s) bsubumnc[:, jmn] = vs.bsubumnc[jmn](s) bsubvmnc[:, jmn] = vs.bsubvmnc[jmn](s) - + theta_pest = np.zeros((ns, nalpha, nl)) phi = np.zeros((ns, nalpha, nl)) - + if theta1d is None: # We are given phi. Compute theta_pest: for js in range(ns): @@ -982,18 +982,23 @@ def vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=F def residual(theta_v, phi0, theta_p_target, jradius): """ - This function is used for computing the an array of values of theta_vmec that + This function is used for computing an array of values of theta_vmec that give a desired theta_pest array. """ return theta_p_target - (theta_v + np.sum(lmns[js, :, None] * np.sin(xm[:, None] * theta_v - xn[:, None] * phi0), axis=0)) - + theta_vmec = np.zeros((ns, nalpha, nl)) for js in range(ns): for jalpha in range(nalpha): theta_guess = theta_pest[js, jalpha, :] - solution = newton(residual, x0=theta_guess - 1.0, args=(phi[js, jalpha, :], theta_pest[js, jalpha, :], js), x1=theta_guess + 1.0) + solution = newton( + residual, + x0=theta_guess, + x1=theta_guess + 0.1, + args=(phi[js, jalpha, :], theta_pest[js, jalpha, :], js), + ) theta_vmec[js, jalpha, :] = solution - + # 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) @@ -1008,16 +1013,16 @@ def residual(theta_v, phi0, theta_p_target, jradius): 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) - + 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) - + 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) - + # Now handle the Nyquist quantities: angle = xm_nyq[:, None, None, None] * theta_vmec[None, :, :, :] - xn_nyq[:, None, None, None] * phi[None, :, :, :] cosangle = np.cos(angle) @@ -1026,25 +1031,25 @@ def residual(theta_v, phi0, theta_p_target, jradius): ncosangle = xn_nyq[:, None, None, None] * cosangle 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) B_sup_theta_pest = iota[:, None, None] * B_sup_phi - + sqrt_g_vmec_alt = R * (d_Z_d_s * d_R_d_theta_vmec - d_R_d_s * d_Z_d_theta_vmec) - + # Note the minus sign. psi in the straight-field-line relation seems to have opposite sign to vmec's phi array. edge_toroidal_flux_over_2pi = -vs.phiedge / (2 * np.pi) - + # ********************************************************************* # Using R(theta,phi) and Z(theta,phi), compute the Cartesian # components of the gradient basis vectors using the dual relations: @@ -1059,56 +1064,56 @@ def residual(theta_v, phi0, theta_p_target, jradius): d_Y_d_theta_vmec = d_R_d_theta_vmec * sinphi d_Y_d_phi = d_R_d_phi * sinphi + R * cosphi d_Y_d_s = d_R_d_s * sinphi - + # Now use the dual relations to get the Cartesian components of grad s, grad theta_vmec, and grad phi: grad_s_X = (d_Y_d_theta_vmec * d_Z_d_phi - d_Z_d_theta_vmec * d_Y_d_phi) / sqrt_g_vmec grad_s_Y = (d_Z_d_theta_vmec * d_X_d_phi - d_X_d_theta_vmec * d_Z_d_phi) / sqrt_g_vmec grad_s_Z = (d_X_d_theta_vmec * d_Y_d_phi - d_Y_d_theta_vmec * d_X_d_phi) / sqrt_g_vmec - + grad_theta_vmec_X = (d_Y_d_phi * d_Z_d_s - d_Z_d_phi * d_Y_d_s) / sqrt_g_vmec grad_theta_vmec_Y = (d_Z_d_phi * d_X_d_s - d_X_d_phi * d_Z_d_s) / sqrt_g_vmec grad_theta_vmec_Z = (d_X_d_phi * d_Y_d_s - d_Y_d_phi * d_X_d_s) / sqrt_g_vmec - + grad_phi_X = (d_Y_d_s * d_Z_d_theta_vmec - d_Z_d_s * d_Y_d_theta_vmec) / sqrt_g_vmec grad_phi_Y = (d_Z_d_s * d_X_d_theta_vmec - d_X_d_s * d_Z_d_theta_vmec) / sqrt_g_vmec grad_phi_Z = (d_X_d_s * d_Y_d_theta_vmec - d_Y_d_s * d_X_d_theta_vmec) / sqrt_g_vmec # End of dual relations. - + # ********************************************************************* # Compute the Cartesian components of other quantities we need: # ********************************************************************* - + grad_psi_X = grad_s_X * edge_toroidal_flux_over_2pi grad_psi_Y = grad_s_Y * edge_toroidal_flux_over_2pi grad_psi_Z = grad_s_Z * edge_toroidal_flux_over_2pi - + # Form grad alpha = grad (theta_vmec + lambda - iota * phi) grad_alpha_X = (d_lambda_d_s - (phi - phi_center) * d_iota_d_s[:, None, None]) * grad_s_X grad_alpha_Y = (d_lambda_d_s - (phi - phi_center) * d_iota_d_s[:, None, None]) * grad_s_Y grad_alpha_Z = (d_lambda_d_s - (phi - phi_center) * d_iota_d_s[:, None, None]) * grad_s_Z - + grad_alpha_X += (1 + d_lambda_d_theta_vmec) * grad_theta_vmec_X + (-iota[:, None, None] + d_lambda_d_phi) * grad_phi_X grad_alpha_Y += (1 + d_lambda_d_theta_vmec) * grad_theta_vmec_Y + (-iota[:, None, None] + d_lambda_d_phi) * grad_phi_Y grad_alpha_Z += (1 + d_lambda_d_theta_vmec) * grad_theta_vmec_Z + (-iota[:, None, None] + d_lambda_d_phi) * grad_phi_Z - + grad_B_X = d_B_d_s * grad_s_X + d_B_d_theta_vmec * grad_theta_vmec_X + d_B_d_phi * grad_phi_X grad_B_Y = d_B_d_s * grad_s_Y + d_B_d_theta_vmec * grad_theta_vmec_Y + d_B_d_phi * grad_phi_Y grad_B_Z = d_B_d_s * grad_s_Z + d_B_d_theta_vmec * grad_theta_vmec_Z + d_B_d_phi * grad_phi_Z - + B_X = edge_toroidal_flux_over_2pi * ((1 + d_lambda_d_theta_vmec) * d_X_d_phi + (iota[:, None, None] - d_lambda_d_phi) * d_X_d_theta_vmec) / sqrt_g_vmec B_Y = edge_toroidal_flux_over_2pi * ((1 + d_lambda_d_theta_vmec) * d_Y_d_phi + (iota[:, None, None] - d_lambda_d_phi) * d_Y_d_theta_vmec) / sqrt_g_vmec B_Z = edge_toroidal_flux_over_2pi * ((1 + d_lambda_d_theta_vmec) * d_Z_d_phi + (iota[:, None, None] - d_lambda_d_phi) * d_Z_d_theta_vmec) / sqrt_g_vmec - + # ********************************************************************* # For gbdrift, we need \vect{B} cross grad |B| dot grad alpha. # For cvdrift, we also need \vect{B} cross grad s dot grad alpha. # Let us compute both of these quantities 2 ways, and make sure the two # approaches give the same answer (within some tolerance). # ********************************************************************* - + B_cross_grad_s_dot_grad_alpha = (B_sub_phi * (1 + d_lambda_d_theta_vmec) \ - B_sub_theta_vmec * (d_lambda_d_phi - iota[:, None, None])) / sqrt_g_vmec - + B_cross_grad_s_dot_grad_alpha_alternate = 0 \ + B_X * grad_s_Y * grad_alpha_Z \ + B_Y * grad_s_Z * grad_alpha_X \ @@ -1116,7 +1121,7 @@ def residual(theta_v, phi0, theta_p_target, jradius): - B_Z * grad_s_Y * grad_alpha_X \ - B_X * grad_s_Z * grad_alpha_Y \ - B_Y * grad_s_X * grad_alpha_Z - + B_cross_grad_B_dot_grad_alpha = 0 \ + (B_sub_s * d_B_d_theta_vmec * (d_lambda_d_phi - iota[:, None, None]) \ + B_sub_theta_vmec * d_B_d_phi * (d_lambda_d_s - (phi - phi_center) * d_iota_d_s[:, None, None]) \ @@ -1124,7 +1129,7 @@ def residual(theta_v, phi0, theta_p_target, jradius): - B_sub_phi * d_B_d_theta_vmec * (d_lambda_d_s - (phi - phi_center) * d_iota_d_s[:, None, None]) \ - B_sub_theta_vmec * d_B_d_s * (d_lambda_d_phi - iota[:, None, None]) \ - B_sub_s * d_B_d_phi * (1 + d_lambda_d_theta_vmec)) / sqrt_g_vmec - + B_cross_grad_B_dot_grad_alpha_alternate = 0 \ + B_X * grad_B_Y * grad_alpha_Z \ + B_Y * grad_B_Z * grad_alpha_X \ @@ -1132,49 +1137,49 @@ def residual(theta_v, phi0, theta_p_target, jradius): - B_Z * grad_B_Y * grad_alpha_X \ - B_X * grad_B_Z * grad_alpha_Y \ - B_Y * grad_B_X * grad_alpha_Z - + grad_alpha_dot_grad_alpha = grad_alpha_X * grad_alpha_X + grad_alpha_Y * grad_alpha_Y + grad_alpha_Z * grad_alpha_Z - + grad_alpha_dot_grad_psi = grad_alpha_X * grad_psi_X + grad_alpha_Y * grad_psi_Y + grad_alpha_Z * grad_psi_Z - + grad_psi_dot_grad_psi = grad_psi_X * grad_psi_X + grad_psi_Y * grad_psi_Y + grad_psi_Z * grad_psi_Z - + B_cross_grad_B_dot_grad_psi = (B_sub_theta_vmec * d_B_d_phi - B_sub_phi * d_B_d_theta_vmec) / sqrt_g_vmec * edge_toroidal_flux_over_2pi - + B_cross_kappa_dot_grad_psi = B_cross_grad_B_dot_grad_psi / modB - + mu_0 = 4 * np.pi * (1.0e-7) B_cross_kappa_dot_grad_alpha = B_cross_grad_B_dot_grad_alpha / modB + mu_0 * d_pressure_d_s[:, None, None] / edge_toroidal_flux_over_2pi - + # stella / gs2 / gx quantities: - + L_reference = vs.Aminor_p B_reference = 2 * abs(edge_toroidal_flux_over_2pi) / (L_reference * L_reference) toroidal_flux_sign = np.sign(edge_toroidal_flux_over_2pi) sqrt_s = np.sqrt(s) - + bmag = modB / B_reference - + gradpar_theta_pest = L_reference * B_sup_theta_pest / modB - + gradpar_phi = L_reference * B_sup_phi / modB - + gds2 = grad_alpha_dot_grad_alpha * L_reference * L_reference * s[:, None, None] - + gds21 = grad_alpha_dot_grad_psi * shat[:, None, None] / B_reference - + gds22 = grad_psi_dot_grad_psi * shat[:, None, None] * shat[:, None, None] / (L_reference * L_reference * B_reference * B_reference * s[:, None, None]) - + # 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 - + # temporary fix. Please see issue #238 and the discussion therein cvdrift = gbdrift - 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 - + # Package results into a structure to return: results = Struct() variables = ['ns', 'nalpha', 'nl', 's', 'iota', 'd_iota_d_s', 'd_pressure_d_s', 'shat', @@ -1194,7 +1199,7 @@ def residual(theta_v, phi0, theta_p_target, jradius): 'grad_alpha_dot_grad_alpha', 'grad_alpha_dot_grad_psi', 'grad_psi_dot_grad_psi', 'L_reference', 'B_reference', 'toroidal_flux_sign', 'bmag', 'gradpar_theta_pest', 'gradpar_phi', 'gds2', 'gds21', 'gds22', 'gbdrift', 'gbdrift0', 'cvdrift', 'cvdrift0'] - + if plot: import matplotlib.pyplot as plt plt.figure(figsize=(13, 7)) @@ -1209,13 +1214,13 @@ def residual(theta_v, phi0, theta_p_target, jradius): plt.plot(phi[0, 0, :], eval(variable + '[0, 0, :]')) plt.xlabel('Standard toroidal angle $\phi$') plt.title(variable) - + plt.figtext(0.5, 0.995, f's={s[0]}, alpha={alpha[0]}', ha='center', va='top') plt.tight_layout() if show: plt.show() - + for v in variables: results.__setattr__(v, eval(v)) - + return results