Merge pull request #1 from landreman/faster-root-solver

vmec_fieldlines: reduced number of iterations for root finding

src/simsopt/mhd/vmec_diagnostics.py

@@ -884,23 +884,23 @@ 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)
     except:
         s = [s]
     s = np.array(s)
     ns = len(s)
-    
+
     # Make sure alpha is an array
     try:
         nalpha = len(alpha)
     except:
         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,45 +909,45 @@ 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
     xn = vs.xn
     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)
     d_iota_d_s = vs.d_iota_d_s(s)
     # 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)
         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)
-    
+
     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,122 +1064,122 @@ 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 \
         + B_Z * grad_s_X * grad_alpha_Y \
         - 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]) \
            + B_sub_phi * d_B_d_s * (1 + d_lambda_d_theta_vmec) \
            - 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 \
         + B_Z * grad_B_X * grad_alpha_Y \
         - 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