code: """
  This module contains the top-level routines for the quasisymmetric
  stellarator construction.
  """
  
  import logging
  import numpy as np
  from scipy.io import netcdf
  #from numba import jit
  
  #logging.basicConfig(level=logging.DEBUG)
  logger = logging.getLogger(__name__)
  
  class Qsc():
      """
      This is the main class for representing the quasisymmetric
      stellarator construction.
      """
      
      # Import methods that are defined in separate files:
      from .init_axis import init_axis, convert_to_spline
      from .calculate_r1 import _residual, _jacobian, solve_sigma_equation, \
          _determine_helicity, r1_diagnostics
      from .grad_B_tensor import calculate_grad_B_tensor, calculate_grad_grad_B_tensor, \
          Bfield_cylindrical, Bfield_cartesian, grad_B_tensor_cartesian, \
          grad_grad_B_tensor_cylindrical, grad_grad_B_tensor_cartesian
      from .calculate_r2 import calculate_r2
      from .calculate_r3 import calculate_r3
      from .mercier import mercier
      from .r_singularity import calculate_r_singularity
      from .plot import plot, plot_boundary, get_boundary, B_fieldline, B_contour, plot_axis
      from .Frenet_to_cylindrical import Frenet_to_cylindrical
      from .to_vmec import to_vmec
      from .util import B_mag
      
      def __init__(self, rc, zs, rs=[], zc=[], nfp=1, etabar=1., sigma0=0., B0=1.,
                   I2=0., sG=1, spsi=1, nphi=61, B2s=0., B2c=0., p2=0., order="r1"):
          """
          Create a quasisymmetric stellarator.
          """
          # First, force {rc, zs, rs, zc} to have the same length, for
          # simplicity.
          nfourier = np.max([len(rc), len(zs), len(rs), len(zc)])
          self.nfourier = nfourier
          self.rc = np.zeros(nfourier)
          self.zs = np.zeros(nfourier)
          self.rs = np.zeros(nfourier)
          self.zc = np.zeros(nfourier)
          self.rc[:len(rc)] = rc
          self.zs[:len(zs)] = zs
          self.rs[:len(rs)] = rs
          self.zc[:len(zc)] = zc
  
          # Force nphi to be odd:
          if np.mod(nphi, 2) == 0:
              nphi += 1
  
          if sG != 1 and sG != -1:
              raise ValueError('sG must be +1 or -1')
          
          if spsi != 1 and spsi != -1:
              raise ValueError('spsi must be +1 or -1')
  
          self.nfp = nfp
          self.etabar = etabar
          self.sigma0 = sigma0
          self.B0 = B0
          self.I2 = I2
          self.sG = sG
          self.spsi = spsi
          self.nphi = nphi
          self.B2s = B2s
          self.B2c = B2c
          self.p2 = p2
          self.order = order
          self.min_R0_threshold = 0.3
          self._set_names()
  
          self.calculate()
  
      def change_nfourier(self, nfourier_new):
          """
          Resize the arrays of Fourier amplitudes. You can either increase
          or decrease nfourier.
          """
          rc_old = self.rc
          rs_old = self.rs
          zc_old = self.zc
          zs_old = self.zs
          index = np.min((self.nfourier, nfourier_new))
          self.rc = np.zeros(nfourier_new)
          self.rs = np.zeros(nfourier_new)
          self.zc = np.zeros(nfourier_new)
          self.zs = np.zeros(nfourier_new)
          self.rc[:index] = rc_old[:index]
          self.rs[:index] = rs_old[:index]
          self.zc[:index] = zc_old[:index]
          self.zs[:index] = zs_old[:index]
          nfourier_old = self.nfourier
          self.nfourier = nfourier_new
          self._set_names()
          # No need to recalculate if we increased the Fourier
          # resolution, only if we decreased it.
          if nfourier_new < nfourier_old:
              self.calculate()
  
      def calculate(self):
          """
          Driver for the main calculations.
          """
          self.init_axis()
          self.solve_sigma_equation()
          self.r1_diagnostics()
          if self.order != 'r1':
              self.calculate_r2()
              if self.order == 'r3':
                  self.calculate_r3()
      
      def get_dofs(self):
          """
          Return a 1D numpy vector of all possible optimizable
          degrees-of-freedom, for simsopt.
          """
          return np.concatenate((self.rc, self.zs, self.rs, self.zc,
                                 np.array([self.etabar, self.sigma0, self.B2s, self.B2c, self.p2, self.I2, self.B0])))
  
      def set_dofs(self, x):
          """
          For interaction with simsopt, set the optimizable degrees of
          freedom from a 1D numpy vector.
          """
          assert len(x) == self.nfourier * 4 + 7
          self.rc = x[self.nfourier * 0 : self.nfourier * 1]
          self.zs = x[self.nfourier * 1 : self.nfourier * 2]
          self.rs = x[self.nfourier * 2 : self.nfourier * 3]
          self.zc = x[self.nfourier * 3 : self.nfourier * 4]
          self.etabar = x[self.nfourier * 4 + 0]
          self.sigma0 = x[self.nfourier * 4 + 1]
          self.B2s = x[self.nfourier * 4 + 2]
          self.B2c = x[self.nfourier * 4 + 3]
          self.p2 = x[self.nfourier * 4 + 4]
          self.I2 = x[self.nfourier * 4 + 5]
          self.B0 = x[self.nfourier * 4 + 6]
          self.calculate()
          logger.info('set_dofs called with x={}. Now iota={}, elongation={}'.format(x, self.iota, self.max_elongation))
          
      def _set_names(self):
          """
          For simsopt, sets the list of names for each degree of freedom.
          """
          names = []
          names += ['rc({})'.format(j) for j in range(self.nfourier)]
          names += ['zs({})'.format(j) for j in range(self.nfourier)]
          names += ['rs({})'.format(j) for j in range(self.nfourier)]
          names += ['zc({})'.format(j) for j in range(self.nfourier)]
          names += ['etabar', 'sigma0', 'B2s', 'B2c', 'p2', 'I2', 'B0']
          self.names = names
  
      @classmethod
      def from_paper(cls, name, **kwargs):
          """
          Get one of the configurations that has been used in our papers.
          Available values for ``name`` are
          ``"r1 section 5.1"``,
          ``"r1 section 5.2"``,
          ``"r1 section 5.3"``,
          ``"r2 section 5.1"``,
          ``"r2 section 5.2"``,
          ``"r2 section 5.3"``,
          ``"r2 section 5.4"``, and
          ``"r2 section 5.5"``.
          These last 5 configurations can also be obtained by specifying an integer 1-5 for ``name``.
          The configurations that begin with ``"r1"`` refer to sections in 
          Landreman, Sengupta, and Plunk, Journal of Plasma Physics 85, 905850103 (2019).
          The configurations that begin with ``"r2"`` refer to sections in 
          Landreman and Sengupta, Journal of Plasma Physics 85, 815850601 (2019).
  
          You can specify any other arguments of the ``Qsc`` constructor
          in ``kwargs``. You can also use ``kwargs`` to override any of
          the properties of the configurations from the papers. For
          instance, you can modify the value of ``etabar`` in the first
          example using
  
          .. code-block::
  
            q = qsc.Qsc.from_paper('r1 section 5.1', etabar=1.1)
          """
  
          def add_default_args(kwargs_old, **kwargs_new):
              """
              Take any key-value arguments in ``kwargs_new`` and treat them as
              defaults, adding them to the dict ``kwargs_old`` only if
              they are not specified there.
              """
              for key in kwargs_new:
                  if key not in kwargs_old:
                      kwargs_old[key] = kwargs_new[key]
  
                      
          if name == "r1 section 5.1":
              """ The configuration from Landreman, Sengupta, Plunk (2019), section 5.1 """
              add_default_args(kwargs, rc=[1, 0.045], zs=[0, -0.045], nfp=3, etabar=-0.9)
                  
          elif name == "r1 section 5.2":
              """ The configuration from Landreman, Sengupta, Plunk (2019), section 5.2 """
              add_default_args(kwargs, rc=[1, 0.265], zs=[0, -0.21], nfp=4, etabar=-2.25)
                  
          elif name == "r1 section 5.3":
              """ The configuration from Landreman, Sengupta, Plunk (2019), section 5.3 """
              add_default_args(kwargs, rc=[1, 0.042], zs=[0, -0.042], zc=[0, -0.025], nfp=3, etabar=-1.1, sigma0=-0.6)
                  
          elif name == "r2 section 5.1" or name == '5.1' or name == 1:
              """ The configuration from Landreman & Sengupta (2019), section 5.1 """
              add_default_args(kwargs, rc=[1, 0.155, 0.0102], zs=[0, 0.154, 0.0111], nfp=2, etabar=0.64, order='r3', B2c=-0.00322)
              
          elif name == "r2 section 5.2" or name == '5.2' or name == 2:
              """ The configuration from Landreman & Sengupta (2019), section 5.2 """
              add_default_args(kwargs, rc=[1, 0.173, 0.0168, 0.00101], zs=[0, 0.159, 0.0165, 0.000985], nfp=2, etabar=0.632, order='r3', B2c=-0.158)
                               
          elif name == "r2 section 5.3" or name == '5.3' or name == 3:
              """ The configuration from Landreman & Sengupta (2019), section 5.3 """
              add_default_args(kwargs, rc=[1, 0.09], zs=[0, -0.09], nfp=2, etabar=0.95, I2=0.9, order='r3', B2c=-0.7, p2=-600000.)
                               
          elif name == "r2 section 5.4" or name == '5.4' or name == 4:
              """ The configuration from Landreman & Sengupta (2019), section 5.4 """
              add_default_args(kwargs, rc=[1, 0.17, 0.01804, 0.001409, 5.877e-05],
                         zs=[0, 0.1581, 0.01820, 0.001548, 7.772e-05], nfp=4, etabar=1.569, order='r3', B2c=0.1348)
                               
          elif name == "r2 section 5.5" or name == '5.5' or name == 5:
              """ The configuration from Landreman & Sengupta (2019), section 5.5 """
              add_default_args(kwargs, rc=[1, 0.3], zs=[0, 0.3], nfp=5, etabar=2.5, sigma0=0.3, I2=1.6, order='r3', B2c=1., B2s=3., p2=-0.5e7)
  
          elif name == "LandremanPaul2021QA" or name == "precise QA":
              """
              A fit of the near-axis model to the quasi-axisymmetric
              configuration in Landreman & Paul, arXiv:2108.03711 (2021).
  
              The fit was performed to the boozmn data using the script
              20200621-01-Extract_B0_B1_B2_from_boozxform
              """
              add_default_args(kwargs,
                               nfp=2,
                               rc=[1.0038581971135636, 0.18400998741139907, 0.021723381370503204, 0.0025968236014410812, 0.00030601568477064874, 3.5540509760304384e-05, 4.102693907398271e-06, 5.154300428457222e-07, 4.8802742243232844e-08, 7.3011320375259876e-09],
                               zs=[0.0, -0.1581148860568176, -0.02060702320552523, -0.002558840496952667, -0.0003061368667524159, -3.600111450532304e-05, -4.174376962124085e-06, -4.557462755956434e-07, -8.173481495049928e-08, -3.732477282851326e-09],
                               B0=1.006541121335688,
                               etabar=-0.6783912804454629,
                               B2c=0.26859318908803137,
                               nphi=99,
                               order='r3')
  
          elif name == "precise QA+well":
              """
              A fit of the near-axis model to the precise quasi-axisymmetric
              configuration from SIMSOPT with magnetic well.
  
              The fit was performed to the boozmn data using the script
              20200621-01-Extract_B0_B1_B2_from_boozxform
              """
              add_default_args(kwargs,
                               nfp=2,
                               rc=[1.0145598919163676, 0.2106377247598754, 0.025469267136340394, 0.0026773601516136727, 0.00021104172568911153, 7.891887175655046e-06, -8.216044358250985e-07, -2.379942694112007e-07, -2.5495108673798585e-08, 1.1679227114962395e-08, 8.961288962248274e-09],
                               zs=[0.0, -0.14607192982551795, -0.021340448470388084, -0.002558983303282255, -0.0002355043952788449, -1.2752278964149462e-05, 3.673356209179739e-07, 9.261098628194352e-08, -7.976283362938471e-09, -4.4204430633540756e-08, -1.6019372369445714e-08],
                               B0=1.0117071561808106,
                               etabar=-0.5064143402495729,
                               B2c=-0.2749140163639202,
                               nphi=99,
                               order='r3')
              
          elif name == "LandremanPaul2021QH" or name == "precise QH":
              """
              A fit of the near-axis model to the quasi-helically symmetric
              configuration in Landreman & Paul, arXiv:2108.03711 (2021).
  
              The fit was performed to the boozmn data using the script
              20211001-02-Extract_B0_B1_B2_from_boozxform
              """
              add_default_args(kwargs,
                               nfp=4,
                               rc=[1.0033608429348413, 0.19993025252481125, 0.03142704185268144, 0.004672593645851904, 0.0005589954792333977, 3.298415996551805e-05, -7.337736061708705e-06, -2.8829857667619663e-06, -4.51059545517434e-07],
                               zs=[0.0, 0.1788824025525348, 0.028597666614604524, 0.004302393796260442, 0.0005283708386982674, 3.5146899855826326e-05, -5.907671188908183e-06, -2.3945326611145963e-06, -6.87509350019021e-07],
                               B0=1.003244143729638,
                               etabar=-1.5002839921360023,
                               B2c=0.37896407142157423,
                               nphi=99,
                               order='r3')
  
          elif name == "precise QH+well":
              """
              A fit of the near-axis model to the precise quasi-helically symmetric
              configuration from SIMSOPT with magnetic well.
  
              The fit was performed to the boozmn data using the script
              20211001-02-Extract_B0_B1_B2_from_boozxform
              """
              add_default_args(kwargs,
                               nfp=4,
                               rc=[1.000474932581454, 0.16345392520298313, 0.02176330066615466, 0.0023779201451133163, 0.00014141976024376502, -1.0595894482659743e-05, -2.9989267970578764e-06, 3.464574408947338e-08],
                               zs=[0.0, 0.12501739099323073, 0.019051257169780858, 0.0023674771227236587, 0.0001865909743321566, -2.2659053455802824e-06, -2.368335337174369e-06, -1.8521248561490157e-08],
                               B0=0.999440074325872,
                               etabar=-1.2115187546668142,
                               B2c=0.6916862277166693,
                               nphi=99,
                               order='r3')
              
          else:
              raise ValueError('Unrecognized configuration name')
  
          return cls(**kwargs)
  
      @classmethod
      def from_cxx(cls, filename):
          """
          Load a configuration from a ``qsc_out.<extension>.nc`` output file
          that was generated by the C++ version of QSC. Almost all the
          data will be taken from the output file, over-writing any
          calculations done in python when the new Qsc object is
          created.
          """
          def to_string(nc_str):
              """ Convert a string from the netcdf binary format to a python string. """
              temp = [c.decode('UTF-8') for c in nc_str]
              return (''.join(temp)).strip()
          
          f = netcdf.netcdf_file(filename, mmap=False)
          nfp = f.variables['nfp'][()]
          nphi = f.variables['nphi'][()]
          rc = f.variables['R0c'][()]
          rs = f.variables['R0s'][()]
          zc = f.variables['Z0c'][()]
          zs = f.variables['Z0s'][()]
          I2 = f.variables['I2'][()]
          B0 = f.variables['B0'][()]
          spsi = f.variables['spsi'][()]
          sG = f.variables['sG'][()]
          etabar = f.variables['eta_bar'][()]
          sigma0 = f.variables['sigma0'][()]
          order_r_option = to_string(f.variables['order_r_option'][()])
          if order_r_option == 'r2.1':
              order_r_option = 'r3'
          if order_r_option == 'r1':
              p2 = 0.0
              B2c = 0.0
              B2s = 0.0
          else:
              p2 = f.variables['p2'][()]
              B2c = f.variables['B2c'][()]
              B2s = f.variables['B2s'][()]
  
          q = cls(nfp=nfp, nphi=nphi, rc=rc, rs=rs, zc=zc, zs=zs,
                  B0=B0, sG=sG, spsi=spsi,
                  etabar=etabar, sigma0=sigma0, I2=I2, p2=p2, B2c=B2c, B2s=B2s, order=order_r_option)
          
          def read(name, cxx_name=None):
              if cxx_name is None: cxx_name = name
              setattr(q, name, f.variables[cxx_name][()])
  
          [read(v) for v in ['R0', 'Z0', 'R0p', 'Z0p', 'R0pp', 'Z0pp', 'R0ppp', 'Z0ppp',
                             'sigma', 'curvature', 'torsion', 'X1c', 'Y1c', 'Y1s', 'elongation']]
          if order_r_option != 'r1':
              [read(v) for v in ['X20', 'X2c', 'X2s', 'Y20', 'Y2c', 'Y2s', 'Z20', 'Z2c', 'Z2s', 'B20']]
              if order_r_option != 'r2':
                  [read(v) for v in ['X3c1', 'Y3c1', 'Y3s1']]
                      
          f.close()
          return q
          
      def min_R0_penalty(self):
          """
          This function can be used in optimization to penalize situations
          in which min(R0) < min_R0_constraint.
          """
          return np.max((0, self.min_R0_threshold - self.min_R0)) ** 2
          
  
apis: numpy.zeros numpy.mod numpy.max numpy.min numpy.array scipy.io.netcdf.netcdf_file