Source code for scimath.units.unit_array

```# (C) Copyright 2005-2021 Enthought, Inc., Austin, TX
#
# This software is provided without warranty under the terms of the BSD
# is also available online at http://www.enthought.com/licenses/BSD.txt
#
# Thanks for using Enthought open source!

# Numeric Library Imports
import numpy

# Enthought library imports
from scimath.units import convert
from scimath.units.unit import unit, dimensionless, IncompatibleUnits
from scimath.units.unit_parser import unit_parser

def __newobj__(cls, *args):
""" Unpickles new-style objects.
"""
return cls.__new__(cls, *args)

_retain_units = [numpy.absolute, numpy.negative,
numpy.floor, numpy.ceil, numpy.rint, numpy.conjugate,
numpy.maximum, numpy.minimum]

_retain_units_if_same = []  # [numpy.add, numpy.subtract]

_retain_units_if_single = []  # [numpy.multiply]

_retain_units_if_only_first = [numpy.remainder]  # , numpy.divide]

[docs]class UnitArray(numpy.ndarray):
""" Define a UnitArray that subclasses from a Numpy array

This class simply adds a "units" object to a numpy array.  It deals
with unit adjustments involving basic arithmetic and comparison
operations between arrays with units.  If incompatible units are
used, an error will be raised.

Note that the "units" may hold some amount of the magnitude, which can
be particularly significant for dimensionless quantities.

You should always convert a UnitArray to the required units before
extracting the numerical value.

This code is shamelessly copied from the numpy matrix class.  It is
of course modified to suite our purposes.
"""

##########################################################################
# numpy.ndarray attributes
##########################################################################

# priority->0.0 results in binary array ops returning UnitArray objects.
__array_priority__ = 10.0

##########################################################################
# UnitArray attributes
#
# Note: These are commented out because we are not using traits.  However,
#       They are left here for documentation.
##########################################################################

# Units specification.  it is a scimath.units.unit object, not a string.
# units = Instance(unit)

# 'type' specifies the type of UnitArray.
# This could be useful for specifying that this is a 'tops' UnitArray or of
# 'flag' or other such types.  I am not sure if this is useful, but I
# think it might be.
# fixme: Not Implemented
# type = Str

##########################################################################
# object interface
##########################################################################
def __repr__(self):
""" String representation using the repr of the unit."""
base_str = self._get_values_base_str()
s = "{klass}({val}, units='{unit!r}')"
klass = type(self).__name__
return s.format(klass=klass, val=base_str, unit=self.units)

def __str__(self):
""" String representation using the label of the unit."""
s = "{klass} ({unit}): {val}"
base_str = self._get_values_base_str()
if self.units.label is not None:
str_unit = self.units.label
else:
str_unit = repr(self.units)

return s.format(klass=type(self).__name__, val=base_str, unit=str_unit)

def __reduce_ex__(self, protocol):
"""
pickling function for classes which inherit from tuple.

__reduce_ex__ must be overloaded for pickling to work. Refer to the docs
in the pickle source code for details as to why.

"""

state = (self.units, super(UnitArray, self).__reduce_ex__(protocol))
return (__newobj__, (self.__class__, ()), state)

def __setstate__(self, state):
"""
unpickling function
"""

super(UnitArray, self).__setstate__(state[1][2])
units = state[0]
self.units = units

def __deepcopy__(self, memo={}):
copy = self.__class__(self.view(numpy.ndarray), copy=True,
units=self.units)
memo[id(self)] = copy
return copy

##########################################################################
# numpy.ndarray interface
##########################################################################

def __new__(cls, data, dtype=None, copy=True, units=None):
""" Called when a new object is created (before __init__).

The default behavior of ndarray is overridden to add units and
family_name to the class.

For more details, see:
http://docs.python.org/ref/customization.html

"""

### Array Setup #######################################################

if isinstance(data, numpy.ndarray):
# Handle the case where we are passed in a numpy array.
if dtype is None:
intype = data.dtype
else:
intype = numpy.dtype(dtype)

new = data.view(cls)

if intype != data.dtype:
res = new.astype(intype)
elif copy:
res = new.copy()
else:
res = new

else:
# Handle other input types (lists, etc.)
arr = numpy.array(data, dtype=dtype, copy=copy)

res = numpy.ndarray.__new__(cls, arr.shape, arr.dtype,
buffer=arr)

### Configure Other Attributes ########################################
if isinstance(units, str):
units = unit_parser.parse_unit(units)

res.units = units

return res

def __array_finalize__(self, obj):
"""
Copy any values that were on the original UnitArray into the output
UnitArray.
"""
units = getattr(obj, 'units', None)
self.units = units

def __array_wrap__(self, obj, context=None):
# Behavior of getting units is
#   to retain units if
#  1) the function is in _retain_units list
#  2) the function is in _retain_units_if_same list
#     and the arguments either have the same units
#     or no units.
#  3) the function is in _retain_units_if_single list
#     and all but one argument does not have units
#     (or is dimensionless)
#  4) the function is in _retain_units_if_only_first
#     and all but the first argument does not have units
#     (or is dimensionless)
#  5) the function has a single argument which is dimensionless
#
# FIXME: To do this "right" would require a ufunc
#      object that could tell how to do unit conversions
#      with its inputs to produce its outputs.
result = obj.view(self.__class__)
theunit = None
if context is not None:
func, args, iout = context
if func in _retain_units:
theunit = getattr(self, 'units', None)
elif func in _retain_units_if_same:
theunit = None
for arg in args:
u = getattr(arg, 'units', None)
if u is None:
continue
if theunit is None:
theunit = u
elif (theunit != u):
theunit = None
elif func in _retain_units_if_single:
theunit = None
for arg in args:
u = getattr(arg, 'units', None)
if u is not None:
if theunit is None:
theunit = u
elif u != dimensionless:  # already a unit
theunit = None
break
elif func in _retain_units_if_only_first:
theunit = getattr(args[0], 'units', None)
for arg in args[1:]:
u = getattr(arg, 'units', None)
if u is not None and u != dimensionless:
theunit = None
break
elif len(args) == 1:
theunit = getattr(self, 'units', None)
if theunit != dimensionless:
theunit = None

result.units = theunit
return result

def __convert_other(self, other):
su = getattr(self, 'units', dimensionless)
ou = getattr(other, 'units', dimensionless)

if su is None and ou is None:
u = None
else:
if isinstance(other, unit):
# Handles 5 * liters
ou = unit(1, other.derivation)
other = convert(other.value, ou, su)
elif isinstance(other, UnitArray):
# Handles UnitArray or UnitScalar
other = convert(numpy.array(other), ou, su)
elif isinstance(other, numpy.ndarray):
if len(other.shape) > 0 and hasattr(
other.item(0), 'derivation'):
# Handles array([1,2,3] * liters)
ou = unit(1, other.item(0).derivation)
other = convert(other / ou, ou, su)
u = su
return other, u

"""
Defines the addition of 2 unitted arrays
"""
other, u = self.__convert_other(other)
result.units = u
return result

"""
Defines the reverse addition of 2 unitted arrays
"""

def __sub__(self, other):
"""
Defines the subtraction of 2 unitted arrays
"""
other, u = self.__convert_other(other)
result = super(UnitArray, self).__sub__(other)
result.units = u
return result

def __rsub__(self, other):
"""
Defines the reverse subtraction of 2 unitted arrays
"""
su = getattr(self, 'units', dimensionless)
ou = getattr(other, 'units', dimensionless)
if su is None and ou is None:
s = self
u = None
else:
s = self.as_units(ou)
u = ou
result = super(UnitArray, s).__rsub__(other)
result.units = u
return result

def __mul__(self, other):
"""
Defines the multiplication of 2 unitted arrays
"""
result = super(UnitArray, self).__mul__(other)
su = getattr(self, 'units', None)
ou = getattr(other, 'units', None)
if su and ou:
# note that there may be a scale factor in the units.
# This may be confusing for otherwise dimensionless
# quantities
result.units = su * ou
elif su:
result.units = su
else:
result.units = ou

# If the units are only a scale factor, apply it to result and set
# to dimensionless
if isinstance(result.units, float):
result = result.as_units(dimensionless)
return result

def __rmul__(self, other):
"""
Defines the reverse multiplication of 2 unitted arrays
"""
# multiplication is commutative
return self.__mul__(other)

def __div__(self, other):
return type(self).__truediv__(self, other)

def __truediv__(self, other):
"""
Defines the division of 2 unitted arrays
"""
result = super(UnitArray, self).__truediv__(other)
su = getattr(self, 'units', None)
ou = getattr(other, 'units', None)
if su and ou:
# note that there may be a scale factor in the units.
# This may be confusing for otherwise dimensionless
# quantities, but the alternative is losing units like
# 'percent'
result.units = su / ou
elif ou:
result.units = 1 / ou
else:
result.units = su

# If the units are only a scale factor, apply it to result and set
# to dimensionless
if isinstance(result.units, float):
result = result.as_units(dimensionless)
return result

def __rdiv__(self, other):
return type(self).__rtruediv__(self, other)

def __rtruediv__(self, other):
"""
Defines the reverse division of 2 unitted arrays
"""
result = super(UnitArray, self).__rtruediv__(other)
su = getattr(self, 'units', None)
ou = getattr(other, 'units', None)
if su and ou:
# note that there may be a scale factor in the units.
# This may be confusing for otherwise dimensionless
# quantities
result.units = ou / su
elif su:
result.units = 1 / su
else:
result.units = ou

# If the units are only a scale factor, apply it to result and set
# to dimensionless
if isinstance(result.units, float):
result = result.as_units(dimensionless)
return result

def __pow__(self, other):
"""
Defines the exponent operator of a unitted array
"""
if isinstance(other, (int, int, float)) or \
(isinstance(other, numpy.ndarray) and other.shape == ()):
if isinstance(other, UnitArray):
if getattr(other, "units", dimensionless) == dimensionless:
other = float(other)
else:
raise IncompatibleUnits("exponent must be dimensionless")
result = super(UnitArray, self).__pow__(other)
su = getattr(self, 'units', None)
if su is not None:
result.units = su**other
return result
else:
raise TypeError("exponent must be an integer, float or 0-d array")

def __rpow__(self, other):
return NotImplemented

def __le__(self, other):
"""
Defines the 'less or equal' operator of 2 unitted arrays
"""
other, u = self.__convert_other(other)
result = super(UnitArray, self).__le__(other)
return result

def __lt__(self, other):
"""
Defines the 'less than' operator of 2 unitted arrays
"""
other, u = self.__convert_other(other)
result = super(UnitArray, self).__lt__(other)
return result

def __ge__(self, other):
"""
Defines the 'greater or equal' operator of 2 unitted arrays
"""
other, u = self.__convert_other(other)
result = super(UnitArray, self).__ge__(other)
return result

def __gt__(self, other):
"""
Defines the 'greater than' operator of 2 unitted arrays
"""
other, u = self.__convert_other(other)
result = super(UnitArray, self).__gt__(other)
return result

def __eq__(self, other):
"""
Defines the 'equal' operator of 2 unitted arrays
"""
try:
other, u = self.__convert_other(other)
result = super(UnitArray, self).__eq__(other)
return result
except:
return False

def __ne__(self, other):
"""
Defines the 'not equal' operator of 2 unitted arrays
"""
try:
other, u = self.__convert_other(other)
result = super(UnitArray, self).__ne__(other)
return result
except:
return True

##########################################################################
# UnitArray interface
##########################################################################

### Unit Conversion ######################################################

def as_units(self, new_units):
""" Convert UnitArray from its current units to a new set of units.

"""
result = self.__class__(convert(self.view(numpy.ndarray),
self.units, new_units))
result.units = new_units

return result

# Unit Conversion ########################################################

def _get_values_base_str(self):
""" Build a string representation of the array values.
"""
# this is a little more complicated than it should be in order
# to be more resilient to changes to numpy ndarray API
base_str = numpy.ndarray.__repr__(self.view(numpy.ndarray))
start = base_str.find('(')
end = base_str.rfind(')')

if start > -1 and end > -1:
base_str = base_str[start+1:end]

return base_str

##########################################################################
# static methods which wrap numpy builtin functions
##########################################################################

@staticmethod
def concatenate(sequences, axis=0):
result = numpy.concatenate(sequences, axis)
result.units = sequences[0].units
return result
```