Repository: PlasmaControl/DESC

Issue:
Utility for copying data between grids that share nodes
EG, when computing iota on specific surfaces we then need to copy those iota values to another grid with the same rho but different theta/zeta.

Right now we do this using `grid.compress` and `grid.expand` but that doesn't really work for the sort of unstructured grids that get created under jit that we need for field line following stuff.

Heres a rough draft:
```python
def copy_data_grids(x, grid1, grid2, match="rpz", tol=1e-14):
    """Copy data from one grid to another at matching nodes.
    
    Given data x corresponding to nodes of grid1, copy data
    to a new array that corresponds to grid2.
    
    Parameters
    ----------
    x : ndarray, shape(grid1.num_nodes,...)
        Data to copy
    grid1, grid2: Grid
        Grid to copy from and copy to
    match {"r", "p", "z"}
        What coordinates to match. For local quantities,
        all three coordinates must match. For flux surface 
        qtys, only rho must match, etc.
    tol : float
        tolerance for considering nodes the same.
        
    Returns
    -------
    y : ndarray, shape(grid2.num_nodes, ...)
        Data copied to grid2
    """
    axes = ()
    if "r" in match:
        axes += (0,)
    if "t" in match:
        axes += (1,)
    if "z" in match:
        axes += (2,)
    y = jnp.zeros((grid2.num_nodes, *x.shape[1:]))
    
    def body(i, y):
        y = jnp.where(jnp.max(jnp.abs((grid2.nodes[:,axes] - grid1.nodes[i,axes])), axis=-1) < tol, x[i], y)
        return y
    
    y = fori_loop(0, grid1.num_nodes, body, y)
    return y
    
```

From what @YigitElma learned doing the zernike radial stuff loops might not be the best option, so vmap or some other tricks might be better. Anybody want to try out a few options?

Discussion:

Patch:
diff --git a/desc/compute/utils.py b/desc/compute/utils.py
index 3ef93e2cc..677b29e66 100644
--- a/desc/compute/utils.py
+++ b/desc/compute/utils.py
@@ -695,24 +695,31 @@ def _get_grid_surface(grid, surface_label):
     """
     assert surface_label in {"rho", "theta", "zeta"}
     if surface_label == "rho":
-        unique_size = grid.num_rho
-        inverse_idx = grid.inverse_rho_idx
         spacing = grid.spacing[:, 1:]
         has_endpoint_dupe = False
+        unique_size = getattr(grid, "num_rho", -1)
+        inverse_idx = getattr(grid, "_inverse_rho_idx", jnp.array([]))
     elif surface_label == "theta":
-        unique_size = grid.num_theta
-        inverse_idx = grid.inverse_theta_idx
         spacing = grid.spacing[:, [0, 2]]
-        has_endpoint_dupe = (grid.nodes[grid.unique_theta_idx[0], 1] == 0) & (
-            grid.nodes[grid.unique_theta_idx[-1], 1] == 2 * np.pi
+        unique_size = getattr(grid, "num_theta", -1)
+        inverse_idx = getattr(grid, "_inverse_theta_idx", jnp.array([]))
+        has_endpoint_dupe = (
+            isinstance(grid, LinearGrid)
+            and hasattr(grid, "_unique_theta_idx")
+            and (grid.nodes[grid.unique_theta_idx[0], 1] == 0)
+            & (grid.nodes[grid.unique_theta_idx[-1], 1] == 2 * np.pi)
         )
     else:
-        unique_size = grid.num_zeta
-        inverse_idx = grid.inverse_zeta_idx
         spacing = grid.spacing[:, :2]
-        has_endpoint_dupe = (grid.nodes[grid.unique_zeta_idx[0], 2] == 0) & (
-            grid.nodes[grid.unique_zeta_idx[-1], 2] == 2 * np.pi / grid.NFP
+        unique_size = getattr(grid, "num_zeta", -1)
+        inverse_idx = getattr(grid, "_inverse_zeta_idx", jnp.array([]))
+        has_endpoint_dupe = (
+            isinstance(grid, LinearGrid)
+            and hasattr(grid, "_unique_zeta_idx")
+            and (grid.nodes[grid.unique_zeta_idx[0], 2] == 0)
+            & (grid.nodes[grid.unique_zeta_idx[-1], 2] == 2 * np.pi / grid.NFP)
         )
+
     return unique_size, inverse_idx, spacing, has_endpoint_dupe
 
 
@@ -722,6 +729,7 @@ def line_integrals(
     line_label="theta",
     fix_surface=("rho", 1.0),
     expand_out=True,
+    tol=1e-14,
 ):
     """Compute line integrals over curves covering the given surface.
 
