[VMEC,SurfaceRZFourier] fix inconsistency regarding mpol

src/simsopt/geo/surfacerzfourier.py

@@ -35,6 +35,10 @@ class SurfaceRZFourier(sopp.SurfaceRZFourier, Surface):
 
     and the same for :math:`z(\theta, \phi)`.
 
+    Note that the `mpol` in this is different than the `mpol` in VMEC:
+    VMEC uses poloidal mode numbers `m=0, 1, ..., (mpol-1)`,
+    but here we use `m=0, 1, ..., mpol`.
+
     Here, :math:`(r,\phi, z)` are standard cylindrical coordinates, and theta
     is any poloidal angle.
 
@@ -190,7 +194,10 @@ def from_wout(cls, filename: str, s: float = 1.0,
             interp = interp1d(s_full_grid, zmnc, kind=interp_kind, axis=0)
             zbc = interp(s)
 
-        mpol = int(np.max(xm))
+        # VMEC only considers poloidal modes up to `m=(mpol-1)`,
+        # but `SurfaceRZFourier` includes `m=mpol`.
+        # Here, `xm` comes from VMEC and `mpol` goes into `SurfaceRZFourier`.
+        mpol = int(np.max(xm)) - 1
         ntor = int(np.max(np.abs(xn)) / nfp)
 
         ntheta = kwargs.pop("ntheta", None)
@@ -640,31 +647,31 @@ def write_nml(self, filename: str):
     def fourier_transform_scalar(self, scalar, mpol=None, ntor=None, normalization=None, **kwargs):
         r"""
         Compute the Fourier components of a scalar on the surface. The scalar
-        is evaluated at the quadrature points on the surface. 
-        The Fourier uses the conventions of the FourierRZSurface series, 
+        is evaluated at the quadrature points on the surface.
+        The Fourier uses the conventions of the FourierRZSurface series,
         with `npol` going from `-ntor` to `ntor` and `mpol` from 0 to `mpol`
-        i.e.: 
+        i.e.:
         :math:`f(\theta, \phi) = \Sum_{m=0}^{mpol} \Sum_{n=-npol}^{npol} A^{mn}_s \sin(m\theta - n*Nfp*\phi)\\
             + A^{mn}_c \cos(m\theta - n*Nfp*\phi)`
         Where the cosine series is only evaluated if the surface is not stellarator
-        symmetric (if the scalar does not adhere to the symmetry of the surface, 
+        symmetric (if the scalar does not adhere to the symmetry of the surface,
         request the cosine series by setting the kwarg stellsym=False)
-        By default, the poloidal and toroidal resolution are the same as those of the surface, but different quantities can be specified in the kwargs. 
+        By default, the poloidal and toroidal resolution are the same as those of the surface, but different quantities can be specified in the kwargs.
         *Arguments*:
             - scalar: 2D array of shape (numquadpoints_phi, numquadpoints_theta).
             - mpol: maximum poloidal mode number of the transform, if None,
                 the mpol attribute of the surface is used.
-            - ntor: maximum toroidal mode number of the transform if None, 
+            - ntor: maximum toroidal mode number of the transform if None,
                 the ntor attribute of the surface is used.
         *Optional keyword arguments*:
-            - normalization: Fourier transform normalization. Can be: 
+            - normalization: Fourier transform normalization. Can be:
               None: forward and back transform are not normalized
               float: forward transform is divided by this number
-            - stellsym: boolean to override the stellsym attribute 
+            - stellsym: boolean to override the stellsym attribute
                 of the surface if you want to force the calculation of the cosine series
         *Returns*:
             - A_mns: 2D array of shape (mpol+1, 2*ntor+1) containing the sine coefficients
-            - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients 
+            - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients
                 (these are zero if the surface is stellarator symmetric)
         """
         assert scalar.shape[0] == self.quadpoints_phi.size, "scalar must be evaluated at the quadrature points on the surface.\n the scalar you passed in has shape {}".format(scalar.shape)
@@ -684,9 +691,9 @@ def fourier_transform_scalar(self, scalar, mpol=None, ntor=None, normalization=N
         factor = 2.0 / (ntheta_grid * nphi_grid)
 
         phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi,
-                                     2 * np.pi * self.quadpoints_theta, 
+                                     2 * np.pi * self.quadpoints_theta,
                                      indexing="ij")
-        
+
         for m in range(mpol + 1):
             nmin = -ntor
             if m == 0: nmin = 1
@@ -701,7 +708,7 @@ def fourier_transform_scalar(self, scalar, mpol=None, ntor=None, normalization=N
                 if not stellsym:
                     cosangle = np.cos(angle)
                     A_mnc[m, n + ntor] = np.sum(scalar * cosangle * factor2)
-        
+
         if not stellsym:
             A_mnc[0, ntor] = np.sum(scalar) / (ntheta_grid * nphi_grid)
         if normalization is not None:
@@ -723,7 +730,7 @@ def inverse_fourier_transform_scalar(self, A_mns, A_mnc, normalization=None, **k
         symmetric.
         *Arguments*:
             - A_mns: 2D array of shape (mpol+1, 2*ntor+1) containing the sine coefficients
-            - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients 
+            - A_mnc: 2D array of shape (mpol+1, 2*ntor+1) containing the cosine coefficients
                 (these are zero if the surface is stellarator symmetric)
         *Optional keyword arguments*:
             - normalization: Fourier transform normalization. Can be:
@@ -752,10 +759,10 @@ def inverse_fourier_transform_scalar(self, A_mns, A_mnc, normalization=None, **k
                 if not stellsym:
                     cosangle = np.cos(angle)
                     scalars = scalars + A_mnc[m, n + ntor] * cosangle
-        
+
         if not stellsym:
             scalars = scalars + A_mnc[0, ntor]
-        if normalization is not None: 
+        if normalization is not None:
             if not isinstance(normalization, float):
                 raise ValueError("normalization must be a float")
             scalars = scalars * normalization

src/simsopt/mhd/vmec.py

@@ -344,17 +344,22 @@ def __init__(self,
             # object, but the mpol/ntor values of either the vmec object
             # or the boundary surface object can be changed independently
             # by the user.
+            # `vi.mpol` comes from VMEC, but VMEC only uses `m=0, 1, ..., (mpol-1)`,
+            # but `SurfaceRZFourier` includes `m=mpol`, so we have to
+            # initialize `SurfaceRZFourier` with the highest poloidal mode
+            # that VMEC uses, which is `vi.mpol-1`.
             self._boundary = SurfaceRZFourier.from_nphi_ntheta(nfp=vi.nfp,
                                                                stellsym=not vi.lasym,
-                                                               mpol=vi.mpol,
+                                                               mpol=vi.mpol - 1,
                                                                ntor=vi.ntor,
                                                                ntheta=ntheta,
                                                                nphi=nphi,
                                                                range=range_surface)
             self.free_boundary = bool(vi.lfreeb)
 
             # Transfer boundary shape data from fortran to the ParameterArray:
-            for m in range(vi.mpol + 1):
+            # The highest mode number relevant to VMEC in its input is `vi.mpol-1`.
+            for m in range(vi.mpol):
                 for n in range(-vi.ntor, vi.ntor + 1):
                     self._boundary.rc[m, n + vi.ntor] = vi.rbc[101 + n, m]
                     self._boundary.zs[m, n + vi.ntor] = vi.zbs[101 + n, m]
@@ -523,7 +528,9 @@ def set_indata(self):
         mpol_capped = np.min([boundary_RZFourier.mpol, 101])
         ntor_capped = np.min([boundary_RZFourier.ntor, 101])
         # Transfer boundary shape data from the surface object to VMEC:
-        for m in range(mpol_capped + 1):
+        # The highest mode number that VMEC can use is `m=100`
+        # if the max resolution is `mpol=101`.
+        for m in range(mpol_capped):
             for n in range(-ntor_capped, ntor_capped + 1):
                 vi.rbc[101 + n, m] = boundary_RZFourier.get_rc(m, n)
                 vi.zbs[101 + n, m] = boundary_RZFourier.get_zs(m, n)
@@ -756,7 +763,7 @@ def run(self):
                     os.remove(filename)
                 except FileNotFoundError:
                     logger.debug(f"Tried to delete the file {filename} but it was not found")
-                    
+
             self.files_to_delete = []
 
             # Record the latest output file to delete if we run again:

tests/geo/test_surface_rzfourier.py

@@ -234,11 +234,13 @@ def test_from_vmec_2_ways(self):
         filename2 = TEST_DIR / 'wout_li383_low_res_reference.nc'
         s1 = SurfaceRZFourier.from_vmec_input(filename1)
         s2 = SurfaceRZFourier.from_wout(filename2)
-        mpol = min(s1.mpol, s2.mpol)
-        ntor = min(s1.ntor, s2.ntor)
-        places = 13
+        self.assertEqual(s1.mpol, s2.mpol)
+        self.assertEqual(s1.ntor, s2.ntor)
         self.assertEqual(s1.nfp, s2.nfp)
         self.assertEqual(s1.stellsym, s2.stellsym)
+        mpol = s1.mpol
+        ntor = s1.ntor
+        places = 13
         for m in range(mpol + 1):
             nmin = 0 if m == 0 else -ntor
             for n in range(nmin, ntor + 1):
@@ -266,8 +268,10 @@ def test_from_vmec_2_ways(self):
         # coordinate-independent properties like the volume and area.
         self.assertAlmostEqual(np.abs(s1.volume()), np.abs(s2.volume()), places=13)
         self.assertAlmostEqual(s1.area(), s2.area(), places=7)
-        mpol = min(s1.mpol, s2.mpol)
-        ntor = min(s1.ntor, s2.ntor)
+        self.assertEqual(s1.mpol, s2.mpol)
+        self.assertEqual(s1.ntor, s2.ntor)
+        mpol = s1.mpol
+        ntor = s1.ntor
         places = 13
         for m in range(mpol + 1):
             nmin = 0 if m == 0 else -ntor
@@ -290,11 +294,13 @@ def test_get_and_write_nml(self):
         with ScratchDir("."):
             s1.write_nml(new_filename)
             s2 = SurfaceRZFourier.from_vmec_input(new_filename)
-        mpol = min(s1.mpol, s2.mpol)
-        ntor = min(s1.ntor, s2.ntor)
-        places = 13
         self.assertEqual(s1.nfp, s2.nfp)
         self.assertEqual(s1.stellsym, s2.stellsym)
+        self.assertEqual(s1.mpol, s2.mpol)
+        self.assertEqual(s1.ntor, s2.ntor)
+        mpol = s1.mpol
+        ntor = s1.ntor
+        places = 13
         for m in range(mpol + 1):
             nmin = 0 if m == 0 else -ntor
             for n in range(nmin, ntor + 1):
@@ -310,11 +316,13 @@ def test_get_and_write_nml(self):
             with open(new_filename, 'w') as f:
                 f.write(nml_str)
             s2 = SurfaceRZFourier.from_vmec_input(new_filename)
-        mpol = min(s1.mpol, s2.mpol)
-        ntor = min(s1.ntor, s2.ntor)
-        places = 13
         self.assertEqual(s1.nfp, s2.nfp)
         self.assertEqual(s1.stellsym, s2.stellsym)
+        self.assertEqual(s1.mpol, s2.mpol)
+        self.assertEqual(s1.ntor, s2.ntor)
+        mpol = s1.mpol
+        ntor = s1.ntor
+        places = 13
         for m in range(mpol + 1):
             nmin = 0 if m == 0 else -ntor
             for n in range(nmin, ntor + 1):
@@ -743,10 +751,10 @@ def test_fourier_transform_scalar(self):
 
         # Create the grid of quadpoints:
         phi2d, theta2d = np.meshgrid(2 * np.pi * s.quadpoints_phi,
-                                     2 * np.pi * s.quadpoints_theta, 
+                                     2 * np.pi * s.quadpoints_theta,
                                      indexing='ij')
 
-        # create a test field where only Fourier elements [m=2, n=3] 
+        # create a test field where only Fourier elements [m=2, n=3]
         # and [m=4,n=5] are nonzero:
         field = 0.8 * np.sin(2*theta2d - 3*s.nfp*phi2d) + 0.2*np.sin(4*theta2d - 5*s.nfp*phi2d)+ 0.7*np.cos(3*theta2d - 3*s.nfp*phi2d)
 

tests/mhd/test_vmec.py

@@ -230,8 +230,10 @@ def test_surface_4_ways(self):
 
         def compare_surfaces_sym(s1, s2):
             logger.debug('compare_surfaces_sym called')
-            mpol = min(s1.mpol, s2.mpol)
-            ntor = min(s1.ntor, s2.ntor)
+            self.assertEqual(s1.mpol, s2.mpol)
+            self.assertEqual(s1.ntor, s2.ntor)
+            mpol = s1.mpol
+            ntor = s1.ntor
             places = 13
             for m in range(mpol + 1):
                 nmin = 0 if m == 0 else -ntor
@@ -243,8 +245,10 @@ def compare_surfaces_asym(s1, s2, places):
             logger.debug('compare_surfaces_asym called')
             self.assertAlmostEqual(np.abs(s1.volume()), np.abs(s2.volume()), places=13)
             self.assertAlmostEqual(s1.area(), s2.area(), places=7)
-            mpol = min(s1.mpol, s2.mpol)
-            ntor = min(s1.ntor, s2.ntor)
+            self.assertEqual(s1.mpol, s2.mpol)
+            self.assertEqual(s1.ntor, s2.ntor)
+            mpol = s1.mpol
+            ntor = s1.ntor
             for m in range(mpol + 1):
                 nmin = 0 if m == 0 else -ntor
                 for n in range(nmin, ntor + 1):