PlasmaControl/DESC

Reminder to update tests/benchmarks with the new Grid API after #2053 is merged into master. (Revert commit 6b68c5679fe53fd83cfcb21871db3ab94f65ab50).

Should pullback a single contangent through the entire desc compute infrastructure to extend factor N memory and performance improvement, where N is number of rows in Jacobian, to full pipeline. All you have to do is put the sparse pullback decorator over that single function. Few lines changed . For meshgrids this is easy. This would require the inputs and outputs the to the registered compute functions to be of shape (rho, theta, zeta, ...) . That's actually the natural shape that np.meshgrid(rho,theta,zeta, ij) returns. Then wrap the full compute pipeline in a single sparse pullback decorator except for Zernike transform. Edit: see below for simpler Before exiting to return data to users use tree_map(meshgrid_reshape, data) so that retuned data to users is same structure as before.

Comments for future consideration/improved usability: angle is a little clunky, it would be nice if it were auto-computed and stored within the Bounce2D functional or within the objective, maybe it is fine for the objective part, but when I compute a quantity outside of an objective it is unintuitive to compute angle first, separately. if angle is nonexistent it should be computed automatically, perhaps inside the Bounce2D creation... Originally posted by @dpanici Reply by @unalmis : Ideally, it would be computable quantity in data_index and in Bounce2D.required_names. Then users can change this code: from desc.integrals import Bounce2D angle = Bounce2D.angle(eq, rho=grid.compress(grid.nodes[: ,0])) datas = eq.compute("effective ripple", grid=grid, angle=angle) to datas = eq.compute("effective ripple", grid=grid)