PR: Get a near-axis configuration from BOOZ_XFORM

.github/workflows/ci.yml

@@ -28,7 +28,7 @@ jobs:
     - name: apt-get stuff needed for libstell and vmec
       run: |
         sudo apt-get update
-        sudo apt-get install -y build-essential gfortran openmpi-bin libopenmpi-dev libnetcdf-dev libnetcdff-dev liblapack-dev libscalapack-mpi-dev libhdf5-dev libhdf5-serial-dev git m4
+        sudo apt-get install -y build-essential gfortran openmpi-bin libopenmpi-dev libnetcdf-dev libnetcdff-dev liblapack-dev libscalapack-mpi-dev libhdf5-dev libhdf5-serial-dev git m4 libfftw3-dev libboost-all-dev libopenblas-dev
 
     - uses: actions/checkout@v2
       # If we want submodules downloaded, uncomment the next 2 lines:
@@ -51,6 +51,9 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install wheel numpy mpi4py scipy matplotlib
+
+    - name: Install booz_xform
+      run: pip install -v git+https://github.com/hiddenSymmetries/booz_xform
 
     - name: env after adding python
       run: env

.gitignore

@@ -8,4 +8,6 @@ jxbout*
 threed*
 wout*
 mercier*
-parvmec*
+parvmec*
+xvmec2000
+timings.txt

qsc/qsc.py

@@ -6,7 +6,7 @@
 import logging
 import numpy as np
 from scipy.io import netcdf
-#from numba import jit
+import matplotlib.pyplot as plt
 
 #logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
