Skip to main content

Quantities in JAX

Project description

unxt

Unitful Quantities in JAX


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

Yes, it supports auto-differentiation (grad, jacobian, hessian) and vectorization (vmap, etc).

Installation

PyPI version PyPI platforms

pip install unxt

Documentation

Documentation Status

Quick example

from unxt import Quantity

x = Quantity(jnp.arange(1, 5, dtype=float), "kpc")
print(x)
# Quantity['length'](Array([1., 2., 3., 4.], dtype=float64), unit='kpc')

# Addition / Subtraction
print(x + x)
# Quantity['length'](Array([2., 4., 6., 8.], dtype=float64), unit='kpc')

# Multiplication / Division
print(2 * x)
# Quantity['length'](Array([2., 4., 6., 8.], dtype=float64), unit='kpc')

y = Quantity(jnp.arange(4, 8, dtype=float), "Gyr")

print(x / y)
# Quantity['speed'](Array([0.25      , 0.4       , 0.5       , 0.57142857], dtype=float64), unit='kpc / Gyr')

# Exponentiation
print(x**2)
# Quantity['area'](Array([0., 1., 4., 9.], dtype=float64), unit='kpc2')

# Unit Checking on operations
try:
    x + y
except Exception as e:
    print(e)
# 'Gyr' (time) and 'kpc' (length) are not convertible

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.

from quaxed import grad, vmap
import quaxed.numpy as jnp

print(jnp.square(x))
# Quantity['area'](Array([ 1.,  4.,  9., 16.], dtype=float64), unit='kpc2')

print(qnp.power(x, 3))
# Quantity['volume'](Array([ 1.,  8., 27., 64.], dtype=float64), unit='kpc3')

print(vmap(grad(lambda x: x**3))(x))
# Quantity['area'](Array([ 3., 12., 27., 48.], dtype=float64), unit='kpc2')

Since Quantity is parametric, it can do runtime dimension checking!

LengthQuantity = Quantity["length"]
print(LengthQuantity(2, "km"))
# Quantity['length'](Array(2, dtype=int64, weak_type=True), unit='km')

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

Citation

DOI

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.

Development

codecov Actions Status

We welcome contributions!

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

unxt-0.20.5.tar.gz (658.9 kB view details)

Uploaded Source

Built Distribution

unxt-0.20.5-py3-none-any.whl (56.0 kB view details)

Uploaded Python 3

File details

Details for the file unxt-0.20.5.tar.gz.

File metadata

  • Download URL: unxt-0.20.5.tar.gz
  • Upload date:
  • Size: 658.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/5.1.1 CPython/3.12.7

File hashes

Hashes for unxt-0.20.5.tar.gz
Algorithm Hash digest
SHA256 1be97a623157fd50a66fb9100caca706dba410243826c52326a757cba4de5c8e
MD5 ca56595d52464d5327f63ec75a9e3a26
BLAKE2b-256 03db4f5fa1760c09ec4a2b981cd96bc64ac3b90014a20567ea0f01181f13ac08

See more details on using hashes here.

Provenance

The following attestation bundles were made for unxt-0.20.5.tar.gz:

Publisher: cd.yml on GalacticDynamics/unxt

Attestations:

File details

Details for the file unxt-0.20.5-py3-none-any.whl.

File metadata

  • Download URL: unxt-0.20.5-py3-none-any.whl
  • Upload date:
  • Size: 56.0 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: twine/5.1.1 CPython/3.12.7

File hashes

Hashes for unxt-0.20.5-py3-none-any.whl
Algorithm Hash digest
SHA256 d7164a523b17af8f763d975adb7095d750d64f2c7330fa870c422a3644900900
MD5 db30eeb1d215d605503a87578ec9cf3e
BLAKE2b-256 86e5b67607e0f0b11a1e4cadba615d5e8d85a764f216961234aefc96f592d055

See more details on using hashes here.

Provenance

The following attestation bundles were made for unxt-0.20.5-py3-none-any.whl:

Publisher: cd.yml on GalacticDynamics/unxt

Attestations:

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page