unxt

Unitful Quantities in JAX

Unxt is unitful quantities and calculations in JAX, built on Equinox and Quax.

Unxt supports JAX's compelling features:

• JIT compilation (jit)
• vectorization (vmap, etc.)
• auto-differentiation (grad, jacobian, hessian)
• GPU/TPU/multi-host acceleration

And best of all, unxt doesn't force you to use special unit-compatible re-exports of JAX libraries. You can use unxt with existing JAX code, and with quax's simple decorator, JAX will work with unxt.Quantity.

Installation

pip install unxt

uv add unxt

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

cd /path/to/parent
git clone https://github.com/GalacticDynamics/unxt.git
cd unxt
pip install -e .  # editable mode

Documentation

For full documentation, including installation instructions, tutorials, and API reference, please see the unxt docs. This README provides a brief overview and some quick examples.

Dimensions

Dimensions represent the physical type of a quantity, such as length, time, or mass.

>>> import unxt as u

Create dimensions from strings:

>>> u.dimension("length")
PhysicalType('length')

Dimensions support mathematical expressions:

>>> u.dimension("length / time")
PhysicalType({'speed', 'velocity'})

Multi-word dimension names require parentheses in expressions:

>>> u.dimension("(amount of substance) / (time)")
PhysicalType('catalytic activity')

Units

Units specify the scale and dimension of measurements.

>>> meter = u.unit("m")
>>> meter
Unit("m")

Units can be combined, either inside the expression string or by arithmetic on Unit objects:

>>> u.unit("km/h")  # in the expression
Unit("km / h")

>>> u.unit("km") / u.unit("h")  # via arithmetic
Unit("km / h")

Get the dimension of a unit:

>>> u.dimension_of(meter)
PhysicalType('length')

Unit Systems

Unit systems define consistent sets of base units for specific domains. unxt provides built-in unit systems and tools for creating custom ones.

Built-in Unit Systems

>>> u.unitsystem("si")  # SI (International System of Units)
unitsystem(m, kg, s, mol, A, K, cd, rad)

>>> u.unitsystem("cgs")  # CGS (centimeter-gram-second)
unitsystem(cm, g, s, dyn, erg, Ba, P, St, rad)

>>> u.unitsystem("galactic")  # galactic (astrophysics)
unitsystem(kpc, Myr, solMass, rad)

Composing Units from a Unit System

Once you have a unit system, you can get units for any physical dimension by indexing the system:

>>> usys = u.unitsystem("si")
>>> usys["length"]
Unit("m")

Custom Unit Systems

Create custom unit systems by specifying base units:

>>> custom_usys = u.unitsystem("km", "h", "tonne", "degree")
>>> custom_usys
unitsystem(km, h, t, deg)

Derived units are then available by dimension:

>>> custom_usys["velocity"]
Unit("km / h")

Dynamical Unit Systems

For domains like gravitational dynamics, use dynamical unit systems where $G = 1$. Specify only 2 of (length, time, mass); the third is computed to make $G = 1$.

>>> from unxt.unitsystems import DynamicalSimUSysFlag

>>> dyn_usys = u.unitsystem(DynamicalSimUSysFlag, "kpc", "Myr")
>>> dyn_usys
LengthMassTimeUnitSystem(length=Unit("kpc"),
                         mass=Unit("1.49828e+10 kpc3 s2 kg / (Myr2 m3)"), time=Unit("Myr"))

The mass unit is the derived one — an exact composite expression, not a rounded label:

>>> dyn_usys["mass"]
Unit("1.49828e+10 kpc3 s2 kg / (Myr2 m3)")

Quantities

Quantities combine values with units, providing type-safe unitful arithmetic.

Quantity (u.Q) is the lightweight, non-parametric default: a single class — and a single JAX pytree type — for all physical dimensions. ParametricQuantity (up.PQ) adds runtime dimension checking and dimension-specific plum dispatch by encoding each dimension in its own on-the-fly class (and pytree type), which grows the type/dispatch surface and adds per-construction overhead. (This is not about jax.jit cache misses: the unit is static, so a jitted function specializes per unit with either class — that part is inherent.) See the Quantity guide for full details; upgrading from an earlier version? See the migration guide.

Basic Quantities

>>> import jax.numpy as jnp

>>> x = u.Q(jnp.arange(1, 5, dtype=float), "km")
>>> x
Quantity(Array([1., 2., 3., 4.], dtype=float32...), unit='km')

The constituent value and unit are accessible as attributes:

>>> x.value
Array([1., 2., 3., 4.], dtype=float32...)

>>> x.unit
Unit("km")

Quantity objects obey the rules of unitful arithmetic — addition, subtraction, multiplication, division, and exponentiation:

>>> x + x
Quantity(Array([2., 4., 6., 8.], dtype=float32...), unit='km')

>>> 2 * x
Quantity(Array([2., 4., 6., 8.], dtype=float32...), unit='km')

>>> y = u.Q(jnp.arange(4, 8, dtype=float), "yr")
>>> x / y
Quantity(
    Array([0.25     , 0.4      , 0.5      , 0.5714286], dtype=float32...), unit='km / yr'
)

>>> x**2
Quantity(Array([ 1.,  4.,  9., 16.], dtype=float32...), unit='km2')

Operations are unit-checked, so mixing incompatible dimensions raises:

