Merge pull request #279 from PlasmaControl/grid_spacing

Grid spacing

desc/grid.py

@@ -425,11 +425,14 @@ def _create_nodes(
             self._L = len(np.atleast_1d(rho))
         if np.isscalar(rho) and (int(rho) == rho) and rho > 0:
             r = np.flipud(np.linspace(1, 0, int(rho), endpoint=axis))
+            # choose dr such that each node has the same weight
             dr = 1 / r.size * np.ones_like(r)
         else:
             r = np.atleast_1d(rho)
             dr = np.zeros_like(r)
             if r.size > 1:
+                # choose dr such that cumulative sums of dr[] are node midpoints
+                # and the total sum is 1
                 dr[0] = (r[0] + r[1]) / 2
                 dr[1:-1] = (r[2:] - r[:-2]) / 2
                 dr[-1] = 1 - (r[-2] + r[-1]) / 2
@@ -449,17 +452,26 @@ def _create_nodes(
             if self.sym:
                 t += t[1] / 2
             dt = 2 * np.pi / t.size * np.ones_like(t)
+
         else:
             t = np.atleast_1d(theta)
             dt = np.zeros_like(t)
             if t.size > 1:
-                dt[0] = (t[0] + t[1]) / 2
-                dt[1:-1] = (t[2:] - t[:-2]) / 2
-                dt[-1] = 2 * np.pi - (t[-2] + t[-1]) / 2
+                # choose dt to be half the cyclic distance of the surrounding two nodes
+                SUP = 2 * np.pi  # supremum
+                dt[0] = t[1] + (SUP - (t[-1] % SUP)) % SUP
+                dt[1:-1] = t[2:] - t[:-2]
+                dt[-1] = t[0] + (SUP - (t[-2] % SUP)) % SUP
+                dt /= 2
+                if t.size == 2:
+                    dt[-1] = dt[0]
             else:
                 dt = np.array([2 * np.pi])
 
         # zeta
+        # note: dz spacing should not depend on NFP
+        # spacing corresponds to a node's weight in an integral --
+        # such as integral = sum(dz * data["B"]) -- not the node's coordinates
         if self.N is not None:
             zeta = 2 * self.N + 1
         else:
@@ -471,9 +483,15 @@ def _create_nodes(
             z = np.atleast_1d(zeta)
             dz = np.zeros_like(z)
             if z.size > 1:
-                dz[0] = (z[0] + z[1]) / 2
-                dz[1:-1] = (z[2:] - z[:-2]) / 2
-                dz[-1] = 2 * np.pi - (z[-2] + z[-1]) / 2
+                # choose dz to be half the cyclic distance of the surrounding two nodes
+                SUP = 2 * np.pi / self.NFP  # supremum
+                dz[0] = z[1] + (SUP - (z[-1] % SUP)) % SUP
+                dz[1:-1] = z[2:] - z[:-2]
+                dz[-1] = z[0] + (SUP - (z[-2] % SUP)) % SUP
+                dz /= 2
+                dz *= self.NFP
+                if z.size == 2:
+                    dz[-1] = dz[0]
             else:
                 dz = np.array([2 * np.pi])
 

devtools/dev-requirements.txt

@@ -12,11 +12,11 @@ sphinx-rtd-theme
 black == 21.7b0
 click == 8.0.4
 flake8 >= 4.0.1
-pylint >= 2.12.2
-pydocstyle >= 2.0.0
 flake8-docstrings
 flake8-eradicate
 flake8-isort
+pydocstyle >= 2.0.0
+pylint >= 2.12.2
 
 # testing and benchmarking
 codecov

devtools/dev-requirements_conda.yml

@@ -17,6 +17,8 @@ dependencies:
     - jax >= 0.2.11, <= 0.2.25
     - jaxlib >= 0.1.69, <= 0.1.76
     - nvgpu
+    # testing and benchmarking
+    - qsc
 
   # building the docs
   - nbsphinx > 0.8.5
@@ -29,12 +31,11 @@ dependencies:
   - black = 21.7b0
   - click = 8.0.4
   - flake8 >= 4.0.1
+  - flake8-docstrings
+  - flake8-eradicate
+  - flake8-isort
+  - pydocstyle >= 2.0.0
   - pylint >= 2.12.2
-  - pip:
-    - pydocstyle >= 2.0.0
-    - flake8-docstrings
-    - flake8-eradicate
-    - flake8-isort
 
   # testing and benchmarking
   - codecov
@@ -47,8 +48,6 @@ dependencies:
   - pytest-mpl = 0.16.1
   - pytest-split
   - shapely >= 1.8.2
-  - pip:
-    - qsc
 
   # building
   - python-build

examples/DESC/ATF

@@ -52,4 +52,4 @@ m:  -2	n:   0	R1 =   0.00000000E+00	Z1 =   6.75000000E-03
 m:   2	n:   1	R1 =  -2.50000000E-02	Z1 =   0.00000000E+00
 m:  -2	n:  -1	R1 =   2.50000000E-02	Z1 =   0.00000000E+00
 m:  -2	n:   1	R1 =   0.00000000E+00	Z1 =  -6.75000000E-03
-m:   2	n:  -1	R1 =   0.00000000E+00	Z1 =  -6.75000000E-03
+m:   2	n:  -1	R1 =   0.00000000E+00	Z1 =  -6.75000000E-03

tests/test_configuration.py

@@ -413,19 +413,22 @@ class TestGetSurfaces:
     def test_get_rho_surface(self):
         """Test getting a constant rho surface."""
         eq = Equilibrium()
-        surf = eq.get_surface_at(rho=0.5)
+        R0 = 10
+        rho = 0.5
+        surf = eq.get_surface_at(rho=rho)
+        assert surf.rho == rho
         np.testing.assert_allclose(
-            surf.compute_surface_area(), 4 * np.pi ** 2 * 10 * 0.5
+            surf.compute_surface_area(), 4 * np.pi ** 2 * R0 * rho
         )
-        assert surf.rho == 0.5
 
     @pytest.mark.unit
     def test_get_zeta_surface(self):
         """Test getting a constant zeta surface."""
         eq = Equilibrium()
         surf = eq.get_surface_at(zeta=np.pi)
-        np.testing.assert_allclose(surf.compute_surface_area(), np.pi * (1.0) ** 2)
         assert surf.zeta == np.pi
+        rho = 1
+        np.testing.assert_allclose(surf.compute_surface_area(), np.pi * rho ** 2)
 
     @pytest.mark.unit
     def test_get_theta_surface(self):

tests/test_grid.py

@@ -68,6 +68,53 @@ def test_linear_grid(self):
         np.testing.assert_allclose(g.spacing.prod(axis=1), g.weights)
         np.testing.assert_allclose(g.weights.sum(), (2 * np.pi) ** 2)
 
+    @pytest.mark.unit
+    def test_linear_grid_spacing(self):
+        """Test linear grid spacing is consistent."""
+
+        def test(endpoint=False, axis=True):
+            nrho = 1
+            ntheta = 5
+            nzeta = 7
+            NFP = 3
+            grid1 = LinearGrid(
+                rho=nrho,
+                theta=ntheta,
+                zeta=nzeta,
+                NFP=NFP,
+                axis=axis,
+                endpoint=endpoint,
+            )
+            grid2 = LinearGrid(
+                rho=np.linspace(1, 0, nrho, endpoint=axis)[::-1],
+                theta=np.linspace(0, 2 * np.pi, ntheta, endpoint=endpoint),
+                zeta=np.linspace(0, 2 * np.pi / NFP, nzeta, endpoint=endpoint),
+                NFP=NFP,
+                axis=axis,
+                endpoint=endpoint,
+            )
+            np.testing.assert_allclose(grid1.nodes, grid2.nodes)
+            np.testing.assert_allclose(grid1.spacing, grid2.spacing)
+
+        test(endpoint=False)
+        test(axis=False)
+        test(axis=True)
+
+    @pytest.mark.unit
+    def test_linear_grid_spacing_two_nodes(self):
+        """Test that 2 node grids assign equal spacing to nodes."""
+        node_count = 2
+        NFP = 7  # any integer > 1 is good candidate for test
+        endpoint = False  # TODO: fix endpoint = True issue later
+        lg = LinearGrid(
+            theta=np.linspace(0, 2 * np.pi, node_count, endpoint=endpoint),
+            zeta=np.linspace(0, 2 * np.pi / NFP, node_count, endpoint=endpoint),
+            NFP=NFP,
+            endpoint=endpoint,
+        )
+        spacing = np.tile([1, np.pi, np.pi], (node_count * node_count, 1))
+        np.testing.assert_allclose(lg.spacing, spacing)
+
     @pytest.mark.unit
     def test_concentric_grid(self):
         """Test node placement in ConcentricGrid."""