Repository: plasmacontrol/desc
Issue: Weights are scaled incorrectly on grids with duplicate nodes

_Originally posted in https://github.com/PlasmaControl/DESC/pull/279#discussion_r987394648_
Link to comment with notebook to test https://github.com/PlasmaControl/DESC/pull/279#issuecomment-1263193774

Our `_scale_weights` does not reduce duplicate node weight correctly

A simple example of this is `nrho=1, ntheta=5, nzeta=1, endpoint=True`. The correct theta spacing for this `[0.78539816 1.57079633 1.57079633 1.57079633 0.78539816]` with drho = 1 always, dzeta = 2pi always. And at this point in the code that is what the spacing is.

`_scale_weights` then detects the duplicate nodes at 0 and 2pi and scales down the node volume drho * dtheta * dzeta. So the end result of the grid spacing is:
```
[[0.82982653 0.65174424 5.21395388]
 [1.04551592 1.64229256 6.56917025]
 [1.04551592 1.64229256 6.56917025]
 [1.04551592 1.64229256 6.56917025]
 [0.82982653 0.65174424 5.21395388]]
```
The reason I think this is wrong is because a rho surface integral uses the dtheta * dzeta portion of the node volume, ignoring drho. Similar for other surface labels. If we wanted the correct rho surface integral we should leave the drho unchanged and scale up /down dtheta * dzeta.

Changed files:
desc/grid.py, tests/test_grid.py

diff --git a/desc/grid.py b/desc/grid.py
index 429b86a5..122e9f0f 100644
--- a/desc/grid.py
+++ b/desc/grid.py
@@ -139,13 +139,36 @@ class Grid(IOAble):
         _, inverse, counts = np.unique(
             nodes, axis=0, return_inverse=True, return_counts=True
         )
-        self._spacing /= np.tile(np.atleast_2d(counts[inverse]).T, 3) ** (1 / 3)
-        # assign weights pretending _enforce_symmetry didn't change theta spacing
+        duplicates = np.tile(np.atleast_2d(counts[inverse]).T, 3)
         temp_spacing = np.copy(self.spacing)
+        temp_spacing /= duplicates ** (1 / 3)
+        # assign weights pretending _enforce_symmetry didn't change theta spacing
         temp_spacing[:, 1] /= dtheta_scale
+        # scale weights sum to full volume
         temp_spacing *= (4 * np.pi**2 / temp_spacing.prod(axis=1).sum()) ** (1 / 3)
         self._weights = temp_spacing.prod(axis=1)