@@ -749,18 +757,9 @@ def line_integrals(
     q : ndarray
         Quantity to integrate.
         The first dimension of the array should have size ``grid.num_nodes``.
-
-        When ``q`` is 1-dimensional, the intention is to integrate,
-        over the domain parameterized by rho, theta, and zeta,
-        a scalar function over the previously mentioned domain.
-
-        When ``q`` is 2-dimensional, the intention is to integrate,
-        over the domain parameterized by rho, theta, and zeta,
-        a vector-valued function over the previously mentioned domain.
-
-        When ``q`` is 3-dimensional, the intention is to integrate,
+        When ``q`` is n-dimensional, the intention is to integrate,
         over the domain parameterized by rho, theta, and zeta,
-        a matrix-valued function over the previously mentioned domain.
+        an n-dimensional function over the previously mentioned domain.
     line_label : str
         The coordinate curve to compute the integration over.
         To clarify, a theta (poloidal) curve is the intersection of a
@@ -774,6 +773,9 @@ def line_integrals(
         shape as the input. Defaults to true so that the output may be
         broadcast in the same way as the input. Setting to false will save
         memory.
+    tol : float
+        Tolerance for considering nodes the same.
+        Only relevant if the grid object doesn't already have this information.
 
     Returns
     -------
@@ -806,10 +808,12 @@ def line_integrals(
     mask = grid.nodes[:, column_id] == fix_surface[1]
     q_prime = (mask * jnp.atleast_1d(q).T / grid.spacing[:, column_id]).T
     (surface_label,) = labels.keys() - {line_label, fix_surface[0]}
-    return surface_integrals(grid, q_prime, surface_label, expand_out)
+    return surface_integrals(grid, q_prime, surface_label, expand_out, tol)
 
 
-def surface_integrals(grid, q=jnp.array([1.0]), surface_label="rho", expand_out=True):
+def surface_integrals(
+    grid, q=jnp.array([1.0]), surface_label="rho", expand_out=True, tol=1e-14
+):
     """Compute a surface integral for each surface in the grid.
 
     Notes
@@ -825,18 +829,9 @@ def surface_integrals(grid, q=jnp.array([1.0]), surface_label="rho", expand_out=
     q : ndarray
         Quantity to integrate.
         The first dimension of the array should have size ``grid.num_nodes``.
-
-        When ``q`` is 1-dimensional, the intention is to integrate,
-        over the domain parameterized by rho, theta, and zeta,
-        a scalar function over the previously mentioned domain.
-
-        When ``q`` is 2-dimensional, the intention is to integrate,
+        When ``q`` is n-dimensional, the intention is to integrate,
         over the domain parameterized by rho, theta, and zeta,
-        a vector-valued function over the previously mentioned domain.
-
-        When ``q`` is 3-dimensional, the intention is to integrate,
-        over the domain parameterized by rho, theta, and zeta,
-        a matrix-valued function over the previously mentioned domain.
+        an n-dimensional function over the previously mentioned domain.
     surface_label : str
         The surface label of rho, theta, or zeta to compute the integration over.
     expand_out : bool
@@ -844,6 +839,9 @@ def surface_integrals(grid, q=jnp.array([1.0]), surface_label="rho", expand_out=
         shape as the input. Defaults to true so that the output may be
         broadcast in the same way as the input. Setting to false will save
         memory.
+    tol : float
+        Tolerance for considering nodes the same.
+        Only relevant if the grid object doesn't already have this information.
 
     Returns
     -------
@@ -852,10 +850,10 @@ def surface_integrals(grid, q=jnp.array([1.0]), surface_label="rho", expand_out=
         By default, the returned array has the same shape as the input.
 
     """
-    return surface_integrals_map(grid, surface_label, expand_out)(q)
+    return surface_integrals_map(grid, surface_label, expand_out, tol)(q)
 
 
-def surface_integrals_map(grid, surface_label="rho", expand_out=True):
+def surface_integrals_map(grid, surface_label="rho", expand_out=True, tol=1e-14):
     """Returns a method to compute any surface integral for each surface in the grid.
 
     Parameters
@@ -869,6 +867,9 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True):
         shape as the input. Defaults to true so that the output may be
         broadcast in the same way as the input. Setting to false will save
         memory.
+    tol : float
+        Tolerance for considering nodes the same.
+        Only relevant if the grid object doesn't already have this information.
 
     Returns
     -------
@@ -888,40 +889,53 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True):
     unique_size, inverse_idx, spacing, has_endpoint_dupe = _get_grid_surface(
         grid, surface_label
     )
+    spacing = jnp.prod(spacing, axis=1)
 
     # Todo: Define masks as a sparse matrix once sparse matrices are no longer
     #       experimental in jax.
-    # The ith row of masks is True only at the indices which correspond to the
-    # ith surface. The integral over the ith surface is the dot product of the
-    # ith row vector and the vector of integrands of all surfaces.
-    masks = inverse_idx == jnp.arange(unique_size)[:, jnp.newaxis]
-    # Imagine a torus cross-section at zeta=π.
-    # A grid with a duplicate zeta=π node has 2 of those cross-sections.
-    #     In grid.py, we multiply by 1/n the areas of surfaces with
-    # duplicity n. This prevents the area of that surface from being
-    # double-counted, as surfaces with the same node value are combined
-    # into 1 integral, which sums their areas. Thus, if the zeta=π
-    # cross-section has duplicity 2, we ensure that the area on the zeta=π
-    # surface will have the correct total area of π+π = 2π.
-    #     An edge case exists if the duplicate surface has nodes with
-    # different values for the surface label, which only occurs when
-    # has_endpoint_dupe is true. If ``has_endpoint_dupe`` is true, this grid
-    # has a duplicate surface at surface_label=0 and
-    # surface_label=max surface value. Although the modulo of these values
-    # are equal, their numeric values are not, so the integration
-    # would treat them as different surfaces. We solve this issue by
-    # combining the indices corresponding to the integrands of the duplicated
-    # surface, so that the duplicate surface is treated as one, like in the
-    # previous paragraph.
-    masks = cond(
-        has_endpoint_dupe,
-        lambda _: put(masks, jnp.array([0, -1]), masks[0] | masks[-1]),
-        lambda _: masks,
-        operand=None,
-    )
-    spacing = jnp.prod(spacing, axis=1)
-
-    def _surface_integrals(q=jnp.array([1.0])):
+    if hasattr(grid, f"num_{surface_label}"):
+        # The ith row of masks is True only at the indices which correspond to the
+        # ith surface. The integral over the ith surface is the dot product of the
+        # ith row vector and the integrand defined over all the surfaces.
+        mask = inverse_idx == jnp.arange(unique_size)[:, jnp.newaxis]
+        # Imagine a torus cross-section at zeta=π.
+        # A grid with a duplicate zeta=π node has 2 of those cross-sections.
+        #     In grid.py, we multiply by 1/n the areas of surfaces with
+        # duplicity n. This prevents the area of that surface from being
+        # double-counted, as surfaces with the same node value are combined
+        # into 1 integral, which sums their areas. Thus, if the zeta=π
+        # cross-section has duplicity 2, we ensure that the area on the zeta=π
+        # surface will have the correct total area of π+π = 2π.
+        #     An edge case exists if the duplicate surface has nodes with
+        # different values for the surface label, which only occurs when
+        # has_endpoint_dupe is true. If ``has_endpoint_dupe`` is true, this grid
+        # has a duplicate surface at surface_label=0 and
+        # surface_label=max surface value. Although the modulo of these values
+        # are equal, their numeric values are not, so the integration
+        # would treat them as different surfaces. We solve this issue by
+        # combining the indices corresponding to the integrands of the duplicated
+        # surface, so that the duplicate surface is treated as one, like in the
+        # previous paragraph.
+        mask = cond(
+            has_endpoint_dupe,
+            lambda _: put(mask, jnp.array([0, -1]), mask[0] | mask[-1]),
+            lambda _: mask,
+            operand=None,
+        )
+    else:
+        expand_out = False
+        nodes = grid.nodes[:, {"rho": 0, "theta": 1, "zeta": 2}[surface_label]]
+        # Converting nodes from numpy.ndarray to jaxlib.xla_extension.ArrayImpl
+        # reduces memory usage by > 400% for the forward computation and Jacobian.
+        nodes = jnp.asarray(nodes)
+        # This branch will execute for custom grids, which don't have a use
+        # case for having duplicate nodes, so we don't bother to modulo nodes
+        # by 2pi or 2pi/NFP.
+        mask = jnp.abs(nodes - nodes[:, jnp.newaxis]) <= tol
+        # The above implementation was benchmarked to be more efficient than
+        # alternatives with explicit loops in GitHub pull request #934.
+
+    def integrate(q=jnp.array([1.0])):
         """Compute a surface integral for each surface in the grid.
 
         Notes
@@ -935,18 +949,9 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True):
         q : ndarray
             Quantity to integrate.
             The first dimension of the array should have size ``grid.num_nodes``.
-
-            When ``q`` is 1-dimensional, the intention is to integrate,
+            When ``q`` is n-dimensional, the intention is to integrate,
             over the domain parameterized by rho, theta, and zeta,
-            a scalar function over the previously mentioned domain.
-
-            When ``q`` is 2-dimensional, the intention is to integrate,
-            over the domain parameterized by rho, theta, and zeta,
-            a vector-valued function over the previously mentioned domain.
-
-            When ``q`` is 3-dimensional, the intention is to integrate,
-            over the domain parameterized by rho, theta, and zeta,
-            a matrix-valued function over the previously mentioned domain.
+            an n-dimensional function over the previously mentioned domain.
 
         Returns
         -------
@@ -954,44 +959,11 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True):
             Surface integral of the input over each surface in the grid.
 
         """
-        axis_to_move = (jnp.ndim(q) == 3) * 2
         integrands = (spacing * jnp.nan_to_num(q).T).T
-        # `integrands` may have shape (g.size, f.size, v.size), where
-        #     g is the grid function depending on the integration variables
-        #     f is a function which may be independent of the integration variables
-        #     v is the vector of components of f (or g).
-        # The intention is to integrate `integrands` which is a
-        #     vector-valued            (with v.size components)
-        #     function-valued          (with image size of f.size)
-        #     function over the grid   (with domain size of g.size = grid.num_nodes)
-        # over each surface in the grid.
-
-        # The distinction between f and v is semantic.
-        # We may alternatively consider an `integrands` of shape (g.size, f.size) to
-        # represent a vector-valued (with f.size components) function over the grid.
-        # Likewise, we may alternatively consider an `integrands` of shape
-        # (g.size, v.size) to represent a function-valued (with image size v.size)
-        # function over the grid. When `integrands` has dimension one, it is a
-        # scalar function over the grid. That is, a
-        #     vector-valued            (with 1 component),
-        #     function-valued          (with image size of 1)
-        #     function over the grid   (with domain size of g.size = grid.num_nodes)
-
-        # The integration is performed by applying `masks`, the surface
-        # integral operator, to `integrands`. This operator hits the matrix formed
-        # by the last two dimensions of `integrands`, for every element along the
-        # previous dimension of `integrands`. Therefore, when `integrands` has three
-        # dimensions, the second must hold g. We may choose which of the first and
-        # third dimensions hold f and v. The choice below transposes `integrands` to
-        # shape (v.size, g.size, f.size). As we expect f.size >> v.size, the
-        # integration is in theory faster since numpy optimizes large matrix
-        # products. However, timing results showed no difference.
-        integrals = jnp.moveaxis(
-            masks @ jnp.moveaxis(integrands, axis_to_move, 0), 0, axis_to_move
-        )
+        integrals = jnp.tensordot(mask, integrands, axes=1)
         return grid.expand(integrals, surface_label) if expand_out else integrals
 
-    return _surface_integrals
+    return integrate
 
 
 def surface_averages(
@@ -1001,6 +973,7 @@ def surface_averages(
     surface_label="rho",
     denominator=None,
     expand_out=True,
+    tol=1e-14,
 ):
     """Compute a surface average for each surface in the grid.
 
@@ -1016,18 +989,9 @@ def surface_averages(
     q : ndarray
         Quantity to average.
         The first dimension of the array should have size ``grid.num_nodes``.
-
-        When ``q`` is 1-dimensional, the intention is to average,
-        over the domain parameterized by rho, theta, and zeta,
-        a scalar function over the previously mentioned domain.
-
-        When ``q`` is 2-dimensional, the intention is to average,
-        over the domain parameterized by rho, theta, and zeta,
-        a vector-valued function over the previously mentioned domain.
-
-        When ``q`` is 3-dimensional, the intention is to average,
+        When ``q`` is n-dimensional, the intention is to average,
         over the domain parameterized by rho, theta, and zeta,
-        a matrix-valued function over the previously mentioned domain.
+        an n-dimensional function over the previously mentioned domain.
     sqrt_g : ndarray
         Coordinate system Jacobian determinant; see ``data_index["sqrt(g)"]``.
     surface_label : str
@@ -1044,6 +1008,9 @@ def surface_averages(
         shape as the input. Defaults to true so that the output may be
         broadcast in the same way as the input. Setting to false will save
         memory.
+    tol : float
+        Tolerance for considering nodes the same.
+        Only relevant if the grid object doesn't already have this information.
 
     Returns
     -------
@@ -1052,10 +1019,12 @@ def surface_averages(
         By default, the returned array has the same shape as the input.
 
     """
-    return surface_averages_map(grid, surface_label, expand_out)(q, sqrt_g, denominator)
+    return surface_averages_map(grid, surface_label, expand_out, tol)(
+        q, sqrt_g, denominator
+    )
 
 
-def surface_averages_map(grid, surface_label="rho", expand_out=True):
+def surface_averages_map(grid, surface_label="rho", expand_out=True, tol=1e-14):
     """Returns a method to compute any surface average for each surface in the grid.
 
     Parameters
@@ -1069,6 +1038,9 @@ def surface_averages_map(grid, surface_label="rho", expand_out=True):
         shape as the input. Defaults to true so that the output may be
         broadcast in the same way as the input. Setting to false will save
         memory.
+    tol : float
+        Tolerance for considering nodes the same.
+        Only relevant if the grid object doesn't already have this information.
 
     Returns
     -------
@@ -1078,7 +1050,9 @@ def surface_averages_map(grid, surface_label="rho", expand_out=True):
         ``function(q, sqrt_g)``.
 
     """
-    compute_surface_integrals = surface_integrals_map(grid, surface_label, False)
+    if not hasattr(grid, f"num_{surface_label}"):
+        expand_out = False  # don't try to expand already expanded output
+    integrate = surface_integrals_map(grid, surface_label, expand_out=False, tol=tol)
 
     def _surface_averages(q, sqrt_g=jnp.array([1.0]), denominator=None):
         """Compute a surface average for each surface in the grid.
@@ -1093,18 +1067,9 @@ def surface_averages_map(grid, surface_label="rho", expand_out=True):
         q : ndarray
             Quantity to average.
             The first dimension of the array should have size ``grid.num_nodes``.
-
-            When ``q`` is 1-dimensional, the intention is to average,
-            over the domain parameterized by rho, theta, and zeta,
-            a scalar function over the previously mentioned domain.
-
-            When ``q`` is 2-dimensional, the intention is to average,
+            When ``q`` is n-dimensional, the intention is to average,
             over the domain parameterized by rho, theta, and zeta,
-            a vector-valued function over the previously mentioned domain.
-
-            When ``q`` is 3-dimensional, the intention is to average,
-            over the domain parameterized by rho, theta, and zeta,
-            a matrix-valued function over the previously mentioned domain.
+            an n-dimensional function over the previously mentioned domain.
         sqrt_g : ndarray
             Coordinate system Jacobian determinant; see ``data_index["sqrt(g)"]``.
         denominator : ndarray
@@ -1123,14 +1088,14 @@ def surface_averages_map(grid, surface_label="rho", expand_out=True):
         """
         q = jnp.atleast_1d(q)
         sqrt_g = jnp.atleast_1d(sqrt_g)
-        numerator = compute_surface_integrals((sqrt_g * q.T).T)
+        numerator = integrate((sqrt_g * q.T).T)
         # memory optimization to call expand() at most once
         if denominator is None:
             # skip integration if constant
             denominator = (
                 (4 * jnp.pi**2 if surface_label == "rho" else 2 * jnp.pi) * sqrt_g
                 if sqrt_g.size == 1
-                else compute_surface_integrals(sqrt_g)
+                else integrate(sqrt_g)
             )
             averages = (numerator.T / denominator).T
             if expand_out:
@@ -1187,28 +1152,16 @@ def surface_integrals_transform(grid, surface_label="rho"):
         The second dimension may discretize some function, f, over the
         codomain, and therefore, have size that matches the desired number of
         points at which the output is evaluated.
-        If the integrand is vector-valued then the third dimension may
-        hold the components of size v.size.
 
-        This method can also be used to compute the output one point at a time.
-        In this case, ``q`` will be at most two-dimensional, and the second
-        dimension may hold the vector components.
-
-        There is technically no difference between the labels f and v, so their
-        roles may be swapped if this is more convenient.
+        This method can also be used to compute the output one point at a time,
+        in which case ``q`` can have shape (``grid.num_nodes``, ).
 
         Input
         -----
-        If ``q`` is one-dimensional, then it should have shape
+        If ``q`` has one-dimension, then it should have shape
         (``grid.num_nodes``, ).
-        If ``q`` is two-dimensional, then either
-            1) g and f are scalar functions, so the input should have shape
-               (``grid.num_nodes``, f.size).
-            2) g (or f) is a vector-valued function, and f has been evaluated at
-               only one point, so the input should have shape
-               (``grid.num_nodes``, v.size).
-        If ``q`` is three-dimensional, then it should have shape
-        (``grid.num_nodes``, f.size, v.size).
+        If ``q`` has multiple dimensions, then it should have shape
+        (``grid.num_nodes``, *f.shape).
 
         Output
         ------
@@ -1216,23 +1169,17 @@ def surface_integrals_transform(grid, surface_label="rho"):
         Tᵤ₁ for a particular surface of constant u₁ in the given grid.
         The order is sorted in increasing order of the values which specify u₁.
 
-        If ``q`` is one-dimensional, the returned array has shape
+        If ``q`` has one dimension, the returned array has shape
         (grid.num_surface_label, ).
-        If ``q`` is two-dimensional, the returned array has shape
-        (grid.num_surface_label, (f or v).size), depending on whether f or v is
-        the relevant label.
-        If ``q`` is three-dimensional, the returned array has shape
-        (grid.num_surface_label, f.size, v.size).
+        If ``q`` has multiple dimensions, the returned array has shape
+        (grid.num_surface_label, *f.shape).
 
     """
-    # Although this method seems to duplicate surface_integrals(), the
-    # intentions of these methods may be to implement different algorithms.
-    # We can rely on surface_integrals() for the computation because its current
-    # implementation is flexible enough to implement both algorithms.
     # Expansion should not occur here. The typical use case of this method is to
     # transform into the computational domain, so the second dimension that
     # discretizes f over the codomain will typically have size grid.num_nodes
     # to broadcast with quantities in data_index.
+    assert hasattr(grid, f"num_{surface_label}")
     return surface_integrals_map(grid, surface_label, expand_out=False)
 
 
@@ -1243,6 +1190,7 @@ def surface_variance(
     bias=False,
     surface_label="rho",
     expand_out=True,
+    tol=1e-14,
 ):
     """Compute the weighted sample variance of ``q`` on each surface of the grid.
 
@@ -1304,6 +1252,9 @@ def surface_variance(
         shape as the input. Defaults to true so that the output may be
         broadcast in the same way as the input. Setting to false will save
         memory.
+    tol : float
+        Tolerance for considering nodes the same.
+        Only relevant if the grid object doesn't already have this information.
 
     Returns
     -------
@@ -1313,7 +1264,7 @@ def surface_variance(
 
     """
     _, _, spacing, _ = _get_grid_surface(grid, surface_label)
-    integrate = surface_integrals_map(grid, surface_label, expand_out=False)
+    integrate = surface_integrals_map(grid, surface_label, expand_out=False, tol=tol)
 
     v1 = integrate(weights)
     v2 = integrate(weights**2 * jnp.prod(spacing, axis=-1))
@@ -1325,9 +1276,13 @@ def surface_variance(
     q = jnp.atleast_1d(q)
     # compute variance in two passes to avoid catastrophic round off error
     mean = (integrate((weights * q.T).T).T / v1).T
-    mean = grid.expand(mean, surface_label)
+    if hasattr(grid, f"num_{surface_label}"):
+        mean = grid.expand(mean, surface_label)
     variance = (correction * integrate((weights * ((q - mean) ** 2).T).T).T / v1).T
-    return grid.expand(variance, surface_label) if expand_out else variance
+    if hasattr(grid, f"num_{surface_label}") and expand_out:
+        return grid.expand(variance, surface_label)
+    else:
+        return variance
 
 
 def surface_max(grid, x, surface_label="rho"):
diff --git a/desc/equilibrium/coords.py b/desc/equilibrium/coords.py
index 7c0b4e044..21ded70e3 100644
--- a/desc/equilibrium/coords.py
+++ b/desc/equilibrium/coords.py
@@ -102,14 +102,15 @@ def map_coordinates(  # noqa: C901
 
     params = setdefault(params, eq.params_dict)
 
+    profiles = get_profiles(inbasis + basis_derivs, eq, None)
+    # do surface average to get iota once
+    if "iota" in profiles and profiles["iota"] is None:
+        profiles["iota"] = eq.get_profile("iota")
+
     @functools.partial(jit, static_argnums=1)
     def compute(y, basis):
         grid = Grid(y, sort=False, jitable=True)
-        profiles = get_profiles(inbasis + basis_derivs, eq, grid, jitable=True)
-        # do surface average to get iota once
-        if "current" in profiles and profiles["current"] is not None:
-            profiles["iota"] = eq.get_profile("iota")
-        transforms = get_transforms(inbasis + basis_derivs, eq, grid, jitable=True)
+        transforms = get_transforms(basis, eq, grid, jitable=True)
         data = compute_fun(eq, basis, params, transforms, profiles)
         x = jnp.array([data[k] for k in basis]).T
         return x
diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py
index 91f96fa87..5a4328659 100644
--- a/desc/equilibrium/equilibrium.py
+++ b/desc/equilibrium/equilibrium.py
@@ -896,9 +896,7 @@ class Equilibrium(IOAble, Optimizable):
             )
             # need to make this data broadcast with the data on the original grid
             data1dr = {
-                key: grid.expand(
-                    grid1dr.compress(val, surface_label="rho"), surface_label="rho"
-                )
+                key: grid.copy_data_from_other(val, grid1dr, surface_label="rho")
                 for key, val in data1dr.items()
                 if key in dep1dr
             }
@@ -925,9 +923,7 @@ class Equilibrium(IOAble, Optimizable):
             )
             # need to make this data broadcast with the data on the original grid
             data1dz = {
-                key: grid.expand(
-                    grid1dz.compress(val, surface_label="zeta"), surface_label="zeta"
-                )
+                key: grid.copy_data_from_other(val, grid1dz, surface_label="zeta")
                 for key, val in data1dz.items()
                 if key in dep1dz
             }
diff --git a/desc/grid.py b/desc/grid.py
index 91e343f85..f7a89d29a 100644
--- a/desc/grid.py
+++ b/desc/grid.py
@@ -5,9 +5,9 @@ from abc import ABC, abstractmethod
 import numpy as np
 from scipy import optimize, special
 
-from desc.backend import jnp, put
+from desc.backend import fori_loop, jnp, put
 from desc.io import IOAble
-from desc.utils import Index
+from desc.utils import Index, errorif
 
 __all__ = [
     "Grid",
@@ -227,32 +227,62 @@ class _Grid(IOAble, ABC):
     @property
     def unique_rho_idx(self):
         """ndarray: Indices of unique rho coordinates."""
-        return self.__dict__.setdefault("_unique_rho_idx", np.array([]))
+        errorif(
+            not hasattr(self, "_unique_rho_idx"),
+            AttributeError,
+            "jit compatible Grid objects do not have unique indices assigned.",
+        )
+        return self._unique_rho_idx
 
     @property
     def unique_theta_idx(self):
         """ndarray: Indices of unique theta coordinates."""
-        return self.__dict__.setdefault("_unique_theta_idx", np.array([]))
+        errorif(
+            not hasattr(self, "_unique_theta_idx"),
+            AttributeError,
+            "jit compatible Grid objects do not have unique indices assigned.",
+        )
+        return self._unique_theta_idx
 
     @property
     def unique_zeta_idx(self):
         """ndarray: Indices of unique zeta coordinates."""
-        return self.__dict__.setdefault("_unique_zeta_idx", np.array([]))
+        errorif(
+            not hasattr(self, "_unique_zeta_idx"),
+            AttributeError,
+            "jit compatible Grid objects do not have unique indices assigned.",
+        )
+        return self._unique_zeta_idx
 
     @property
     def inverse_rho_idx(self):
         """ndarray: Indices of unique_rho_idx that recover the rho coordinates."""
-        return self.__dict__.setdefault("_inverse_rho_idx", np.array([]))
+        errorif(
+            not hasattr(self, "_inverse_rho_idx"),
+            AttributeError,
+            "jit compatible Grid objects do not have inverse indices assigned.",
+        )
+        return self._inverse_rho_idx
 
     @property
     def inverse_theta_idx(self):
         """ndarray: Indices of unique_theta_idx that recover the theta coordinates."""
-        return self.__dict__.setdefault("_inverse_theta_idx", np.array([]))
+        errorif(
+            not hasattr(self, "_inverse_theta_idx"),
+            AttributeError,
+            "jit compatible Grid objects do not have inverse indices assigned.",
+        )
+        return self._inverse_theta_idx
 
     @property
     def inverse_zeta_idx(self):
         """ndarray: Indices of unique_zeta_idx that recover the zeta coordinates."""
-        return self.__dict__.setdefault("_inverse_zeta_idx", np.array([]))
+        errorif(
+            not hasattr(self, "_inverse_zeta_idx"),
+            AttributeError,
+            "jit compatible Grid objects do not have inverse indices assigned.",
+        )
+        return self._inverse_zeta_idx
 
     @property
     def axis(self):
@@ -299,7 +329,7 @@ class _Grid(IOAble, ABC):
             The array to compress.
             Should usually represent a surface function (constant over a surface)
             in an array that matches the grid's pattern.
-        surface_label : str
+        surface_label : {"rho", "theta", "zeta"}
             The surface label of rho, theta, or zeta.
 
         Returns
@@ -313,10 +343,25 @@ class _Grid(IOAble, ABC):
         assert surface_label in {"rho", "theta", "zeta"}
         assert len(x) == self.num_nodes
         if surface_label == "rho":
+            errorif(
+                not hasattr(self, "_unique_rho_idx"),
+                AttributeError,
+                "compress operation undefined for jit compatible grids",
+            )
             return x[self.unique_rho_idx]
         if surface_label == "theta":
+            errorif(
+                not hasattr(self, "_unique_theta_idx"),
+                AttributeError,
+                "compress operation undefined for jit compatible grids",
+            )
             return x[self.unique_theta_idx]
         if surface_label == "zeta":
+            errorif(
+                not hasattr(self, "_unique_zeta_idx"),
+                AttributeError,
+                "compress operation undefined for jit compatible grids",
+            )
             return x[self.unique_zeta_idx]
 
     def expand(self, x, surface_label="rho"):
@@ -332,7 +377,7 @@ class _Grid(IOAble, ABC):
             that the first element corresponds to the value associated with the
             smallest surface, and the last element corresponds to the value
             associated with the largest surface.
-        surface_label : str
+        surface_label : {"rho", "theta", "zeta"}
             The surface label of rho, theta, or zeta.
 
         Returns
@@ -343,15 +388,74 @@ class _Grid(IOAble, ABC):
         """
         assert surface_label in {"rho", "theta", "zeta"}
         if surface_label == "rho":
+            errorif(
+                not hasattr(self, "_inverse_rho_idx"),
+                AttributeError,
+                "expand operation undefined for jit compatible grids",
+            )
             assert len(x) == self.num_rho
             return x[self.inverse_rho_idx]
         if surface_label == "theta":
+            errorif(
+                not hasattr(self, "_inverse_theta_idx"),
+                AttributeError,
+                "expand operation undefined for jit compatible grids",
+            )
             assert len(x) == self.num_theta
             return x[self.inverse_theta_idx]
         if surface_label == "zeta":
+            errorif(
+                not hasattr(self, "_inverse_zeta_idx"),
+                AttributeError,
+                "expand operation undefined for jit compatible grids",
+            )
             assert len(x) == self.num_zeta
             return x[self.inverse_zeta_idx]
 
+    def copy_data_from_other(self, x, other_grid, surface_label="rho", tol=1e-14):
+        """Copy data x from other_grid to this grid at matching surface label.
+
+        Given data x corresponding to nodes of other_grid, copy data
+        to a new array that corresponds to this grid.
+
+        Parameters
+        ----------
+        x : ndarray, shape(other_grid.num_nodes,...)
+            Data to copy
+        other_grid: Grid
+            Grid to copy from.
+        surface_label : {"rho", "theta", "zeta"}
+            The surface label of rho, theta, or zeta.
+        tol : float
+            tolerance for considering nodes the same.
+
+        Returns
+        -------
+        y : ndarray, shape(grid2.num_nodes, ...)
+            Data copied to grid2
+        """
+        x = jnp.asarray(x)
+        assert surface_label in {"rho", "theta", "zeta"}
+        try:
+            xc = other_grid.compress(x, surface_label)
+            y = self.expand(xc, surface_label)
+        except AttributeError:
+            self_nodes = jnp.asarray(self.nodes)
+            other_nodes = jnp.asarray(other_grid.nodes)
+            axis = {"rho": 0, "theta": 1, "zeta": 2}[surface_label]
+            y = jnp.zeros((self.num_nodes, *x.shape[1:]))
+
+            def body(i, y):
+                y = jnp.where(
+                    jnp.abs(self_nodes[:, axis] - other_nodes[i, axis]) <= tol,
+                    x[i],
+                    y,
+                )
+                return y
+
+            y = fori_loop(0, other_grid.num_nodes, body, y)
+        return y
+
     def replace_at_axis(self, x, y, copy=False, **kwargs):
         """Replace elements of ``x`` with elements of ``y`` at the axis of grid.
 
@@ -403,7 +507,7 @@ class Grid(_Grid):
         etc may be wrong if grid contains duplicate nodes.
     """
 
-    def __init__(self, nodes, sort=False, jitable=False):
+    def __init__(self, nodes, sort=False, jitable=False, **kwargs):
         # Python 3.3 (PEP 412) introduced key-sharing dictionaries.
         # This change measurably reduces memory usage of objects that
         # define all attributes in their __init__ method.
@@ -420,14 +524,19 @@ 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 = 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)
+            # allow for user supplied indices/inverse indices for special cases
+            for attr in [
+                "_unique_rho_idx",
+                "_unique_theta_idx",
+                "_unique_zeta_idx",
+                "_inverse_rho_idx",
+                "_inverse_theta_idx",
+                "_inverse_zeta_idx",
+            ]:
+                if attr in kwargs:
+                    setattr(self, attr, jnp.asarray(kwargs.pop(attr)))
         else:
             self._enforce_symmetry()
             self._axis = self._find_axis()
@@ -441,9 +550,10 @@ class Grid(_Grid):
             ) = self._find_unique_inverse_nodes()
             self._weights = self._scale_weights()
 
-        self._L = self.num_rho
-        self._M = self.num_theta
-        self._N = self.num_zeta
+        self._L = self.num_nodes
+        self._M = self.num_nodes
+        self._N = self.num_nodes
+        errorif(len(kwargs), ValueError, f"Got unexpected kwargs {kwargs.keys()}")
 
     def _create_nodes(self, nodes):
         """Allow for custom node creation.
diff --git a/desc/input_reader.py b/desc/input_reader.py
index 0bf71f3e0..8789eec6a 100644
--- a/desc/input_reader.py
+++ b/desc/input_reader.py
@@ -874,7 +874,7 @@ class InputReader:
 
         # fit profiles to power series
         grid = LinearGrid(L=eq.L_grid, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP)
-        rho = grid.nodes[grid._unique_rho_idx, 0]
+        rho = grid.nodes[grid.unique_rho_idx, 0]
         if not isinstance(eq.pressure, PowerSeriesProfile):
             pressure = grid.compress(eq.compute("p", grid=grid)["p"])
             pres_profile = PowerSeriesProfile.from_values(
diff --git a/desc/profiles.py b/desc/profiles.py
index 5b9ade5bf..65e413989 100644
--- a/desc/profiles.py
+++ b/desc/profiles.py
@@ -645,9 +645,9 @@ class PowerSeriesProfile(_Profile):
         if self.sym:
             # need to pad with odd numbered modes
             params = jnp.array([params, jnp.zeros_like(params)]).flatten(order="F")
-        r = grid.nodes[grid.unique_rho_idx, 0]
+        r = grid.nodes[:, 0]
         f = polyval_vec(polyder_vec(jnp.atleast_2d(params[::-1]), dr, False), r)[0]
-        return f[grid.inverse_rho_idx]
+        return f
 
     @classmethod
     def from_values(cls, x, y, order=6, rcond=None, w=None, sym="auto", name=""):
@@ -775,11 +775,11 @@ class SplineProfile(_Profile):
             params = self.params
         if dt != 0 or dz != 0:
             return jnp.zeros_like(grid.nodes[:, 0])
-        xq = grid.nodes[grid.unique_rho_idx, 0]
         x = self.knots
         f = params
+        xq = grid.nodes[:, 0]
         fq = interp1d(xq, x, f, method=self._method, derivative=dr, extrap=True)
-        return fq[grid.inverse_rho_idx]
+        return fq
 
 
 class MTanhProfile(_Profile):
@@ -920,16 +920,14 @@ class MTanhProfile(_Profile):
         if dt != 0 or dz != 0:
             return jnp.zeros_like(grid.nodes[:, 0])
 
-        xq = grid.nodes[grid.unique_rho_idx, 0]
-
         ped = params[0]
         offset = params[1]
         sym = params[2]
         width = params[3]
         core_poly = params[4:]
-
+        xq = grid.nodes[:, 0]
         y = MTanhProfile._mtanh(xq, ped, offset, sym, width, core_poly, dx=dr)
-        return y[grid.inverse_rho_idx]
+        return y
 
     @classmethod
     def from_values(


Test patch:
diff --git a/tests/test_compute_utils.py b/tests/test_compute_utils.py
index 02a6947b8..b8a69c20f 100644
--- a/tests/test_compute_utils.py
+++ b/tests/test_compute_utils.py
@@ -118,6 +118,27 @@ class TestComputeUtils:
                 # theta integrals are poorly defined on concentric grids
                 test_b_theta(label, cg_sym, eq)
 
+    @pytest.mark.unit
+    def test_unknown_unique_grid_integral(self):
+        """Test that averages are invariant to whether grids have unique_idx."""
+        lg = LinearGrid(L=L, M=M, N=N, NFP=NFP, endpoint=False)
+        q = jnp.arange(lg.num_nodes) ** 2
+        result = surface_integrals(lg, q, surface_label="rho")
+        del lg._unique_rho_idx
+        np.testing.assert_allclose(
+            surface_integrals(lg, q, surface_label="rho"), result
+        )
+        result = surface_averages(lg, q, surface_label="theta")
+        del lg._unique_theta_idx
+        np.testing.assert_allclose(
+            surface_averages(lg, q, surface_label="theta"), result
+        )
+        result = surface_variance(lg, q, surface_label="zeta")
+        del lg._unique_zeta_idx
+        np.testing.assert_allclose(
+            surface_variance(lg, q, surface_label="zeta"), result
+        )
+
     @pytest.mark.unit
     def test_surface_integrals_transform(self):
         """Test surface integral of a kernel function."""
diff --git a/tests/test_grid.py b/tests/test_grid.py
index 0f773afdd..62309419a 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -785,30 +785,53 @@ def test_find_least_rational_surfaces():
 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
-        )
+    eq.change_resolution(3, 3, 3, 6, 6, 6)
+    # field lines on two surfaces
+    rho = np.concatenate([0.5 * np.ones(10), 0.7 * np.ones(10)])
+    theta = np.concatenate([np.linspace(0, 1, 10), np.linspace(0, 1, 10)]) * 2 * np.pi
+    zeta = np.concatenate([np.linspace(0, 1, 10), np.linspace(0, 1, 10)]) * 2 * np.pi
+    grid1 = Grid(np.array([rho, theta, zeta]).T, jitable=False)
+    grid2 = Grid(np.array([rho, theta, zeta]).T, jitable=True)
+    grid3 = Grid(np.array([rho, theta, zeta]).T, jitable=True, _unique_rho_idx=[0, 10])
+    np.testing.assert_allclose(grid1.nodes, grid2.nodes)
+
+    x = np.random.random(grid1.num_nodes)
+
+    # these shouldn't error
+    _ = grid1.unique_rho_idx
+    _ = grid1.unique_theta_idx
+    _ = grid1.unique_zeta_idx
+    _ = grid1.inverse_rho_idx
+    _ = grid1.inverse_theta_idx
+    _ = grid1.inverse_zeta_idx
+    _ = grid3.unique_rho_idx
+
+    with pytest.raises(AttributeError):
+        _ = grid2.unique_rho_idx
+    with pytest.raises(AttributeError):
+        _ = grid2.unique_theta_idx
+    with pytest.raises(AttributeError):
+        _ = grid2.unique_zeta_idx
+    with pytest.raises(AttributeError):
+        _ = grid2.inverse_rho_idx
+    with pytest.raises(AttributeError):
+        _ = grid2.inverse_theta_idx
+    with pytest.raises(AttributeError):
+        _ = grid2.inverse_zeta_idx
+
+    y1 = grid1.copy_data_from_other(x, grid2, "rho")
+    y2 = grid2.copy_data_from_other(x, grid1, "rho")
+    y3 = grid3.copy_data_from_other(x, grid1, "rho")
+
+    np.testing.assert_allclose(y1, y2)
+    np.testing.assert_allclose(y1, y3)
 
     # 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)
+    with pytest.raises(AttributeError):
+        _ = eq.compute(["|B|"], grid=grid2, override_grid=True)["|B|"]
+    b3 = eq.compute(["|B|"], grid=grid3, override_grid=True)["|B|"]
+    np.testing.assert_allclose(b1, b3)
 
 
 @pytest.mark.unit