numeric arrays with units

John Hunter jdhunter at ace.bsd.uchicago.edu
Thu May 22 13:24:34 EDT 2003


Recently I asked about using Numeric arrays in combination with a
units package.  Fernando Perez pointed me to ScientificPython's
PhysicalQuantity module, which works with ufuncs, but not array funcs.
So I followed the lead of Numeric's UserArray and wrote a Numeric
array work-alike that incorporates units information.

You can combine operations of PhysicalQuantities and UnitArrays in the
same way you can combine scalar and array operations in Numeric, with
automatic unit conversion if the units are compatible, as in

    >>> a = UnitArray([1., 2., 3.], 'g')
    >>> b = UnitArray([0.0001, 0.1, 3.], 'kg')
    >>> c = PhysicalQuantity(6, 'mg')
    >>> print a+c
    UnitArray([ 1.006,  2.006,  3.006]) g
    >>> print a+b
    UnitArray([  1.10000000e+00,   1.02000000e+02,   3.00300000e+03]) g
    >>> print a>b
    [1 0 0]


You can use ufuncs on UnitArrays if the dimensions are correct

    >>> t = UnitArray(arange(0.0, 5.0, 0.2), 's')
    >>> f = UnitArray(arange(5.0, 10.0, 0.2), 'Hz')
    >>> sin(2*pi*f*t)
    array([  0.00000000e+00,   2.48689887e-01,   8.44327926e-01,   7.70513243e-01,
      ...snip...

or
    >>> t = UnitArray(arange(0.0, 1.0, 0.2), 's')
    >>> tau = PhysicalQuantity(50, 'ms')
    >>> exp(-t/tau)     # OK
    array([  1.00000000e+00,   1.83156389e-02,   3.35462628e-04,   
             6.14421235e-06, 1.12535175e-07])
    >>> exp(-t)  # bad, exp arg must be dimensionless
    Traceback (most recent call last):
      File "<stdin>", line 1, in ?
    AttributeError: PhysicalQuantity instance has no attribute 'exp'


as well as array funcs

   >>> a = UnitArray([1., 2., 3.], 's')
   >>> print sum(a)
   6.0 s
   >>> print cumsum(a)
   [1.0 s  3.0 s  6.0 s ]
   >>> print cumproduct(a)
   [1.0 s  2.0 s**2  6.0 s**3 ]

There are some unit tests at the end of the UnitArray module, which
pass, but let me know if you encounter any bugs because this is not
extensively tested.  

And I have a question: all the Numeric array funcs start with

def somefunc(x):  
   x = array(x, copy=0)
   blah_blah_blah

which converts an array like 

  >>> a = UnitArray([1., 2., 3.], 's')

to 
  >>> array(a, copy=0)
  array([1.0 s , 2.0 s , 3.0 s ],'O')

ie, an array of PhysicalQuantities.  I would rather this return an
UnitArray (which stores an array of scalars and a unit attribute.  I
know that you can overload the __array__ method of UnitArray, but I
don't know how to do it to achieve the desired behavior.  So if you
have any suggestions, let me know.

John Hunter


"""
This code is a melding of UserArray from Numeric and PhysicalQuantity
from ScientificPython, and requires Numeric and ScientificPython.  If
you want a standalone version which doesn't require ScientificPython,
email me at jdhunter at ace.bsd.uchicago.edu and I can send you one (that
has all the requisite ScientificPython in one file with the UnitArray
code).

The class UnitArray is based on Numeric's UserArray, which is
distributed under the python license


"""
from Numeric import *
from Scientific.Physics.PhysicalQuantities import PhysicalQuantity, \
     isPhysicalQuantity, PhysicalUnit, _findUnit

class UnitArray:
    "A Numeric array with units"
    def __init__(self, a, unit):
        "a is a Numeric array, unit is a unit string"
        try: a.shape
        except AttributeError: a = array(a)
        self.value = a
        self.shape = self.value.shape
        self._typecode = self.value.typecode()
        self.name = str(self.__class__).split()[0]
        self.unit = _findUnit(unit)


    def _sum(self, other, sign1, sign2):  
	if not isPhysicalQuantity(other):
	    raise TypeError, 'Incompatible types'
	new_value = sign1*self.value + \
		    sign2*other.value*other.unit.conversionFactorTo(self.unit)
	return self._rc(new_value)

    def __repr__(self): 
        if len(self.value.shape) > 0:
            return self.__class__.__name__+repr(self.value)[len("array"):] + ' ' + self.unit.name()
        else:
            return self.__class__.__name__+"("+repr(self.value)+")"  + ' ' + self.unit.name()


    def __float__(self): 
        return self.__class__(float(asarray(self.value)), self.unit)

    # Array as sequence
    def __len__(self): return len(self.value)

    def __getitem__(self, index):
        return self._rc(self.value[index])

    def __getslice__(self, i, j): 
        return self._rc(self.value[i:j])


    def __setitem__(self, index, value):
        self.value[index] = asarray(value,self._typecode)

    def __setslice__(self, i, j, value):
        self.value[i:j] = asarray(value,self._typecode)

    def __del__(self):
        # necessary?
        for att in self.__dict__.keys():
            delattr(self,att)

    def __abs__(self): return self._rc(absolute(self.value))
    def __neg__(self): return self._rc(-self.value)

    def __add__(self, other): 
	return self._sum(other, 1, 1)

    __radd__ = __add__

    def __iadd__(self, other):  
	if not isPhysicalQuantity(other):
	    raise TypeError, 'Incompatible types'

        add(self.value, other.value*other.unit.conversionFactorTo(self.unit),
            self.value)
        return self
        
    def __sub__(self, other): 
	return self._sum(other, 1, -1)
    
    def __rsub__(self, other):  
        return self._sum(other, -1, 1)

    def __isub__(self, other): 
        if not isPhysicalQuantity(other):
	    raise TypeError, 'Incompatible types'
        subtract(self.value,
                 other.value*other.unit.conversionFactorTo(self.unit),
                 self.value)
        return self

    def __mul__(self, other): 
	if not isPhysicalQuantity(other):
	    return self._rc(self.value*other)

	a = self.value*other.value
	unit = self.unit*other.unit
	if unit.isDimensionless():
	    return a*unit.factor
	else:
	    return self._rc(a, unit)

    __rmul__ = __mul__


    def __imul__(self, other): 

	if isPhysicalQuantity(other):
            multiply(self.value, other.value, self.value)
            self.unit = self.unit*other.unit
            if self.unit.isDimensionless():
                return self.value*unit.factor
        else:
            self.value *= other
        return self

    def __div__(self, other): 
	if not isPhysicalQuantity(other):
	    return self._rc(divide(self.value,other))

	a = divide(self.value, other.value)
	unit = self.unit/other.unit
	if unit.isDimensionless():
	    return a*unit.factor    # yes multiply
	else:
	    return self._rc(a, unit)

    def __rdiv__(self, other): 
	if not isPhysicalQuantity(other):
	    return self._rc(divide(other, self.value), unit=1/self.unit)

	a = divide(other.value, self.value)
	unit = other.unit/self.unit
	if unit.isDimensionless():
	    return a*unit.factor    # yes multiply
	else:
	    return self._rc(a, unit)

    def __idiv__(self, other): 
	if isPhysicalQuantity(other):
            divide(self.value, other.value, self.value)
            self.unit = self.unit/other.unit
            if self.unit.isDimensionless():
                return self.value*unit.factor  # yes multiply
        else:
            divide(self.value, other, self.value)
        return self
    

    def __pow__(self,other):
        if isPhysicalQuantity(other):
            raise TypeError, 'Powers should be dimensionless'

        # PhysicalQuantity will only let us use integer powers as
        # exponents, so we'll have to do this in a loop.

        try: iter(other)
        except TypeError:
            unit = self.unit**other
            return self._rc(self.value**other, unit)
        else:
            a = self.value**other
            units = ['']*len(a)
            ret = [None]*len(a)
            for i in xrange(len(other)):
                unit = self.unit**other[i]
                ret[i] = PhysicalQuantity(a[i], unit)
        
            return array(ret)

    def __neg__(self):
        return self._rc(-self.value)

    def __pos__(self):
        return self._rc(self.value)

    def __abs__(self):
        return self._rc(abs(self.value))

    def __invert__(self):
        unit = 1/self.unit
        return self._rc(invert(self.value), unit)

    def _scalarfunc(a, func):
        if len(a.shape) == 0:
            return self._rc(func(a[0]))
        else:
            raise TypeError, "only rank-0 arrays can be converted to Python scalars."
        
    def __complex__(self): return self._scalarfunc(complex)
    def __float__(self): return self._scalarfunc(float)
    def __int__(self): return self._scalarfunc(int)
    def __long__(self): return self._scalarfunc(long)
    def __hex__(self): return self._scalarfunc(hex)
    def __oct__(self): return self._scalarfunc(oct)

    def __lt__(self,other):
	if not isPhysicalQuantity(other):
	    raise TypeError, 'Incompatible types'
	return less(self.value,
                    other.value*other.unit.conversionFactorTo(self.unit))

    def __le__(self,other):
	if not isPhysicalQuantity(other):
	    raise TypeError, 'Incompatible types'
	return less_equal(self.value,
                    other.value*other.unit.conversionFactorTo(self.unit))

    def __eq__(self,other):
	if not isPhysicalQuantity(other):
	    return 0
        try: factor = other.unit.conversionFactorTo(self.unit)
        except TypeError: return 0
	return equal(self.value, other.value*factor)

    def __ne__(self,other): 
	if not isPhysicalQuantity(other):
            return 1
        try: factor = other.unit.conversionFactorTo(self.unit)
        except TypeError: return 0
	return not_equal(self.value, other.value*factor)

    def __gt__(self,other):
	if not isPhysicalQuantity(other):
	    raise TypeError, 'Incompatible types'
	return greater(self.value,
                    other.value*other.unit.conversionFactorTo(self.unit))

    def __ge__(self,other): 
	if not isPhysicalQuantity(other):
	    raise TypeError, 'Incompatible types'
	return greater_equal(self.value,
                    other.value*other.unit.conversionFactorTo(self.unit))
        
    def copy(self):
        return self._rc(self.value.copy())

    def tostring(self): return self.value.tostring()

    def byteswapped(self): return self._rc(self.value.byteswapped())

    def astype(self, typecode):
        return self._rc(self.value.astype(typecode))
   
    def typecode(self):
        return self._typecode

    def itemsize(self): return self.value.itemsize()

    def iscontiguous(self): return self.value.iscontiguous()

    def _rc(self, a, unit=None):
        if unit is None: unit = self.unit.name()
        if len(shape(a)) == 0:
            return PhysicalQuantity(a, unit)
        return self.__class__(a, unit)

    def __setattr__(self,attr,value):
        if attr=='shape':
            self.value.shape=value
        self.__dict__[attr]=value

    def __getattr__(self,attr):
        # for .attributes for example, and any future attributes
        if attr == 'real':
            return self._rc(self.value.real)
        elif attr == 'imag':
            return self._rc(self.value.imag)
        elif attr == 'flat':
            return self._rc(self.value.flat)
        return getattr(self.value, attr)
            

if __name__ == '__main__':
    # add two arrays with compatible units
    # add two arrays with compatible units
    a = UnitArray([1., 2., 3.], 'g')
    b = UnitArray([1., 2., 3.], 'mg')
    c = PhysicalQuantity(6, 's')

    assert(a+b==UnitArray([1.001, 2.002, 3.003], 'g'))
    assert(a*b==UnitArray([1., 4., 9.], 'g*mg'))
    # floating point margin
    assert(a/b-array([1000.0, 1000.0, 1000.0])<array([1e-10, 1e-10, 1e-10]))

    try: a+c
    except TypeError: pass
    else:
        raise TypeError, 'Added incompatible types'


    assert(a*c == UnitArray([6., 12., 18.], 'g*s'))
    assert(c/a == UnitArray([6., 3., 2.], 's/g'))

    c = PhysicalQuantity(6, 'kg')
    assert(a+c == UnitArray([6001., 6002., 6003.], 'g'))
    assert(c/a == array([6000., 3000., 2000.]))
    assert(c-a == UnitArray([5999., 5998., 5997.], 'g'))

    a += c
    a *= b
    a /= b
    assert(a==UnitArray([6001., 6002., 6003.], 'g'))




    a = UnitArray([1., 2., 3., 4.], 'Hz')
    b = UnitArray([0.0001, 0.1, 1.0, 0.004], 'kHz')
    assert(sum(a<b)==2)
    assert(sum(a>b)==1)
    assert(sum(a==b)==1)
    assert(sum(a>=b)==2)
    assert(sum(a!=b)==3)

    c = UnitArray([0.0001, 0.1, 1.0, 1.0], 'l')

    try: a<c
    except TypeError: pass
    else:
        raise TypeError, 'Compared incompatible types'

    a = UnitArray([1., 2., 3., 4.], 'g')

    assert( sum(a) == PhysicalQuantity(10, 'g'))
    print cumsum(a) 
    print cumproduct(a)

    tau = PhysicalQuantity(0.05, 's')
    t = UnitArray(arange(0.0, 1.0, 0.2), 's')
    i = UnitArray(arange(0.0, 1.0, 0.2), 'C')
    r = UnitArray(exp(-t/tau), 'V')
    convolve(i,r)

    t = UnitArray(arange(0.0, 5.0, 0.2), 's')
    f = UnitArray(arange(5.0, 10.0, 0.2), 'Hz')
    sin(2*pi*f*t)

    try: sin(2*pi*f)
    except TypeError: pass
    else:
        raise 'Took a sin of Hz, should be illegal'

    t = UnitArray(arange(0.0, 5.0, 0.2), 's')
    t -= PhysicalQuantity(1.0, 'h')
    assert(t[0]==PhysicalQuantity('-3600s'))

    x = sort( UnitArray([1, -1, 2], 'cal'))
    print type(x), type(t)
    print array(UnitArray([1, -1, 2], 'cal'), copy=0)


    print 'all tests passed'





More information about the Python-list mailing list