WIP added FT normalization and some tests

src/simsopt/field/normal_field.py

@@ -302,11 +302,10 @@ def get_vns_asarray(self, mpol=None, ntor=None):
         elif ntor > self.ntor:
             raise ValueError('ntor out of bound')
 
-
         vns = np.zeros((mpol + 1, 2 * ntor + 1))
         for mm in range(0, mpol + 1):
             for nn in range(-ntor, ntor + 1):
-                if mm == 0 and nn < 0: continue
+                if mm == 0 and nn <= 0: continue
                 vns[mm, ntor + nn] = self.get_vns(mm, nn)
 
         return vns
@@ -367,7 +366,7 @@ def set_vns_asarray(self, vns, mpol=None, ntor=None):
 
         for mm in range(0, mpol + 1):
             for nn in range(-ntor, ntor + 1):
-                if mm == 0 and nn < 0: continue
+                if mm == 0 and nn <= 0: continue
                 self.set_vns(mm, nn, vns[mm, ntor + nn])
 
     def set_vnc_asarray(self, vnc, mpol=None, ntor=None):

src/simsopt/geo/surface.py

@@ -679,7 +679,7 @@ def interpolate_on_arclength_grid(self, function, theta_evaluate):
 
         return function_interpolated
 
-    def fourier_transform_field(self, field, mpol=None, ntor=None, **kwargs):
+    def fourier_transform_field(self, field, mpol=None, ntor=None, normalization=None, **kwargs):
         r"""
         Compute the Fourier components of a field on the surface. The field
         is evaluated at the quadrature points on the surface. 
@@ -698,14 +698,17 @@ def fourier_transform_field(self, field, mpol=None, ntor=None, **kwargs):
             - 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: 
+              None: forward and back transform are not normalized
+              float: forward transform is divided by this number
             - 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 
                 (these are zero if the surface is stellarator symmetric)
         """
-        assert(field.shape == [self.quadpoints_phi.size, self.quadpoints_theta.size],
+        assert (field.shape == [self.quadpoints_phi.size, self.quadpoints_theta.size],
                "Field must be evaluated at the quadrature points on the surface. \
                the field you passed in has shape {}".format(field.shape))
         stellsym = kwargs.pop('stellsym', self.stellsym)
@@ -715,38 +718,44 @@ def fourier_transform_field(self, field, mpol=None, ntor=None, **kwargs):
         if ntor is None:
             try: ntor = self.ntor
             except AttributeError: raise ValueError("ntor must be specified")
-        A_mns = np.zeros((int(mpol + 1), int(2 * ntor + 1))) # sine coefficients
-        A_mnc = np.zeros((int(mpol + 1), int(2 * ntor + 1))) # cosine coefficients
+        A_mns = np.zeros((int(mpol + 1), int(2 * ntor + 1)))  # sine coefficients
+        A_mnc = np.zeros((int(mpol + 1), int(2 * ntor + 1)))  # cosine coefficients
         ntheta_grid = len(self.quadpoints_theta)
         nphi_grid = len(self.quadpoints_phi)
 
         factor = 2.0 / (ntheta_grid * nphi_grid)
 
-
-        phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi, 
-                                     2 * np.pi * self.quadpoints_theta)
+        phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi,
+                                     2 * np.pi * self.quadpoints_theta, 
+                                     indexing="ij")
 
         for m in range(mpol + 1):
             nmin = -ntor
-            if m==0: nmin = 1
+            if m == 0: nmin = 1
             for n in range(nmin, ntor+1):
                 angle = m * theta2d - n * self.nfp * phi2d
                 sinangle = np.sin(angle)
                 if not stellsym: 
                     cosangle = np.cos(angle)
                 factor2 = factor
                 # The next 2 lines ensure inverse Fourier transform(Fourier transform) = identity
-                if np.mod(ntheta_grid, 2) == 0 and m  == (ntheta_grid/2): factor2 = factor2 / 2
-                if np.mod(nphi_grid,2) == 0 and abs(n) == (nphi_grid/2): factor2 = factor2 / 2
+                if np.mod(ntheta_grid, 2) == 0 and m == (ntheta_grid/2): factor2 = factor2 / 2
+                if np.mod(nphi_grid, 2) == 0 and abs(n) == (nphi_grid/2): factor2 = factor2 / 2
                 A_mns[m, n + ntor] = np.sum(field * sinangle * factor2)
                 if not stellsym:
                     A_mnc[m, n + ntor] = np.sum(field * cosangle * factor2)
 
         if not stellsym:
             A_mnc[0, ntor] = np.sum(field) / (ntheta_grid * nphi_grid)
