Add gather function; refine math, linalg, autograd, and constants (#44)

* Bump version to 0.0.18, add gather function, and enhance round function with decimals parameter

* Update Publish.yml to enhance package build and publishing process for SAIUnit and BrainUnit

* Add @set_module_as decorator to activation functions and update Publish.yml for Python setup and dependency installation

* Refactor unit handling in functions and improve type assertions for clarity

* Fix unit constants for parsec, barrel, and degree Fahrenheit for accuracy

* Fix input dimension handling in FFT functions for improved accuracy

* Refactor dynamic slicing functions to improve handling of Quantity objects and maintain unit consistency

* Fix references to 'brainunit' in jacobian functions and add allow_int parameter for flexibility

* Refactor linalg module to improve QR decomposition handling and clean up code formatting

* Remove unused allow_int parameter from jacfwd function and update related documentation

* Fix barrel unit scale adjustment in _unit_constants.py and update related test assertions

* Update CI configuration to remove Python versions 3.11 and 3.12 from testing matrix

Author: chaoming0625

.github/workflows/CI-daily.yml

@@ -37,7 +37,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.13"]
         jax-version: ["0.4.38", "0.5.2", "0.6.0", ""]
         # Optional: Exclude incompatible combinations if needed
         # exclude:

.github/workflows/CI.yml

@@ -28,7 +28,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.10", "3.11", "3.12", "3.13" ]
+        python-version: [ "3.10", "3.13" ]
 
     steps:
       - name: Cancel Previous Runs
@@ -58,7 +58,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.10", "3.11", "3.12", "3.13" ]
+        python-version: [ "3.10", "3.13" ]
 
     steps:
       - name: Cancel Previous Runs
@@ -124,7 +124,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ "3.10", "3.11", "3.12", "3.13" ]
+        python-version: [ "3.10", "3.13" ]
 
     steps:
       - name: Cancel Previous Runs

.github/workflows/Publish.yml

@@ -1,7 +1,7 @@
 name: Publish to PyPI.org
 on:
   release:
-    types: [published]
+    types: [ published ]
 jobs:
   pypi:
     runs-on: ubuntu-latest
@@ -10,8 +10,29 @@ jobs:
         uses: actions/checkout@v5
         with:
           fetch-depth: 0
-      - run: python setup.py bdist_wheel
-      - name: Publish package
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.13'
+      - name: Install build dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install build setuptools wheel
+      - name: build SAIUnit package
+        run: |
+          python -m build
+      - name: Publish SAIUnit package
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}
+      - name: build BrainUnit package
+        run: |
+          python -m pip install dist/saiunit* jax jaxlib numpy typing_extensions
+          python make_brainunit_setup.py
+          cd ./brainunit
+          python -m build
+      - name: Publish BrainUnit package
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           password: ${{ secrets.PYPI_API_TOKEN }}
+          packages-dir: brainunit/dist/

pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools", "numpy", 'jax', 'jaxlib']
+requires = ["setuptools"]
 build-backend = "setuptools.build_meta"
 
 

saiunit/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.0.16"
+__version__ = "0.0.18"
 
 from . import _matplotlib_compat
 from . import autograd

saiunit/_unit_constants.py

@@ -84,7 +84,7 @@
 au = astronomical_unit = Unit.create(meter.dim, name="astronomical unit", dispname="AU", scale=meter.scale + 11,
                                      factor=1.495978707)
 light_year = Unit.create(meter.dim, name="light year", dispname="ly", scale=meter.scale + 15, factor=9.460730777119564)
-parsec = Unit.create(meter.dim, name="parsec", dispname="pc", scale=meter.scale + 16, factor=3.085677581491367e16)
+parsec = Unit.create(meter.dim, name="parsec", dispname="pc", scale=meter.scale + 16, factor=3.085677581491367)
 
 # ----- Pressure -----
 atm = atmosphere = Unit.create(pascal.dim, name="atmosphere", dispname="atm", scale=pascal.scale + 5, factor=1.013249966)
@@ -104,7 +104,7 @@
                                            factor=2.95735295625)
 fluid_ounce_imp = Unit.create(meter3.dim, name="imperial fluid ounce", dispname="fl oz imp", scale=meter3.scale - 5,
                               factor=2.84130742)
-bbl = barrel = Unit.create(meter3.dim, name="barrel", dispname="bbl", scale=meter3.scale + 2, factor=1.5898729493)
+bbl = barrel = Unit.create(meter3.dim, name="barrel", dispname="bbl", scale=meter3.scale + 2, factor=1.58987294928)
 
 # ----- Speed -----
 speed_unit = meter / second
