Merge branch 'dp/docs-update-data' of github.com:PlasmaControl/DESC i…

…nto dp/docs-update-data

desc/grid.py

@@ -422,12 +422,12 @@ def __init__(self, nodes, sort=False, jitable=False):
             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 @@ def _create_nodes(self, nodes):
         # 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
 

tests/test_grid.py

@@ -5,6 +5,7 @@
 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)