diff --git a/src/simsopt/mhd/vmec_diagnostics.py b/src/simsopt/mhd/vmec_diagnostics.py index d0671b548..5d7296a06 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 @@ -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 = False + for jmn in range(mnmax): rmnc[:, jmn] = vs.rmnc[jmn](s) zmns[:, jmn] = vs.zmns[jmn](s) @@ -974,29 +982,22 @@ 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 value of theta_vmec that - gives a desired theta_pest. - """ - - """ - 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 + 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[jradius, :] * np.sin(xm * theta_v - xn * phi0))) + return theta_p_target - (theta_v + np.sum(lmns[js, :, None] * np.sin(xm[:, None] * theta_v - xn[:, None] * phi0), axis=0)) - # Solve for theta_vmec corresponding to theta_pest: 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 = theta_pest[js, jalpha, :] + 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, :, :, :] @@ -1169,13 +1170,13 @@ 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 \ - / (modB * modB * modB) * toroidal_flux_sign + # 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] \ - * toroidal_flux_sign / (edge_toroidal_flux_over_2pi * modB * modB) + # 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 diff --git a/tests/mhd/test_vmec_diagnostics.py b/tests/mhd/test_vmec_diagnostics.py index 9fabf509d..e18fcd120 100755 --- a/tests/mhd/test_vmec_diagnostics.py +++ b/tests/mhd/test_vmec_diagnostics.py @@ -546,7 +546,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) @@ -602,9 +602,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):