@@ -371,3 +371,144 @@ def min_R0_penalty(self):
         """
         return np.max((0, self.min_R0_threshold - self.min_R0)) ** 2
 
+    @classmethod
+    def from_boozxform(cls, booz_xform_file, order='r2', max_s_for_fit = 0.4, N_phi = 200, max_n_to_plot = 2, show=False,
+                         vmec_file=None, rc=[], rs=[], zc=[], zs=[], sigma0=0, I2=0, p2=0):
+        """
+        Load a configuration from a VMEC and a BOOZ_XFORM output files
+        """
+        # Read properties of BOOZ_XFORM output file
+        f = netcdf.netcdf_file(booz_xform_file,'r',mmap=False)
+        bmnc = f.variables['bmnc_b'][()]
+        ixm = f.variables['ixm_b'][()]
+        ixn = f.variables['ixn_b'][()]
+        jlist = f.variables['jlist'][()]
+        ns = f.variables['ns_b'][()]
+        nfp = f.variables['nfp_b'][()]
+        Psi = f.variables['phi_b'][()]
+        Psi_a = np.abs(Psi[-1])
+        iotaVMECt=f.variables['iota_b'][()][1]
+        f.close()
+
+        if vmec_file!=None:
+            # Read axis-shape from VMEC output file
+            f = netcdf.netcdf_file(vmec_file,'r',mmap=False)
+            am = f.variables['am'][()]
+            rc = f.variables['raxis_cc'][()]
+            zs = -f.variables['zaxis_cs'][()]
+            try:
+                rs = -f.variables['raxis_cs'][()]
+                zc = f.variables['zaxis_cc'][()]
+                logger.info('Non stellarator symmetric configuration')
+            except:
+                rs=[]
+                zc=[]
+                logger.info('Stellarator symmetric configuration')
+            f.close()
+        elif rc!=[]:
+            # Read axis-shape from input parameters
+            rc=rc
+            rs=rs
+            zc=zc
+            zs=zs
+        else:
+            raise Exception("Axis shape not specified")
+
+        # Calculate nNormal
+        stel = Qsc(rc=rc, rs=rs, zc=zc, zs=zs, nfp=nfp)
+        nNormal = stel.iotaN - stel.iota
+
+        # Prepare coordinates for fit
+        s_full = np.linspace(0,1,ns)
+        ds = s_full[1] - s_full[0]
+        #s_half = s_full[1:] - 0.5*ds
+        s_half = s_full[jlist-1] - 0.5*ds
+        mask = s_half < max_s_for_fit
+        s_fine = np.linspace(0,1,400)
+        sqrts_fine = s_fine
+        phi = np.linspace(0,2*np.pi / nfp, N_phi)
+        B0  = np.zeros(N_phi)
+        B1s = np.zeros(N_phi)
+        B1c = np.zeros(N_phi)
+        B20 = np.zeros(N_phi)
+        B2s = np.zeros(N_phi)
+        B2c = np.zeros(N_phi)
+
+        # Perform fit
+        numRows=3
+        numCols=max_n_to_plot*2+1
+        fig=plt.figure(num=None, figsize=(16, 9), dpi=80, facecolor='w', edgecolor='k')
+        for jmn in range(len(ixm)):
+            m = ixm[jmn]
+            n = ixn[jmn] / nfp
+            if m>2:
+                continue
+            doplot = (np.abs(n) <= max_n_to_plot) & show
+            row = m
+            col = n+max_n_to_plot
+            if doplot:
+                plt.subplot(int(numRows),int(numCols),int(row*numCols + col + 1))
+                plt.plot(np.sqrt(s_half), bmnc[:,jmn],'.-')
+                # plt.xlabel(r'$\sqrt{s}$')
+                plt.title('bmnc(m='+str(m)+' n='+str(n)+')')
+            if m==0:
+                # For m=0, fit a polynomial in s (not sqrt(s)) that does not need to go through the origin.
+                degree = 4
+                p = np.polyfit(s_half[mask], bmnc[mask,jmn], degree)
+                B0 += p[-1] * np.cos(n*nfp*phi)
+                B20 += p[-2] * np.cos(2*n*nfp*phi)
+                if doplot:
+                    plt.plot(np.sqrt(s_fine), np.polyval(p, s_fine),'r')
+            if m==1:
+                # For m=1, fit a polynomial in sqrt(s) to an odd function
+                x1 = np.sqrt(s_half[mask])
+                y1 = bmnc[mask,jmn]
+                x2 = np.concatenate((-x1,x1))
+                y2 = np.concatenate((-y1,y1))
+                degree = 5
+                p = np.polyfit(x2,y2, degree)
+                B1c += p[-2] * (np.sin(n*nfp*phi) * np.sin(nNormal*phi) + np.cos(n*nfp*phi) * np.cos(nNormal*phi))
+                B1s += p[-2] * (np.sin(n*nfp*phi) * np.cos(nNormal*phi) - np.cos(n*nfp*phi) * np.sin(nNormal*phi))
+                if doplot:
+                    plt.plot(sqrts_fine, np.polyval(p, sqrts_fine),'r')
+            if m==2:
+                # For m=2, fit a polynomial in s (not sqrt(s)) that does need to go through the origin.
+                x1 = s_half[mask]
+                y1 = bmnc[mask,jmn]
+                degree = 4
+                p = np.polyfit(x1,y1, degree)
+                B2c += p[-2] * (np.sin(2*n*nfp*phi) * np.sin(2*nNormal*phi) + np.cos(2*n*nfp*phi) * np.cos(2*nNormal*phi))
+                B2s += p[-2] * (np.sin(2*n*nfp*phi) * np.cos(2*nNormal*phi) - np.cos(2*n*nfp*phi) * np.sin(2*nNormal*phi))
+                if doplot:
+                    plt.plot(np.sqrt(s_fine), np.polyval(p, s_fine),'r')
+        if show:
+            plt.show()
+        # Convert expansion in sqrt(s) to an expansion in r
+        BBar = np.mean(B0)
+        sqrt_s_over_r = np.sqrt(np.pi * BBar / Psi_a)
+        B1s *= -sqrt_s_over_r
+        B1c *= -sqrt_s_over_r
+        B20 *= sqrt_s_over_r*sqrt_s_over_r
+        B2c *= sqrt_s_over_r*sqrt_s_over_r
+        B2s *= sqrt_s_over_r*sqrt_s_over_r
+        eta_bar = np.mean(B1c) / BBar
+
+        # NEEDS A WAY TO READ I2 FROM VMEC OR BOOZ_XFORM
+
+        if p2==0 and vmec_file!=None:
+            r  = np.sqrt(Psi_a/(np.pi*BBar))
+            p2 = - am[0]/r/r
+
+        if order=='r1':
+            q = cls(rc=rc,rs=rs,zc=zc,zs=zs,etabar=eta_bar,nphi=N_phi,nfp=nfp,B0=BBar,sigma0=sigma0, I2=I2)
+        else:
+            q = cls(rc=rc,rs=rs,zc=zc,zs=zs,etabar=eta_bar,nphi=N_phi,nfp=nfp,B0=BBar,sigma0=sigma0, I2=I2, B2c=np.mean(B2c), B2s=np.mean(B2s), order=order, p2=p2)
+
+        q.B0_boozxform_array=B0
+        q.B1c_boozxform_array=B1c
+        q.B1s_boozxform_array=B1s
+        q.B20_boozxform_array=B20
+        q.B2c_boozxform_array=B2c
+        q.B2s_boozxform_array=B2s
+
+        return q

qsc/tests/test_from_boozxform.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+
+from qsc.util import to_Fourier
+import unittest
+import os
+from scipy.io import netcdf
+import numpy as np
+import logging
+from qsc.qsc import Qsc
+from mpi4py import MPI
+import vmec
+import booz_xform as bx
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def compare_from_boozxform(name, r=0.008, nphi=151):
+    # Add the directory of this file to the specified filename:
+    inputFile="input."+str(name).replace(" ","")
+    abs_filename = os.path.join(os.path.dirname(__file__), inputFile)
+    woutFile="wout_"+str(name).replace(" ","")+".nc"
+    boozFile="boozmn_"+str(name).replace(" ","")+".nc"
+    # Run pyQsc and create a VMEC input file
+    logger.info('Creating pyQSC configuration')
+    py = Qsc.from_paper(name, nphi=nphi)
+    logger.info('Outputing to VMEC')
+
+    py.to_vmec(inputFile,r,params={"ftol_array": [2e-17,5e-17,3e-16,5e-16], "ns_array": [16,49,101,151], "niter_array": [2000,2000,2000,2000]})
+
+    # Run VMEC
+    fcomm = MPI.COMM_WORLD.py2f()
+    logger.info("Calling runvmec. comm={}".format(fcomm))
+    ictrl=np.array([15,0,0,0,0], dtype=np.int32)
+    vmec.runvmec(ictrl, inputFile, True, fcomm, '')
+    # Check that VMEC converged
+    assert ictrl[1] == 11
+    vmec.cleanup(True)
+
+    # Run VMEC locally
+    # bashCommand = "./xvmec2000 input."+str(name).replace(" ","")
+    # from subprocess import run
+    # run(bashCommand.split())
+
+    # Run BOOZ_XFORM
+    b1 = bx.Booz_xform()
+
+    b1.read_wout(woutFile)
+    b1.compute_surfs = [0,5,10,15,20,25,30,40,50,60,75,90,105,125,149]
+    b1.mboz = 120
+    b1.nboz = 40
+    b1.run()
+    b1.write_boozmn(boozFile)
+
+    # Read local BOOZ_XFORM file
+    # b1.read_boozmn(boozFile)
+    # Check results
+    stel = Qsc.from_boozxform(boozFile, vmec_file=woutFile, order=py.order, sigma0=py.sigma0, I2=py.I2)
+    logger.info('Initial iota on axis = '+str(py.iota)+'for case '+name)
+    logger.info('Final   iota on axis = '+str(stel.iota)+'for case '+name)
+    assert np.isclose(py.iota,stel.iota,rtol=1e-2)
+    logger.info('Initial B0 = '+str(py.B0)+'for case '+name)
+    logger.info('Final   B0 = '+str(stel.B0)+'for case '+name)
+    assert np.isclose(py.B0,stel.B0,rtol=1e-2)
+    logger.info('Initial etabar = '+str(py.etabar)+'for case '+name)
+    logger.info('Final   etabar = '+str(stel.etabar)+'for case '+name)
+    assert np.isclose(py.etabar,stel.etabar,rtol=2e-2)
+    if stel.order != 'r1':
+        logger.info('Initial B20 mean = '+str(py.B20_mean)+'for case '+name)
+        logger.info('Final   B20 mean = '+str(stel.B20_mean)+'for case '+name)
+        logger.info('Initial B20 variation = '+str(py.B20_variation)+'for case '+name)
+        logger.info('Final   B20 variation = '+str(stel.B20_variation)+'for case '+name)
+        logger.info('Initial B2c = '+str(py.B2c)+'for case '+name)
+        logger.info('Final   B2c = '+str(stel.B2c)+'for case '+name)
+        logger.info('Initial B2s = '+str(py.B2s)+'for case '+name)
+        logger.info('Final   B2s = '+str(stel.B2s)+'for case '+name)
+        assert np.isclose(py.B20_mean,stel.B20_mean,rtol=1e-2,atol=5e-3)
+        assert np.isclose(py.B20_variation,stel.B20_variation,rtol=1e-2,atol=5e-3)
+        assert np.isclose(py.B2c,stel.B2c,rtol=2e-2,atol=5e-3)
+        assert np.isclose(py.B2s,stel.B2s,rtol=1e-3,atol=1e-3)
+
+class FromBoozxformTests(unittest.TestCase):
+
+    def __init__(self, *args, **kwargs):
+        super(FromBoozxformTests, self).__init__(*args, **kwargs)
+        logger = logging.getLogger('qsc.qsc')
+        logger.setLevel(1)
+        # self.cases=["r1 section 5.1","r1 section 5.2","r1 section 5.3",\
+        #             "r2 section 5.1","r2 section 5.2","r2 section 5.3","r2 section 5.4","r2 section 5.5"]
+        self.cases=["r2 section 5.1"]
+
+    def test_from_boozxform(self):
+        """
+        Verify that the we can read the Booz_Xform and VMEC files and
+        that iota on axis match the predicted values when using the to_vmec
+        and from_boozxform functions one after the other.
+        """
+        for case in self.cases:
+            logger.info('Going through case '+case)
+            compare_from_boozxform(case)

qsc/tests/test_to_vmec.py

@@ -24,8 +24,7 @@ def compare_to_vmec(name, r=0.005, nphi=151):
     abs_filename = os.path.join(os.path.dirname(__file__), inputFile)
     # Run pyQsc and create a VMEC input file
     logger.info('Creating pyQSC configuration')
-    order = 'r2' if name[1] == '2' else 'r1'
-    py = Qsc.from_paper(name, nphi=nphi, order=order)
+    py = Qsc.from_paper(name, nphi=nphi)
     logger.info('Outputing to VMEC')
     py.to_vmec(inputFile,r)
     # Run VMEC