diff --git a/brainunit/__init__.py b/brainunit/__init__.py index c78f24b..fce3ee9 100644 --- a/brainunit/__init__.py +++ b/brainunit/__init__.py @@ -15,8 +15,16 @@ __version__ = "0.0.1.1" +import importlib + from . import math +from . import _base +from . import _unit_common +from . import _unit_constants +from . import _unit_shortcuts + from ._base import * +from ._base import _default_magnitude, _siprefixes from ._base import __all__ as _base_all from ._unit_common import * from ._unit_common import __all__ as _common_all @@ -27,3 +35,71 @@ __all__ = ['math'] + _common_all + _std_units_all + _constants_all + _base_all del _common_all, _std_units_all, _constants_all, _base_all + + +def set_default_magnitude( + magnitude: int | dict[str, int], + # unit: Unit = None, +): + """ + Set the default magnitude for units. + + Parameters + ---------- + magnitude : int | dict[str, int] + The magnitude to set. If an int is given, it will be set for all + dimensions. If a dict is given, it will be set for the specified + dimensions. + + Examples + -------- + >>> set_default_magnitude('n') # Sets the default magnitude to 'nano' (10e-9) + >>> set_default_magnitude(-9) # Alternatively, use an integer to represent the exponent of 10 + >>> set_default_magnitude({'m': -3, 'kg': -9}) # Set the default magnitude for 'metre' to 'milli' and 'kilogram' to 'nano' + >>> set_default_magnitude({'m': 'm', 'kg': 'n'}) # Alternatively, use a string to represent the magnitude + >>> set_default_magnitude(-3, unit=volt) # Set the default magnitude for the 'volt' unit to 'milli' + """ + global _default_magnitude + if isinstance(magnitude, int): + # if isinstance(unit, Unit): + # for key, dim in zip(_default_magnitude.keys(), unit.dim._dims): + # _default_magnitude[key] = magnitude / abs(dim) if dim != 0 else 0 + # else: + _default_magnitude.update((key, magnitude) for key in _default_magnitude) + elif isinstance(magnitude, str): + # if isinstance(unit, Unit): + # for key, dim in zip(_default_magnitude.keys(), unit.dim._dims): + # _default_magnitude[key] = _siprefixes[magnitude] / abs(dim) if dim != 0 else 0 + # else: + _default_magnitude.update((key, _siprefixes[magnitude]) for key in _default_magnitude) + elif isinstance(magnitude, dict): + _default_magnitude.update((key, 0) for key in _default_magnitude) + for key, value in magnitude.items(): + if isinstance(value, int): + _default_magnitude[key] = value + elif isinstance(value, str): + _default_magnitude[key] = _siprefixes[value] + else: + raise ValueError(f"Invalid magnitude value: {value}") + else: + raise ValueError(f"Invalid magnitude: {magnitude}") + + global _unit_common + global _unit_constants + global _unit_shortcuts + # Reload modules + importlib.reload(_unit_common) + importlib.reload(_unit_constants) + importlib.reload(_unit_shortcuts) + + from ._base import __all__ as _base_all + from ._unit_common import __all__ as _common_all + from ._unit_constants import __all__ as _constants_all + from ._unit_shortcuts import __all__ as _std_units_all + globals().update({k: getattr(_unit_common, k) for k in _common_all}) + globals().update({k: getattr(_unit_constants, k) for k in _constants_all}) + globals().update({k: getattr(_unit_shortcuts, k) for k in _std_units_all}) + + global __all__ + __all__ = ['math'] + _common_all + _std_units_all + _constants_all + _base_all + del _common_all, _std_units_all, _constants_all, _base_all \ No newline at end of file diff --git a/brainunit/_base.py b/brainunit/_base.py index 070c439..93d859e 100644 --- a/brainunit/_base.py +++ b/brainunit/_base.py @@ -48,7 +48,7 @@ 'check_units', 'is_scalar_type', 'fail_for_dimension_mismatch', - 'assert_quantity' + 'assert_quantity', ] _all_slice = slice(None, None, None) @@ -56,6 +56,17 @@ _allow_python_scalar_value = False _auto_register_unit = True +# for setting the default magnitude of the unit +_default_magnitude = { + "m": 0, + "kg": 0, + "s": 0, + "A": 0, + "K": 0, + "mol": 0, + "cd": 0, +} + @contextmanager def turn_off_auto_unit_register(): @@ -584,6 +595,31 @@ def get_dim(obj) -> Dimension: raise TypeError(f"Object of type {type(obj)} does not have dimensions") +def get_unit(obj) -> Unit: + """ + Return the unit of any object that has them. + + Parameters + ---------- + obj : `object` + The object to check. + + Returns + ------- + unit : Unit + The physical unit of the `obj`. + """ + try: + return obj.unit + except AttributeError: + if isinstance(obj, (numbers.Number, jax.Array, np.number, np.ndarray)): + return Unit(1, name='1', dispname='1') + try: + return Quantity(obj).unit + except TypeError: + raise TypeError(f"Object of type {type(obj)} does not have a unit") + + def have_same_unit(obj1, obj2) -> bool: """Test if two values have the same dimensions. @@ -969,12 +1005,15 @@ def check_units_and_collect_values(lst): def _get_dim(dim: Dimension, unit: 'Unit'): if dim != DIMENSIONLESS and unit is not None: - raise ValueError("Cannot specify both a dimension and a unit") + return None, dim if dim == DIMENSIONLESS: if unit is None: return None, DIMENSIONLESS else: - return unit.value, unit.dim + try: + return unit.value, unit.dim + except: + return None, unit.dim else: return None, dim @@ -987,7 +1026,7 @@ class Quantity(object): """ __module__ = "brainunit" - __slots__ = ('_value', '_dim') + __slots__ = ('_value', '_dim', '_unit') _value: Union[jax.Array, numbers.Number] _dim: Dimension __array_priority__ = 1000 @@ -1005,6 +1044,8 @@ def __init__( if isinstance(value, numbers.Number): self._dim = dim self._value = (value if scale is None else (value * scale)) + if dim is not DIMENSIONLESS: + self._unit = unit return if isinstance(value, (list, tuple)): @@ -1023,6 +1064,7 @@ def __init__( # array value if isinstance(value, Quantity): self._dim = value.dim + self._unit = Unit(1, name='1', dispname='1') if unit is None else unit self._value = jnp.array(value.value, dtype=dtype) return @@ -1044,6 +1086,12 @@ def __init__( # dimension self._dim = dim + # unit + if unit is None: + self._unit = Unit(1, name='1', dispname='1') + else: + self._unit = unit + @property def value(self) -> jax.Array | numbers.Number: # return the value @@ -1101,7 +1149,8 @@ def dim(self, *args): @property def unit(self) -> 'Unit': - return Unit(1., self.dim, register=False) + return self._unit + # return Unit(1., self.dim, register=False) @unit.setter def unit(self, *args): @@ -1214,14 +1263,14 @@ def get_best_unit(self, *regs) -> 'Quantity': The best unit for this `Array`. """ if self.is_unitless: - return Unit(1) + return Unit(1, name='1', dispname='1') if len(regs): for r in regs: try: return r[self] except KeyError: pass - return Quantity(1, dim=self.dim) + return self.unit else: return self.get_best_unit(standard_unit_register, user_unit_register, additional_unit_register) @@ -1556,6 +1605,7 @@ def _binary_operation( """ other = _to_quantity(other) other_dim = None + other_unit = None if fail_for_mismatch: if inplace: @@ -1568,9 +1618,13 @@ def _binary_operation( if other_dim is None: other_dim = get_dim(other) + if other_unit is None: + other_unit = get_unit(other) + new_dim = unit_operation(self.dim, other_dim) + new_unit = unit_operation(self.unit, other_unit) result = value_operation(self.value, other.value) - r = Quantity(result, dim=new_dim) + r = Quantity(result, dim=new_dim, unit=new_unit) if inplace: self.update_value(r.value) return self @@ -2688,11 +2742,18 @@ def __init__( # Whether this unit is a combination of other units self.iscompound = iscompound - super().__init__(value, dtype=dtype, dim=dim) + if dim == DIMENSIONLESS: + super().__init__(value, dtype=dtype, dim=dim) + else: + super().__init__(value, dtype=dtype, dim=dim, unit=self) if _auto_register_unit and register: register_new_unit(self) + @property + def unit(self) -> 'Unit': + return self + @staticmethod def create(unit: Dimension, name: str, dispname: str, scale: int = 0): """ @@ -2718,6 +2779,8 @@ def create(unit: Dimension, name: str, dispname: str, scale: int = 0): name = str(name) dispname = str(dispname) + scale = calculate_scale(unit=unit, scale=scale) + u = Unit( 10.0 ** scale, dim=unit, @@ -3274,3 +3337,37 @@ def new_f(*args, **kwds): return new_f return do_check_units + + +def calculate_scale( + unit: Dimension, + scale: int = 0 +) -> int: + """ + Calculate the scale for a unit. + + Parameters + ---------- + unit : Dimension + The unit to determine the scale for. + + Returns + ------- + scale : int + The scale for the unit. + + Examples + -------- + >>> set_default_magnitude({'m': -3, 'kg': -9}) + >>> calculate_scale(get_or_create_dimension(m=1)) + -3 + >>> calculate_scale(get_or_create_dimension(kg=1)) + -9 + >>> calculate_scale(get_or_create_dimension(m=1, kg=1)) + -12 + >>> calculate_scale(get_or_create_dimension(m=2, kg=-1)) + 3 + """ + for dim, magnitude in zip(unit._dims, _default_magnitude.values()): + scale -= dim * magnitude + return scale diff --git a/brainunit/_unit_test.py b/brainunit/_unit_test.py index c7401a4..152b9f2 100644 --- a/brainunit/_unit_test.py +++ b/brainunit/_unit_test.py @@ -608,7 +608,6 @@ def test_addition_subtraction(): with pytest.raises(DimensionMismatchError): np.array([5], dtype=np.float64) - q - # Check that operations with 0 work with pytest.raises(DimensionMismatchError): assert_quantity(q + 0, q.value, volt) @@ -1473,3 +1472,123 @@ def f(b): return b f(a) + + +def test_set_default_magnitude_1(): + bu.set_default_magnitude(-3) + # from brainunit import second, ms + q1 = 3 * bu.second + q2 = 3 * bu.ms + + assert q1.to_value() == 3e3 + assert q2.to_value() == 3. + + bu.set_default_magnitude(0) + # from brainunit import second, ms + q1 = 3 * bu.second + q2 = 3 * bu.ms + assert q1.to_value() == 3. + assert q2.to_value() == 3e-3 + + bu.set_default_magnitude(3) + # from brainunit import second, ms + q1 = 3 * bu.second + q2 = 3 * bu.ms + assert q1.to_value() == 3e-3 + assert q2.to_value() == 3e-6 + +def test_set_default_magnitude_2(): + bu.set_default_magnitude(-3) + # from brainunit import second, ms, meter, kmeter + q1 = 3 * bu.second * bu.meter + q2 = 3 * bu.ms * bu.kmeter + + value = str(q1) + + assert q1.to_value() == 3e6 + assert q2.to_value() == 3e6 + + bu.set_default_magnitude(0) + # from brainunit import second, ms, meter, kmeter + q1 = 3 * bu.second * bu.meter + q2 = 3 * bu.ms * bu.kmeter + assert q1.to_value() == 3. + assert q2.to_value() == 3. + + bu.set_default_magnitude(3) + # from brainunit import second, ms, meter, kmeter + q1 = 3 * bu.second * bu.meter + q2 = 3 * bu.ms * bu.kmeter + assert q1.to_value() == 3e-6 + assert q2.to_value() == 3e-6 + +def test_set_default_magnitude_3(): + bu.set_default_magnitude({'s': 'm'}) + from brainunit import second, ms, meter, kmeter + q1 = 3 * bu.second * bu.meter + q2 = 3 * bu.ms * bu.kmeter + + assert q1.to_value() == 3e3 + assert q2.to_value() == 3e3 + + bu.set_default_magnitude(0) + # from brainunit import second, ms, meter, kmeter + q1 = 3 * bu.second * bu.meter + q2 = 3 * bu.ms * bu.kmeter + assert q1.to_value() == 3. + assert q2.to_value() == 3. + + bu.set_default_magnitude({'s': 3}) + # from brainunit import second, ms, meter, kmeter + q1 = 3 * bu.second * bu.meter + q2 = 3 * bu.ms * bu.kmeter + assert q1.to_value() == 3e-3 + assert q2.to_value() == 3e-3 + +def test_set_default_magnitude_4(): + # volt: m=2, kg=1, s=-3, A=-1 + bu.set_default_magnitude(-3) + # from brainunit import volt, mV + q1 = 3 * bu.volt + q2 = 3 * bu.mV + + assert q1.to_value() == 3e-3 + assert q2.to_value() == 3e-6 + + bu.set_default_magnitude(0) + # from brainunit import volt, mV + q1 = 3 * bu.volt + q2 = 3 * bu.mV + assert q1.to_value() == 3. + assert q2.to_value() == 3e-3 + + bu.set_default_magnitude(3) + # from brainunit import volt, mV + q1 = 3 * bu.volt + q2 = 3 * bu.mV + assert q1.to_value() == 3e3 + assert q2.to_value() == 3 + +# def test_set_default_magnitude_5(): +# # volt: m=2, kg=1, s=-3, A=-1 --> (if set to milivolt) _default_magnitude = {'m': -1.5, 'kg': -3, 's': -1, 'A': -3} (WRONG!) +# bu.set_default_magnitude(-3, unit=bu.volt) +# from brainunit import volt, mV +# q1 = 3 * volt +# q2 = 3 * mV +# +# assert q1.to_value() == 3e3 +# assert q2.to_value() == 3 +# +# bu.set_default_magnitude(0) +# from brainunit import volt, mV +# q1 = 3 * volt +# q2 = 3 * mV +# assert q1.to_value() == 3. +# assert q2.to_value() == 3e-3 +# +# bu.set_default_magnitude(3) +# from brainunit import volt, mV +# q1 = 3 * volt +# q2 = 3 * mV +# assert q1.to_value() == 3e-3 +# assert q2.to_value() == 3e-6 \ No newline at end of file