jpu

JAX + единицы

Построен с JAX и PINT!

Этот модуль обеспечивает интерфейс между JAX и Pint, чтобы JAX поддерживать операции с единицами. Распространение единиц происходит во время трассировки, поэтому пишеткие функции не должны видеть стоимости времени выполнения. Эта библиотека экспериментальна, поэтому ожидайте некоторых острых краев.

Например:

 >> > import jax
>> > import jax . numpy as jnp
>> > import jpu
>> >
>> > u = jpu . UnitRegistry ()
>> >
>> > @ jax . jit
... def add_two_lengths ( a , b ):
...     return a + b
...
>> > add_two_lengths ( 3 * u . m , jnp . array ([ 4.5 , 1.2 , 3.9 ]) * u . cm )
< Quantity ([ 3.045 3.012 3.039 ], 'meter' ) > 

Установка

Для установки используйте pip :

python -m pip install jpu

Единственные зависимости - jax и pint , и они также будут установлены, если еще не в вашей среде. Посмотрите на документы JAX для получения дополнительной информации об установке JAX в разных системах.

Использование

Вот немного более полный пример:

 >> > import jax
>> > import numpy as np
>> > from jpu import UnitRegistry , numpy as jnpu
>> >
>> > u = UnitRegistry ()
>> >
>> > @ jax . jit
... def projectile_motion ( v_init , theta , time , g = u . standard_gravity ):
...     """Compute the motion of a projectile with support for units"""
...     x = v_init * time * jnpu . cos ( theta )
...     y = v_init * time * jnpu . sin ( theta ) - 0.5 * g * jnpu . square ( time )
...     return x . to ( u . m ), y . to ( u . m )
...
>> > x , y = projectile_motion (
...     5.0 * u . km / u . h , 60 * u . deg , np . linspace ( 0 , 1 , 50 ) * u . s
... )

Технические детали и ограничения

Наиболее значительным ограничением этой библиотеки является тот факт, что пользователи должны использовать функции jpu.numpy при взаимодействии с «величинами» с единицами вместо интерфейса jax.numpy . Это потому, что JAX еще не (еще?) Предоставляет общий интерфейс для отправки UFUNCS на пользовательских классах массива. Я играл с незарегистрированным интерфейсом __jax_array__ , но он не очень гибкий, и в настоящее время он не совместим с объектами Pytree.

До настоящего времени реализовано только подмножество интерфейса numpy / jax.numpy . Запросы на привлечение добавления более широкой поддержки (включая подмодули) будут приветствовать!