Skip to content

Commit

Permalink
[math] add math operations (#1)
Browse files Browse the repository at this point in the history
* Fix numpy array creation methods

* Update

* Update test cases

* update codes

* upgrade __init__.py

---------

Co-authored-by: Chaoming Wang <adaduo@outlook.com>
  • Loading branch information
Routhleck and chaoming0625 authored Jun 8, 2024
1 parent bc5ba2a commit 06ce267
Show file tree
Hide file tree
Showing 6 changed files with 1,185 additions and 825 deletions.
3 changes: 2 additions & 1 deletion brainunit/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 *
Expand All @@ -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
21 changes: 13 additions & 8 deletions brainunit/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
'Quantity',
'Unit',
'UnitRegistry',
'DIMENSIONLESS',
'DimensionMismatchError',
'get_or_create_dimension',
'get_unit',
Expand All @@ -49,7 +50,6 @@
]

_all_slice = slice(None, None, None)

random = None
_unit_checking = True
_automatically_register_units = True
Expand Down Expand Up @@ -242,6 +242,7 @@ class Dimension:
indices, allowing for a very fast dimensionality check with ``is``.
"""

__module__ = "brainunit"
__slots__ = ["_dims"]
__array_priority__ = 1000

Expand Down Expand Up @@ -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':
"""
Expand All @@ -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.
Expand All @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 #
Expand Down Expand Up @@ -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__(
Expand Down Expand Up @@ -2706,6 +2709,8 @@ class UnitRegistry:
__getitem__
"""

__module__ = "brainunit"

def __init__(self):
self.units = collections.OrderedDict()
self.units_for_dimensions = collections.defaultdict(dict)
Expand Down
Loading

0 comments on commit 06ce267

Please sign in to comment.