grid-dependent coordinate Jacobian in compute functions

[unalmis]

As we add quantities that require a grid generated from other coordinates, we will need to use the custom Grid class to compute more often. An inconvenience of these grids is that they do not automatically compute differential area and volume elements required for surface and volume integrals. (It's difficult since the structure of the grid is not known apriori). So we need to compute them on other grids and interpolate or broadcast. It can be cumbersome to do this and will require more memory.

We should consider generalizing the compute functions to feed in coordinate Jacobians based on grid.coordinates, (or just both Jacobians, one for the computational grid desc_grid and another for the source grid of that compute fun, accessible via desc_grid.source_grid after #1024). Then we should be able to compute these quadratures on the source grid that has the differential area elements defined. Also the source grids may typically be tensor product grids while the desc grid is not, so you would get better convergence properties of quadratures on the source grid.

[unalmis]

I'm suggesting to generalize the compute functions to something like

@register_compute_fun(...)
def _thing(data, transforms, profiles, **kwargs):
   grid = getattr(transforms["grid"], "source_grid", transforms["grid"])
   data["thing"] = surface_integrals(grid, thing, sqrt_g=get_Jacobian(data, grid.coordinates))
   return data

[f0uriest]

I could see something like this working, but I'd vote for moving the extra logic to inside the surface integral function, or a higher order wrapper for it.

There may also be some complications with dependencies, since now what qty are needed to compute a given thing could depend on what grid coordinates are used?

[unalmis]

I have updated this issue as the old comments were out of date.

With the change in the above comment, stuff like this

data = compute_fun(
    eq,
    self._keys_1dr, # flux surface quantities like iota, iota_r, etc.
    params,
    constants["transforms_1dr"],
    constants["profiles"],
)
grid = eq.get_rtz_grid(
    rho,
    vartheta,
    phi,
    coordinates="rvp",
    period=(np.inf, 2 * np.pi, 2 * np.pi),
    params=params,
)
data = {
    key: grid.copy_data_from_other(data[key], self._grid_1dr)
    for key in self._keys_1dr
}
data = compute_fun(
    eq,
    self._keys, # stuff to compute along field lines
    params,
    get_transforms(self._keys, eq, grid, jitable=True),
    constants["profiles"],
    data=data,
)

could be simplified to code below. In particular the additional transforms_1dr built for a uniformly spaced DESC coordinate grid would not be necessary as the vartheta and phi coordinates are usually uniformly spaced.

grid = eq.get_rtz_grid(
    rho,
    vartheta,
    phi,
    coordinates="rvp",
    period=(np.inf, 2 * np.pi, 2 * np.pi),
    params=params,
)
data = compute_fun(
    eq,
    self._keys,
    params,
    get_transforms(self._keys, eq, grid, jitable=True),
    constants["profiles"],
)