Repository: PlasmaControl/DESC

Issue:
Fix weights and indexing for `Grid` when `jitable=True`
Like the cause of https://github.com/PlasmaControl/DESC/pull/854#discussion_r1504948317

Basically, when we set `jitable=True` the resulting weights aren't scaled correctly which can throw off any averaging done with that grid. This is usually fine because any averaging is done on a structured grid, but we should still fix it (and using a structured grid doesn't always fix it, see below).

The other problem is that `Grid.unique_rho_idx` etc are set as `np.arange(num_nodes)` when `jitable=True`, but it uses `np.unique` when `jitable=False`. `np.unique` returns the _sorted_ unique values, which means that when we do `grid1.expand(grid2.compress(x))` things get all jumbled up because `grid1` has sorted indices and `grid2` doesn't.

Discussion:

Patch:
diff --git a/desc/grid.py b/desc/grid.py
index 5578911a8..0a5051036 100644
--- a/desc/grid.py
+++ b/desc/grid.py
@@ -422,12 +422,12 @@ class Grid(_Grid):
             r = jnp.where(r == 0, 1e-12, r)
             self._nodes = jnp.array([r, t, z]).T
             self._axis = np.array([], dtype=int)
-            self._unique_rho_idx = np.arange(self._nodes.shape[0])
-            self._unique_theta_idx = np.arange(self._nodes.shape[0])
-            self._unique_zeta_idx = np.arange(self._nodes.shape[0])
-            self._inverse_rho_idx = np.arange(self._nodes.shape[0])
-            self._inverse_theta_idx = np.arange(self._nodes.shape[0])
-            self._inverse_zeta_idx = np.arange(self._nodes.shape[0])
+            self._unique_rho_idx = jnp.argsort(self._nodes[:, 0])
+            self._unique_theta_idx = jnp.argsort(self._nodes[:, 1])
+            self._unique_zeta_idx = jnp.argsort(self._nodes[:, 2])
+            self._inverse_rho_idx = jnp.argsort(self._unique_rho_idx)
+            self._inverse_theta_idx = jnp.argsort(self._unique_theta_idx)
+            self._inverse_zeta_idx = jnp.argsort(self._unique_zeta_idx)
             # don't do anything fancy with weights
             self._weights = self._spacing.prod(axis=1)
         else:
@@ -470,7 +470,9 @@ class Grid(_Grid):
         # surfaces outside the interval [0, 2pi] as duplicates. However, most
         # surface integral computations are done with LinearGrid anyway.
         spacing = (  # make weights sum to 4pi^2
-            jnp.ones_like(nodes) * jnp.array([1, 2 * np.pi, 2 * np.pi]) / nodes.shape[0]
+            jnp.ones_like(nodes)
+            * jnp.array([1, 2 * np.pi, 2 * np.pi])
+            / nodes.shape[0] ** (1 / 3)
         )
         return nodes, spacing
 


Test patch:
diff --git a/tests/test_grid.py b/tests/test_grid.py
index b2fcb679d..52b1c8572 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -5,6 +5,7 @@ import pytest
 from scipy import special
 
 from desc.equilibrium import Equilibrium
+from desc.examples import get
 from desc.grid import (
     ConcentricGrid,
     Grid,
@@ -774,3 +775,48 @@ def test_find_least_rational_surfaces():
     max_rational = max(lior)
 
     assert np.all(np.array(lio) > max_rational)
+
+
+@pytest.mark.unit
+def test_custom_jitable_grid_indexing():
+    """Test that unique/inverse indices are set correctly when jitable=True."""
+    eq = get("NCSX")
+    rho = np.random.random(100)
+    theta = np.random.random(100) * 2 * np.pi
+    zeta = np.random.random(100) * 2 * np.pi / eq.NFP
+    grid1 = Grid(np.array([rho, theta, zeta]).T, jitable=True)
+    grid2 = Grid(np.array([rho, theta, zeta]).T, jitable=False)
+
+    attrs = [
+        "unique_rho_idx",
+        "inverse_rho_idx",
+        "unique_theta_idx",
+        "inverse_theta_idx",
+        "unique_zeta_idx",
+        "inverse_zeta_idx",
+    ]
+
+    for attr in attrs:
+        np.testing.assert_array_equal(
+            getattr(grid1, attr), getattr(grid2, attr), err_msg=attr
+        )
+
+    # make sure compress/expand done when override_grid=True works as expected
+    b1 = eq.compute(["|B|"], grid=grid1, override_grid=True)["|B|"]
+    b2 = eq.compute(["|B|"], grid=grid2, override_grid=True)["|B|"]
+    np.testing.assert_allclose(b1, b2)
+
+
+@pytest.mark.unit
+def test_custom_jitable_grid_weights():
+    """Test that grid weights are set correctly when jitable=True."""
+    rho = np.random.random(100)
+    theta = np.random.random(100) * 2 * np.pi
+    zeta = np.random.random(100) * 2 * np.pi
+    grid1 = Grid(np.array([rho, theta, zeta]).T, jitable=True)
+    grid2 = Grid(np.array([rho, theta, zeta]).T, jitable=False)
+
+    np.testing.assert_allclose(grid1.spacing, grid2.spacing)
+    np.testing.assert_allclose(grid1.weights, grid2.weights)
+    np.testing.assert_allclose(grid1.weights.sum(), 4 * np.pi**2)
+    np.testing.assert_allclose(grid2.weights.sum(), 4 * np.pi**2)