Skip to content

Commit

Permalink
fix unit and quantity disaply
Browse files Browse the repository at this point in the history
  • Loading branch information
chaoming0625 committed Aug 23, 2024
1 parent f6456ab commit f48ac46
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 21 deletions.
91 changes: 70 additions & 21 deletions brainunit/_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -1146,34 +1146,33 @@ def _assert_same_base(u1, u2):

def _create_name(dim: Dimension, base, scale) -> str:
if dim == DIMENSIONLESS:
name = f"Unit({base}^{scale})"
return f"Unit({base}^{scale})"
else:
if _is_tracer(scale) or scale == 0.:
name = f"{dim}"
else:
name = f"{base}^{scale} * {dim}"
return name
return name


def _find_name(dim: Dimension, base, scale) -> Tuple[str, str]:
def _find_a_name(dim: Dimension, base, scale) -> Optional[str]:
if dim == DIMENSIONLESS:
name = f"Unit({base}^{scale})"
return name, name
return name

if isinstance(base, (int, float)):
if isinstance(scale, (int, float)):
key = (dim, scale, base)
if key in _standard_units:
name = _standard_units[key].name
return name, name
else:
key = (dim, 0, base)
if key in _standard_units:
name = _standard_units[key].name
return name, name

name = _create_name(dim, base, scale)
return name, name
return name
key = (dim, 0, base)
if key in _standard_units:
name = _standard_units[key].name
if _is_tracer(scale):
return name
else:
return f"{base}^{scale} * {name}"


_standard_units: Dict[Tuple, 'Unit'] = {}
Expand Down Expand Up @@ -1547,8 +1546,15 @@ def __mul__(self, other) -> 'Unit' | Quantity:
_assert_same_base(self, other)
scale = self.scale + other.scale
dim = self.dim * other.dim
name, dispname = _find_name(dim, self.base, scale)
return Unit(dim, scale=scale, base=self.base, name=name, dispname=dispname, iscompound=True)
name = _find_a_name(dim, self.base, scale)
dispname = name
if name is None:
name = f"{self.name} * {other.name}"
dispname = f"{self.dispname} * {other.dispname}"
iscompound = True
else:
iscompound = False
return Unit(dim, scale=scale, base=self.base, name=name, dispname=dispname, iscompound=iscompound)

elif isinstance(other, Quantity):
return Quantity(other._mantissa, unit=(self * other.unit))
Expand Down Expand Up @@ -1577,8 +1583,27 @@ def __div__(self, other) -> 'Unit':
scale = self.scale - other.scale
dim = self.dim / other.dim
_assert_same_base(self, other)
name, dispname = _find_name(dim, self.base, scale)
return Unit(dim, base=self.base, scale=scale, name=name, dispname=dispname, iscompound=True)
name = _find_a_name(dim, self.base, scale)
dispname = name
if name is None:
if self.iscompound:
dispname = f"({self.dispname})"
name = f"({self.name})"
else:
dispname = self.dispname
name = self.name
dispname += "/"
name += " / "
if other.iscompound:
dispname += f"({other.dispname})"
name += f"({other.name})"
else:
dispname += other.dispname
name += other.name
iscompound = True
else:
iscompound = False
return Unit(dim, base=self.base, scale=scale, name=name, dispname=dispname, iscompound=iscompound)
else:
raise TypeError(f"unit {self} cannot divide by a non-unit {other}")

Expand All @@ -1587,8 +1612,19 @@ def __rdiv__(self, other) -> 'Unit' | Quantity:
if is_scalar_type(other) and other == 1:
dim = self.dim ** -1
scale = -self.scale
name, dispname = _find_name(dim, self.base, scale)
return Unit(dim, base=self.base, scale=scale, name=name, dispname=dispname, iscompound=True)
name = _find_a_name(dim, self.base, scale)
dispname = name
if name is None:
if self.iscompound:
dispname = f"({self.dispname})"
name = f"({self.name})"
else:
dispname = self.dispname
name = self.name
iscompound = True
else:
iscompound = False
return Unit(dim, base=self.base, scale=scale, name=name, dispname=dispname, iscompound=iscompound)

elif isinstance(other, Unit):
return other.__div__(self)
Expand Down Expand Up @@ -1627,8 +1663,21 @@ def __pow__(self, other):
if is_scalar_type(other):
dim = self.dim ** other
scale = self.scale * other
name, dispname = _find_name(dim, self.base, scale)
return Unit(dim, base=self.base, scale=scale, name=name, dispname=dispname, iscompound=True)
name = _find_a_name(dim, self.base, scale)
dispname = name
if name is None:
if self.iscompound:
dispname = f"({self.dispname})"
name = f"({self.name})"
else:
dispname = self.dispname
name = self.name
dispname += f"^{str(other)}"
name += f" ** {repr(other)}"
iscompound = True
else:
iscompound = False
return Unit(dim, base=self.base, scale=scale, name=name, dispname=dispname, iscompound=iscompound)
else:
raise TypeError(f"unit cannot perform an exponentiation (unit ** other) with a non-scalar, "
f"since one unit cannot contain multiple units. \n"
Expand Down
10 changes: 10 additions & 0 deletions brainunit/_base_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,6 +245,16 @@ def test_display(self):
display_in_unit(10 * nS, ohm)
assert_equal(display_in_unit(10.0, Unit(scale=1)), "1. * Unit(10.0^1)")
assert_equal(display_in_unit(3 * bu.kmeter / bu.meter), '3. * Unit(10.0^3)')
assert_equal(str(bu.mS / bu.cm ** 2), 'mS/cmeter2')

assert_equal(display_in_unit(10. * bu.mV), '10. * mvolt')
assert_equal(display_in_unit(10. * bu.ohm * bu.amp), '10. * volt')
assert_equal(display_in_unit(120. * (bu.mS / bu.cm ** 2)), '120. * msiemens / cmeter2')
assert_equal(display_in_unit(3.0 * bu.kmeter / 130.51 * bu.meter), '0.02298675 * 10.0^3 * meter2')
assert_equal(display_in_unit(3.0 * bu.kmeter / (130.51 * bu.meter)), '0.02298675 * Unit(10.0^3)')
assert_equal(display_in_unit(3.0 * bu.kmeter / 130.51 * bu.meter * bu.cm ** -2), '0.02298675 * Unit(10.0^7)')
assert_equal(display_in_unit(3.0 * bu.kmeter / 130.51 * bu.meter * bu.cm ** -1), '0.02298675 * 10.0^5 * meter')
assert_equal(display_in_unit(1. * bu.joule / bu.kelvin), '1. * joule / kelvin')

def test_display2(self):

Expand Down

0 comments on commit f48ac46

Please sign in to comment.