minor tweaks to mgrid test and docstrings

src/simsopt/field/magneticfield.py

@@ -112,7 +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 magnetic potential, defaults to false.
+            include_potential: Boolean to include the vector potential A. Defaults to false.
         """
 
         rs = np.linspace(rmin, rmax, nr, endpoint=True)
@@ -138,7 +138,7 @@ def to_mgrid(self, filename, nr=10, nphi=4, nz=12, rmin=1.0, rmax=2.0, zmin=-0.5
 
         if include_potential:
             A = self.A_cyl()
-            # shape the Potential components
+            # shape the potential components
             ar, ap, az = A.T
             ar_3 = ar.reshape((nphi, nz, nr))
             ap_3 = ap.reshape((nphi, nz, nr))

src/simsopt/field/mgrid.py

@@ -75,15 +75,15 @@ def __init__(self,  # fname='temp', #binary=False,
 
     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.
@@ -92,9 +92,9 @@ def add_field_cylindrical(self, br, bp, bz, ar=None, ap=None, az=None, 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 magnetic potential
-            ap: the azimuthal component of the magnetic potential
-            az: the axial component of the magnetic potential
+            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
         '''
 

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, include_potential=True):
+    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, include_potential=True)
+            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:
@@ -1062,18 +1062,19 @@ def test_to_mgrid(self, include_potential=True):
                 Br = f.variables["br_001"][()]
                 Bphi = f.variables["bp_001"][()]
                 Bz = f.variables["bz_001"][()]
-                Ar = f.variables["ar_001"][()]
-                Aphi = f.variables["ap_001"][()]
-                Az = f.variables["az_001"][()]
                 assert nr == f.dimensions["rad"]
                 assert nphi == f.dimensions["phi"]
                 assert nz == f.dimensions["zee"]
                 assert Br.shape == (nphi, nz, nr)
                 assert Bphi.shape == (nphi, nz, nr)
                 assert Bz.shape == (nphi, nz, nr)
-                assert Ar.shape == (nphi, nz, nr)
-                assert Aphi.shape == (nphi, nz, nr)
-                assert Az.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)
@@ -1084,9 +1085,15 @@ def test_to_mgrid(self, include_potential=True):
                             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])
-                            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])
+                            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

@@ -56,7 +56,7 @@ def test_add_field_cylinder(self):
         assert mgrid.coil_names[0] == '__________test_coil___________'
         assert np.allclose(mgrid.br_arr[0], br)
 
-    def test_write(self): 
+    def test_write(self):
         mgrid = MGrid.from_file(test_file)
         with tempfile.TemporaryDirectory() as tmpdir:
             filename = Path(tmpdir) / 'mgrid.test.nc'