Authors: Dmitrii Kochkov and Stephan Hoyer
Coordax makes it easy to associate array dimensions with coordinates in the context of scientific simulation codes written in JAX. This allows for efficient and expressive manipulation of data defined on structured grids, enabling operations like differentiation and interpolation with respect to physical coordinates.
Coordax was designed to meet the needs of NeuralGCM, but we hope it will be useful more broadly!
- Compute on locally-positional axes via coordinate map (
cmap) - Coordinate objects that carry discretization details and custom methods
- Lossless conversion to and from Xarray data structures (e.g., for serialization)
Coordax is particularly well-suited for scientific simulations where it is crucial to propagate discretization details and associated objects throughout the computation, such as Earth system modeling of fluid dynamics. The approach to labeled dimensions was originally forked from Daniel Johnson's Penzai, which may be a better fit for simpler use-cases.
Xarray does indeed support putting JAX arrays into Xarray data structures, an
approach used by the
GraphCast codebase.
This works reasonably well, but wrapping JAX in Xarray will always be at least a
little bit awkward, because Xarray was designed for the needs of data analysis
rather than modeling, and cannot build core functionality on top of JAX power
features such as vmap. For JAX-native simulations, we believe Coordax is a
better choice.
Coming soon!