-        self._spacing *= (4 * np.pi**2 / self.spacing.prod(axis=1).sum()) ** (1 / 3)
+
+        # Spacing is the differential element used for integration over surfaces.
+        # For this, 2 columns of the matrix are used.
+        # Spacing is rescaled below to get the correct double product for each pair
+        # of columns in grid.spacing.
+        # The reduction of weight on duplicate nodes should be accounted for
+        # by the 2 columns of spacing which span the surface.
+        self._spacing /= duplicates ** (1 / 2)
+        # Note we rescale 3 columns by the factor that 'should' rescale 2 columns,
+        # so grid.spacing is valid for integrals over all surface labels.
+        # Because a surface integral always ignores 1 column, with this approach,
+        # duplicates nodes are scaled down properly regardless of which two columns
+        # span the surface.
+
+        # The following operation is not a general solution to return the weight
+        # removed from the duplicate nodes back to the unique nodes.
+        # For this reason, duplicates should typically be deleted rather that rescaled.
+        # Note we multiply each column by duplicates^(1/6) to account for the extra
+        # division by duplicates^(1/2) in one of the columns above.
+        self._spacing *= (
+            4 * np.pi**2 / (self.spacing * duplicates ** (1 / 6)).prod(axis=1).sum()
+        ) ** (1 / 3)
 
     def _create_nodes(self, nodes):
         """Allow for custom node creation.
@@ -452,7 +475,11 @@ class LinearGrid(Grid):
             if self.sym:
                 t += t[1] / 2
             dt = 2 * np.pi / t.size * np.ones_like(t)
-
+            if endpoint and t.size > 1:
+                # increase node weight to account for duplicate node
+                dt *= t.size / (t.size - 1)
+                # scale_weights() will reduce endpoint (dt[0] and dt[-1])
+                # duplicate node weight
         else:
             t = np.atleast_1d(theta)
             dt = np.zeros_like(t)
@@ -465,13 +492,20 @@ class LinearGrid(Grid):
                 dt /= 2
                 if t.size == 2:
                     dt[-1] = dt[0]
+                if endpoint:
+                    # The cyclic distance algorithm above correctly weights
+                    # the duplicate node spacing at theta = 0 and 2pi.
+                    # However, scale_weights() is not aware of this, so we multiply
+                    # by 2 to counteract the reduction that will be done there.
+                    dt[0] *= 2
+                    dt[-1] *= 2
             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
+        # such as integral = sum(dt * dz * data["B"]) -- not the node's coordinates
         if self.N is not None:
             zeta = 2 * self.N + 1
         else:
@@ -479,6 +513,11 @@ class LinearGrid(Grid):
         if np.isscalar(zeta) and (int(zeta) == zeta) and zeta > 0:
             z = np.linspace(0, 2 * np.pi / self.NFP, int(zeta), endpoint=endpoint)
             dz = 2 * np.pi / z.size * np.ones_like(z)
+            if endpoint and z.size > 1:
+                # increase node weight to account for duplicate node
+                dz *= z.size / (z.size - 1)
+                # scale_weights() will reduce endpoint (dz[0] and dz[-1])
+                # duplicate node weight
         else:
             z = np.atleast_1d(zeta)
             dz = np.zeros_like(z)
@@ -492,6 +531,13 @@ class LinearGrid(Grid):
                 dz *= self.NFP
                 if z.size == 2:
                     dz[-1] = dz[0]
+                if endpoint:
+                    # The cyclic distance algorithm above correctly weights
+                    # the duplicate node spacing at zeta = 0 and 2pi / NFP.
+                    # However, _scale_weights() is not aware of this, so we multiply
+                    # by 2 to counteract the reduction that will be done there.
+                    dz[0] *= 2
+                    dz[-1] *= 2
             else:
                 dz = np.array([2 * np.pi])
 
diff --git a/tests/test_grid.py b/tests/test_grid.py
index 7b111a89..e5bbb457 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -64,9 +64,11 @@ class TestGrid:
         ).T
 
         np.testing.assert_allclose(g.nodes, nodes)
-        # spacing.prod == weights for linear grids (not true for concentric)
-        np.testing.assert_allclose(g.spacing.prod(axis=1), g.weights)
         np.testing.assert_allclose(g.weights.sum(), (2 * np.pi) ** 2)
+        # spacing.prod == weights for linear grids
+        # this is not true for concentric or any grid with duplicates
+        if not endpoint:
+            np.testing.assert_allclose(g.spacing.prod(axis=1), g.weights)
 
     @pytest.mark.unit
     def test_linear_grid_spacing(self):
@@ -97,6 +99,7 @@ class TestGrid:
             np.testing.assert_allclose(grid1.spacing, grid2.spacing)
 
         test(endpoint=False)
+        test(endpoint=True)
         test(axis=False)
         test(axis=True)
 
@@ -105,15 +108,137 @@ class TestGrid:
         """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,
