coordinax

🚀 Get Started

coordinax enables working with coordinates and reference frames with JAX.

coordinax supports JAX’s main features:

• JIT compilation (jit())

• vectorization (vmap(), etc.)

• auto-differentiation (grad(), jacobian(), jax.hessian())

• GPU/TPU/multi-host acceleration

And best of all, coordinax doesn’t force you to use special unit-compatible re-exports of JAX libraries. You can use coordinax with existing JAX code, and with one simple decorator (quax.quaxify()), JAX will work with coordinax objects.

---

Installation

pip

pip install coordinax

uv

uv add coordinax

source, via pip

To install the latest development version of coordinax directly from the GitHub repository, use pip:

uv add git+https://https://github.com/GalacticDynamics/coordinax.git@main

You can customize the branch by replacing main with any other branch name.

building from source

To build coordinax from source, clone the repository and install it with uv:

cd /path/to/parent
git clone https://https://github.com/GalacticDynamics/coordinax.git
cd coordinax
uv pip install -e .  # editable mode

Quickstart

The coordinax package has powerful tools for representing, using, and transforming coordinate objects, such as:

• specific Quantity subclasses like Angle and Distance

• and more!

This functionality is organized into submodules available under the top-level coordinax namespace. You can import them directly, or for many objects use the coordinax.main namespace to access them.

>>> import coordinax
>>> import sys

>>> sorted(name.removeprefix("coordinax.")
...        for name in sys.modules
...        if name.startswith("coordinax.") and name.count(".") == 1)[1:]
['angles', 'api', 'astro', 'charts', 'curveframes', 'distances', 'frames', 'hypothesis', 'internal', 'interop', 'main', 'manifolds', 'representations', 'transforms', 'vectors']

We recommend importing as needed:

• coordinax.main as cx : probably everything you need!

• coordinax.angles as cxa : further angle-specific functionality.

• coordinax.distances as cxd : further distance-specific functionality.

• coordinax.charts as cxc : chart-specific functionality.

• coordinax.frames as cxf : frame-specific functionality.

• coordinax.manifolds as cxm : manifold-specific functionality.

• coordinax.representations as cxr : representation-specific functionality.

• coordinax.transforms as cxt : transform-specific functionality.

• coordinax.vectors as cxv : vector-specific functionality.

• coordinax.astro as cxastro : astronomy-specific functionality. Note that this package is an optional extra, so you may need to install it separately.

• coordinax.hypothesis as cxst : property-based testing strategies for coordinax. Note that this package is an optional extra, so you may need to install it separately.

• coordinax.interop.astropy as cxapy : interoperability with astropy. Note that this package is an optional extra, so you may need to install it separately.

Angles and Distances

coordinax is built on top of unxt, which provides support for quantity objects that represent a data array with an associated unit with the unxt.quantity.Quantity class. These Quantity objects can be used throughout coordinax, but coordinax also provides specific classes that offer additional functionality.

Let’s start with angles, which are represented by the Angle class. This class enforces that the inputted units have angular dimensions and provides some other useful utilities for working with angles. For example, the resulting Angle (a re-export of unxt.Angle) object can be wrapped to a specific range to conform to a branch cut (e.g., 0 to \(2\pi\) or \(-180^\circ\) to \(180^\circ\)).

>>> import coordinax.main as cx
>>> import unxt as u

>>> a = cx.Angle(370, "deg")
>>> a
Angle(370, 'deg')

>>> a.wrap_to(u.Q(0, "deg"), u.Q(360, "deg"))
Angle(10, 'deg')

Similarly, the Distance class represents distances in coordinax:

>>> d = cx.Distance(10, "kpc")
>>> d
Distance(10, 'kpc')

but other distance-like objects can be represented with the Parallax and DistanceModulus classes. These classes check that the units have distance dimensions, and they provide useful properties for converting between different distance representations.

>>> import coordinax.astro as cxastro
>>> import plum

>>> plum.convert(d, cxastro.Parallax)
Parallax(4.84813681e-10, 'rad')

>>> plum.convert(d, cxastro.DistanceModulus)
DistanceModulus(15., 'mag')

Ecosystem

coordinax’s Dependencies

• unxt: Quantities in JAX.

• Equinox: one-stop JAX library, for everything that isn’t already in core JAX.

• Quax: JAX + multiple dispatch + custom array-ish objects.

• Quaxed: pre-quaxifyed Jax.

• plum: multiple dispatch in python

coordinax’s Dependents

• galax: Galactic dynamics in JAX.