@@ -118,8 +118,8 @@
 # ----- Temperature -----
 # TODO: The relationship between Celsius and Kelvin should be linear, but the current implementation is not.
 # zero_Celsius = Unit.create(kelvin.dim, name="zero Celsius", dispname="0°C", scale=kelvin.scale, factor=273.15)
-degree_Fahrenheit = Unit.create(kelvin.dim, name="degree Fahrenheit", dispname="°F", scale=kelvin.scale + 2,
-                                factor=2.55927778)
+degree_Fahrenheit = Unit.create(kelvin.dim, name="degree Fahrenheit", dispname="°F", scale=kelvin.scale,
+                                factor=5/9)
 
 # ----- Energy -----
 eV = electron_volt = Unit.create(joule.dim, name="electronvolt", dispname="eV", scale=joule.scale - 19, factor=1.602176565)
@@ -130,7 +130,7 @@
 Btu = Btu_IT = Unit.create(joule.dim, name="British thermal unit (International Table)", dispname="Btu IT",
                            scale=joule.scale + 3, factor=1.05505585262)
 Btu_th = Unit.create(joule.dim, name="British thermal unit (thermochemical)", dispname="Btu th", scale=joule.scale + 3,
-                     factor=1.0543499999744)
+                     factor=1.05435026444)
 ton_TNT = Unit.create(joule.dim, name="ton of TNT", dispname="ton TNT", scale=joule.scale + 9, factor=4.184)
 
 # ----- Power -----

saiunit/autograd/_jacobian.py

@@ -29,7 +29,6 @@
     _check_output_dtype_jacfwd
 )
 from jax.api_util import argnums_partial
-from jax.extend import linear_util
 
 from saiunit._base import Quantity, maybe_decimal, get_magnitude, get_unit
 from saiunit._compatible_import import safe_map, wrap_init
