https://github.com/galacticdynamics/coordinax
Coordinates in JAX
https://github.com/galacticdynamics/coordinax
coordinate-systems coordinate-transformation coordinates jax multiple-dispatch quax reference-frames vectors
Last synced: 29 days ago
JSON representation
Coordinates in JAX
- Host: GitHub
- URL: https://github.com/galacticdynamics/coordinax
- Owner: GalacticDynamics
- License: mit
- Created: 2024-02-10T21:21:00.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-09-04T23:23:12.000Z (about 1 month ago)
- Last Synced: 2025-09-05T01:15:19.276Z (about 1 month ago)
- Topics: coordinate-systems, coordinate-transformation, coordinates, jax, multiple-dispatch, quax, reference-frames, vectors
- Language: Python
- Homepage: https://coordinax.readthedocs.io/en/latest/
- Size: 2.27 MB
- Stars: 25
- Watchers: 3
- Forks: 4
- Open Issues: 15
-
Metadata Files:
- Readme: README.md
- Contributing: .github/CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Security: SECURITY.md
Awesome Lists containing this project
README
coordinax
Coordinates in JAX
---
Coordinax enables calculations with coordinates in
[JAX](https://jax.readthedocs.io/en/latest/). Built on
[Equinox](https://docs.kidger.site/equinox/) and
[Quax](https://github.com/patrick-kidger/quax).
## Installation
[![PyPI platforms][pypi-platforms]][pypi-link]
[![PyPI version][pypi-version]][pypi-link]
```bash
pip install coordinax
```
## Documentation
[](https://coordinax.readthedocs.io/en/)
## Quick example
```python
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)
#
q2 = cx.vconvert(cx.SphericalPos, q)
print(q2)
#
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)
#
p2 = cx.vconvert(cx.SphericalVel, p, q)
print(p2)
#
# 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")
#
#
coord = cx.Coordinate({"length": q, "speed": p}, frame=icrs_frame)
print(coord)
# Coordinate(
# KinematicSpace({
# 'length': ,
# 'speed':
# }),
# frame=ICRS()
# )
print(coord.to_frame(gc_frame))
# Coordinate(
# KinematicSpace({
# 'length': ,
# 'speed':
# }),
# frame=Galactocentric( ... )
# )
```
## Citation
[![DOI][zenodo-badge]][zenodo-link]
If you found this library to be useful in academic work, then please cite.
## Development
[![Actions Status][actions-badge]][actions-link]
[![Documentation Status][rtd-badge]][rtd-link]
[![codecov][codecov-badge]][codecov-link]
[![SPEC 0 — Minimum Supported Dependencies][spec0-badge]][spec0-link]
[![pre-commit][pre-commit-badge]][pre-commit-link]
[![ruff][ruff-badge]][ruff-link]
We welcome contributions!
[actions-badge]: https://github.com/GalacticDynamics/coordinax/workflows/CI/badge.svg
[actions-link]: https://github.com/GalacticDynamics/coordinax/actions
[codecov-badge]: https://codecov.io/gh/GalacticDynamics/unxt/graph/badge.svg
[codecov-link]: https://codecov.io/gh/GalacticDynamics/unxt
[pre-commit-badge]: https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit
[pre-commit-link]: https://pre-commit.com
[pypi-link]: https://pypi.org/project/coordinax/
[pypi-platforms]: https://img.shields.io/pypi/pyversions/coordinax
[pypi-version]: https://img.shields.io/pypi/v/coordinax
[rtd-badge]: https://readthedocs.org/projects/coordinax/badge/?version=latest
[rtd-link]: https://coordinax.readthedocs.io/en/latest/?badge=latest
[ruff-badge]: https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/charliermarsh/ruff/main/assets/badge/v2.json
[ruff-link]: https://docs.astral.sh/ruff/
[spec0-badge]: https://img.shields.io/badge/SPEC-0-green?labelColor=%23004811&color=%235CA038
[spec0-link]: https://scientific-python.org/specs/spec-0000/
[zenodo-badge]: https://zenodo.org/badge/DOI/10.5281/zenodo.15320465.svg
[zenodo-link]: https://zenodo.org/doi/10.5281/zenodo.10850557