+        if normalization is not None:
+            if type(normalization) is not float:
+                raise ValueError("normalization must be a float")
+            A_mns = A_mns / normalization
+            A_mnc = A_mnc / normalization
+
         return A_mns, A_mnc
 
-    def inverse_fourier_transform_field(self, A_mns, A_mnc, **kwargs):
+    def inverse_fourier_transform_field(self, A_mns, A_mnc, normalization=None, **kwargs):
         r"""
         Compute the inverse Fourier transform of a field on the surface. The field
         is evaluated at the quadrature points on the surface. The Fourier
@@ -755,22 +764,30 @@ def inverse_fourier_transform_field(self, A_mns, A_mnc, **kwargs):
             + A^{mn}_c \cos(m\theta - n*Nfp*\phi)`
         Where the cosine series is only evaluated if the surface is not stellarator
         symmetric.
-
-        UNTESTED
+        *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 
+                (these are zero if the surface is stellarator symmetric)
+        *Optional keyword arguments*:
+            - normalization: Fourier transform normalization. Can be:
+                None: forward and back transform are not normalized
+                float: inverse transform is multiplied by this number
+            - stellsym: boolean to override the stellsym attribute of the surface
         """
         mpol = A_mns.shape[0] - 1
         ntor = int((A_mns.shape[1] - 1) / 2)
         stellsym = kwargs.pop('stellsym', self.stellsym)
         ntheta_grid = len(self.quadpoints_theta)
         nphi_grid = len(self.quadpoints_phi)
 
-        phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi, 
-                                     2 * np.pi * self.quadpoints_theta)
+        phi2d, theta2d = np.meshgrid(2 * np.pi * self.quadpoints_phi,
+                                     2 * np.pi * self.quadpoints_theta,
+                                     indexing="ij")
 
-        field = np.zeros((ntheta_grid, nphi_grid))
+        field = np.zeros((nphi_grid, ntheta_grid))
         for m in range(mpol + 1):
             nmin = -ntor
-            if m==0: nmin = 1
+            if m == 0: nmin = 1
             for n in range(nmin, ntor+1):
                 angle = m * theta2d - n * self.nfp * phi2d
                 sinangle = np.sin(angle)
@@ -782,6 +799,10 @@ def inverse_fourier_transform_field(self, A_mns, A_mnc, **kwargs):
 
         if not stellsym:
             field = field + A_mnc[0, ntor]
+        if normalization is not none: 
+            if type(normalization) is not float:
+                raise ValueError("normalization must be a float")
+            field = field * normalization
         return field
 
 
@@ -812,8 +833,6 @@ def signed_distance_from_surface(xyz, surface):
     return signed_dists
 
 
-
-
 class SurfaceClassifier():
     r"""
     Takes in a toroidal surface and constructs an interpolant of the signed distance function

src/simsopt/mhd/spec.py

@@ -245,6 +245,33 @@ def __init__(self,
                                                             m + si.mmpol]
         self._boundary.local_full_x = self._boundary.get_dofs()
 
+        # If the equilibrium is freeboundary, we need to read the computational
+        # boundary as well. Otherwise set the outermost boundary as the
+        # computational boundary. 
+        if si.lfreebound:
+            self._computational_boundary = SurfaceRZFourier(nfp=si.nfp,
+                                                            stellsym=stellsym,
+                                                            mpol=si.mpol,
+                                                            ntor=si.ntor)
+            for m in range(si.mpol + 1):
+                for n in range(-si.ntor, si.ntor + 1):
+                    self._computational_boundary.rc[m,
+                                                    n + si.ntor] = si.rwc[n + si.mntor,
+                                                                          m + si.mmpol]
+                    self._computational_boundary.zs[m,
+                                                    n + si.ntor] = si.zws[n + si.mntor,
+                                                                          m + si.mmpol]
+                    if not stellsym:
+                        self._computational_boundary.rs[m,
+                                                        n + si.ntor] = si.rws[n + si.mntor,
+                                                                              m + si.mmpol]
+                        self._computational_boundary.zc[m,
+                                                        n + si.ntor] = si.zwc[n + si.mntor,
+                                                                              m + si.mmpol]
+            self._computational_boundary.fix_all()  # computational boundary is not moved!
+        else:  # in non-free-boundary case, computational boundary is the same as the plasma boundary
+            self._computational_boundary = self._boundary 
+
         self.need_to_run_code = True
         self.counter = -1
 
@@ -291,6 +318,18 @@ def boundary(self):
         """
         return self._boundary
 
