Repository: PlasmaControl/DESC

Issue:
Create surface with specified shape parameters
Similar to the model we have for approximately QI surfaces, it would be neat to have a simple option for creating a surface based on simple scalar shape parameters:

- [ ] Major radius
- [ ] minor radius / aspect ratio
- [ ] elongation
- [ ] triangularity
- [ ] "twist" (something with n=1)?
- [ ] torsion

Discussion:

Patch:
diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py
index ffcdf40a7..3a4cf0c13 100644
--- a/desc/geometry/surface.py
+++ b/desc/geometry/surface.py
@@ -11,7 +11,7 @@ from desc.grid import Grid, LinearGrid
 from desc.io import InputReader
 from desc.optimizable import optimizable_parameter
 from desc.transform import Transform
-from desc.utils import check_nonnegint, check_posint, copy_coeffs, setdefault
+from desc.utils import check_nonnegint, check_posint, copy_coeffs, errorif, setdefault
 
 from .core import Surface
 
@@ -519,6 +519,109 @@ class FourierRZToroidalSurface(Surface):
         )
         return surf
 
+    @classmethod
+    def from_shape_parameters(
+        cls,
+        major_radius=10.0,
+        aspect_ratio=10.0,
+        elongation=1.0,
+        triangularity=0.0,
+        squareness=0.0,
+        eccentricity=0.0,
+        torsion=0.0,
+        twist=0.0,
+        NFP=1,
+        sym=True,
+    ):
+        """Create a surface using a generalized Miller parameterization.
+
+        Parameters
+        ----------
+        major_radius : float > 0
+            Average major radius.
+        aspect_ratio : float > 0
+            Ratio of major radius / minor radius.
+        elongation : float > 0
+            Elongation of the cross section = major axis / minor axis. Value of 1 gives
+            a circular cross section. Value > 1 gives vertical elongated cross section,
+            value < 1 gives horizontally elongated section.
+        triangularity : float
+            Positive triangularity makes a "▷" like cross section, negative
+            triangularity makes a "◁" like cross section. Surface may self-intersect
+            if abs(triangularity) > 1
+        squareness : float
+            Positive squareness makes a "□" type cross section. Negative squareness
+            makes a "+" like cross section. Surface may self-intersect if
+            abs(squareness) > 0.5
+        eccentricity : float in [0, 1)
+            Eccentricity of the magnetic axis. Value of 0 gives circular axis, value of
+            1 gives infinitely elongated axis.
+        torsion : float
+            How non-planar the magnetic axis is.
+        twist : float
+            How many times the cross section rotates per field period. For integer
+            values it is a rigid rotation in phi, for non-integer values it also deforms
+            as it goes around toroidally. Values > 0 give a CCW rotation,
+            values < 0 give CW rotation. Cross section has zero volume if
+            abs(twist)%1 == 1/2.
+        NFP : int
+            Number of field periods.
+        sym : bool (optional)
+            Whether to enforce stellarator symmetry.
+
+        Returns
+        -------
+        surface : FourierRZToroidalSurface
+            Surface with given geometric properties.
+
+        """
+        errorif(major_radius <= 0, ValueError, "major_radius must be positive")
+        errorif(aspect_ratio <= 0, ValueError, "aspect_ratio must be positive")
+        errorif(elongation <= 0, ValueError, "elongation should be positive")
+        errorif(eccentricity < 0, ValueError, "eccentricity must be in [0,1)")
+        errorif(eccentricity >= 1, ValueError, "eccentricity must be in [0,1)")
+        errorif(
+            abs(twist) % 1 == 0.5,
+            ValueError,
+            "surface has no volume for abs(twist)%1 == 0.5",
+        )
+        grid = LinearGrid(L=0, M=30, N=30, NFP=NFP, endpoint=True)
+        theta = grid.nodes[:, 1]
+        zeta = grid.nodes[:, 2]
+
+        # area = pi*a*b = pi*r^2
+        # r ~ sqrt(a*b) -> a = r^2/b
+        # elongation = a/b = r^2/b^2 -> b = r/sqrt(elongation)
+        r = major_radius / aspect_ratio
+        b = r / np.sqrt(elongation)
+        a = r * np.sqrt(elongation)
+
+        # create cross section shape using miller parameterization
+        Rp = b * np.cos(
+            theta + triangularity * np.sin(theta) - squareness * np.sin(2 * theta)
+        )
+        Zp = -a * np.sin(theta + squareness * np.sin(2 * theta))
+
+        # create axis shape using ellipse + torsion
+        a = 2 * major_radius / (1 + np.sqrt(1 - eccentricity**2))
+        b = a * np.sqrt(1 - eccentricity**2)
+        Ra = (a * b) / np.sqrt(
+            b**2 * np.cos(NFP * zeta) ** 2 + a**2 * np.sin(NFP * zeta) ** 2
+        )
+        # max(2, NFP) ensures that nonzero torsion gives nonplanar axis even for NFP=1
+        # otherwise it just tilts the whole axis.
+        Za = torsion * np.sin(max(2, NFP) * zeta)
+        # combine axis + cross section with twist
+        R = Ra + Rp * np.cos(twist * NFP * zeta) - Zp * np.sin(twist * NFP * zeta)
+        Z = Za + Zp * np.cos(twist * NFP * zeta) + Rp * np.sin(twist * NFP * zeta)
+
+        # self._compute_orientation falsely returns -1 when twist = 1 since there is
+        # no m=1, n=0 mode, but by construction this should be right handed
+        # so we can skip that check.
+        return cls.from_values(
+            np.array([R, zeta, Z]).T, theta, NFP=NFP, sym=sym, check_orientation=False
+        )
+
     def constant_offset_surface(
         self, offset, grid=None, M=None, N=None, full_output=False
     ):


