Merge branch 'master' of github.com:hiddenSymmetries/simsopt

.github/workflows/extensive_test.yml

@@ -98,7 +98,7 @@ jobs:
     - name: Install python dependencies
       run: |
            sudo apt-get install graphviz graphviz-dev
-           pip install wheel numpy scipy f90nml h5py scikit-build cmake qsc sympy pyevtk matplotlib ninja plotly networkx pygraphviz
+           pip install wheel numpy scipy f90nml h5py scikit-build cmake qsc sympy pyevtk matplotlib ninja plotly networkx pygraphviz ground bentley_ottmann
 
     - name: Install booz_xform
       if: contains(matrix.packages, 'vmec') || contains(matrix.packages, 'all')

pyproject.toml

@@ -71,7 +71,7 @@ dynamic = ["version"]
 [project.optional-dependencies]
 SPEC = ["py_spec>=3.0.1", "pyoculus>=0.1.1", "h5py>=3.1.0"]
 MPI = ["mpi4py>=3.0.3"]
-VIS = ["vtk >= 8.1.2", "PyQt5", "mayavi", "plotly", "networkx"]
+VIS = ["vtk >= 8.1.2", "PyQt5", "plotly", "networkx"]
 DOCS = ["sphinx", "sphinx-rtd-theme"]
 
 [project.urls]

src/simsopt/field/magneticfield.py

@@ -92,7 +92,8 @@ def to_vtk(self, filename, nr=10, nphi=10, nz=10, rmin=1.0, rmax=2.0, zmin=-0.5,
         contig = np.ascontiguousarray
         gridToVTK(filename, X, Y, Z, pointData={"B": (contig(vals[..., 0]), contig(vals[..., 1]), contig(vals[..., 2]))})
 
-    def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5, zmax=0.5, nfp=1):
+    def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5, zmax=0.5, nfp=1, 
+                 include_potential=False):
         """Export the field to the mgrid format for free boundary calculations.
 
         The field data is represented as a single "current group". For
@@ -111,6 +112,7 @@ def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5
             zmin: Minimum value of z for the grid.
             zmax: Maximum value of z for the grid.
             nfp: Number of field periods.
+            include_potential: Boolean to include the vector potential A. Defaults to false.
         """
 
         rs = np.linspace(rmin, rmax, nr, endpoint=True)
@@ -133,11 +135,23 @@ def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5
         br_3 = br.reshape((nphi, nz, nr))
         bp_3 = bp.reshape((nphi, nz, nr))
         bz_3 = bz.reshape((nphi, nz, nr))
+
+        if include_potential:
+            A = self.A_cyl()
+            # shape the potential components
+            ar, ap, az = A.T
+            ar_3 = ar.reshape((nphi, nz, nr))
+            ap_3 = ap.reshape((nphi, nz, nr))
+            az_3 = az.reshape((nphi, nz, nr))
 
         mgrid = MGrid(nfp=nfp,
                       nr=nr, nz=nz, nphi=nphi,
                       rmin=rmin, rmax=rmax, zmin=zmin, zmax=zmax)
-        mgrid.add_field_cylindrical(br_3, bp_3, bz_3, name='simsopt_coils')  
+        if include_potential:
+            mgrid.add_field_cylindrical(br_3, bp_3, bz_3, ar=ar_3, ap=ap_3, az=az_3, name='simsopt_coils')  
+        else:
+            mgrid.add_field_cylindrical(br_3, bp_3, bz_3, name='simsopt_coils')  
+
 
         mgrid.write(filename)
 

src/simsopt/field/mgrid.py