+    @property
+    def computational_boundary(self):
+        """
+        Getter for the computational boundary. 
+        Same as the plasma boundary in the non-free-boundary case, 
+        gives the surface on which Vns and Vnc are defined in the free-boundary case.
+
+        Returns:
+            SurfaceRZFourier instance representing the plasma boundary
+        """
+        return self._computational_boundary
+
     @property
     def pressure_profile(self):
         """

tests/geo/test_surface_rzfourier.py

@@ -672,7 +672,7 @@ def test_surface_from_other_surface_with_different_resolution(self):
 
         s2 = SurfaceRZFourier.from_other_surface(s, mpol=3, ntor=2, range='field period')
         s3 = SurfaceRZFourier.from_other_surface(s2, nfp=4, ntheta=100, nphi=100, range='half period')
-        s4 = SurfaceRZFourier.from_other_surface(s3, nfp=s.nfp,  range='full torus')
+        s4 = SurfaceRZFourier.from_other_surface(s3, nfp=s.nfp, range='full torus')
         self.assertAlmostEqual(s.area(), s4.area(), places=4)
         self.assertAlmostEqual(s.volume(), s4.volume(), places=3)
 
@@ -713,31 +713,41 @@ def test_fourier_transform_field(self):
         s.rc[1, 0] = 0.4
         s.zs[1, 0] = 0.2
 
-        # Create a field evaluated on the quadpoints:
-        phi2d, theta2d = np.meshgrid(2 * np.pi * s.quadpoints_phi, 
-                                    2 * np.pi * s.quadpoints_theta)
+        # Create the grid of quadpoints:
+        phi2d, theta2d = np.meshgrid(2 * np.pi * s.quadpoints_phi,
+                                     2 * np.pi * s.quadpoints_theta, 
+                                     indexing='ij')
 
         # create a test field where only Fourier elements [m=2, n=3] 
         # and [m=4,n=5] are nonzero:
         field = []
         for phi, theta in zip(phi2d.flatten(), theta2d.flatten()):
-            field.append( 0.8* np.sin(2*theta - 3*s.nfp*phi) + 0.2*np.sin(4*theta - 5*s.nfp*phi)
-                         + 0.7*np.cos(3*theta - 3*s.nfp*phi) )
+            field.append(0.8 * np.sin(2*theta - 3*s.nfp*phi) + 0.2*np.sin(4*theta - 5*s.nfp*phi)
+                         + 0.7*np.cos(3*theta - 3*s.nfp*phi))
         field = np.array(field).reshape(phi2d.shape)
 
-
         # Transform the field to Fourier space:
         ft_sines, ft_cosines = s.fourier_transform_field(field, stellsym=False)
         self.assertAlmostEqual(ft_sines[2, 3+s.ntor], 0.8)
         self.assertAlmostEqual(ft_sines[4, 5+s.ntor], 0.2)
         self.assertAlmostEqual(ft_cosines[3, 3+s.ntor], 0.7)
 
+        # Test that all other elements are close to zero
+        sines_mask = np.ones_like(ft_sines, dtype=bool)
+        cosines_mask = sines_mask
+        sines_mask[2, 3 + s.ntor] = False
+        sines_mask[4, 5 + s.ntor] = False
+        cosines_mask[3, 3 + s.ntor] = False
+        self.assertEqual(np.all(np.abs(ft_sines[sines_mask]) < 1e-10), True)
+        self.assertEqual(np.all(np.abs(ft_cosines[cosines_mask]) < 1e-10), True)
+
         # Transform back to real space:
         field2 = s.inverse_fourier_transform_field(ft_sines, ft_cosines, stellsym=False)
 
         # Check that the result is the same as the original field:
         np.testing.assert_allclose(field, field2)
-
+
+
 class SurfaceRZPseudospectralTests(unittest.TestCase):
     def test_names(self):
         """

tests/mhd/test_spec.py

@@ -439,7 +439,7 @@ def test_integrated_stellopt_scenarios_2dof(self):
             self.assertAlmostEqual(equil.iota(), -0.4114567, places=3)
             self.assertAlmostEqual(prob.objective(), 7.912501330E-04, places=3)
 
-    def test_freeboundary_vs_coil_field(self)
+    def test_freeboundary_vs_coil_field(self):
         """
         This test evaluates the Vns components from a coil set on the computational 
         boundary and creates a 1-volume SPEC vacuum equilibrium.