jpu

وحدات Jax +

بنيت مع jax و pint!

توفر هذه الوحدة واجهة بين Jax و Pint للسماح Jax بدعم العمليات مع الوحدات. يحدث انتشار الوحدات في وقت التتبع ، لذلك يجب أن لا ترى وظائف Jated أي تكلفة وقت تشغيل. هذه المكتبة تجريبية ، لذا توقع بعض الحواف الحادة.

على سبيل المثال:

 >> > 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 . سحب طلبات إضافة دعم أوسع (بما في ذلك النسيج الفرعي) سيكون موضع ترحيب!