[math] add math operations (#1)

* Fix numpy array creation methods * Update * Update test cases * update codes * upgrade __init__.py --------- Co-authored-by: Chaoming Wang <[email protected]>
chaobrain · Jun 8, 2024 · 06ce267 · 06ce267
1 parent bc5ba2a
commit 06ce267
Show file tree

Hide file tree

Showing 6 changed files with 1,185 additions and 825 deletions.
diff --git a/brainunit/__init__.py b/brainunit/__init__.py
@@ -15,6 +15,7 @@
 
 __version__ = "0.0.1"
 
+from . import math
 from ._base import *
 from ._base import __all__ as _base_all
 from ._unit_common import *
@@ -24,5 +25,5 @@
 from ._unit_shortcuts import *
 from ._unit_shortcuts import __all__ as _std_units_all
 
-__all__ = _common_all + _std_units_all + _constants_all + _base_all
+__all__ = ['math'] + _common_all + _std_units_all + _constants_all + _base_all
 del _common_all, _std_units_all, _constants_all, _base_all
diff --git a/brainunit/_base.py b/brainunit/_base.py
@@ -34,6 +34,7 @@
   'Quantity',
   'Unit',
   'UnitRegistry',
+  'DIMENSIONLESS',
   'DimensionMismatchError',
   'get_or_create_dimension',
   'get_unit',
@@ -49,7 +50,6 @@
 ]
 
 _all_slice = slice(None, None, None)
-
 random = None
 _unit_checking = True
 _automatically_register_units = True
@@ -242,6 +242,7 @@ class Dimension:
   indices, allowing for a very fast dimensionality check with ``is``.
   """
 
+  __module__ = "brainunit"
   __slots__ = ["_dims"]
   __array_priority__ = 1000
 
@@ -748,9 +749,9 @@ def in_best_unit(x, precision=None):
   return x.repr_in_unit(u, precision=precision)
 
 
-def array_with_units(
+def array_with_unit(
     floatval,
-    units: Dimension,
+    unit: Dimension,
     dtype: bst.typing.DTypeLike = None
 ) -> 'Quantity':
   """
@@ -764,7 +765,7 @@ def array_with_units(
   ----------
   floatval : `float`
       The floating point value of the array.
-  units: Dimension
+  unit: Dimension
       The unit dimensions of the array.
   dtype: `dtype`, optional
       The data type of the array.
@@ -777,10 +778,10 @@ def array_with_units(
   Examples
   --------
   >>> from brainunit import *
-  >>> array_with_units(0.001, volt.unit)
+  >>> array_with_unit(0.001, volt.unit)
   1. * mvolt
   """
-  return Quantity(floatval, unit=get_or_create_dimension(units._dims), dtype=dtype)
+  return Quantity(floatval, unit=get_or_create_dimension(unit._dims), dtype=dtype)
 
 
 def is_unitless(obj) -> bool:
@@ -936,6 +937,7 @@ class Quantity(object):
   unit. It is used to represent all physical quantities in ``BrainCore``.
 
   """
+  __module__ = "brainunit"
   __slots__ = ('_value', '_unit')
   _value: Union[jax.Array, numbers.Number]
   _unit: Dimension
@@ -1701,7 +1703,7 @@ def __round__(self, ndigits: int = None) -> 'Quantity':
     return Quantity(self.value.__round__(ndigits), unit=self.unit)
 
   def __reduce__(self):
-    return array_with_units, (self.value, self.unit, self.value.dtype)
+    return array_with_unit, (self.value, self.unit, self.value.dtype)
 
   # ----------------------- #
   #      NumPy methods      #
@@ -2438,8 +2440,9 @@ class Unit(Quantity):
    3. * joule
 
   """
-  __slots__ = ["_value", "_unit", "scale", "_dispname", "_name", "iscompound"]
 
+  __module__ = "brainunit"
+  __slots__ = ["_value", "_unit", "scale", "_dispname", "_name", "iscompound"]
   __array_priority__ = 1000
 
   def __init__(
@@ -2706,6 +2709,8 @@ class UnitRegistry:
   __getitem__
   """
 
+  __module__ = "brainunit"
+
   def __init__(self):
     self.units = collections.OrderedDict()
     self.units_for_dimensions = collections.defaultdict(dict)