coordinax

Coordinates in JAX

Coordinax enables calculations with coordinates in JAX. Built on Equinox and Quax.

Installation

pip install coordinax

Documentation

Quick example

import jax.numpy as jnp
import unxt as u
import coordinax as cx

q = cx.CartesianPos3D(
    x=u.Q(jnp.arange(0, 10.0), "kpc"),
    y=u.Q(jnp.arange(5, 15.0), "kpc"),
    z=u.Q(jnp.arange(10, 20.0), "kpc"),
)
print(q)
# <CartesianPos3D: (x, y, z) [kpc]
#     [[ 0.  5. 10.]
#      [ 1.  6. 11.]
#      ...
#      [ 8. 13. 18.]
#      [ 9. 14. 19.]]>

q2 = cx.vconvert(cx.SphericalPos, q)
print(q2)
# <SphericalPos: (r[kpc], theta[rad], phi[rad])
#     [[11.18   0.464  1.571]
#      [12.57   0.505  1.406]
#      ...
#      [23.601  0.703  1.019]
#      [25.259  0.719  0.999]]>

p = cx.CartesianVel3D(
    x=u.Q(jnp.arange(0, 10.0), "km/s"),
    y=u.Q(jnp.arange(5, 15.0), "km/s"),
    z=u.Q(jnp.arange(10, 20.0), "km/s"),
)
print(p)
# <CartesianVel3D: (x, y, z) [km / s]
#     [[ 0.  5. 10.]
#      [ 1.  6. 11.]
#      ...
#      [ 8. 13. 18.]
#      [ 9. 14. 19.]]>

p2 = cx.vconvert(cx.SphericalVel, p, q)
print(p2)
# <SphericalVel: (r[km / s], theta[km rad / (kpc s)], phi[km rad / (kpc s)])
#     [[ 1.118e+01  2.980e-08  0.000e+00]
#      [ 1.257e+01  1.192e-07  1.490e-08]
#      ...
#      [ 2.360e+01  8.941e-08  0.000e+00]
#      [ 2.526e+01  2.086e-07  2.980e-08]]>


# Transforming between frames
icrs_frame = cx.frames.ICRS()
gc_frame = cx.frames.Galactocentric()
op = cx.frames.frame_transform_op(icrs_frame, gc_frame)
q_gc, p_gc = op(q, p)
print(q_gc, p_gc, sep="\n")
# <CartesianPos3D: (x, y, z) [kpc]
#     [[-1.732e+01  5.246e+00  3.614e+00]
#      ...
#      [-3.004e+01  1.241e+01 -1.841e+00]]>
# <CartesianVel3D: (x, y, z) [km / s]
#      [[  3.704 250.846  11.373]
#       ...
#       [ -9.02  258.012   5.918]]>

coord = cx.Coordinate({"length": q, "speed": p}, frame=icrs_frame)
print(coord)
# Coordinate(
#   {
#     'length':
#     <CartesianPos3D: (x, y, z) [kpc]
#         [[ 0.  5. 10.]
#          ...
#          [ 9. 14. 19.]]>,
#     'speed':
#     <CartesianVel3D: (x, y, z) [km / s]
#         [[ 0.  5. 10.]
#          ...
#          [ 9. 14. 19.]]>
#   },
#   frame=ICRS()
# )

print(coord.to_frame(gc_frame))
# Coordinate(
#   {
#     'length':
#     <CartesianPos3D: (x, y, z) [kpc]
#         [[-1.732e+01  5.246e+00  3.614e+00]
#          ...
#          [-3.004e+01  1.241e+01 -1.841e+00]]>,
#     'speed':
#     <CartesianVel3D: (x, y, z) [km / s]
#         [[  3.704 250.846  11.373]
#          ...
#          [ -9.02  258.012   5.918]]>
#   },
#   frame=Galactocentric( ... )
# )

Citation

If you found this library to be useful in academic work, then please cite.

Development

We welcome contributions!

Metadata

Release files for coordinax 0.24.2

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for coordinax 0.24.2
File	Size	Uploaded
coordinax-0.24.2.tar.gz	729.8 kB	Oct 11, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for coordinax 0.24.2
File	Interpreter	ABI	Platform	Reset
coordinax-0.24.2-py3-none-any.whl 195.4 kB Oct 11, 2026	Python 3	none	any	Details

Total release size: 925.2 kB

Release files / coordinax-0.24.2.tar.gz

Download URL	coordinax-0.24.2.tar.gz
Size	729.8 kB
Tags	Source
SHA-256 checksum How to use checksums	`319448c79e2ddd62eb88e75a211470307ea857f34d761db27df4670dfcacd1be`
BLAKE2b-256 checksum How to use checksums	`a5315cfadf65d9d6c4eb103cf97b7649789a68a3311503bcbaec2f0f6cfc600e`
Upload date	Oct 11, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/GalacticDynamics/coordinax
Publishing commit: github.com/GalacticDynamics/coordinax/tree/908ba8aaa88c33c0f8a1c895897eaed61a2133a0

Workflow

Publishing configuration: github.com/GalacticDynamics/coordinax/blob/908ba8aaa88c33c0f8a1c895897eaed61a2133a0/.github/workflows/cd.yml
Publishing logs: github.com/GalacticDynamics/coordinax/actions/runs/38098305306/attempts/1

Artifact

Subject: coordinax-0.24.2.tar.gz
SHA-256 checksum: 319448c79e2ddd62eb88e75a211470307ea857f34d761db27df4670dfcacd1be Verifying attestations

Release files / coordinax-0.24.2-py3-none-any.whl

Download URL	coordinax-0.24.2-py3-none-any.whl
Size	195.4 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`1b16cf655a1bf5573c6c217a85db3c69a4c17cb7b812412310afb094eb232e8b`
BLAKE2b-256 checksum How to use checksums	`35b948e53284c5e6f0d6a66c7662dee4ca9630787d4f73c97d98f6eacc605e28`
Upload date	Oct 11, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	Yes
Uploaded via	`twine/7.0.0 CPython/3.13.14`

Provenance

Identity

Publishing platform: GitHub Actions
Publishing repository: github.com/GalacticDynamics/coordinax
Publishing commit: github.com/GalacticDynamics/coordinax/tree/908ba8aaa88c33c0f8a1c895897eaed61a2133a0

Workflow

Publishing configuration: github.com/GalacticDynamics/coordinax/blob/908ba8aaa88c33c0f8a1c895897eaed61a2133a0/.github/workflows/cd.yml
Publishing logs: github.com/GalacticDynamics/coordinax/actions/runs/38098305306/attempts/1

Artifact

Subject: coordinax-0.24.2-py3-none-any.whl
SHA-256 checksum: 1b16cf655a1bf5573c6c217a85db3c69a4c17cb7b812412310afb094eb232e8b Verifying attestations