"""numarray: The big enchilada numeric module $Id: numarray.py,v 1.7 2002/03/11 15:34:56 jaytmiller Exp $ """ import types, math, os.path import memory import ndarray import _bytes import _numarray import ufunc from ufunc import * import arrayprint from numerictypes import * from ndarray import * from math import pi # rename built-in function type so not to conflict with keyword _type = type MAX_LINE_WIDTH = None PRECISION = None SUPPRESS_SMALL = None _FPE_DIVIDE_BY_ZERO = 1 _FPE_OVERFLOW = 2 _FPE_UNDERFLOW = 4 _FPE_INVALID = 8 class NumError: """Defines how numeric errors should be handled.""" def __init__(self, all="warn", overflow=None, underflow=None, dividebyzero=None, invalid=None): self._errorset = ("ignore", "warn", "raise") self.setMode(all=all, underflow=underflow, overflow=overflow, dividebyzero=dividebyzero, invalid=invalid) def setMode(self, all="warn", overflow=None, underflow=None, dividebyzero=None, invalid=None): if all in self._errorset: self.overflowMode = all self.underflowMode = all self.dividebyzeroMode = all self.invalidMode = all if overflow in self._errorset: self.overflowMode = overflow if underflow in self._errorset: self.underflowMode = underflow if dividebyzero in self._errorset: self.dividebyzeroMode = dividebyzero if invalid in self._errorset: self.invalidMode = invalid Error = NumError() class MathDomainError(ArithmeticError): pass class UnderflowError(ArithmeticError): pass def handleError(errorStatus, sourcemsg): """Take error status and use error mode to handle it.""" if errorStatus & _FPE_INVALID: if Error.invalidMode == "warn": print "Warning: Encountered invalid numarray result(s)", sourcemsg if Error.invalidMode == "raise": raise MathDomainError(sourcemsg) if errorStatus & _FPE_DIVIDE_BY_ZERO: if Error.dividebyzeroMode == "warn": print "Warning: Encountered divide by zero(s)", sourcemsg if Error.dividebyzeroMode == "raise": raise ZeroDivisionError(sourcemsg) if errorStatus & _FPE_OVERFLOW: if Error.overflowMode == "warn": print "Warning: Encountered overflow(s)", sourcemsg if Error.overflowMode == "raise": raise OverflowError(sourcemsg) if errorStatus & _FPE_UNDERFLOW: if Error.underflowMode == "warn": print "Warning: Encountered underflow(s)", sourcemsg if Error.underflowMode == "raise": raise UnderflowError(sourcemsg) PyINT_TYPES = { types.IntType: 1, types.LongType: 1, } PyREAL_TYPES = { types.IntType: 1, types.LongType: 1, types.FloatType: 1, } # Python numeric types with values indicating level in type hierarchy PyNUMERIC_TYPES = { types.IntType: 0, types.LongType: 1, types.FloatType: 2, types.ComplexType: 3 } # Mapping back from level to type PyLevel2Type = {} for key, value in PyNUMERIC_TYPES.items(): PyLevel2Type[value] = key del key, value # Mapping from Python to Numeric types Py2NumType = { types.IntType: Int32, types.LongType: Int32, types.FloatType: Float64, types.ComplexType: Complex128 } def array2list(arr): #XXX Probably should make this a tolist method of NumArray (and write in C) retlist = [] if len(arr._shape) <= 1: for item in arr: retlist.append(item) else: for subarray in arr: retlist.append(array2list(subarray)) return retlist # array factory functions def _flattenseq(seq): """Turn a nested set of sequences into a list of non sequence values. Returns list [seq] if argument is not a sequence. """ #XXX Should be written in C for efficiency #XXX Use Tkinter._flatten as model -- it is vastly faster try: stypes = map(type, seq) except TypeError: return [seq] if types.ListType not in stypes and types.TupleType not in stypes: return seq retlist = [] tt = (types.TupleType, types.ListType) for i in xrange(len(seq)): if stypes[i] in tt: retlist.extend(_flattenseq(seq[i])) else: retlist.append(seq[i]) return retlist def _maxtype(args): """Find the maximum scalar numeric type in the arguments. An exception is raised if the types are not python numeric types. """ if not len(args): return None tlevel = map(PyNUMERIC_TYPES.get, map(type, args)) if None in tlevel: i = tlevel.index(None) raise TypeError("Expecting a python numeric type, got a %s" % _type(args[i]).__name__) return PyLevel2Type[max(tlevel)] def _storePyValueInBuffer(buffer, Ctype, index, value): """Store a python value in a buffer, index is in element units, not bytes""" # Do not use for complex scalars! Ctype._conv.fromPyValue(value, buffer._data, index*Ctype.bytes, Ctype.bytes, 0) def _storePyValueListInBuffer(buffer, Ctype, valuelist): # Do not use for complex values! for i in xrange(len(valuelist)): _storePyValueInBuffer(buffer, Ctype, i, valuelist[i]) def _fillarray(size, start, delta, type=None): ptype = _maxtype((start, delta)) if PyINT_TYPES.has_key(ptype): ptype = Int32 # this should change to Int64 when available elif PyREAL_TYPES.has_key(ptype): ptype = Float64 else: ptype = Complex128 if type: outtype = _getType(type) if _isComplexType(ptype) and not _isComplexType(outtype): raise TypeError("outtype must be a complex type") else: outtype = ptype if _isComplexType(outtype): # Not memory efficient at the moment real = _fillarray(size, complex(start).real, complex(delta).real, type = _realtype(ptype)) imag = _fillarray(size, complex(start).imag, complex(delta).imag, type = _realtype(ptype)) outarr = ComplexArray((size,), outtype) outarr.getreal()[:] = real outarr.getimag()[:] = imag else: # save parameters in a buffer parbuffer = ufunc._bufferPool.getBuffer() _storePyValueListInBuffer(parbuffer, ptype, [start, delta]) cfunction = _numarray.functionDict['fillarray'+ptype.name] outarr = NumArray((size,), outtype) if ptype == outtype: # no conversion necessary, simple case _numarray.CheckFPErrors() cfunction(size, 1, 1, ((outarr._data, 0), (parbuffer._data, 0))) errorstatus = _numarray.CheckFPErrors() if errorstatus: handleError(errorstatus, " in fillarray") else: # use buffer loop convbuffer = ufunc._bufferPool.getBuffer() convfunction = ptype._conv.astype[outtype.name] bsize = len(convbuffer._data)/ptype.bytes iters, lastbsize = divmod(size, bsize) _numarray.CheckFPErrors() outoff = 0 for i in xrange(iters + (lastbsize>0)): if i == iters: bsize = lastbsize cfunction(bsize, 1, 1, ((convbuffer._data, 0), (parbuffer._data, 0))) convfunction(bsize, 1, 1, ((convbuffer._data, 0), (outarr._data, outoff))) outoff += bsize*outtype.bytes start += delta * bsize _storePyValueListInBuffer(parbuffer, ptype, [start, delta]) errorstatus = _numarray.CheckFPErrors() if errorstatus: handleError(errorstatus, " in fillarray") return outarr def zeros(shape, type=None): shape = ndarray.getShape(shape) retarr = _fillarray(ndarray.product(shape), 0, 0, type) retarr.setshape(shape) return retarr def ones(shape, type=None): shape = ndarray.getShape(shape) retarr = _fillarray(ndarray.product(shape), 1, 0, type) retarr.setshape(shape) return retarr def arange(a1, a2=None, stride=1, type=None): if a2 == None: start = 0 + (0*a1) # to make it same type as stop stop = a1 else: start = a1 +(0*a2) stop = a2 delta = (stop-start)/stride ## xxx int divide issue if _type(delta) == types.ComplexType: # xxx What make sense here? size = int(math.ceil(delta.real)) else: size = int(math.ceil((stop-start)/float(stride))) return _fillarray(size, start, stride, type) arrayrange = arange # alias arange as arrayrange. def identity(n, type=None): a = zeros(shape=(n,n), type=type) i = arange(n) a.put(i, i, 1) return a def dot(a,b): """dot(a,b) matrix-multiplies a by b. """ return innerproduct(a, swapaxes(inputarray(b), -1, -2)) matrixmultiply = dot # Deprecated in Numeric def outerproduct(a,b): """outerproduct(a,b) computes the NxM outerproduct of N vector 'a' and M vector 'b', where result[i,j] = a[i]*b[j]. """ a=reshape(inputarray(a), (-1,1)) # ravel a into an Nx1 b=reshape(inputarray(b), (1,-1)) # ravel b into a 1xM return matrixmultiply(a,b) # return NxM result def sum(a, axis=0): """sum(a,axis=0) is a synonym for add.reduce(a, axis). """ a = inputarray(a) if axis == 0: return ufunc.add.reduce(a) else: return swapaxes(ufunc.add.reduce(swapaxes(a, 0, axis)), 0, axis) def cumsum(a, axis=0): """cumsum(a,axis=0) is a synonym for add.accumulate(a, axis). """ a = inputarray(a) if axis == 0: return ufunc.add.accumulate(a) else: return swapaxes(ufunc.add.accumulate(swapaxes(a, 0, axis)), 0, axis) def product(a, axis=0): """product(a,axis=0) is a synonym for multiply.reduce(a, axis). """ a = inputarray(a) if axis == 0: return ufunc.multiply.reduce(a) else: return swapaxes(ufunc.multiply.reduce(swapaxes(a, 0, axis)), 0, axis) def cumproduct(a, axis=0): """cumproduct(a,axis=0) is a synonym for multiply.accumulate(a,axis). """ a = inputarray(a) if axis == 0: return ufunc.multiply.accumulate(a) else: return swapaxes(ufunc.multiply.accumulate(swapaxes(a, 0, axis)), 0, axis) def alltrue(a, axis=0): """alltrue(a,axis=0) is a synonym for logical_and.reduce(a,axis). """ a = inputarray(a) if axis == 0: return ufunc.logical_and.reduce(a) else: return swapaxes(ufunc.logical_and(swapaxes(a, 0, axis)), 0, axis) def sometrue(a, axis=0): """sometrue(a,axis=0) is a synonym for logical_or.reduce(a,axis). """ a = inputarray(a) if axis == 0: return ufunc.logical_or.reduce(a) else: return swapaxes(ufunc.logical_or(swapaxes(a, 0, axis)), 0, axis) def _frontseqshape(seq): """Find the length of all the first elements, return as a list""" if not len(seq): return (0,) try: shape = [] while 1: shape.append(len(seq)) seq = seq[0] except TypeError: return shape def _setarray(arr, seq, indices): """Populate array with sequence values starting at specified partial indices""" try: depth = len(indices) if len(arr._shape) == depth+1: # Must be at lowest level for i in xrange(arr._shape[depth]): arr[tuple(indices+[i])] = seq[i] else: for i in xrange(arr._shape[depth]): _setarray(arr, seq[i], indices+[i]) except IndexError: raise IndexError("Supplied sequence not shape consistent.") def fromlist(seq, type=None, shape=None): """fromlist(seq, type=None, shape=None) creates a NumArray from the sequence 'seq' which must be a list or tuple of python numeric types. If type is specified, it is as the type of the resulting NumArray. If shape is specified, it becomes the shape of the result and must have the same number of elements as seq. """ highest_type = _maxtype(_flattenseq(seq)) tshape = _frontseqshape(seq) if shape is not None and ndarray.product(shape) != ndarray.product(tshape): raise ValueError("shape incompatible with sequence") ndim = len(tshape) if ndim <= 0: raise TypeError("Argument must be a sequence") if type is None: type = Py2NumType.get(highest_type) if type is None: raise TypeError("Cannot create array of type %s" % highest_type.__name__) tshape = tuple(tshape) if _isComplexType(type): arr = ComplexArray(tshape, type) else: arr = NumArray(tuple(tshape), type) _setarray(arr, seq, []) if shape is not None: arr.setshape(shape) return arr def array(buffer=None, shape=None, type=None): """array(buffer=None, shape=None, type=None) constructs a NumArray by calling NumArray, one of its factory functions (fromstring, fromfile, fromlist), or by making a copy. """ if isinstance(shape, types.IntType): shape = (shape,) if buffer is None and (shape is None or type is None): raise ValueError("Must define shape and type if buffer is None") elif buffer is None or ndarray.SuitableBuffer(buffer): return NumArray(buffer=buffer, shape=shape, type=type) elif isinstance(buffer, types.StringType): return fromstring(buffer, shape=shape, type=type) elif (isinstance(buffer, types.ListType) or isinstance(buffer, types.TupleType)): return fromlist(buffer, type, shape) elif isinstance(buffer, NDArray): a = buffer.astype(type) # make a re-typed copy if shape is not None: a.setshape(shape) return a elif isinstance(buffer, types.FileType): return fromfile(buffer, type=type, shape=shape) else: raise ValueError("Unknown input type") def inputarray(seq, type=None): """inputarray(seq, type=None) converts scalars, lists and tuples to arrays if possible, passes NumArrays thru making copies only to convert types, and raises a TypeError for everything else. """ if __builtins__["type"](seq) in [types.IntType, types.FloatType, types.LongType]: return fromlist([seq], type=type) elif isinstance(seq, types.ListType) or isinstance(seq, types.TupleType): return fromlist(seq, type=type) elif isinstance(seq, NDArray): if type is None or seq._type == type or seq._type.name == type: return seq else: return seq.astype(type) else: raise TypeError("seq is not a scalar, list, tuple, or NumArray") def _getType(type): """Return the numeric type object for type type may be the name of a type object or the actual object """ if isinstance(type, NumericType): return type try: return typeDict[type] except KeyError: raise TypeError("Not a numeric type") def fromstring(datastring, type, shape=None): """Create an array from binary data contained in a string (by copying)""" type = _getType(type) if not shape: size, rem = divmod(len(datastring), type.bytes) if rem: raise ValueError("Type size inconsistent with string length") else: shape = (size,) # default to a 1-d array elif _type(shape) is types.IntType: shape = (shape,) if len(datastring) != (ndarray.product(shape)*type.bytes): raise ValueError("Specified shape and type inconsistent with string length") arr = NumArray(shape, type) strbuff = buffer(datastring) nelements = arr.nelements() # Currently uses only the byte-by-byte copy, should be optimized to use # larger element copies when possible. cfunc = _bytes.functionDict["copyNbytes"] cfunc((nelements,), strbuff, 0, (type.bytes,), arr._data, 0, (type.bytes,), type.bytes) if arr._type == Bool: arr = ufunc.not_equal(arr, 0) return arr def fromfile(file, type, shape=None): """Create an array from binary file data If file is a string then that file is opened, else it is assumed to be a file object. No options at the moment, all file positioning must be done prior to this function call with a file object """ type = _getType(type) name = 0 if _type(file) == _type(""): name = 1 file = open(file, 'rb') size = os.path.getsize(file.name) - file.tell() if not shape: nelements, rem = divmod(size, type.bytes) if rem: raise ValueError( "Type size inconsistent with shape or remaining bytes in file") shape = (nelements,) elif _type(shape) is types.IntType: shape=(shape,) nbytes = ndarray.product(shape)*type.bytes if nbytes > size: raise ValueError( "Not enough bytes left in file for specified shape and type") # create the array arr = NumArray(shape, type) # Use of undocumented file method! XXX nbytesread = memory.file_readinto(file, arr._data) if nbytesread != nbytes: raise IOError("Didn't read as many bytes as expected") if name: file.close() if arr._type == Bool: arr = ufunc.not_equal(arr, 0) return arr class NumArray(NDArray): """Fundamental Numeric Array Class""" def __init__(self, shape, type, buffer=None, byteoffset=0, bytestride=None, byteswap=0, aligned=1): # need to add a whole bunch of checking here (or then again, # maybe not.) type = _getType(type) itemsize = type.bytes NDArray.__init__(self, shape, itemsize, buffer=buffer, byteoffset=byteoffset, bytestride=bytestride, aligned=aligned) self._type = type self._byteswap = byteswap and not (itemsize == 1) # This attribute is used by subclasses to indicate to ufunc that # they should be called to perform the ufunc (e.g., complex numbers). self._myufuncs = None def astype(self, type=None): """Return a copy of the array converted to the given type""" if type is None: type = self._type type = _getType(type) if type == self._type: # always return a copy even if type is same return self.copy() if type._conv: retarr = NumArray(self._shape, type) if (self._contiguous and self._aligned and not self._byteswap): _numarray.CheckFPErrors() cfunc = self._type._conv.astype[type.name] cfunc(self.nelements(), 1, 1, ((self._data, 0), (retarr._data, 0))) errorstatus = _numarray.CheckFPErrors() if errorstatus: handleError(errorstatus, " during type conversion") else: ufunc._copyFromAndConvert(self, retarr) elif type.fromtype: retarr = type.fromtype(self) else: raise TypeError("Don't know how to convert from %s to %s" % (self._type.name, type.name)) return retarr def new(self, type, byteswap=0): """Return a new array of given type with same shape as this array Note this only creates the array; it does not copy the data. """ type = _getType(type) arr = self.view() arr._itemsize = type.bytes arr._type = type arr._byteswap = byteswap and not (arr._itemsize == 1) arr._byteoffset = 0 arr._bytestride = arr._itemsize arr._contiguous = 1 arr._aligned = 1 arr._data = memory.memory_buffer(arr._itemsize * arr.nelements()) arr._strides = arr._stridesFromShape() return arr def type(self): """Return the type object for the array""" return self._type def tolist(self): """Convert array to list(s) of Python scalars""" return array2list(self) def isbyteswapped(self): """True if data in buffer is byteswapped relative to native form""" return self._byteswap def byteswap(self): """Byteswap the data in place""" if self._itemsize != 1: fname = "byteswap"+str(self._itemsize)+"bytes" cfunc = _bytes.functionDict[fname] cfunc(self._shape, self._data, self._byteoffset, self._strides, self._data, self._byteoffset, self._strides) def _getitem(self, key): """Returns value for single element in array using byte offset Does no validity checking on key (assumes it is valid). """ byteoffset = self._getByteOffset(key) return self._type._conv.asPyValue(self._data, byteoffset, self._itemsize, self._byteswap) def _setitem(self, key, value): """Sets a single value in an array using byte offset Does no validity checking on key (assumes it is valid). """ byteoffset = self._getByteOffset(key) self._type._conv.fromPyValue(value, self._data, byteoffset, self._itemsize, self._byteswap) def _copyFrom(self, arr, subShape=None, offset=None): """Copy elements from another array. This overrides the generic NDArray version """ if subShape == None: subShape = self._shape if offset == None: offset = self._byteoffset # Test for simple case first if isinstance(arr, NumArray): if (arr._type == self._type and arr._shape == subShape and arr._byteswap == self._byteswap and ndarray.product(arr._strides) != 0 and arr._aligned and self._aligned): cfunc = _bytes.functionDict['copy'+`self._itemsize`+'bytes'] cfunc(subShape, arr._data, arr._byteoffset, arr._strides, self._data, offset, self._strides) return elif PyNUMERIC_TYPES.has_key(_type(arr)): # Convert scalar to a one element array for broadcasting arr = array([arr]) else: raise TypeError('argument is not array or number') barr = self._broadcast(arr, raiseOnFail=1) ufunc._copyFromAndConvert(barr, self) def copy(self): c = ndarray.NDArray.copy(self) if self._byteswap: c.byteswap() c._byteswap = 0 return c def __str__(self): return arrayprint.array2string(self, MAX_LINE_WIDTH, PRECISION, SUPPRESS_SMALL, ' ', 0) def __repr__(self): return arrayprint.array2string(self, MAX_LINE_WIDTH, PRECISION, SUPPRESS_SMALL, ', ', 1) def __neg__(self): return ufunc.minus(self) def __add__(self, operand): return ufunc.add(self, operand) def __radd__(self, operand): return ufunc.add(operand, self) def __sub__(self, operand): return ufunc.subtract(self, operand) def __rsub__(self, operand): return ufunc.subtract(operand, self) def __mul__(self, operand): return ufunc.multiply(self, operand) def __rmul__(self, operand): return ufunc.multiply(operand, self) def __div__(self, operand): return ufunc.divide(self, operand) def __rdiv__(self, operand): return ufunc.divide(operand, self) def __mod__(self, operand): return ufunc.remainder(self, operand) def __rmod__(self, operand): return ufunc.remainder(operand, self) def __pow__(self, operand): return ufunc.power(self, operand) def __rpow__(self, operand): return ufunc.power(operand, self) def __and__(self, operand): return ufunc.bitwise_and(self, operand) def __rand__(self, operand): return ufunc.bitwise_and(operand, self) def __or__(self, operand): return ufunc.bitwise_or(self, operand) def __ror__(self, operand): return ufunc.bitwise_or(operand, self) def __xor__(self, operand): return ufunc.bitwise_xor(self, operand) def __rxor__(self, operand): return ufunc.bitwise_xor(operand, self) def __invert__(self): return ufunc.bitwise_not(self) def __rshift__(self, operand): return ufunc.rshift(self, operand) def __rrshift__(self, operand): return ufunc.rshift(operand, self) def __lshift__(self, operand): return ufunc.lshift(self, operand) def __rlshift__(self, operand): return ufunc.lshift(operand, self) def __pos__(self): return self # augmented assignment operators def __iadd__(self, operand): return ufunc.add(self, operand, self) def __isub__(self, operand): return ufunc.subtract(self, operand, self) def __imul__(self, operand): return ufunc.multiply(self, operand, self) def __idiv__(self, operand): return ufunc.divide(self, operand, self) def __imod__(self, operand): return ufunc.remainder(self, operand, self) def __ipow__(self, operand): return ufunc.power(self, operand, self) def __iand__(self, operand): return ufunc.bitwise_and(self, operand, self) def __ior__(self, operand): return ufunc.bitwise_or(self, operand, self) def __ixor__(self, operand): return ufunc.bitwise_xor(self, operand, self) def __irshift(self, operand): return ufunc.rshift(self, operand, self) def __ilshift(self, operand): return ufunc.lshift(self, operand, self) # rich comparisons (only works in Python 2.1 and later) def __lt__(self, operand): return ufunc.less(self, operand) def __gt__(self, operand): return ufunc.greater(self, operand) def __le__(self, operand): return ufunc.less_equal(self, operand) def __ge__(self, operand): return ufunc.greater_equal(self, operand) def __eq__(self, operand): if operand is None: return 0 else: return ufunc.equal(self, operand) def __ne__(self, operand): if operand is None: return 1 else: return ufunc.not_equal(self, operand) def __nonzero__(self): b = self.view() b.ravel() return bitwise_or.reduce(b!=0) def sort(self, axis=-1): if axis==-1: ufunc._sortN(self) else: self.swapaxes(axis,-1) ufunc._sortN(self) self.swapaxes(axis,-1) def argsort(self, axis=-1): a = self.copy() w = arange(a.nelements()) w = reshape(w, a._shape) if axis==-1: ufunc._argsortN(a,w) else: a.swapaxes(axis,-1) w.swapaxes(axis,-1) ufunc._argsortN(a,w) w.swapaxes(axis,-1) return w class ComplexArray(NumArray): """Complex Arrays Although it inherits from Numarray, this class inherits virtually none of its methods! The inheritance is solely so that complex arrays can be tested as an instance of NumArrays and thus be properly considered as such. """ def __init__(self, shape, type, buffer=None, byteoffset=0, bytestride=None, byteswap=0, aligned=1, real=None, imag=None): type = _getType(type) if not _isComplexType(type): raise TypeError("Type must be a Complex number type") itemsize = type.bytes NDArray.__init__(self, shape, itemsize, buffer=buffer, byteoffset=byteoffset, bytestride=bytestride, aligned=aligned) self._type = type if not bytestride: rbytestride = itemsize else: rbytestride = bytestride self._byteswap = byteswap if real is not None: self.getreal()[:] = real if imag is not None: self.getimag()[:] = imag self._myufuncs = 1 def getreal(self): t = _realtype(self._type) arr = NumArray(self._shape, t, buffer=self._data, byteoffset=self._byteoffset, bytestride=self._bytestride, aligned=self._aligned) arr._contiguous = self._contiguous return arr def setreal(self, value): self.getreal()[:] = value def getimag(self): t = _realtype(self._type) arr = NumArray(self._shape, t, buffer=self._data, byteoffset=self._byteoffset+t.bytes, bytestride=self._bytestride, aligned=self._aligned) arr._contiguous = self._contiguous return arr def setimag(self, value): self.getimag()[:] = value imag = property(getimag, setimag, doc="imaginary component of complex array") real = property(getreal, setreal, doc="real component of complex array") setimaginary = setimag getimaginary = getimag imaginary = imag def _getitem(self, key): """Returns value for single element in array using byte offset Does no validity checking on byteoffset (assumes it is valid). """ return complex(self.getreal()._getitem(key), self.getimag()._getitem(key)) def _setitem(self, key, value): """Sets a single value in an array using byte offset Does no validity checking on byteoffset (assumes it is valid). """ cvalue = complex(value) self.getreal()._setitem(key, cvalue.real) self.getimag()._setitem(key, cvalue.imag) def astype(self, type = None): """Return a copy of the array converted to the given type""" if type is None: type = self._type type = _getType(type) if type == self._type: # always return a copy even if type is same return self.copy() if not _isComplexType(type): raise TypeError("Can only convert complex arrays to another complex array") realtype = _realtype(type) retarr = ComplexArray(self._shape, type) retarr.getreal()[:] = self.getreal().astype(realtype) retarr.getimag()[:] = self.getimag().astype(realtype) return retarr def byteswap(self): """Byteswap the data in place""" self.getreal().byteswap() self.getimag().byteswap() def _copyFrom(self, arr, subShape=None, offset=None): """Copy elements from another array. This overrides the NumArray version """ if isinstance(arr, ComplexArray): real, imag = arr.getreal(), arr.getimag() elif _type(arr)==types.ComplexType: real, imag = arr.real, arr.imag else: real, imag = arr, 0. self.getreal()._copyFrom(real, subShape, offset) self.getimag()._copyFrom(imag, subShape, offset) def _ufbinarysetup(self, arr, outarr): if not isinstance(arr, NumArray): arr = array([arr]) # Figure out output shape. broadcast, dummy = self._dualbroadcast(arr) if broadcast is None: raise ValueError("Operands have incompatible shapes") if outarr is None: # Determine type of output. if isinstance(arr, NumArray) and arr._type == Float64: outtype = Complex128 else: outtype = self._type outarr = ComplexArray(broadcast._shape, outtype) else: ComplexArray._ufoutcheck(broadcast, outarr) # Split second operand into appropriate real and imag components. if isinstance(arr, ComplexArray): real, imag = arr.getreal(), arr.getimag() elif _type(arr) == types.ComplexType: real, imag = arr.real, arr.imag else: real, imag = arr, 0. return real, imag, outarr def _ufunarysetup(self, outarr): if outarr is None: outarr = ComplexArray(self._shape, self._type) else: ComplexArray._ufoutcheck(self, outarr) return outarr def _ufunary_real_setup(self, outarr): if outarr is None: outarr = NumArray(self._shape, self.getreal()._type) else: if self._shape != outarr._shape: raise ValueError("Supplied output array has an incompatible shape") return outarr def _ufoutcheck(inarr, outarr): if not isinstance(outarr, ComplexArray): raise TypeError("Supplied output array must be of complex type") if inarr._shape != outarr._shape: raise ValueError("Supplied output array has an incompatible shape") def conjugate(self): ufunc.minus(self.getimag(), self.getimag()) def _abs_sq(self, outarr=None): outarr = self._ufunary_real_setup(outarr) sr, si = self.getreal(), self.getimag() ufunc.multiply(sr, sr, outarr) temp = si * si return ufunc.add(outarr, temp, outarr) def _abs(self, outarr=None): return ufunc.sqrt(self._abs_sq(outarr)) def _minus(self, outarr): outarr = self._ufunarysetup(outarr) ufunc.minus(self.getreal(), outarr.getreal()) ufunc.minus(self.getimag(), outarr.getimag()) return outarr def _add(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) ufunc.add(self.getreal(), real, outarr.getreal()) ufunc.add(self.getimag(), imag, outarr.getimag()) return outarr def _right_add(self, arr, outarr): return self._add(arr, outarr) def _subtract(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) ufunc.subtract(self.getreal(), real, outarr.getreal()) ufunc.subtract(self.getimag(), imag, outarr.getimag()) return outarr def _right_subtract(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) ufunc.subtract(real, self.getreal(), outarr.getreal()) ufunc.subtract(real, self.getimag(), outarr.getimag()) return outarr def _multiply(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() temp1 = sr*real temp2 = si*real ufunc.multiply(si, imag, tr) ufunc.subtract(temp1, tr, tr) ufunc.multiply(sr, imag, temp1) ufunc.add(temp1, temp2, ti) return outarr def _right_multiply(self, arr, outarr): return self._multiply(arr, outarr) def _divide(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() temp1 = sr*real temp2 = si*real if not isinstance(arr, NumArray): d = array([arr+0j])._abs_sq() else: d= arr.astype(self._type)._abs_sq() ufunc.multiply(si, imag, tr) ufunc.add(tr, temp1, tr) ufunc.multiply(sr, imag, temp1) ufunc.subtract(temp2, temp1, ti) ufunc.divide(tr, d, tr) ufunc.divide(ti, d, ti) return outarr def _right_divide(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() temp1 = sr*real temp2 = si*real d = self._abs_sq() ufunc.multiply(si, imag, tr) ufunc.add(tr, temp1, tr) ufunc.multiply(sr, imag, temp1) ufunc.subtract(temp1, temp2, ti) ufunc.divide(tr, d, tr) ufunc.divide(ti, d, ti) return outarr def _exp(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() realexp = ufunc.exp(sr) ufunc.cos(si, tr) ufunc.sin(si, ti) ufunc.multiply(tr, realexp, tr) ufunc.multiply(ti, realexp, ti) return outarr def _log(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() radius = self._abs(None) ufunc.log(radius, tr) ufunc.arctan2(si, sr, ti) return outarr def _log10(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() radius = self._abs(None) ufunc.log10(radius, tr) ufunc.divide(ufunc.arctan2(si, sr, ti), ufunc.log(10.), ti) return outarr def _power(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) ufunc.log(self, outarr) ufunc.multiply(arr, outarr, outarr) ufunc.exp(outarr, outarr) return outarr def _right_power(self, arr, outarr): real, imag, outarr = self._ufbinarysetup(arr, outarr) outarr[:] = 0.+0.j if not isinstance(arr, NumArray): arr = array([arr]) ufunc.add(arr, outarr, outarr) ufunc.log(outarr, outarr) ufunc.multiply(self, outarr, outarr) ufunc.exp(outarr, outarr) return outarr def _sqrt(self, outarr): outarr = self._ufunarysetup(outarr) return self._power(0.5, outarr) def _sin(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() temp = ufunc.sin(sr) ufunc.multiply(ufunc.cosh(si, tr), temp, tr) ufunc.cos(sr, temp) ufunc.multiply(ufunc.sinh(si, ti), temp, ti) return outarr def _cos(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() temp = ufunc.cos(sr) ufunc.multiply(ufunc.cosh(si, tr), temp, tr) ufunc.sin(sr, temp) ufunc.multiply(ufunc.sinh(si, ti), temp, ti) return outarr def _tan(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() sinr = ufunc.sin(sr) cosr = ufunc.cos(sr) sinhi, coshi = tr, ti # alias ufunc.sinh(si, sinhi) ufunc.cosh(si, coshi) temp1 = sinr*sinr ufunc.multiply(temp1, sinhi, temp1) ufunc.multiply(temp1, sinhi, temp1) temp2 = cosr*cosr ufunc.multiply(temp2, coshi, temp2) ufunc.multiply(temp2, coshi, temp2) ufunc.add(temp1, temp2, temp1) ufunc.multiply(sinhi, coshi, ti) ufunc.multiply(sinr, cosr, tr) ufunc.divide(tr, temp1, tr) ufunc.divide(ti, temp1, ti) return outarr def _sinh(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() temp = cos(si) ufunc.sinh(sr, tr) ufunc.multiply(tr, temp, tr) temp = sin(si) ufunc.cosh(sr, ti) ufunc.multiply(ti, temp, ti) return outarr def _cosh(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() temp = cos(si) ufunc.cosh(sr, tr) ufunc.multiply(tr, temp, tr) temp = sin(si) ufunc.sinh(sr, ti) ufunc.multiply(ti, temp, ti) return outarr def _tanh(self, outarr): outarr = self._ufunarysetup(outarr) sr, si = self.getreal(), self.getimag() tr, ti = outarr.getreal(), outarr.getimag() sini = ufunc.sin(si) cosi = ufunc.cos(si) sinhr, coshr = tr, ti # alias ufunc.sinh(sr, sinhr) ufunc.cosh(sr, coshr) temp1 = sini*sini ufunc.multiply(temp1, sinhr, temp1) ufunc.multiply(temp1, sinhr, temp1) temp2 = cosi*cosi ufunc.multiply(temp2, coshr, temp2) ufunc.multiply(temp2, coshr, temp2) ufunc.add(temp1, temp2, temp1) ufunc.multiply(sinhr, coshr, tr) ufunc.multiply(sini, cosi, ti) ufunc.divide(tr, temp1, tr) ufunc.divide(ti, temp1, ti) return outarr def dummy(): """ remainder" arcsin", 1 arccos", 1 arctan", 1 arctan2 sqrt", 1 equal", not_equal" greater", greater_equal less", less_equal", logical_and", logical_or", logical_xor", logical_not", bitwise_and", bitwise_or", bitwise_xor", bitwise_not", rshift", 2, lshift", 2, floor", 1, 1, floor", 1, 1, ceil", 1, 1, ceil", 1, 1, maximum", 2, minimum", 2, """ pass def Complex64_fromtype(arr): """Used for converting other types to Complex64. This is used to set an fromtype attribute of the ComplexType objects""" rarr = arr.astype(Float32) retarr = ComplexArray(arr._shape, Complex64) retarr.getreal()[:] = rarr retarr.getimag()[:] = 0. return retarr def Complex128_fromtype(arr): """Used for converting other types to Complex128. This is used to set an fromtype attribute of the ComplexType objects""" rarr = arr.astype(Float64) retarr = ComplexArray(arr._shape, Complex128) retarr.getreal()[:] = rarr retarr.getimag()[:] = 0. return retarr # Check whether byte order is big endian or little endian. from sys import byteorder isBigEndian = (byteorder == "big") del byteorder # Add fromtype function to Complex types Complex64.fromtype = Complex64_fromtype Complex128.fromtype = Complex128_fromtype # Return type of complex type components def _isComplexType(type): return type in [Complex64, Complex128] def _realtype(complextype): if complextype == Complex64: return Float32 else: return Float64