🚀 Get Started

`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:

pip install git+https://https://github.com/GalacticDynamics/coordinax.git

building from source

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

cd /path/to/parent
git clone https://https://github.com/GalacticDynamics/coordinax.git
cd coordinax
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
vector objects
- 1D, 2D, 3D and N-dimensional vector classes in many different representation types (e.g., CartesianPos1D, CartesianPos2D, CartesianPos3D, SphericalPos, ProlateSpheroidalPos, etc.)
- time-differential vector objects, like velocities and accelerations (e.g., CartesianVel3D, SphericalVel, etc.)
- collections of vector objects like KinematicSpace
transformations of vector types (vconvert())
operations on vectors (ops)
reference frames and coordinate systems (frames)
coordinates that combine vectors and frames (Coordinate)
and more!

This functionality is organized into submodules, which are imported into the top-level coordinax namespace. You can import them directly, or use the coordinax namespace to access them.

>>> import coordinax as cx

>>> from inspect import ismodule
>>> [name for name in cx.__all__ if ismodule(getattr(cx, name))]
['angle', 'distance', 'vecs', 'ops', 'frames']

We recommend importing as needed:

coordinax as cx
coordinax.vecs as cxv
coordinax.ops as cxo
coordinax.frames as cxf

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π or -180º to 180º).

>>> import unxt as u

>>> a = u.Angle(370, "deg")
>>> a
Angle(Array(370, dtype=int32, weak_type=True), unit='deg')

>>> a.wrap_to(u.Quantity(0, "deg"), u.Quantity(360, "deg"))
Angle(Array(10, dtype=int32, weak_type=True), unit='deg')

Similarly, the Distance class represents distances in coordinax:

>>> d = cx.distance.Distance(10, "kpc")
>>> d
Distance(Array(10, dtype=int32, weak_type=True), unit='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.

>>> d.parallax
Parallax(Array(4.848137e-10, dtype=float32, weak_type=True), unit='rad')

>>> d.distance_modulus
DistanceModulus(Array(15., dtype=float32), unit='mag')

Creating and Working with Vector Objects#

Vector objects in coordinax represent positions, velocities, and accelerations in a variety of coordinate systems and dimensions. Here are some common operations:

Constructing Vector Objects#

You can create a vector by specifying its components and units:

>>> import coordinax.vecs as cxv

>>> q = cxv.CartesianPos3D.from_([1, 2, 3], "kpc")
>>> print(q)
<CartesianPos3D: (x, y, z) [kpc]
    [1 2 3]>

The from_() method is a flexible constructor that allows you to create vectors from various input formats, such as lists, tuples, or NumPy arrays. Direct construction is also possible by specifying values for all components:

>>> q = cxv.CartesianPos3D(x=u.Quantity(1, "kpc"), y=u.Quantity(2, "kpc"), z=u.Quantity(3, "kpc"))
>>> print(q)
<CartesianPos3D: (x, y, z) [kpc]
    [1 2 3]>

The most flexible way to create vectors is to use the vector() method, which infers the appropriate vector class based on the provided inputs:

>>> q = cxv.vector([1, 2, 3], "kpc")
>>> print(q)
<CartesianPos3D: (x, y, z) [kpc]
    [1 2 3]>

Vector Conversion#

Vectors can be converted between different coordinate representations using the vconvert() method. For example, to convert a Cartesian position vector to spherical coordinates:

>>> sph = q.vconvert(cxv.SphericalPos)
>>> print(sph)
<SphericalPos: (r[kpc], theta[rad], phi[rad])
    [3.742 0.641 1.107]>

Transforming Velocities#

Velocity vectors can also be converted to other representations, but require specifying the corresponding position:

>>> v = cxv.CartesianVel3D.from_([4, 5, 6], "kpc/Myr")
>>> v_sph = v.vconvert(cxv.SphericalVel, q)
>>> print(v_sph)
<SphericalVel: (r[kpc / Myr], theta[rad / Myr], phi[rad / Myr])
    [ 8.552  0.383 -0.6  ]>

Creating a `KinematicSpace` Object#

A KinematicSpace object collects related vectors (e.g., position, velocity, acceleration) into a single container:

>>> import coordinax as cx

>>> space = cx.KinematicSpace(length=q, speed=v)
>>> print(space)
KinematicSpace({
   'length': <CartesianPos3D: (x, y, z) [kpc]
       [1 2 3]>,
   'speed': <CartesianVel3D: (x, y, z) [kpc / Myr]
       [4 5 6]>
})

You can convert all vectors in a KinematicSpace to a different representation at once:

>>> space_sph = space.vconvert(cxv.SphericalPos)
>>> print(space_sph)
KinematicSpace({
    'length': <SphericalPos: (r[kpc], theta[rad], phi[rad])
                [3.742 0.641 1.107]>,
    'speed': <SphericalVel: (r[kpc / Myr], theta[rad / Myr], phi[rad / Myr])
                [ 8.552  0.383 -0.6  ]>
})

Operators on Vectors#

The coordinax.ops module (shorthand cxo) provides a framework for and set of vector operations that work seamlessly with all coordinax vector types.

>>> import coordinax.ops as cxo

>>> op = cxo.GalileanSpatialTranslation.from_([10, 10, 10], "kpc")

>>> print(op(q))
<CartesianPos3D: (x, y, z) [kpc]
    [11 12 13]>

Reference Frames and Coordinates#

coordinax.frames (shorthand cxf) provides a framework for defining and working with reference frames and coordinate systems.

>>> import coordinax.frames as cxf

>>> alice = cxf.Alice()
>>> alice
Alice()

>>> bob = cxf.Bob()
>>> bob
Bob()

Frames can be used to define coordinate transformations. For example, you can transform a position vector from the Alice frame to the Bob frame:

>>> op = cxf.frame_transform_op(alice, bob)
>>> t = u.Quantity(1, "yr")
>>> print(op(t, q)[1])
<CartesianPos3D: (x, y, z) [kpc]
    [1. 2. 3.]>

Coordinate objects can also be created to represent positions in a specific frame:

>>> coord = cxf.Coordinate(q, frame=alice)
>>> print(coord)
Coordinate(
    { 'length': <CartesianPos3D: (x, y, z) [kpc]
                    [1 2 3]> },
    frame=Alice()
)

>>> coord.to_frame(bob, t)
Coordinate(
    KinematicSpace({ 'length': CartesianPos3D(
        x=Quantity(f32[], unit='kpc'),
        y=Quantity(f32[], unit='kpc'),
        z=Quantity(f32[], unit='kpc')
    ) }),
    frame=Bob()
)

>>> coord.vconvert(cxv.SphericalPos)
Coordinate(
    KinematicSpace({ 'length': SphericalPos(
          r=Distance(f32[], unit='kpc'),
          theta=Angle(f32[], unit='rad'),
          phi=Angle(f32[], unit='rad')
        ) }),
    frame=Alice()
)

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.