>>> try:
...     x + y
... except Exception as e:
...     print(e)
'yr' (time) and 'km' (length) are not convertible

Quantities can be converted to different units, by function or by method:

>>> u.uconvert("m", x)  # via function
Quantity(Array([1000., 2000., 3000., 4000.], dtype=float32...), unit='m')

>>> x.uconvert("m")  # via method
Quantity(Array([1000., 2000., 3000., 4000.], dtype=float32...), unit='m')

ParametricQuantity — from the separate unxts.parametric package (pip install unxts.parametric, imported below as up) — adds runtime dimension checking on construction. Use up.PQ["length"] to create a parametric type that raises if the unit's physical type does not match:

>>> import unxts.parametric as up

>>> LengthQuantity = up.PQ["length"]
>>> LengthQuantity(2, "km")
ParametricQuantity(Array(2, dtype=int32...), unit='km')

A unit whose physical type does not match the parameter is rejected:

>>> try:
...     LengthQuantity(2, "s")
... except ValueError as e:
...     print(e)
Physical type mismatch.

By contrast, the default u.Q["length"] accepts the subscript but does not check dimensions — it silently builds a Quantity with the mismatched unit:

>>> u.Q["length"](2, "s")
Quantity(Array(2, dtype=int32...), unit='s')

Use up.PQ["length"] when you need the runtime guard. See the unxts.parametric guide for the full API.

Quantity

Quantity (aliased as u.Q) is the default lightweight class. It does not do runtime dimension checking on construction, which makes it the fastest option for performance-critical code:

>>> bq = u.quantity.Quantity(jnp.array([1.0, 2.0, 3.0]), "m")
>>> bq
Quantity(Array([1., 2., 3.], dtype=float32...), unit='m')

>>> bq * 2
Quantity(Array([2., 4., 6.], dtype=float32...), unit='m')

Angle

Angle is a specialized quantity with wrapping support for angular values:

>>> theta = u.Angle(jnp.array([0, 90, 180, 270, 360]), "deg")
>>> theta
Angle(Array([  0,  90, 180, 270, 360], dtype=int32...), unit='deg')

Angles can optionally be wrapped into a specified range:

>>> angle = u.Angle(jnp.array([370, -10]), "deg")
>>> angle.wrap_to(u.Q(0, "deg"), u.Q(360, "deg"))
Angle(Array([ 10, 350], dtype=int32...), unit='deg')

StaticQuantity

For static configuration values (e.g., JAX static arguments), use StaticQuantity, which stores NumPy values and rejects JAX arrays:

>>> import numpy as np
>>> from functools import partial
>>> import jax

>>> cfg = u.StaticQuantity(np.array([1.0, 2.0]), "m")

>>> @partial(jax.jit, static_argnames=("q",))
... def add(x, q):
...     return x + jnp.asarray(q.value)

>>> add(1.0, cfg)
Array([2., 3.], dtype=float32...)

StaticValue

If you want a Quantity that keeps a static value but still participates in regular arithmetic, wrap the value with StaticValue. Arithmetic behaves like the wrapped array, and StaticValue + StaticValue returns a StaticValue. Equality between two StaticValues (== / !=) returns a scalar bool — which is what makes a StaticValue-backed quantity hashable and usable as a jax.jit static argument. Ordering (<, <=, >, >=), and == / != against a raw array, return element-wise NumPy boolean arrays:

>>> sv = u.quantity.StaticValue(np.array([1.0, 2.0]))
>>> q_static = u.Q(sv, "m")
>>> q = u.Q(jnp.array([3.0, 4.0]), "m")

>>> q_static + q
Quantity(Array([4., 6.], dtype=float32...), unit='m')

Equality between two StaticValues is a scalar bool:

>>> sv2 = u.quantity.StaticValue(np.array([2.0, 1.0]))
>>> sv == sv2
False
>>> sv == u.quantity.StaticValue(np.array([1.0, 2.0]))
True

Ordering, and equality against a raw array, are element-wise NumPy boolean arrays:

>>> sv < sv2
array([ True, False])

>>> sv == np.array([1.0, 2.0])
array([ True,  True])

JAX Integration

unxt is built on quax, which enables custom array-ish objects in JAX. For convenience we use the quaxed library, which is just a quax.quaxify wrapper around jax to avoid boilerplate code.

[!NOTE]

Using quaxed is optional. You can directly use quaxify, and even apply it to the top-level function instead of individual functions.

Using the x quantity from the earlier examples:

>>> from quaxed import grad, vmap
>>> import quaxed.numpy as qnp

>>> qnp.square(x)
Quantity(Array([ 1.,  4.,  9., 16.], dtype=float32...), unit='km2')

>>> qnp.power(x, 3)
Quantity(Array([ 1.,  8., 27., 64.], dtype=float32...), unit='km3')

>>> vmap(grad(lambda x: x**3))(x)
Quantity(Array([ 3., 12., 27., 48.], dtype=float32...), unit='km2')

See the documentation for more examples and details of JIT and AD

Citation

If you found this library to be useful and want to support the development and maintenance of lower-level code libraries for the scientific community, please consider citing this work.

Contributing and Development

We welcome contributions! Contributions are how open source projects improve and grow.

To contribute to unxt, please fork the repository, make a development branch, develop on that branch, then open a pull request from the branch in your fork to main.

To report bugs, request features, or suggest other ideas, please open an issue.

For more information, see CONTRIBUTING.md.