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 platforms PyPI version

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.16.1.tar.gz (540.4 kB view details)

Uploaded Source

Built Distribution

unxt-0.16.1-py3-none-any.whl (49.9 kB view details)

Uploaded Python 3

File details

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

File metadata

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

File hashes

Hashes for unxt-0.16.1.tar.gz
Algorithm Hash digest
SHA256 d27e8101c4a270527eac11f165ee144c3a9df7a40f4eecd882295fafe181272d
MD5 b531d4ddeea902a41f97a92675b1c296
BLAKE2b-256 6e44571c8c22984362e87d83aea05eb06305611b491272c0e4798459d49fbd7f

See more details on using hashes here.

File details

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

File metadata

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

File hashes

Hashes for unxt-0.16.1-py3-none-any.whl
Algorithm Hash digest
SHA256 2cb7b82d226e73ca58e42352ff920d96ad2f863ab9dc30a5c73daf6aa27d85df
MD5 17880f67b2f308d7791e04781da5c667
BLAKE2b-256 28a4fd1aa4e20bb3dc6afde7f6b4e84ffc86e155c5515aafc942fa5c8033a7c3

See more details on using hashes here.

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