coordinax

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.Quantity(jnp.arange(0, 10.0), "kpc"),
    y=u.Quantity(jnp.arange(5, 15.0), "kpc"),
    z=u.Quantity(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.Quantity(jnp.arange(0, 10.0), "km/s"),
    y=u.Quantity(jnp.arange(5, 15.0), "km/s"),
    z=u.Quantity(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 / (km s)], phi[km rad / (km s)])
#     [[ 1.118e+01 -3.886e-16  0.000e+00]
#      [ 1.257e+01 -1.110e-16  0.000e+00]
#      ...
#      [ 2.360e+01  0.000e+00  0.000e+00]
#      [ 2.526e+01 -2.776e-16  0.000e+00]]>


# Transforming between frames
icrs_frame = cx.frames.ICRS()
gc_frame = cx.frames.Galactocentric()
op = cxf.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(cx.Space(length=q, speed=p), frame=icrs_frame)
print(coord)
# Coordinate(
#     data=Space({
#        '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(
#     data=Space({
#        '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!