Add 'sin' poloidal distance function

doc/whats-new.md

@@ -37,6 +37,9 @@ Release history
   By [Ben Dudson](https://github.com/bendudson)
 - Add flag for `tokamak_example.py` to generate grids with reversed current. This
   prevents negative `dx`/`J` and allows them to be run in BOUT++.
+  By [Mike Kryjak](https://github.com/mikekryjak)
+- Add `sin` poloidal distance function (#207).
+  By [Ben Dudson](https://github.com/bendudson)
 
 0.5.2 (13th March 2023)
 -------------------------

hypnotoad/core/equilibrium.py

@@ -1524,7 +1524,9 @@ def func(psival, position, eps=1e-10):
             return [dpsidr * norm, dpsidz * norm]
 
         result = solve_ivp(
-            func, (psi(*p), self.psival), [p.R, p.Z]  # Range of psi
+            func,
+            (psi(*p), self.psival),
+            [p.R, p.Z],  # Range of psi
         )  # Starting location
         if not result.success:
             raise SolutionError("refinePointIntegrate failed to converge")
@@ -2095,15 +2097,16 @@ class EquilibriumRegion(PsiContour):
         # Input parameters for poloidal spacing functions
         #################################################
         poloidal_spacing_method=WithMeta(
-            "sqrt",
+            "sin",
             doc=(
                 "Method to use for poloidal spacing function: 'sqrt' for "
                 "getSqrtPoloidalSpacingFunction; 'monotonic' for "
                 "getMonotonicPoloidalDistanceFunc; 'linear' for "
-                "getLinearPoloidalDistanceFunc"
+                "getLinearPoloidalDistanceFunc; 'sin' for "
+                "getSinPoloidalDistanceFunc"
             ),
             value_type=str,
-            allowed=["sqrt", "monotonic", "linear"],
+            allowed=["sqrt", "monotonic", "linear", "sin"],
         ),
         xpoint_poloidal_spacing_length=WithMeta(
             lambda options: 5.0e-2 if options.orthogonal else 4.0,
@@ -2669,7 +2672,7 @@ def getSpacings(self):
                 self.nonorthogonal_options.nonorthogonal_xpoint_poloidal_spacing_range_outer  # noqa: E501
             )
         else:
-            raise ValueError(f"Unrecognized value before '.' in " f"kind={self.kind}")
+            raise ValueError(f"Unrecognized value before '.' in kind={self.kind}")
         if self.kind.split(".")[1] == "wall":
             sqrt_a_upper = None
             sqrt_b_upper = self.getTargetParameter("target_poloidal_spacing_length")
@@ -2913,6 +2916,19 @@ def getSfuncFixedSpacing(
                 npoints - 1,
             )
             self._checkMonotonic([(sfunc, "linear")], total_distance=distance)
+        elif method == "sin":
+            if spacing_lower is None:
+                spacing_lower = spacings["monotonic_d_lower"]
+            if spacing_upper is None:
+                spacing_upper = spacings["monotonic_d_upper"]
+            sfunc = self.getSinPoloidalDistanceFunc(
+                distance,
+                npoints - 1,
+                self.user_options.N_norm_prefactor * self.ny_total,
+                d_lower=spacing_lower,
+                d_upper=spacing_upper,
+            )
+            self._checkMonotonic([(sfunc, "sin")], total_distance=distance)
         elif method == "nonorthogonal":
             if (
                 self.nonorthogonal_options.nonorthogonal_spacing_method
@@ -3924,6 +3940,38 @@ def getLinearPoloidalDistanceFunc(self, length, N):
         """
         return lambda i: i / N * length
 
+    def getSinPoloidalDistanceFunc(self, length, N, N_norm, d_lower=None, d_upper=None):
+        """
+        Fit to dist = a*i + b*i^2 + c*[i - sin(2pi*i/n) * n/(2pi)]
+        so grid spacing ddist / di = a + 2b * i + c * [1 - cos(2pi * i / n)]
+        i = 0 => 0
+        i = N => length
+
+        This method adapted from the IDL hypnotoad
+        """
+        dl0 = length * N_norm / N
+
+        if d_lower is None:
+            d_lower = dl0
+        if d_upper is None:
+            d_upper = dl0
+
+        # Modify so that spacing remains monotonic
+        if d_upper + d_lower > 2 * dl0:
+            fact = (d_upper + d_lower) / (2 * dl0)
+            d_upper /= fact
+            d_lower /= fact
+
+        a = d_lower / N_norm
+        b = (d_upper / N_norm - a) / (2.0 * N)
+        c = length / N - a - b * N
+
+        return (
+            lambda i: a * i
+            + b * i**2
+            + c * (i - numpy.sin(2 * numpy.pi * i / N) * N / (2 * numpy.pi))
+        )
+
 
 class Equilibrium:
     """
@@ -4423,7 +4471,7 @@ def findRoots_1d(
             lucky_roots = numpy.where(interval_f == 0.0)
             if len(lucky_roots[0]) > 0:
                 raise NotImplementedError(
-                    "Don't handle interval points that happen to land " "on a root yet!"
+                    "Don't handle interval points that happen to land on a root yet!"
                 )
             intervals_with_roots = numpy.where(
                 numpy.sign(interval_f[:-1]) != numpy.sign(interval_f[1:])

hypnotoad/gui/gui.py

@@ -33,7 +33,6 @@
 from ..core.equilibrium import SolutionError
 from ..__init__ import __version__
 
-
 COLOURS = {
     "red": "#aa0000",
 }

hypnotoad/scripts/hypnotoad_circular.py

@@ -14,11 +14,9 @@
 def get_arg_parser():
     from argparse import ArgumentParser
 
-    parser = ArgumentParser(
-        description="""
+    parser = ArgumentParser(description="""
         Create a grid file for a configuration with circular, concentric flux surfaces.
-        """
-    )
+        """)
     parser.add_argument("inputfile", nargs="?", default=None)
     parser.add_argument("--pdb", action="store_true", default=False)
     parser.add_argument("--plot-regions", action="store_true", default=False)

hypnotoad/scripts/hypnotoad_plot_equilibrium.py

@@ -4,11 +4,9 @@
 def get_arg_parser():
     from argparse import ArgumentParser
 
-    parser = ArgumentParser(
-        description="""
+    parser = ArgumentParser(description="""
         Plot the equilibrium stored in a geqdsk file
-        """
-    )
+        """)
     parser.add_argument(
         "equilibrium_file", help="Path to equilibrium file in geqdsk format"
     )

hypnotoad/scripts/hypnotoad_plot_grid_cells.py

@@ -4,16 +4,14 @@
 def get_arg_parser():
     from argparse import ArgumentParser
 
-    parser = ArgumentParser(
-        description="""
+    parser = ArgumentParser(description="""
         Script to plot grid cells for a BOUT++ grid, using the corner positions saved by
         hypnotoad.
 
         The CELL_CENTRE positions are shown in the plot as black dots, the cells are
         shown by joining the corners with black lines. The branch cuts (if shown) are
         thick red lines.
-        """
-    )
+        """)
     parser.add_argument(
         "gridfile", help="Path to the grid file to plot grid cells from"
     )

hypnotoad/scripts/hypnotoad_recreate_inputs.py

@@ -9,12 +9,10 @@
 
 
 def get_arg_parser():
-    parser = AP(
-        description="""
+    parser = AP(description="""
         Recreate input files that were used to create the grid from a hypnotoad grid
         file
-        """
-    )
+        """)
     parser.add_argument("grid_file", type=str)
     parser.add_argument(
         "-g",

hypnotoad/scripts/hypnotoad_torpex.py

@@ -35,11 +35,9 @@
 def get_arg_parser():
     from argparse import ArgumentParser
 
-    parser = ArgumentParser(
-        description="""
+    parser = ArgumentParser(description="""
         Create grid for TORPEX X-point configuration
-        """
-    )
+        """)
     parser.add_argument("filename")
     parser.add_argument("--noplot", action="store_true")
 

hypnotoad/test_suite/test_equilibrium.py

@@ -1326,6 +1326,87 @@ def test_getSqrtPoloidalDistanceFuncBothBoth(self, eqReg):
             abs=1.0e-5,
         )
 
+    def test_getSinPoloidalDistanceFuncLinear(self, eqReg):
+        L = 2.0
+        N = 10.0
+        N_norm = 1
+        f = eqReg.getSinPoloidalDistanceFunc(L, N, N_norm)
+        # f(0) = 0
+        assert f(0.0) == tight_approx(0.0)
+        # f(N) = L
+        assert f(10.0) == tight_approx(2.0)
+        # f(i) = i/N*L
+        assert f(3.0) == tight_approx(0.6)
+
+        # test gradients at upper and lower bounds
+        dfdi = L / N
+        # i=0, interior
+        itest = 1.0e-3
+        assert (f(itest) - f(0.0)) / itest == tight_approx(dfdi)
+        # i=0, extrapolating
+        itest = -1.0e-3
+        assert (f(itest) - f(0.0)) / itest == tight_approx(dfdi)
+        # i=N, interior
+        itest = N - 1.0e-3
+        assert (f(N) - f(itest)) / (N - itest) == tight_approx(dfdi)
+        # i=N, extrapolating
+        itest = N + 1.0e-3
+        assert (f(N) - f(itest)) / (N - itest) == tight_approx(dfdi)
+
+    def test_getSinPoloidalDistanceFuncLower(self, eqReg):
+        d_lower = 0.1
+        L = 2.0
+        N = 10.0
+        N_norm = 20.0
+        f = eqReg.getSinPoloidalDistanceFunc(L, N, N_norm, d_lower=d_lower)
+        # f(0) = 0
+        assert f(0.0) == tight_approx(0.0)
+        # f(N) = L
+        assert f(N) == tight_approx(L)
+        # for i<<1, f ~ d_lower*i/N_norm
+        itest = 0.01
+        assert f(itest) == pytest.approx(d_lower * itest / N_norm, abs=1.0e-5)
+
+    def test_getSinPoloidalDistanceFuncUpper(self, eqReg):
+        d_upper = 0.1
+        L = 2.0
+        N = 10.0
+        N_norm = 20.0
+        f = eqReg.getSinPoloidalDistanceFunc(L, N, N_norm, d_upper=d_upper)
+        # f(0) = 0
+        assert f(0.0) == tight_approx(0.0)
+        # f(N) = L
+        assert f(N) == tight_approx(L)
+        # for (N-i)<<1, f ~ L - d_upper*(N-i)/N_norm
+        itest = N - 0.01
+        assert f(itest) == pytest.approx(L - d_upper * (N - itest) / N_norm, abs=1.0e-5)
+
+    def test_getSinPoloidalDistanceFuncBoth(self, eqReg):
+        d_lower = 0.1
+        d_upper = 0.2
+        L = 2.0
+        N = 10.0
+        N_norm = 40.0
+        f = eqReg.getSinPoloidalDistanceFunc(
+            L, N, N_norm, d_lower=d_lower, d_upper=d_upper
+        )
+        # f(0) = 0
+        assert f(0.0) == tight_approx(0.0)
+        # f(N) = L
+        assert f(N) == tight_approx(L)
+        # for i<<1, f ~ b_lower*i/N_norm
+        itest = 0.01
+        assert f(itest) == pytest.approx(d_lower * itest / N_norm, abs=1.0e-5)
+        # Check we can extrapolate at lower end
+        itest = -0.01
+        assert f(itest) == pytest.approx(d_lower * itest / N_norm, abs=1.0e-5)
+        # for (N-i)<<1, f ~ L - b_upper*(N-i)/N_norm
+        itest = N - 0.01
+        assert f(itest) == pytest.approx(L - d_upper * (N - itest) / N_norm, abs=1.0e-5)
+        # Check we can extrapolate at upper end
+        itest = N + 0.01
+        assert f(itest) == pytest.approx(L - d_upper * (N - itest) / N_norm, abs=1.0e-5)
+
     def test_combineSfuncsPoloidalSpacing(self, eqReg):
         n = len(eqReg)
         L = eqReg.totalDistance(psi=eqReg.psi, equilibrium=eqReg.equilibrium)