@@ -134,15 +133,15 @@ def jacrev(
     In particular, an array is produced (with no pytrees involved) when the
     function input ``x`` and output ``fun(x)`` are each a single array, as in the
     ``simple_function`` example above. If ``fun(x)`` has shape ``(out1, out2, ...)`` and ``x``
-    has shape ``(in1, in2, ...)`` then ``saiunit.autograd.jacrec(fun)(x)`` has shape
+    has shape ``(in1, in2, ...)`` then ``saiunit.autograd.jacrev(fun)(x)`` has shape
     ``(out1, out2, ..., in1, in2, ..., in1, in2, ...)``. To flatten pytrees into
     1D vectors, consider using :py:func:`jax.flatten_util.flatten_pytree`.
     """
     _check_callable(fun)
 
     @wraps(fun)
     def jacfun(*args, **kwargs):
-        f = wrap_init(fun, args, kwargs, 'brainunit.autograd.jacrev')
+        f = wrap_init(fun, args, kwargs, 'saiunit.autograd.jacrev')
         f_partial, dyn_args = argnums_partial(f, argnums, args, require_static_args_hashable=False)
         jax.tree.map(partial(_check_input_dtype_jacrev, holomorphic, allow_int), dyn_args)
         if not has_aux:
@@ -238,7 +237,7 @@ def jacfwd(
     fun: Callable,
     argnums: int | Sequence[int] = 0,
     has_aux: bool = False,
-    holomorphic: bool = False
+    holomorphic: bool = False,
 ) -> Callable:
     """
     Physical unit-aware version of `jax.jacfwd <https://jax.readthedocs.io/en/latest/_autosummary/jax.jacfwd.html>`_.
@@ -276,7 +275,7 @@ def jacfwd(
      [[3., 0.],
       [0., 4.]] * ohm)
 
-    `jacfwd` is a generalization of the usual definition of the JacFwd(Jacobian Reverse Mode).
+    `jacfwd` is a generalization of the usual definition of the JacFwd(Jacobian Forward Mode).
     that supports nested Python containers (i.e. pytrees) as inputs and outputs.
     The tree structure of ``saiunit.autograd.jacfwd(fun)(x)`` is given by forming a tree
     product of the structure of ``fun(x)`` with a tree product of two copies of
@@ -324,7 +323,7 @@ def jacfwd(
     In particular, an array is produced (with no pytrees involved) when the
     function input ``x`` and output ``fun(x)`` are each a single array, as in the
     ``simple_function`` example above. If ``fun(x)`` has shape ``(out1, out2, ...)`` and ``x``
-    has shape ``(in1, in2, ...)`` then ``saiunit.autograd.jacrec(fun)(x)`` has shape
+    has shape ``(in1, in2, ...)`` then ``saiunit.autograd.jacfwd(fun)(x)`` has shape
     ``(out1, out2, ..., in1, in2, ..., in1, in2, ...)``. To flatten pytrees into
     1D vectors, consider using :py:func:`jax.flatten_util.flatten_pytree`.
     """
@@ -333,7 +332,7 @@ def jacfwd(
 
     @wraps(fun)
     def jacfun(*args, **kwargs):
-        f = wrap_init(fun, args, kwargs, 'brainunit.autograd.jacfwd')
+        f = wrap_init(fun, args, kwargs, 'saiunit.autograd.jacfwd')
         f_partial, dyn_args = argnums_partial(f, argnums, args, require_static_args_hashable=False)
         jax.tree.map(partial(_check_input_dtype_jacfwd, holomorphic), dyn_args)
         if not has_aux:

saiunit/constants.py

@@ -14,12 +14,12 @@
 # ==============================================================================
 
 r"""
-A module providing some physical units as `Quantity` objects. Note that these
-units are not imported by wildcard imports, they
+A module providing some physical constants as `Quantity` objects. Note that these
+constants are not imported by wildcard imports, they
 have to be imported explicitly. You can use ``import ... as ...`` to import them
 with shorter names, e.g.::
 
-    from saiunit.constants import faraday_constant as F
+    from saiunit.constants import faraday as F
 
 The available constants are:
 
@@ -61,18 +61,18 @@
 from ._unit_constants import speed_unit
 
 __all__ = [
-    'arcmin', 'arcminute', 'arcsec', 'arcsecond', 'atomic_mass', 'au', 'astronomical_unit',
-    'angstrom', 'atm', 'atmosphere', 'avogadro', 'bar', 'blob', 'boltzmann', 'Btu', 'Btu_IT',
+    'acre', 'arcmin', 'arcminute', 'arcsec', 'arcsecond', 'atomic_mass', 'au', 'astronomical_unit',
+    'angstrom', 'atm', 'atmosphere', 'avogadro', 'bar', 'barrel', 'bbl', 'blob', 'boltzmann', 'Btu', 'Btu_IT',
     'Btu_th', 'carat', 'calorie', 'calorie_IT', 'calorie_th', 'day', 'degree', 'degree_Fahrenheit',
     'dyn', 'dyne', 'eV', 'electron_mass', 'electric', 'electronvolt', 'elementary_charge', 'erg',
     'faraday', 'fermi', 'fluid_ounce', 'fluid_ounce_US', 'fluid_ounce_imp', 'foot', 'gas', 'grain',
     'gallon', 'gallon_US', 'gallon_imp', 'gram', 'hectare', 'hour', 'hp', 'horsepower', 'IMF',
-    'inch', 'julian_year', 'kelvin', 'kgf', 'kilogram_force', 'knot', 'lb', 'lbf', 'light_year',
-    'long_ton', 'mach', 'magnetic', 'meter', 'metric_ton', 'micron', 'mil', 'mile', 'minute',
+    'inch', 'julian_year', 'kelvin', 'kgf', 'kilogram_force', 'kmh', 'knot', 'lb', 'lbf', 'light_year',
+    'long_ton', 'mach', 'magnetic', 'meter', 'metric_ton', 'micron', 'mil', 'mile', 'minute', 'mmHg',
     'molar_mass', 'month', 'mph', 'nautical_mile', 'newton', 'ounce', 'oz', 'parsec', 'pica',
-    'point', 'pound', 'psi', 'radian', 'second', 'short_ton', 'slug', 'slinch', 'speed_unit',
+    'point', 'pound', 'pound_force', 'psi', 'radian', 'second', 'short_ton', 'slug', 'slinch', 'speed_unit',
     'stone', 'survey_foot', 'survey_mile', 'torr', 'troy_ounce', 'troy_pound', 'ton_TNT', 'week',
-    'watt', 'year', 'zero_celsius'
+    'watt', 'yard', 'year', 'zero_celsius'
 ]
 
 #: Avogadro constant (http://physics.nist.gov/cgi-bin/cuu/Value?na)
@@ -90,9 +90,9 @@
 #: gas constant (http://physics.nist.gov/cgi-bin/cuu/Value?r)
 gas = np.asarray(8.3144598) * (joule / mole / kelvin)
 #: Magnetic constant (http://physics.nist.gov/cgi-bin/cuu/Value?mu0)
-magnetic = np.asarray(4 * np.pi * 1e-7) * (newton / amp ** 2)
+magnetic = np.asarray(1.25663706212e-6) * (newton / amp ** 2)
 #: Molar mass constant (http://physics.nist.gov/cgi-bin/cuu/Value?mu)
-molar_mass = np.asarray(1.) * (gram / mole)
+molar_mass = np.asarray(1e-3) * (kilogram / mole)
 #: zero degree Celsius
 zero_celsius = np.asarray(273.15) * kelvin
 
@@ -161,7 +161,9 @@
 bbl = barrel = np.asarray(1.58987294928e2) * meter3  # Barrel (oil)
 
 # ----- Temperature -----
-degree_Fahrenheit = np.asarray(2.55927778e2) * kelvin  # Fahrenheit
+# Note: Fahrenheit is a temperature scale, not a unit. Use conversion functions instead.
+# This constant represents the conversion factor from Fahrenheit to Kelvin degrees
+degree_Fahrenheit = np.asarray(5/9) * kelvin  # Fahrenheit degree size in Kelvin
 
 # ----- Speed -----
 kmh = np.asarray(2.77777778e-1) * speed_unit  # Kilometer per hour

saiunit/constants_test.py

@@ -79,9 +79,9 @@ def test_quantity_constants_and_unit_constants(self):
         import saiunit.constants as quantity_constants
         import saiunit._unit_constants as unit_constants
         for c in constants_list:
+            print(c)
             q_c = getattr(quantity_constants, c)
             u_c = getattr(unit_constants, c)
             assert u.math.isclose(
-                q_c.to_decimal(q_c.unit),
-                (1. * u_c).to_decimal(q_c.unit)
+                q_c.to_decimal(q_c.unit), (1. * u_c).to_decimal(q_c.unit)
             ), f"Mismatch between {c} in quantity_constants and unit_constants"

saiunit/fft/_fft_change_unit.py

@@ -625,7 +625,8 @@ def fftn(
         >>> u.math.allclose(x, u.fft.ifftn(x_fftn))
         Array(True, dtype=bool)
     """
-    n = _calculate_fftn_dimension(a.ndim, axes)
+    input_ndim = a.ndim if hasattr(a, 'ndim') else jnp.asarray(a).ndim
+    n = _calculate_fftn_dimension(input_ndim, axes)
     _unit_change_fun = lambda u: u * (second ** n)
     # TODO: may cause computation overhead?
     fftn._unit_change_fun = _unit_change_fun
@@ -730,7 +731,8 @@ def rfftn(
         >>> u.fft.rfftn(x1)
         ArrayImpl([10.+0.j, -2.+2.j, -2.+0.j], dtype=complex64) * meter * second
     """
-    n = _calculate_fftn_dimension(a.ndim, axes)
+    input_ndim = a.ndim if hasattr(a, 'ndim') else jnp.asarray(a).ndim
+    n = _calculate_fftn_dimension(input_ndim, axes)
     _unit_change_fun = lambda u: u * (second ** n)
     # TODO: may cause computation overhead?
     rfftn._unit_change_fun = _unit_change_fun
@@ -978,7 +980,8 @@ def ifftn(
         ArrayImpl([[ 2.5 +0.j  ,  0.  -0.58j,  0.  +0.58j],
                    [ 0.17+0.j  , -0.83-0.29j, -0.83+0.29j]], dtype=complex64) * meter / second2
     """
-    n = _calculate_fftn_dimension(a.ndim, axes)
+    input_ndim = a.ndim if hasattr(a, 'ndim') else jnp.asarray(a).ndim
+    n = _calculate_fftn_dimension(input_ndim, axes)
     _unit_change_fun = lambda u: u / (second ** n)
     # TODO: may cause computation overhead?
     ifftn._unit_change_fun = _unit_change_fun
@@ -1067,7 +1070,8 @@ def irfftn(
                    [[-2., -2., -2.],
                     [-2., -2., -2.]]], dtype=float32) * meter / second
     """
-    n = _calculate_fftn_dimension(a.ndim, axes)
+    input_ndim = a.ndim if hasattr(a, 'ndim') else jnp.asarray(a).ndim
+    n = _calculate_fftn_dimension(input_ndim, axes)
     _unit_change_fun = lambda u: u / (second ** n)
     # TODO: may cause computation overhead?
     irfftn._unit_change_fun = _unit_change_fun
@@ -1188,7 +1192,7 @@ def rfftfreq(
 ) -> Union[Quantity, jax.typing.ArrayLike]:
     """Return sample frequencies for the discrete Fourier transform.
 
-    saiunit implementation of :func:`numpy.fft.fftfreq`. Returns frequencies appropriate
+    saiunit implementation of :func:`numpy.fft.rfftfreq`. Returns frequencies appropriate
     for use with the outputs of :func:`~saiunit.fft.rfft` and
     :func:`~saiunit.fft.irfft`.
 
@@ -1210,8 +1214,8 @@ def rfftfreq(
     Example:
         >>> import saiunit as u
         >>> import jax.numpy as jnp
-        >>> x = jnp.array([1, 2, 3, 4]) * u.second
-        >>> u.fft.rfftfreq(4, x)
+        >>> d = 1 * u.second
+        >>> u.fft.rfftfreq(4, d)
         ArrayImpl([0.  , 0.25, 0.5 ], dtype=float32) * hertz
     """
     if isinstance(d, Quantity):