+
+        def test(endpoint):
+            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,
+            )
+            # when endpoint is true, rho nodes spacing should be scaled down
+            # so that theta / zeta surface integrals weigh duplicate nodes less
+            spacing = np.tile(
+                [1 / 2 if endpoint else 1, np.pi, np.pi], (node_count**2, 1)
+            )
+            np.testing.assert_allclose(lg.spacing, spacing)
+
+        test(endpoint=False)
+        test(endpoint=True)
+
+    @pytest.mark.unit
+    def test_duplicate_node_areas(self):
+        """Test surface areas on grids with a duplicate node from endpoint=True."""
+
+        def test(grid):
+            surface_area = compress(
+                grid, surface_integrals(grid, surface_label="rho"), "rho"
+            )
+            np.testing.assert_allclose(4 * np.pi**2, surface_area)
+
+            surface_area = compress(
+                grid, surface_integrals(grid, surface_label="theta"), "theta"
+            )
+            np.testing.assert_allclose(2 * np.pi, surface_area[1:-1])
+            if ntheta > 1:
+                # theta=0 and theta=2pi surface should have half weight
+                np.testing.assert_allclose(np.pi, surface_area[0])
+                np.testing.assert_allclose(np.pi, surface_area[-1])
+
+            surface_area = compress(
+                grid, surface_integrals(grid, surface_label="zeta"), "zeta"
+            )
+            np.testing.assert_allclose(2 * np.pi, surface_area[1:-1])
+            if nzeta > 1:
+                # zeta=0 and zeta=2pi/NFP surface should have half weight
+                np.testing.assert_allclose(np.pi, surface_area[0])
+                np.testing.assert_allclose(np.pi, surface_area[-1])
+
+        nrho = 4
+        ntheta = 5
+        nzeta = 7
+        NFP = 3
+        endpoint = True
+        test(LinearGrid(rho=nrho, theta=ntheta, zeta=nzeta, NFP=NFP, endpoint=endpoint))
+        test(
+            LinearGrid(
+                rho=np.linspace(1, 0, nrho)[::-1],
+                theta=np.linspace(0, 2 * np.pi, ntheta, endpoint=endpoint),
+                zeta=np.linspace(0, 2 * np.pi / NFP, nzeta, endpoint=endpoint),
+                NFP=NFP,
+                endpoint=endpoint,
+            )
+        )
+
+    @pytest.mark.unit
+    def test_duplicate_node_spacing(self):
+        """Test surface differential element weight on grid with endpoint=True."""
+        nrho = 1
+        ntheta = 8  # unique theta count
+        nzeta = 13  # unique zeta count
+        NFP = 3
+        axis = True
+        endpoint = True
+
+        def test(grid):
+            is_theta_dupe = endpoint & (grid.nodes[:, 1] % (2 * np.pi) == 0)
+            is_zeta_dupe = endpoint & (grid.nodes[:, 2] % (2 * np.pi / NFP) == 0)
+
+            def test_surface(actual_ds, desired_ds):
+                for index, ds in enumerate(actual_ds):
+                    if is_theta_dupe[index] and is_zeta_dupe[index]:
+                        # the grid has 4 of these nodes
+                        np.testing.assert_allclose(ds, desired_ds / 4)
+                    elif is_theta_dupe[index] or is_zeta_dupe[index]:
+                        # the grid has 2 of these nodes
+                        np.testing.assert_allclose(ds, desired_ds / 2)
+                    else:
+                        # unique node
+                        np.testing.assert_allclose(ds, desired_ds)
+
+            # rho surface
+            test_surface(
+                grid.spacing[:, 1:].prod(axis=1),
+                (2 * np.pi / ntheta) * (2 * np.pi / nzeta),
+            )
+            # theta surface
+            test_surface(
+                grid.spacing[:, [0, 2]].prod(axis=1),
+                (1 / nrho) * (2 * np.pi / nzeta),
+            )
+            # zeta surface
+            test_surface(
+                grid.spacing[:, :2].prod(axis=1),
+                (1 / nrho) * (2 * np.pi / ntheta),
+            )
+
+        test(
+            LinearGrid(
+                rho=nrho,
+                theta=ntheta + endpoint,
+                zeta=nzeta + endpoint,
+                NFP=NFP,
+                axis=axis,
+                endpoint=endpoint,
+            )
+        )
+        # The test below might fail for theta and zeta surfaces only if nrho > 1.
+        # This is unrelated to how duplicate node spacing is handled.
+        # The cause is because grid construction does not always
+        # compute drho as constant (even when the rho nodes are linearly spaced),
+        # and this test assumes drho to be a constant for grids without duplicates.
+        test(
+            LinearGrid(
+                rho=np.linspace(1, 0, nrho, endpoint=axis)[::-1],
+                theta=np.linspace(0, 2 * np.pi, ntheta + endpoint, endpoint=endpoint),
+                zeta=np.linspace(
+                    0, 2 * np.pi / NFP, nzeta + endpoint, endpoint=endpoint
+                ),
+                NFP=NFP,
+                axis=axis,
+                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):