Merge pull request #185 from hiddenSymmetries/Vmec_without_vmec

Initialize Vmec from a wout file without vmec or mpi

src/simsopt/mhd/boozer.py

@@ -26,12 +26,7 @@
     booz_xform = None
     logger.debug(str(e))
 
-if MPI is not None:
-    try:
-        from .vmec import Vmec
-    except ImportError as e:
-        Vmec = None
-        logger.debug(str(e))
+from .vmec import Vmec
 
 from .._core.graph_optimizable import Optimizable
 

src/simsopt/mhd/vmec.py

@@ -177,13 +177,6 @@ def __init__(self,
                  keep_all_files: bool = False,
                  ntheta=50,
                  nphi=50):
-        if MPI is None:
-            raise RuntimeError("mpi4py needs to be installed for running VMEC")
-        if vmec is None:
-            raise RuntimeError(
-                "Running VMEC from simsopt requires VMEC python extension. "
-                "Install the VMEC python extension from "
-                "https://https://github.com/hiddenSymmetries/VMEC2000")
 
         if filename is None:
             # Read default input file, which should be in the same
@@ -202,23 +195,34 @@ def __init__(self,
         else:
             raise ValueError('Invalid filename')
 
+        self.wout = Struct()
+
         # Get MPI communicator:
-        if mpi is None:
+        if (mpi is None and MPI is not None):
             self.mpi = MpiPartition(ngroups=1)
         else:
             self.mpi = mpi
-        comm = self.mpi.comm_groups
-        self.fcomm = comm.py2f()
-
-        self.ictrl = np.zeros(5, dtype=np.int32)
-        self.iter = -1
-        self.keep_all_files = keep_all_files
-        self.files_to_delete = []
-        self.wout = Struct()
-        self.indata = vmec.vmec_input  # Shorthand
-        vi = vmec.vmec_input  # Shorthand
 
         if self.runnable:
+            if MPI is None:
+                raise RuntimeError("mpi4py needs to be installed for running VMEC")
+            if vmec is None:
+                raise RuntimeError(
+                    "Running VMEC from simsopt requires VMEC python extension. "
+                    "Install the VMEC python extension from "
+                    "https://https://github.com/hiddenSymmetries/VMEC2000")
+
+            comm = self.mpi.comm_groups
+            self.fcomm = comm.py2f()
+
+            self.ictrl = np.zeros(5, dtype=np.int32)
+            self.iter = -1
+            self.keep_all_files = keep_all_files
+            self.files_to_delete = []
+
+            self.indata = vmec.vmec_input  # Shorthand
+            vi = vmec.vmec_input  # Shorthand
+
             self.ictrl[0] = restart_flag + readin_flag
             self.ictrl[1] = 0  # ierr
             self.ictrl[2] = 0  # numsteps
@@ -264,15 +268,14 @@ def __init__(self,
                         self._boundary.zc[m, n + vi.ntor] = vi.zbc[101 + n, m]
             self._boundary.local_full_x = self._boundary.get_dofs()
 
-            # Handle a few variables that are not Parameters:
             self.need_to_run_code = True
-
         else:
             # Initialized from a wout file, so not runnable.
             self._boundary = SurfaceRZFourier.from_wout(filename)
             self.output_file = filename
             self.load_wout()
 
+        # Handle a few variables that are not Parameters:
         x0 = self.get_dofs()
         fixed = np.full(len(x0), True)
         names = ['delt', 'tcon0', 'phiedge', 'curtor', 'gamma']
@@ -299,17 +302,22 @@ def boundary(self, boundary):
             self.need_to_run_code = True
 
     def get_dofs(self):
-        return np.array([self.indata.delt, self.indata.tcon0,
-                         self.indata.phiedge, self.indata.curtor,
-                         self.indata.gamma])
+        if not self.runnable:
+            # Use default values from vmec_input
+            return np.array([1, 1, 1, 0, 0])
+        else:
+            return np.array([self.indata.delt, self.indata.tcon0,
+                             self.indata.phiedge, self.indata.curtor,
+                             self.indata.gamma])
 
     def set_dofs(self, x):
-        self.need_to_run_code = True
-        self.indata.delt = x[0]
-        self.indata.tcon0 = x[1]
-        self.indata.phiedge = x[2]
-        self.indata.curtor = x[3]
-        self.indata.gamma = x[4]
+        if self.runnable:
+            self.need_to_run_code = True
+            self.indata.delt = x[0]
+            self.indata.tcon0 = x[1]
+            self.indata.phiedge = x[2]
+            self.indata.curtor = x[3]
+            self.indata.gamma = x[4]
 
     def recompute_bell(self, parent=None):
         self.need_to_run_code = True
@@ -585,7 +593,6 @@ def vacuum_well(self):
         half mesh, we extrapolate by half of a radial grid point to s
         = 0 and 1.
         """
-
         self.run()
 
         # gmnc is on the half mesh, so drop the 0th radial entry:

src/simsopt/mhd/vmec_diagnostics.py

@@ -351,7 +351,8 @@ def J(self):
         """
         Computes the quantity :math:`J` described in the class definition.
         """
-        self.vmec.need_to_run_code = True
+        # if self.vmec.runnable:
+        #     self.vmec.need_to_run_code = True
         self.vmec.run()
         return 0.5 * np.sum((self.vmec.wout.iotas[1::]
                              - self.iota_target(self.vmec.s_half_grid))**2) * self.vmec.ds
@@ -461,7 +462,6 @@ def J(self):
         """
         Computes the quantity :math:`J` described in the class definition.
         """
-        self.vmec.need_to_run_code = True
         self.vmec.run()
         return np.sum(self.weight_function(self.vmec.s_half_grid) * self.vmec.wout.iotas[1:]) \
             / np.sum(self.weight_function(self.vmec.s_half_grid))
@@ -580,8 +580,6 @@ def J(self):
         """
         Computes the quantity :math:`J` described in the class definition.
         """
-
-        self.vmec.need_to_run_code = True
         self.vmec.run()
         return np.sum((self.weight_function1(self.vmec.s_half_grid)-self.weight_function2(self.vmec.s_half_grid)) * self.vmec.wout.vp[1:]) \
             / np.sum((self.weight_function1(self.vmec.s_half_grid)+self.weight_function2(self.vmec.s_half_grid)) * self.vmec.wout.vp[1:])

tests/mhd/test_vmec.py

@@ -17,14 +17,99 @@
 from simsopt.objectives.graph_least_squares import LeastSquaresProblem
 from simsopt.geo.surfacerzfourier import SurfaceRZFourier
 
-if (MPI is not None) and vmec_found:
-    from simsopt.mhd.vmec import Vmec
+from simsopt.mhd.vmec import Vmec
+
 from . import TEST_DIR
 
 logger = logging.getLogger(__name__)
 #logging.basicConfig(level=logging.DEBUG)
 
 
+class InitializedFromWout(unittest.TestCase):
+    def test_vacuum_well(self):
+        """
+        Test the calculation of magnetic well. This is done by comparison
+        to a high-aspect-ratio configuration, in which case the
+        magnetic well can be computed analytically by the method in
+        Landreman & Jorge, J Plasma Phys 86, 905860510 (2020). The
+        specific configuration considered is the one of section 5.4 in
+        Landreman & Sengupta, J Plasma Phys 85, 815850601 (2019), also
+        considered in the 2020 paper. We increase the mean field B0 to
+        2T in that configuration to make sure all the factors of B0
+        are correct.
+
+        """
+        filename = os.path.join(TEST_DIR, 'wout_LandremanSengupta2019_section5.4_B2_A80_reference.nc')
+        vmec = Vmec(filename)
+
+        well_vmec = vmec.vacuum_well()
+        # Let psi be the toroidal flux divided by (2 pi)
+        abs_psi_a = np.abs(vmec.wout.phi[-1]) / (2 * np.pi)
+
+        # Data for this configuration from the near-axis construction code qsc:
+        # https://github.com/landreman/pyQSC
+        # or
+        # https://github.com/landreman/qsc
+        B0 = 2.0
+        G0 = 2.401752071286676
+        d2_volume_d_psi2 = 25.3041656424299
+
+        # See also "20210504-01 Computing magnetic well from VMEC.docx" by MJL
+        well_analytic = -abs_psi_a * B0 * B0 * d2_volume_d_psi2 / (4 * np.pi * np.pi * np.abs(G0))
+
+        logger.info(f'well_vmec: {well_vmec}  well_analytic: {well_analytic}')
+        np.testing.assert_allclose(well_vmec, well_analytic, rtol=2e-2, atol=0)
+
+    def test_iota(self):
+        """
+        Test the functions related to iota.
+        """
+        filename = os.path.join(TEST_DIR, 'wout_LandremanSengupta2019_section5.4_B2_A80_reference.nc')
+        vmec = Vmec(filename)
+
+        iota_axis = vmec.iota_axis()
+        iota_edge = vmec.iota_edge()
+        mean_iota = vmec.mean_iota()
+        mean_shear = vmec.mean_shear()
+        # These next 2 lines are different ways the mean iota and
+        # shear could be defined. They are not mathematically
+        # identical to mean_iota() and mean_shear(), but they should
+        # be close.
+        mean_iota_alt = (iota_axis + iota_edge) * 0.5
+        mean_shear_alt = iota_edge - iota_axis
+        logger.info(f"iota_axis: {iota_axis}, iota_edge: {iota_edge}")
+        logger.info(f"    mean_iota: {mean_iota},     mean_shear: {mean_shear}")
+        logger.info(f"mean_iota_alt: {mean_iota_alt}, mean_shear_alt: {mean_shear_alt}")
+
+        self.assertAlmostEqual(mean_iota, mean_iota_alt, places=3)
+        self.assertAlmostEqual(mean_shear, mean_shear_alt, places=3)
+
+    def test_error_on_rerun(self):
+        """
+        If a vmec object is initialized from a wout file, and if the dofs
+        are then changed, vmec output functions should raise an
+        exception.
+        """
+        filename = os.path.join(TEST_DIR, 'wout_li383_low_res_reference.nc')
+        vmec = Vmec(filename)
+        iota = vmec.mean_iota()
+        vmec.boundary.set_rc(1, 0, 2.0)
+        with self.assertRaises(RuntimeError):
+            iota2 = vmec.mean_iota()
+
+
+@unittest.skipIf((MPI is not None) and (vmec_found), "Interface to MPI and VMEC found")
+class VmecTestsWithoutMPIorvmec(unittest.TestCase):
+    def test_runnable_raises(self):
+        """
+        Test that RuntimeError is raised if Vmec initialized from
+        input file without vmec or MPI
+        """
+        from simsopt.mhd.vmec import Vmec
+        with self.assertRaises(RuntimeError):
+            v = Vmec()
+
+
 @unittest.skipIf((MPI is None) or (not vmec_found), "Valid Python interface to VMEC not found")
 class VmecTests(unittest.TestCase):
     def test_init_defaults(self):
@@ -178,19 +263,6 @@ def test_2_init_methods(self):
             np.testing.assert_allclose(iota1, iota2, atol=1e-10)
             np.testing.assert_allclose(bmnc1, bmnc2, atol=1e-10)
 
-    def test_error_on_rerun(self):
-        """
-        If a vmec object is initialized from a wout file, and if the dofs
-        are then changed, vmec output functions should raise an
-        exception.
-        """
-        filename = os.path.join(TEST_DIR, 'wout_li383_low_res_reference.nc')
-        vmec = Vmec(filename)
-        iota = vmec.mean_iota()
-        vmec.boundary.set_rc(1, 0, 2.0)
-        with self.assertRaises(RuntimeError):
-            iota2 = vmec.mean_iota()
-
     def test_vmec_failure(self):
         """
         Verify that failures of VMEC are correctly caught and represented
@@ -256,64 +328,6 @@ def test_vmec_failure(self):
             print(f)
             np.testing.assert_allclose(f, correct_f, rtol=0.1)
 
-    def test_vacuum_well(self):
-        """
-        Test the calculation of magnetic well. This is done by comparison
-        to a high-aspect-ratio configuration, in which case the
-        magnetic well can be computed analytically by the method in
-        Landreman & Jorge, J Plasma Phys 86, 905860510 (2020). The
-        specific configuration considered is the one of section 5.4 in
-        Landreman & Sengupta, J Plasma Phys 85, 815850601 (2019), also
-        considered in the 2020 paper. We increase the mean field B0 to
-        2T in that configuration to make sure all the factors of B0
-        are correct.
-
-        """
-        filename = os.path.join(TEST_DIR, 'wout_LandremanSengupta2019_section5.4_B2_A80_reference.nc')
-        vmec = Vmec(filename)
-
-        well_vmec = vmec.vacuum_well()
-        # Let psi be the toroidal flux divided by (2 pi)
-        abs_psi_a = np.abs(vmec.wout.phi[-1]) / (2 * np.pi)
-
-        # Data for this configuration from the near-axis construction code qsc:
-        # https://github.com/landreman/pyQSC
-        # or
-        # https://github.com/landreman/qsc
-        B0 = 2.0
-        G0 = 2.401752071286676
-        d2_volume_d_psi2 = 25.3041656424299
-
-        # See also "20210504-01 Computing magnetic well from VMEC.docx" by MJL
-        well_analytic = -abs_psi_a * B0 * B0 * d2_volume_d_psi2 / (4 * np.pi * np.pi * np.abs(G0))
-
-        logger.info(f'well_vmec: {well_vmec}  well_analytic: {well_analytic}')
-        np.testing.assert_allclose(well_vmec, well_analytic, rtol=2e-2, atol=0)
-
-    def test_iota(self):
-        """
-        Test the functions related to iota.
-        """
-        filename = os.path.join(TEST_DIR, 'wout_LandremanSengupta2019_section5.4_B2_A80_reference.nc')
-        vmec = Vmec(filename)
-
-        iota_axis = vmec.iota_axis()
-        iota_edge = vmec.iota_edge()
-        mean_iota = vmec.mean_iota()
-        mean_shear = vmec.mean_shear()
-        # These next 2 lines are different ways the mean iota and
-        # shear could be defined. They are not mathematically
-        # identical to mean_iota() and mean_shear(), but they should
-        # be close.
-        mean_iota_alt = (iota_axis + iota_edge) * 0.5
-        mean_shear_alt = iota_edge - iota_axis
-        logger.info(f"iota_axis: {iota_axis}, iota_edge: {iota_edge}")
-        logger.info(f"    mean_iota: {mean_iota},     mean_shear: {mean_shear}")
-        logger.info(f"mean_iota_alt: {mean_iota_alt}, mean_shear_alt: {mean_shear_alt}")
-
-        self.assertAlmostEqual(mean_iota, mean_iota_alt, places=3)
-        self.assertAlmostEqual(mean_shear, mean_shear_alt, places=3)
-
     #def test_stellopt_scenarios_1DOF_circularCrossSection_varyR0_targetVolume(self):
         """
         This script implements the "1DOF_circularCrossSection_varyR0_targetVolume"