80 check boozer routines (#83)

.gitignore

@@ -61,3 +61,4 @@ xsol.dat
 yvec.dat
 include/
 lib/
+.venv

python/libneo/boozer.py

@@ -17,64 +17,76 @@
 length_cgs_to_si = 1e-2
 debug = False
 
-def get_boozer_transform(stor, nth):
-  from scipy.interpolate import CubicSpline
 
-  efit_to_boozer.efit_to_boozer.init()
+def get_boozer_transform(stor, num_theta):
+    from scipy.interpolate import CubicSpline
 
-  R_axis, Z_axis = get_magnetic_axis()
+    efit_to_boozer.efit_to_boozer.init()
 
-  ns = stor.shape[0]
+    ns = stor.shape[0]
 
-  G_of_thb = []
-  dth_of_thb = []
-  for ks in np.arange(ns - 1):
+    G_of_thb = []
+    dth_of_thb = []
+    for ks in np.arange(ns - 1):
+        th_boozers, th_geoms, G_s = get_angles_and_transformation(stor[ks], num_theta)
+        dth, th_boozers, th_geoms = get_sign_dependent_thetas(th_geoms, th_boozers)
+
+        G_of_thb.append(CubicSpline(th_boozers, G_s, bc_type="periodic"))
+        dth_of_thb.append(CubicSpline(th_boozers, dth, bc_type="periodic"))
+
+    return dth_of_thb, G_of_thb, th_boozers, th_geoms
+
+
+def get_angles_and_transformation(s, num_theta):
     psi = np.array(0.0)
-    s = np.array(stor[ks])
+    R_axis, Z_axis = get_magnetic_axis()
 
     th_boozers = []
     th_geoms = []
     Gs = []
-    for th_symflux in 2 * np.pi * np.arange(2 * nth) / (2 * nth):
-      inp_label = 1
-      (
-        q,
-        _,
-        _,
-        _,
-        _,
-        _,
-        _,
-        R,
-        _,
-        _,
-        Z,
-        _,
-        _,
-        G,
-        _,
-        _,
-      ) = efit_to_boozer.efit_to_boozer.magdata(
-        inp_label, s, psi, th_symflux)
-      Gs.append(G)
-      th_boozers.append(th_symflux + G / q)
-      th_geoms.append(np.arctan2(
+
+    for th_symflux in 2 * np.pi * np.arange(2 * num_theta) / (2 * num_theta):
+        inp_label = 1
+        (q, _, _, _, _, _, _,
+        R, _, _, 
+        Z, _, _, 
+        G, _, _,
+        ) = efit_to_boozer.efit_to_boozer.magdata(inp_label, s, psi, th_symflux)
+        Gs.append(G)
+        th_boozers.append(th_symflux + G / q)
+        th_geoms.append(np.arctan2(
           Z*length_cgs_to_si - Z_axis,
           R*length_cgs_to_si - R_axis))
+
     Gs = np.array(Gs)
     Gs = np.append(Gs, Gs[0])
-    th_boozers = np.unwrap(th_boozers)
-    th_boozers = np.append(th_boozers, th_boozers[0] + 2 * np.pi)
-    th_geoms = np.unwrap(th_geoms)
-    th_geoms = np.append(th_geoms, th_geoms[0] + 2 * np.pi)
 
-    G_of_thb.append(CubicSpline(th_boozers, Gs, bc_type="periodic"))
-    dth_of_thb.append(CubicSpline(th_boozers, th_geoms - th_boozers, bc_type="periodic"))
+    return th_boozers, th_geoms, Gs
+
+
+def get_sign_dependent_thetas(th_geoms, th_boozers):
 
-  return dth_of_thb, G_of_thb
+  th_boozers = np.unwrap(th_boozers)
+  monotone_sign_boozers = np.sign(np.mean(np.sign(np.diff(th_boozers))))
 
+  th_geoms = np.unwrap(th_geoms)
+  monotone_sign_geoms = np.sign(np.mean(np.sign(np.diff(th_geoms))))
 
-def get_B0_of_s_theta_boozer(stor, nth):
+  th_boozers = np.append(th_boozers, th_boozers[0] + monotone_sign_boozers * 2 * np.pi)
+  th_geoms = np.append(th_geoms, th_geoms[0] + monotone_sign_geoms * 2 * np.pi)
+
+  if monotone_sign_boozers == monotone_sign_geoms:
+    sign = -1
+  else:
+    sign = 1 
+
+  dth= th_geoms + sign * th_boozers
+
+  return dth, th_boozers, th_geoms
+
+
+
+def get_B0_of_s_theta_boozer(stor, num_theta):
   from scipy.interpolate import CubicSpline
 
   efit_to_boozer.efit_to_boozer.init()
@@ -88,7 +100,7 @@ def get_B0_of_s_theta_boozer(stor, nth):
 
     th_boozers = []
     B0mod = []
-    for th_symflux in 2 * np.pi * np.arange(2 * nth) / (2 * nth):
+    for th_symflux in 2 * np.pi * np.arange(2 * num_theta) / (2 * num_theta):
       inp_label = 1
       (
         q,
@@ -113,7 +125,8 @@ def get_B0_of_s_theta_boozer(stor, nth):
       B0mod.append(B0mod_temp)
 
     th_boozers = np.unwrap(th_boozers)
-    th_boozers = np.append(th_boozers, th_boozers[0] + 2 * np.pi)
+    monotone_sign_boozers = np.sign(np.mean(np.sign(np.diff(th_boozers))))
+    th_boozers = np.append(th_boozers, th_boozers[0] + monotone_sign_boozers * 2 * np.pi)
 
     B0mod = np.array(B0mod)
     B0mod = np.append(B0mod, B0mod[0])
@@ -122,90 +135,145 @@ def get_B0_of_s_theta_boozer(stor, nth):
   return B0
 
 
-
-def get_boozer_harmonics(f, stor, nth, nph, m0b, n, dth_of_thb, G_of_thb):
+def get_boozer_harmonics(f, stor, num_theta, num_phi, m0b, n, dth_of_thb, G_of_thb):
   """
   f(stor, th_geom, ph) of normalized toroidal flux and geometric angles
 
   Returns: Fourier harmonics fmn in terms of Boozer angles
+  Set num_phi=1 for axisymmetric case.
   """
 
   ns = stor.shape[0]
 
   fmn = np.zeros((ns - 1, 2 * m0b + 1), dtype=complex)
 
-  # Loop over flux surface index
-  th_geoms = np.zeros((ns - 1, nth))
-  phs = np.zeros((ns - 1, nth, nph))
+  th_geoms, phs = get_thgeoms_phs(ns, num_theta, num_phi, dth_of_thb, G_of_thb)
+
 
+  for kth in np.arange(num_theta):
+    if debug: print(f"kth = {kth}/{num_theta-1}")
+    thb = kth * 2 * np.pi / (num_theta)
+    for kph in np.arange(num_phi):
+      phb = kph * 2 * np.pi / (num_phi)
+      f_values = f(stor[:-1], th_geoms[:, kth], phs[:, kth, kph])
+
+      for m in np.arange(-m0b, m0b + 1):
+        # Take a sum over all theta values here
+        fmn[:, m + m0b] += f_values * np.exp(-1j * (m * thb + n * phb))
+
+  return fmn / ((num_theta) * (num_phi))
+
+def get_thgeoms_phs(ns, num_theta, num_phi, dth_of_thb, G_of_thb):
+
+  th_geoms = np.zeros((ns - 1, num_theta))
+  phs = np.zeros((ns - 1, num_theta, num_phi))
+
+  # Loop over flux surface index
   for ks in np.arange(ns - 1):
     if debug: print(f"ks = {ks}/{ns-2}")
-    for kth in np.arange(nth):
-      thb = kth * 2 * np.pi / nth
+    for kth in np.arange(num_theta):
+      thb = kth * 2 * np.pi / num_theta
       th_geoms[ks, kth] = thb + dth_of_thb[ks](thb)
       G = G_of_thb[ks](thb)
-      for kph in np.arange(nph):
-        phb = kph * 2 * np.pi / nph
+      for kph in np.arange(num_phi):
+        phb = kph * 2 * np.pi / num_phi
         phs[ks, kth, kph] = phb - G
 
-  for kth in np.arange(nth):
-    if debug: print(f"kth = {kth}/{nth-1}")
-    thb = kth * 2 * np.pi / nth
-    for kph in np.arange(nph):
-      phb = kph * 2 * np.pi / nph
+  return th_geoms, phs
 
-      f_values = f(stor[:-1], th_geoms[:, kth], phs[:, kth, kph])
+def get_boozer_harmonics_divide_f_by_B0(f, stor, num_theta, num_phi, m0b, n, dth_of_thb, G_of_thb):
+  """
+  f(stor, th_geom, ph) of normalized toroidal flux and geometric angles
+
+  Returns: Fourier harmonics fmn in terms of Boozer angles, i.e. fmn(s,m) for a given
+  toroidal mode number n
+  """
+
+  ns = stor.shape[0]
+  fmn = np.zeros((ns - 1, 2 * m0b + 1), dtype=complex)
+
+  th_geoms, phs = get_thgeoms_phs(ns, num_theta, num_phi, dth_of_thb, G_of_thb)
+  B0 = get_B0_of_s_theta_boozer(stor, num_theta)
+
+  for kth in np.arange(num_theta):
+    if debug: print(f"kth = {kth}/{num_theta-1}")
+    thb = kth * 2 * np.pi / num_theta
+    B0_arr = np.array([spline(thb) for spline in B0])
+
+    for kph in np.arange(num_phi):
+      phb = kph * 2 * np.pi / num_phi
+
+      f_values = np.array(f(stor[:-1], th_geoms[:, kth], phs[:, kth, kph]), dtype=complex) \
+        / B0_arr
 
       for m in np.arange(-m0b, m0b + 1):
         # Take a sum over all theta values here
         fmn[:, m + m0b] += f_values * np.exp(-1j * (m * thb + n * phb))
 
-  return fmn / (nth * nph)
-
+  return fmn / (num_theta * num_phi)
 
-def get_boozer_harmonics_divide_f_by_B0(f, stor, nth, nph, m0b, n, dth_of_thb, G_of_thb):
+def get_boozer_harmonics_divide_f_by_B0_1D(f, stor, num_theta, num_phi, m0b, n, dth_of_thb, G_of_thb):
   """
-  f(stor, th_geom, ph) of normalized toroidal flux and geometric angles
+  f_n(stor, th_geom) of normalized toroidal flux and geometric poloidal angle for given 
+  toroidal mode number n (also in geometric angle). Takes also the transformation from 
+  geometric to Boozer angle in toroidal direction into account (additional phase factor
+  depending on G(\vatheta_B)).
 
   Returns: Fourier harmonics fmn in terms of Boozer angles, i.e. fmn(s,m) for a given
   toroidal mode number n
   """
 
   ns = stor.shape[0]
-
   fmn = np.zeros((ns - 1, 2 * m0b + 1), dtype=complex)
 
-  # Loop over flux surface index
-  th_geoms = np.zeros((ns - 1, nth))
-  phs = np.zeros((ns - 1, nth, nph))
+  th_geoms, phs = get_thgeoms_phs(ns, num_theta, num_phi, dth_of_thb, G_of_thb)
+  B0 = get_B0_of_s_theta_boozer(stor, num_theta)
 
-  for ks in np.arange(ns - 1):
-    if debug: print(f"ks = {ks}/{ns-2}")
-    for kth in np.arange(nth):
-      thb = kth * 2 * np.pi / nth
-      th_geoms[ks, kth] = thb + dth_of_thb[ks](thb)
-      G = G_of_thb[ks](thb)
-      for kph in np.arange(nph):
-        phb = kph * 2 * np.pi / nph
-        phs[ks, kth, kph] = phb - G
+  for kth in np.arange(num_theta):
+    if debug: print(f"kth = {kth}/{num_theta-1}")
+    thb = kth * 2 * np.pi / num_theta
+    B0_arr = np.array([spline(thb) for spline in B0])
 
-  B0 = get_B0_of_s_theta_boozer(stor, nth)
+    G = np.array([G_of_thb[ks](thb) for ks in np.arange(ns - 1)])
 
-  for kth in np.arange(nth):
-    if debug: print(f"kth = {kth}/{nth-1}")
-    thb = kth * 2 * np.pi / nth
-    for kph in np.arange(nph):
-      phb = kph * 2 * np.pi / nph
+    f_values = np.array(f(stor[:-1], th_geoms[:, kth], phs[:, kth, 0]), dtype=complex) \
+      / B0_arr * np.exp(-1j * n * G)
 
-      f_values = np.array(f(stor[:-1], th_geoms[:, kth], phs[:, kth, kph])) \
-        / np.array([spline(thb) for spline in B0])
+    for m in np.arange(-m0b, m0b + 1):
+      # Take a sum over all theta values here
+      fmn[:, m + m0b] += f_values * np.exp(-1j * (m * thb))
 
-      for m in np.arange(-m0b, m0b + 1):
-        # Take a sum over all theta values here
-        fmn[:, m + m0b] += f_values * np.exp(-1j * (m * thb + n * phb))
+  return fmn / (num_theta * num_phi)
+
+def get_boozer_harmonics_divide_f_by_B0_1D_fft(f, stor, num_theta, n, dth_of_thb=None, G_of_thb=None):
+  """
+  f_n(ind_stor, th_geom) of normalized toroidal flux (index is used as argument!) and 
+  geometric poloidal angle for given 
+  toroidal mode number n (also in geometric angle). Takes also the transformation from 
+  geometric to Boozer angle in toroidal direction into account (additional phase factor
+  depending on G(\vatheta_B)).
+
+  Returns: Fourier harmonics fmn in terms of Boozer angles, i.e. fmn(s,m) for a given
+  toroidal mode number n
+  """
+
+  ns = stor.shape[0]
+  fmn = np.zeros((ns - 1, num_theta), dtype=complex)
+
+  if dth_of_thb==None or G_of_thb==None:  
+    dth_of_thb, G_of_thb, _, _ = get_boozer_transform(stor, num_theta)
+
+  th_geoms, phs = get_thgeoms_phs(ns, num_theta, 1, dth_of_thb, G_of_thb)
+  theta_boozers = np.linspace(0, 2*np.pi, num_theta, endpoint=False)
 
-  return fmn / ((2 * np.pi) ** 2 * nth * nph)
+  B0 = get_B0_of_s_theta_boozer(stor, num_theta)
+  #breakpoint()
 
+  fmn = np.zeros((ns -1, num_theta), dtype=complex)
+  for js, _ in enumerate(stor[:-1]):
+    FF = f(js, th_geoms[js,:]) / np.abs(B0[js](theta_boozers)) * np.exp(-1j * n * G_of_thb[js](theta_boozers))
+    fmn[js, :] = np.fft.fft(FF)
+  return fmn / num_theta
 
 def get_magnetic_axis():
   psi = np.array(0.0)
@@ -1341,18 +1409,18 @@ def get_rho_poloidal(self):
 
     from math import sqrt
     from numpy import array
-    from scipy.integrate import simps
+    from scipy.integrate import simpson
 
     iota = array(self.iota)
 
-    psi_pol_a = simps(iota, self.s)# Order of arguments is y, x.
+    psi_pol_a = simpson(iota, x=self.s)# Order of arguments is y, x.
 
     rho_poloidal = []
 
     rho_poloidal.append(0.0)
 
     for k in range(1+1, len(self.iota)+1):
-      rho_poloidal.append(sqrt(simps(iota[0:k], self.s[0:k])/ psi_pol_a))
+      rho_poloidal.append(sqrt(simpson(iota[0:k], x=self.s[0:k])/ psi_pol_a))
 
     # Interpolate first point
     rho_poloidal[0] = sqrt(self.s[0]/self.s[1])*rho_poloidal[1]

test/python/.gitignore

@@ -0,0 +1,2 @@
+out.*
+*.dat

test/python/efit_to_boozer.inp

@@ -0,0 +1 @@
+/Users/markusmarkl/plasma/6codes/0projects/2024-mastu-rmps/MHD_current/efit_to_boozer.inp

test/python/field_divB0.inp

@@ -0,0 +1 @@
+/Users/markusmarkl/plasma/6codes/0projects/2024-mastu-rmps/MHD_current/field_divB0.inp