coordax — tutorials
Pedagogical tutorials for coordax, the coordinate-aware JAX array library
coordax notebooks¶
Showcase notebooks ported from jejjax.numpy
arrays and the full xarray Hoyer & Hamman (2017) stack.
Each notebook is executed end-to-end with outputs embedded, so everything (prints, tables, numeric checks) renders inline in the MyST docs site without re-execution.
Each sub-section is a curated landing page that leads with the math, numerics, and references before pointing at the notebooks themselves.
Foundations¶
| Notebook | Topic |
|---|---|
foundations/01_create_datasets.ipynb | Wrapping arrays as Field objects with LabeledAxis / SizedAxis |
foundations/02_ops_unary_binary.ipynb | Arithmetic, comparison, and unary ops on Field; broadcasting rules |
foundations/03_ops_coordinates.ipynb | isel, sel, reindexing, CartesianProduct, coordinate composition |
foundations/04_reductions.ipynb | Coordinate-aware reductions: sum, mean, max over named dims |
Derivatives¶
| Notebook | Topic |
|---|---|
derivatives/05_finite_difference.ipynb | Periodic + non-uniform finite-difference derivatives; cmap pattern |
derivatives/06_spherical_harmonics_derivatives.ipynb | Spherical-harmonic derivatives on a Gauss-Legendre lat-lon grid (no FD pole singularity); Laplacian; vorticity and divergence |
derivatives/07_finite_volume.ipynb | Cell-centred FV operators; flux divergence; conservative schemes |
Dynamics¶
| Notebook | Topic |
|---|---|
dynamics/08_ode_integration.ipynb | Integrating ODEs (advection-diffusion) with diffrax; state as Field |
dynamics/09_ode_parameter_state_estimation.ipynb | Joint parameter/state estimation via optax + jax.value_and_grad |
dynamics/10_pde_parameter_estimation.ipynb | Learning PDE parameters from data; coordinate-aware residuals |
Running locally¶
These notebooks depend on coordax and its ML stack (JAX, diffrax, optax,
equinox). The committed .ipynb files carry their cell outputs, so MyST
renders them inline without needing the kernel installed — no extra setup
is required to view the site.
Re-executing the notebooks is not part of this repo’s tooling. The source
environment lives in jej_vc_snippets/jax/coordax,
which ships a pixi environment with coordax, JAX, diffrax, optax, and
equinox pinned together.
- Hoyer, S., & Hamman, J. (2017). xarray: N-D Labeled Arrays and Datasets in Python. Journal of Open Research Software, 5(1). 10.5334/jors.148