2X faster vmec_fieldlines

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
 

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):