Mgrid improvements

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 @@
         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 @@
             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 @@
         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 @@
                 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 @@
             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 @@
             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 @@
                     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

tests/field/test_magneticfields.py

@@ -1041,14 +1041,14 @@ def test_to_vtk(self):
         bs = BiotSavart(coils)
         bs.to_vtk('/tmp/bfield')
 
-    def test_to_mgrid(self):
+    def subtest_to_mgrid(self, include_potential):
         curves, currents, ma = get_ncsx_data()
         nfp = 3
         coils = coils_via_symmetries(curves, currents, nfp, True)
         bs = BiotSavart(coils)
         with tempfile.TemporaryDirectory() as tmpdir:
             filename = Path(tmpdir) / "mgrid.bfield.nc"
-            bs.to_mgrid(filename, nfp=nfp)
+            bs.to_mgrid(filename, nfp=nfp, include_potential=include_potential)
 
             # Compare the B data in the file to a separate evaluation here
             with netcdf_file(filename, mmap=False) as f:
@@ -1068,6 +1068,13 @@ def test_to_mgrid(self):
                 assert Br.shape == (nphi, nz, nr)
                 assert Bphi.shape == (nphi, nz, nr)
                 assert Bz.shape == (nphi, nz, nr)
+                if include_potential:
+                    Ar = f.variables["ar_001"][()]
+                    Aphi = f.variables["ap_001"][()]
+                    Az = f.variables["az_001"][()]
+                    assert Ar.shape == (nphi, nz, nr)
+                    assert Aphi.shape == (nphi, nz, nr)
+                    assert Az.shape == (nphi, nz, nr)
                 r = np.linspace(rmin, rmax, nr)
                 phi = np.linspace(0, 2 * np.pi / nfp, nphi, endpoint=False)
                 z = np.linspace(zmin, zmax, nz)
@@ -1078,6 +1085,15 @@ def test_to_mgrid(self):
                             np.testing.assert_allclose(Br[jphi, jz, jr], bs.B_cyl()[0, 0])
                             np.testing.assert_allclose(Bphi[jphi, jz, jr], bs.B_cyl()[0, 1])
                             np.testing.assert_allclose(Bz[jphi, jz, jr], bs.B_cyl()[0, 2])
+                            if include_potential:
+                                np.testing.assert_allclose(Ar[jphi, jz, jr], bs.A_cyl()[0, 0])
+                                np.testing.assert_allclose(Aphi[jphi, jz, jr], bs.A_cyl()[0, 1])
+                                np.testing.assert_allclose(Az[jphi, jz, jr], bs.A_cyl()[0, 2])
+
+    def test_to_mgrid(self):
+        for include_potential in [True, False]:
+            with self.subTest(include_potential=include_potential):
+                self.subtest_to_mgrid(include_potential)
 
     def test_poloidal_field(self):
         B0 = 1.1

tests/field/test_mgrid.py

@@ -18,7 +18,7 @@
 
 TEST_DIR = (Path(__file__).parent / ".." / "test_files").resolve()
 test_file = TEST_DIR / 'mgrid.pnas-qa-test-lowres-standard.nc'
-
+test_file2 = TEST_DIR / 'mgrid_ncsx_lowres_test.nc'
 
 class Testing(unittest.TestCase):
 
@@ -36,6 +36,13 @@ def test_from_file(self):
         assert mgrid.br_arr.shape == (1, 6, 11, 11)
         assert mgrid.br[0, 0, 0] == -1.0633399551863771  # -0.9946816978184079
 
+        mgrid = MGrid.from_file(test_file2)
+        assert mgrid.rmin == 1.0
+        assert mgrid.bvec.shape == (10, 12, 4, 3)
+        assert mgrid.bz[1, 1, 1] == -1.012339153040808
+        assert mgrid.ap[1, 1, 1] == -0.3719177477496187
+
+
     def test_add_field_cylinder(self):
         N_points = 5
         br = np.ones(N_points)