Coordinax
Coordinates in JAX
Install / Use
/learn @GalacticDynamics/CoordinaxREADME
<h1 align='center'> coordinax </h1>
<h3 align="center">Coordinates in JAX</h3>
<p align="center">
<a href="https://pypi.org/project/coordinax/"> <img alt="PyPI: coordinax" src="https://img.shields.io/pypi/v/coordinax?style=flat" /> </a>
<a href="https://pypi.org/project/coordinax/"> <img alt="PyPI versions: coordinax" src="https://img.shields.io/pypi/pyversions/coordinax" /> </a>
<a href="https://coordinax.readthedocs.io/en/"> <img alt="ReadTheDocs" src="https://img.shields.io/badge/read_docs-here-orange" /> </a>
<a href="https://pypi.org/project/coordinax/"> <img alt="coordinax license" src="https://img.shields.io/github/license/GalacticDynamics/coordinax" /> </a>
</p>
<p align="center">
<a href="https://github.com/GalacticDynamics/coordinax/actions"> <img alt="CI status" src="https://github.com/GalacticDynamics/coordinax/workflows/CI/badge.svg" /> </a>
<a href="https://coordinax.readthedocs.io/en/"> <img alt="ReadTheDocs" src="https://readthedocs.org/projects/coordinax/badge/?version=latest" /> </a>
<a href="https://codecov.io/gh/GalacticDynamics/coordinax"> <img alt="codecov" src="https://codecov.io/gh/GalacticDynamics/coordinax/graph/badge.svg" /> </a>
<a href="https://scientific-python.org/specs/spec-0000/"> <img alt="ruff" src="https://img.shields.io/badge/SPEC-0-green?labelColor=%23004811&color=%235CA038" /> </a>
<a href="https://docs.astral.sh/ruff/"> <img alt="ruff" src="https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v2.json" /> </a>
<a href="https://pre-commit.com"> <img alt="pre-commit" src="https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit" /> </a>
</p>
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 / (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 = 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(
# KinematicSpace({
# '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(
# KinematicSpace({
# '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!
<!-- prettier-ignore-start --> <!-- prettier-ignore-end -->