Merge pull request #233 from hiddenSymmetries/tq/mgrid

Tq/mgrid

docs/source/simsopt._core.rst

@@ -14,7 +14,7 @@ simsopt.\_core.derivative module
    :private-members:
 
 simsopt.\_core.descriptor module
---------------------------
+--------------------------------
 
 .. automodule:: simsopt._core.descriptor
    :members:

docs/source/simsopt.field.rst

@@ -44,6 +44,14 @@ simsopt.field.magneticfieldclasses module
    :undoc-members:
    :show-inheritance:
 
+simsopt.field.mgrid module
+--------------------------
+
+.. automodule:: simsopt.field.mgrid
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 simsopt.field.normal\_field module
 ----------------------------------
 

examples/2_Intermediate/free_boundary_vmec.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+
+"""
+This example shows how to take coils in simsopt, write an mgrid file, and then
+run free-boundary vmec. No optimization of the plasma or coil shapes is
+performed in this example.
+
+You can run this example with one or multiple MPI processes.
+"""
+
+from pathlib import Path
+import numpy as np
+from simsopt.configs import get_w7x_data
+from simsopt.field import BiotSavart, coils_via_symmetries
+from simsopt.mhd import Vmec
+from simsopt.util import MpiPartition
+
+nfp = 5
+
+# Load in some coils
+curves, currents, magnetic_axis = get_w7x_data()
+coils = coils_via_symmetries(curves, currents, nfp, True)
+bs = BiotSavart(coils)
+
+# Number of grid points in the toroidal angle:
+nphi = 24
+
+mgrid_file = "mgrid.w7x.nc"
+bs.to_mgrid(
+    mgrid_file,
+    nr=64,
+    nz=65,
+    nphi=nphi,
+    rmin=4.5,
+    rmax=6.3,
+    zmin=-1.0,
+    zmax=1.0,
+    nfp=nfp,
+)
+
+# Create a VMEC object from an input file:
+TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve()
+input_file = str(TEST_DIR / "input.W7-X_standard_configuration")
+
+mpi = MpiPartition(1)
+vmec = Vmec(input_file, mpi=mpi)
+
+# That input file was for fixed-boundary. We need to change some of 
+# the vmec input parameters for a free-boundary calculation:
+vmec.indata.lfreeb = True
+vmec.indata.mgrid_file = mgrid_file
+vmec.indata.nzeta = nphi
+# All the coils are written into a single "current group", so we only need to
+# set a single entry in vmec's "extcur" array:
+vmec.indata.extcur[0] = 1.0
+
+# Lower the resolution, so the example runs faster:
+vmec.indata.mpol = 6
+vmec.indata.ntor = 6
+vmec.indata.ns_array[2] = 0
+ftol = 1e-10
+vmec.indata.ftol_array[1] = ftol
+
+vmec.run()
+
+assert vmec.wout.fsql < ftol
+assert vmec.wout.fsqr < ftol
+assert vmec.wout.fsqz < ftol
+assert vmec.wout.ier_flag == 0
+np.testing.assert_allclose(vmec.wout.volume_p, 28.6017247168422, rtol=0.01)
+np.testing.assert_allclose(vmec.wout.rmnc[0, 0], 5.561878306096512, rtol=0.01)
+np.testing.assert_allclose(vmec.wout.bmnc[0, 1], 2.78074392223658, rtol=0.01)

examples/run_vmec_examples

@@ -39,6 +39,7 @@ mpiexec $MPI_OPTIONS -n 2 ./2_Intermediate/constrained_optimization.py
 
 mpiexec $MPI_OPTIONS -n 2 ./2_Intermediate/QH_fixed_resolution.py
 mpiexec $MPI_OPTIONS -n 2 ./2_Intermediate/QH_fixed_resolution_boozer.py
+mpiexec $MPI_OPTIONS -n 2 ./2_Intermediate/free_boundary_vmec.py
 
 mpiexec $MPI_OPTIONS -n 2 ./3_Advanced/single_stage_optimization.py
 mpiexec $MPI_OPTIONS -n 2 ./3_Advanced/single_stage_optimization_finite_beta.py

src/simsopt/field/__init__.py

@@ -3,6 +3,7 @@
 from .coil import *
 from .magneticfield import *
 from .magneticfieldclasses import *
+from .mgrid import *
 from .normal_field import *
 from .tracing import *
 
@@ -12,6 +13,7 @@
     + coil.__all__
     + magneticfield.__all__
     + magneticfieldclasses.__all__
+    + mgrid.__all__
     + normal_field.__all__
     + tracing.__all__
 )

src/simsopt/field/magneticfield.py

@@ -3,6 +3,7 @@
 import simsoptpp as sopp
 from .._core.optimizable import Optimizable
 from .._core.json import GSONDecoder
+from .mgrid import MGrid
 
 __all__ = ['MagneticField', 'MagneticFieldSum', 'MagneticFieldMultiply']
 
@@ -91,6 +92,58 @@ 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):
+        """Export the field to the mgrid format for free boundary calculations.
+
+        The field data is represented as a single "current group". For
+        free-boundary vmec, the "extcur" array should have a single nonzero
+        element, set to 1.0.
+
+        In the future, we may want to implement multiple current groups.
+
+        Args:
+            filename: Name of the NetCDF file to save.
+            nr: Number of grid points in the major radius dimension.
+            nphi: Number of planes in the toroidal angle.
+            nz: Number of grid points in the z coordinate.
+            rmin: Minimum value of major radius for the grid.
+            rmax: Maximum value of major radius for the grid.
+            zmin: Minimum value of z for the grid.
+            zmax: Maximum value of z for the grid.
+            nfp: Number of field periods.
+        """
+
+        rs = np.linspace(rmin, rmax, nr, endpoint=True)
+        phis = np.linspace(0, 2 * np.pi / nfp, nphi, endpoint=False)
+        zs = np.linspace(zmin, zmax, nz, endpoint=True)
+
+        Phi, Z, R = np.meshgrid(phis, zs, rs, indexing='ij')
+        X = R * np.cos(Phi)
+        Y = R * np.sin(Phi)
+        Z = Z
+
+        RPhiZ = np.zeros((R.size, 3))
+        RPhiZ[:, 0] = R.flatten()
+        RPhiZ[:, 1] = Phi.flatten()  
+        RPhiZ[:, 2] = Z.flatten()
+
+        # get field on the grid
+        self.set_points_cyl(RPhiZ)
+        B = self.B_cyl()
+
+        # shape the components
+        br, bp, bz = B.T
+        br_3 = br.reshape((nphi, nz, nr))
+        bp_3 = bp.reshape((nphi, nz, nr))
+        bz_3 = bz.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')  
+
+        mgrid.write(filename)
+
 
 class MagneticFieldMultiply(MagneticField):
     """