Test patch:
diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py
index 32e3307ef..c6d3a2348 100644
--- a/tests/test_surfaces.py
+++ b/tests/test_surfaces.py
@@ -354,6 +354,52 @@ class TestFourierRZToroidalSurface:
                 sym=True,
             )
 
+    @pytest.mark.unit
+    def test_surface_from_shape_parameters(self):
+        """Test that making a surface with specified R0,a etc gives correct shape."""
+        R0 = 8
+        a = 2
+        e = 2.1
+        # basic rotating ellipse, parameters should be ~exact
+        surf = FourierRZToroidalSurface.from_shape_parameters(
+            major_radius=R0,
+            aspect_ratio=R0 / a,
+            elongation=e,
+            triangularity=0.0,
+            squareness=0,
+            eccentricity=0,
+            torsion=0,
+            twist=1,
+            NFP=2,
+            sym=True,
+        )
+        eq = Equilibrium(surface=surf)
+        np.testing.assert_allclose(R0, eq.compute("R0")["R0"], rtol=1e-8)
+        np.testing.assert_allclose(a, eq.compute("a")["a"], rtol=1e-8)
+        np.testing.assert_allclose(
+            e, eq.compute("a_major/a_minor")["a_major/a_minor"], rtol=1e-4
+        )
+
+        # slightly more complex shape, parameters are only approximate
+        surf = FourierRZToroidalSurface.from_shape_parameters(
+            major_radius=R0,
+            aspect_ratio=R0 / a,
+            elongation=e,
+            triangularity=0.3,
+            squareness=0.1,
+            eccentricity=0.1,
+            torsion=0.2,
+            twist=1,
+            NFP=2,
+            sym=True,
+        )
+        eq = Equilibrium(surface=surf)
+        np.testing.assert_allclose(R0, eq.compute("R0")["R0"], rtol=1e-2)
+        np.testing.assert_allclose(a, eq.compute("a")["a"], rtol=5e-2)
+        np.testing.assert_allclose(
+            e, eq.compute("a_major/a_minor")["a_major/a_minor"], rtol=2e-2
+        )
+
 
 class TestZernikeRZToroidalSection:
     """Tests for ZernikeRZToroidalSection class."""