@@ -69,17 +69,21 @@ def __init__(self,  # fname='temp', #binary=False,
         self.bz_arr = [] 
         self.bp_arr = [] 
 
-    def add_field_cylindrical(self, br, bp, bz, name=None):
+        self.ar_arr = []
+        self.az_arr = []
+        self.ap_arr = []
+
+    def add_field_cylindrical(self, br, bp, bz, ar=None, ap=None, az=None, name=None):
         '''
-        This function saves the vector field B.
-        B is defined by cylindrical components.
+        This function saves the vector field B, and (optionally) the vector potential A, to the Mgrid object.
+        B and A are provided on a tensor product grid in cylindrical components.
 
         The Mgrid array assumes B is sampled linearly first in r, then z, and last phi.
         Python arrays use the opposite convention such that B[0] gives a (r,z) square at const phi
         and B[0,0] gives a radial line and const phi and z.
 
-        It is assumed that the (br,bp,bz) inputs for this function is already in a
-        (nphi, nz, nr) shaped array.
+        It is assumed that each of the inputs ``br``, ``bp``, and ``bz`` for this function are already
+        arrays of shape ``(nphi, nz, nr)``. The same is true for ``ar``, ``ap``, and ``az`` if they are provided.
 
         This function may be called once for each coil group, 
         to save sets of fields that can be scaled using EXTCUR in VMEC.
@@ -88,13 +92,24 @@ def add_field_cylindrical(self, br, bp, bz, name=None):
             br: the radial component of B field
             bp: the azimuthal component of B field
             bz: the axial component of B field
+            ar: the radial component of the vector potential A
+            ap: the azimuthal component of the vector potential A
+            az: the axial component of the vector potential A
             name: Name of the coil group
         '''
 
         # appending B field to an array for all coil groups.
         self.br_arr.append(br)
         self.bz_arr.append(bz)
         self.bp_arr.append(bp)
+
+        # add potential
+        if ar is not None:
+            self.ar_arr.append(ar)
+        if ap is not None:
+            self.az_arr.append(az)
+        if az is not None:
+            self.ap_arr.append(ap)
 
         # add coil label
         if (name is None):
@@ -104,9 +119,10 @@ def add_field_cylindrical(self, br, bp, bz, name=None):
         self.coil_names.append(label)
         self.n_ext_cur = self.n_ext_cur + 1
 
+
         # TO-DO: this function could check for size consistency, between different fields, and for the (nr,nphi,nz) settings of a given instance
 
-    def write(self, filename):
+    def write(self, filename): 
         '''
         Export class data as a netCDF binary.
 
@@ -176,6 +192,17 @@ def write(self, filename):
                 var_br_001[:, :, :] = self.br_arr[j]
                 var_bz_001[:, :, :] = self.bz_arr[j]
                 var_bp_001[:, :, :] = self.bp_arr[j]
+
+                #If the potential value is not an empty cell, then include it
+                if len(self.ar_arr) > 0 :
+                    var_ar_001 = ds.createVariable('ar'+tag, 'f8', ('phi', 'zee', 'rad'))
+                    var_ar_001[:, :, :] = self.ar_arr[j]
+                if len(self.ap_arr) > 0 :
+                    var_ap_001 = ds.createVariable('ap'+tag, 'f8', ('phi', 'zee', 'rad'))
+                    var_ap_001[:, :, :] = self.ap_arr[j]
+                if len(self.az_arr) > 0 :
+                    var_az_001 = ds.createVariable('az'+tag, 'f8', ('phi', 'zee', 'rad'))
+                    var_az_001[:, :, :] = self.az_arr[j]
 
     @classmethod 
     def from_file(cls, filename):
@@ -203,7 +230,10 @@ def from_file(cls, filename):
             mgrid.filename = filename
             mgrid.n_ext_cur = int(f.variables['nextcur'].getValue())
             coil_data = f.variables['coil_group'][:]
-            mgrid.coil_names = [_unpack(coil_data[j]) for j in range(mgrid.n_ext_cur)] 
+            if len(f.variables['coil_group'].dimensions) == 2:
+                mgrid.coil_names = [_unpack(coil_data[j]) for j in range(mgrid.n_ext_cur)] 
+            else:
+                mgrid.coil_names = [_unpack(coil_data)]
 
             mgrid.mode = f.variables['mgrid_mode'][:][0].decode()
             mgrid.raw_coil_current = np.array(f.variables['raw_coil_cur'][:])
@@ -212,6 +242,14 @@ def from_file(cls, filename):
             bp_arr = []
             bz_arr = []
 
+            ar_arr = []
+            ap_arr = []
+            az_arr = []
+
+            ar = []
+            ap = []
+            az = []
+
             nextcur = mgrid.n_ext_cur
             for j in range(nextcur):
                 idx = '{:03d}'.format(j+1)
@@ -228,24 +266,67 @@ def from_file(cls, filename):
                     bp_arr.append(bp)
                     bz_arr.append(bz)
 
+                #check for potential in mgrid file
+                if 'ar_'+idx in f.variables:
+                    ar = f.variables['ar_'+idx][:]
+                    if mgrid.mode == 'S':
+                        ar_arr.append(ar * mgrid.raw_coil_current[j])
+                    else:
+                        ar_arr.append(ar)
+
+                if 'ap_'+idx in f.variables:
+                    ap = f.variables['ap_'+idx][:]
+                    if mgrid.mode == 'S':
+                        ap_arr.append(ap * mgrid.raw_coil_current[j])
+                    else:
+                        ap_arr.append(ap)
+
+                if 'az_'+idx in f.variables:
+                    az = f.variables['az_'+idx][:]
+                    if mgrid.mode == 'S':
+                        az_arr.append(az * mgrid.raw_coil_current[j])
+                    else:
+                        az_arr.append(az)
+
             mgrid.br_arr = np.array(br_arr)
             mgrid.bp_arr = np.array(bp_arr)
             mgrid.bz_arr = np.array(bz_arr)
 
+            mgrid.ar_arr = np.array(ar_arr)
+            mgrid.ap_arr = np.array(ap_arr)
+            mgrid.az_arr = np.array(az_arr)
+
             # sum over coil groups
             if nextcur > 1:
                 br = np.sum(br_arr, axis=0)
                 bp = np.sum(bp_arr, axis=0)
                 bz = np.sum(bz_arr, axis=0)
+
+                if len(mgrid.ar_arr) > 0:
+                    ar = np.sum(ar_arr, axis=0)
+                if len(mgrid.ap_arr) > 0:
+                    ap = np.sum(ap_arr, axis=0)
+                if len(mgrid.az_arr) > 0:
+                    az = np.sum(az_arr, axis=0)
             else:
                 br = br_arr[0]
                 bp = bp_arr[0]
                 bz = bz_arr[0]
+                if len(mgrid.ar_arr) > 0:
+                    ar = ar_arr[0]
+                if len(mgrid.ap_arr) > 0:
+                    ap = ap_arr[0]
+                if len(mgrid.az_arr) > 0:
+                    az = az_arr[0]
 
             mgrid.br = br
             mgrid.bp = bp
             mgrid.bz = bz
 
+            mgrid.ar = ar
+            mgrid.ap = ap
+            mgrid.az = az
+
             mgrid.bvec = np.transpose([br, bp, bz])
 
         return mgrid