Implemented make_rotating_ellipse

src/simsopt/geo/surfacerzfourier.py

@@ -637,6 +637,43 @@ def write_nml(self, filename: str):
         with open(filename, 'w') as f:
             f.write(self.get_nml())
 
+    def make_rotating_ellipse(self, major_radius, minor_radius, elongation, torsion=0):
+        """
+        Set the surface shape to be a rotating ellipse with the given
+        parameters.
+
+        Values of ``elongation`` larger than 1 will result in the elliptical
+        cross-section at :math:`\phi=0` being taller than it is wide.
+        Values of ``elongation`` less than 1 will result in the elliptical
+        cross-section at :math:`\phi=0` being wider than it is tall.
+
+        The sign convention is such that both the rotating elongation and
+        positive ``torsion`` will contribute positively to iota according to
+        VMEC's sign convention.
+
+        Args:
+            major_radius: Average major radius of the surface.
+            minor_radius: Average minor radius of the surface.
+            elongation: Elongation of the elliptical cross-section.
+            torsion: Value to use for the (m,n)=(0,1) mode of RC and -ZS, which
+                controls the torsion of the magnetic axis.
+        """
+
+        self.local_full_x = np.zeros_like(self.local_full_x)
+        self.set_rc(0, 0, major_radius)
+        self.set_rc(0, 1, torsion)
+        self.set_zs(0, 1, -torsion)
+
+        sqrt_elong = np.sqrt(elongation)
+        amplitude = 0.5 * minor_radius * (1 / sqrt_elong - sqrt_elong)
+        self.set_rc(1, 1, amplitude)
+        self.set_zs(1, 1, -amplitude)
+
+        amplitude = 0.5 * minor_radius * (1 / sqrt_elong + sqrt_elong)
+        self.set_rc(1, 0, amplitude)
+        self.set_zs(1, 0, amplitude)
+
+
     return_fn_map = {'area': sopp.SurfaceRZFourier.area,
                      'volume': sopp.SurfaceRZFourier.volume,
                      'aspect-ratio': Surface.aspect_ratio}

tests/geo/test_surface_rzfourier.py

@@ -10,6 +10,13 @@
 from simsopt.geo.surface import Surface
 from simsopt._core.optimizable import Optimizable
 
+try:
+    import vmec
+except ImportError:
+    vmec = None
+
+from simsopt.mhd import Vmec
+
 TEST_DIR = Path(__file__).parent / ".." / "test_files"
 
 stellsym_list = [True, False]
@@ -631,6 +638,98 @@ def test_shared_dof_serialization(self):
             self.assertAlmostEqual(s2.area(), surf_objs[1].area())
             self.assertIs(surf_objs[0].dofs, surf_objs[1].dofs)
 
+    def test_make_rotating_ellipse(self):
+        major_radius = 8.4
+        minor_radius = 2.3
+        elongation = 2.7
+        torsion = 0.6
+        nfp = 3
+        sqrt_elong = np.sqrt(elongation)
+        surf = SurfaceRZFourier.from_nphi_ntheta(ntor=2, mpol=3, nphi=2, ntheta=4, range="field period", nfp=nfp)
+        surf.make_rotating_ellipse(major_radius, minor_radius, elongation, torsion)
+
+        xyz = surf.gamma()
+        R = np.sqrt(xyz[:, :, 0]**2 + xyz[:, :, 1]**2)
+        Z = xyz[:, :, 2]
+
+        # Check phi=0 plane:
+        np.testing.assert_allclose(
+            R[0, :],
+            [major_radius + torsion + minor_radius / sqrt_elong,
+             major_radius + torsion,
+             major_radius + torsion - minor_radius / sqrt_elong,
+             major_radius + torsion]
+        )
+        np.testing.assert_allclose(
+            Z[0, :],
+            [0,
+             minor_radius * sqrt_elong,
+             0,
+             -minor_radius * sqrt_elong],
+            atol=1e-14,
+        )
+
+        # Check phi=pi/nfp plane:
+        np.testing.assert_allclose(
+            R[1, :],
+            [major_radius - torsion + minor_radius * sqrt_elong,
+             major_radius - torsion,
+             major_radius - torsion - minor_radius * sqrt_elong,
+             major_radius - torsion]
+        )
+        np.testing.assert_allclose(
+            Z[1, :],
+            [0,
+             minor_radius / sqrt_elong,
+             0,
+             -minor_radius / sqrt_elong],
+            atol=1e-14,
+        )
+
+        # Now make the same surface shape with more quadpoints:
+        surf = SurfaceRZFourier.from_nphi_ntheta(ntor=1, mpol=1, nphi=64, ntheta=65, range="field period", nfp=nfp)
+        surf.make_rotating_ellipse(major_radius, minor_radius, elongation, torsion)
+        np.testing.assert_allclose(surf.major_radius(), major_radius)
+        np.testing.assert_allclose(surf.minor_radius(), minor_radius)
+        np.testing.assert_allclose(surf.aspect_ratio(), major_radius / minor_radius)
+
+        # Check that the cross-sectional area is correct at every phi:
+        gamma = surf.gamma()
+        R = np.sqrt(gamma[:, :, 0]**2 + gamma[:, :, 1]**2)
+        gammadash2 = surf.gammadash2()
+        dZdtheta = gammadash2[:, :, 2]
+        dtheta = surf.quadpoints_theta[1] - surf.quadpoints_theta[0]
+        area_vs_phi = np.abs(np.sum(R * dZdtheta, axis=1) * dtheta)
+        np.testing.assert_allclose(area_vs_phi, np.pi * minor_radius**2)
+
+    @unittest.skipIf(vmec is None, "vmec python extension is not installed")
+    def test_make_rotating_ellipse_iota(self):
+        """make_rotating_ellipse() should give positive iota."""
+        filename = str(TEST_DIR / 'input.LandremanPaul2021_QH_reactorScale_lowres')
+        with ScratchDir("."):
+            eq = Vmec(filename)
+            eq.indata.mpol = 4  # Lower resolution to expedite test
+            eq.indata.ntor = 4
+            eq.indata.ftol_array[:2] = [1e-8, 1e-10]
+
+            # Try the case of elongation=1 with positive torsion:
+            major_radius = 8.4
+            minor_radius = 1.3
+            elongation = 1.0
+            torsion = 0.9
+            eq.boundary.make_rotating_ellipse(major_radius, minor_radius, elongation, torsion)
+            eq.run()
+            np.testing.assert_array_less(0, eq.wout.iotaf)
+
+            # Try the case of zero torsion with rotating elongation:
+            major_radius = 8.4
+            minor_radius = 1.3
+            elongation = 2.1
+            torsion = 0.0
+            eq.boundary.make_rotating_ellipse(major_radius, minor_radius, elongation, torsion)
+            eq.run()
+            np.testing.assert_array_less(0, eq.wout.iotaf)
+
 
 class SurfaceRZPseudospectralTests(unittest.TestCase):
     def test_names(self):