"""ndarray: Base multidimensional array class This is the base multidimensional array class which implements all structural array operations but treats the array contents as opaque objects $Id: ndarray.py,v 1.8 2002/03/11 15:34:56 jaytmiller Exp $ """ import types as _types import math as _math import operator import numerictypes as _numerictypes import sys import _bytes import memory NewAxis = None _IOBLOCKSIZE = 1024**2 def product(shape): n = 1 for s in shape: n *= s return n def _isIntegerSequence(seq): """Determine if all the elements of a sequence are integers""" return _isTypeSequence(seq, [_types.IntType, _types.LongType]) def _isTypeSequence(seq, typelist): """Determine if all the elements of a sequence of types in typelist""" # This is about twice as fast as the naive looping algorithm titem = map(type, seq) sum = 0 for ttype in typelist: sum += titem.count(ttype) return (sum == len(titem)) def getShape(shape, *args): """Verifies that this is a legal shape specification and returns tuple Shape can be an integer or a sequence of integers. Also can pass several integer arguments. Raises an exception on problems. """ try: if type(shape) in [_types.IntType, _types.LongType]: shape = (shape,) + args else: if args: raise TypeError shape = tuple(shape) if _isIntegerSequence(shape): return shape except TypeError: pass raise TypeError("Shape must be sequence of integers") def _broadcast(arr, sshape, raiseOnFail=0): """Return broadcast view of arr, else return None.""" ashape = arr._shape # Just return arr if they have the same shape if sshape == ashape: return arr srank = len(sshape) arank = len(ashape) if arank > srank: return None if arank == srank: astrides = list(arr._strides) else: astrides = [0]*(srank-arank) + list(arr._strides) ashape = sshape[0:srank-arank] + ashape if ashape != sshape: for i in range(srank): if sshape[i] != ashape[i]: if ashape[i] == 1: astrides[i] = 0 else: if raiseOnFail: raise ValueError("Arrays have incompatible shapes") else: return None tarr = arr.view() tarr._shape = sshape tarr._contiguous = 0 tarr._strides = tuple(astrides) return tarr def info(arr): print "shape:", arr._shape print "strides:", arr._strides print "byteoffset:", arr._byteoffset print "bytestride:", arr._bytestride print "itemsize:", arr._itemsize print "aligned:", arr._aligned print "contiguous:", arr._contiguous print "data:", repr(arr._data) def SuitableBuffer(b): """SuitableBuffer(b) determines whether 'b' can be used as an NDArray buffer. """ return ((type(b) is _types.BufferType) or ((type(b) is _types.InstanceType) and ("__buffer__" in dir(b.__class__)) and ("resize" in dir(b.__class__)))) # The following adds support for properties (e.g. .shape and .flat) for recent # versions of python. With properties, numarray works more like Numeric, # so it is possible to say: a.shape = (10,10) instead of a.setshape((10,10)). # For python versions earlier than (2,2,0) it is necessary to use # the equivalent methods (e.g. getshape(), setshape(), and getflat()) if sys.version_info < (2,2,0): # Make the compile-time function "property" do nothing and return None # since properties are not supported in this python version. def property(fget, fset=None, fdel=None, doc=None): return None # Can't subclass "object" since it hasn't been invented yet. Still need # an EmptyClass for view(). class _NDArray: pass else: # sys.version_info >= (2,2,0) # Python now provides property as a builtin. # Subclass _NDArray from object so that properties can be supported. class _NDArray(object): pass class NDArray(_NDArray): """Multi-dimensional array abstract base class This class defines the structural operations common to arrays. Subclasses must provide the semantical interpretation of elements, including the __str__, __repr__, _getitem, and _setitem methods. Given an array index of arr[k,j,i] it is always true that the byte offset of the element in the array is computed thusly: with shape[0] --> dimension of current view strides[0] --> bytestride for k index dimension element_byte_offset = byteoffset + ( i*strides[2] + j*strides[1] + k*strides[0]) where 0 <= i < shape[2], 0 <= j < shape[1], 0 <= k < shape[0] For contiguous arrays strides[i] = shape[i+1]*strides[i+1] Summary of attribute meanings: _data buffer with data for the array _shape dimensions of the array _aligned True only if _itemsize = _bytestride. It doesn't matter if they happen to be aligned by accident, if there is any stuff between items, arrays are always treated as nonaligned. _contiguous True only if there are no gaps of items in arrays. An array can be contiguous but not aligned. For example a record array may have a bytestride of 5 but a numeric array contained within may be considered contiguous if no items are skipped (even though there are gaps between items). In general, contiguous is true if strides[i] = shape[i+1]*strides[i+1]. _byteoffset The byte offset of the first element from the beginning of the buffer _bytestride The separation between items in bytes. _itemsize The size of items in bytes """ def __init__(self, shape, itemsize, buffer=None, byteoffset=0, bytestride=None, aligned=1): self._itemsize = itemsize self._aligned = aligned or (self._itemsize == 1) self._byteoffset = byteoffset self._shape = getShape(shape) if bytestride is None: self._bytestride = itemsize else: # note that the `aligned' parameter may be overridden if other # parameters indicate array is not aligned if bytestride > itemsize: self._aligned = 0 or (self._itemsize == 1) elif bytestride < itemsize: raise ValueError("bytestride must be >= itemsize") self._bytestride = bytestride if buffer: if not SuitableBuffer(buffer): raise TypeError("buffer won't work as NDArray buffer") self._data = buffer else: size = self._bytestride * self.nelements() self._data = memory.memory_buffer(size) self._contiguous = 1 # Now ya see it... self._strides = self._stridesFromShape() self._contiguous = self.iscontiguous() # Now ya don't... def __len__(self): if len(self._shape): return self._shape[0] else: return 1 def itemsize(self): """Size (in bytes) of an array element""" return self._itemsize def isaligned(self): return self._aligned #XXX Note from pg: nelements() does not work for broadcast arrays #XXX Does that matter? def nelements(self): # Could cache this value and change it when shape changes. # Do that if we eventually allow direct assignment to shape # (which will require implementing __setitem__) return product(self._shape) def _stridesFromShape(self): """Compute the strides from shape for a contiguous array""" if not self._contiguous: raise ValueError("Cannot determine strides from shape for a non contiguous array.") ndim = len(self._shape) if ndim: strides = [self._bytestride]*ndim for i in xrange(ndim-2, -1, -1): strides[i] = strides[i+1] * self._shape[i+1] return tuple(strides) else: return (0,) # scalar def _getByteOffset(self, indices): """Compute the byte offset for the given indices (which may be only partial)""" offset = self._byteoffset for i in xrange(len(indices)): offset += indices[i]*self._strides[i] return offset # These are used by __getitem__ and __setitem__. They should be defined # by the subclass. #XXX RLW: the base class could return a rank-0 (rank-1?) array with #XXX a single element for these methods def _getitem(self, byteoffset): pass def _setitem(self, byteoffset, value): pass def __getitem__(self, key): return self._universalIndexing(key) def __setitem__(self, key, value): self._universalIndexing(key, value) def __delitem__(self, key): raise TypeError("Can't delete array elements.") def _universalIndexing(self, key, value=None): """Handles both getting (value == None) and setting (value != None)""" if value is not None and \ type(value) in [_types.ListType, _types.TupleType]: value = self.array(value) tkey = key if type(key) in [_types.SliceType, _types.EllipsisType, _types.IntType, _types.LongType, _types.ListType] or \ isinstance(key, _numarray.NumArray): # make it a tuple and fall through to the next block tkey = (tkey,) if type(tkey) == _types.TupleType: # most complex ones fall in this category tkey = list(tkey) if _isIntegerSequence(tkey): return self._simpleIndexing(tkey, value) # more complex if self._isSlice(tkey): # i.e., doesn't contain arrays of any kind return self._slicedIndexing(tkey, value) else: return self._arrayIndexing(tkey, value) else: raise IndexError("Illegal index") def _simpleIndexing(self, tkey, value): if tuple(tkey) == (0,) and self._strides == (0,): # scalar case return self._getitem(tkey) if len(tkey) > len(self._shape): raise IndexError("Too many indices") newshape = self._shape[len(tkey):] #XXX RLW: Should the next loop be moved into computeByteOffset? for i in range(len(tkey)): if tkey[i] < 0: tkey[i] += self._shape[i] if not (0 <= tkey[i] < self._shape[i]): raise IndexError("Index out of range") offset = self._getByteOffset(tkey) if len(tkey) == len(self._shape): # single values if value is None: return self._getitem(tkey) else: self._setitem(tkey, value) else: # subarrays retarr = self.view() retarr._shape = newshape retarr._strides = self._strides[len(tkey):] retarr._byteoffset = offset if value is None: return retarr else: retarr._copyFrom(value) def _slicedIndexing(self, tkey, value): # strip NewAxis values ipos = 0 newaxispos = [] while ipos < len(tkey): if tkey[ipos] is NewAxis: del tkey[ipos] newaxispos.append(ipos) else: ipos += 1 # handle ellipsis ellipsispos = None for i in xrange(len(tkey)): if type(tkey[i]) == _types.EllipsisType: # Only the first one means expand, all the rest are same as ':' if ellipsispos: tkey[i] = slice(None, None, None) else: ellipsispos = i if ellipsispos is not None: # insert extra dims for the first ellipsis idiff = len(self._shape) - len(tkey) + 1 if idiff >= 1: tkey[ellipsispos: ellipsispos+1] = [slice(None, None, None)]*idiff # gotta adjust new axis positions accordingly else: # ignore ellipsis after key "full" del tkey[ellipsispos:ellipsispos+1] for j in range(len(newaxispos)): if newaxispos[j] > ellipsispos: newaxispos[j] += idiff-1 if len(tkey) > len(self._shape): raise IndexError("too many indices") # Do our thing, checking indices and slices as we go retarr = self.view() retshape = [] retstrides = [] firstelement = [] for i in xrange(len(tkey)): shape = self._shape[i] index = tkey[i] if type(index) in [_types.IntType, _types.LongType]: # single valued index if index < 0: index += shape if 0 <= index < shape: firstelement.append(index) else: raise IndexError("Index out of range") elif type(index) == _types.SliceType: # slice case stride = index.step if stride is None: stride = 1 elif stride == 0: raise IndexError("slice step of zero not allowed") start = index.start if start is None: if stride>0: start = 0 else: start = shape-1 elif start < 0: start += shape if start<0: start = 0 elif stride>0: if start > shape: start = shape else: if start > shape-1: start = shape-1 stop = index.stop if stop is None: if stride>0: stop = shape else: stop = -1 elif stop < 0: stop += shape if stop<0: stop = 0 elif stop > shape: stop = shape firstelement.append(start) retstrides.append(self._strides[i]*stride) newdim = int(_math.ceil(float(stop-start)/stride)) if newdim<0: newdim = 0 retshape.append(newdim) else: raise IndexError("Illegal value for index") retarr._byteoffset = self._getByteOffset(firstelement) if newaxispos: # insert all the new axes newaxispos.reverse() for axispos in newaxispos: retshape.insert(axispos, 1) if axispos < len(retstrides): retstrides.insert(axispos, retstrides[axispos]) else: retstrides.insert(axispos, retarr._itemsize) retarr._shape = tuple(retshape + list(self._shape[len(firstelement):])) retarr._strides = tuple(retstrides + list(self._strides[len(firstelement):])) # determine if new array is contiguous if retarr._strides[-1:] != (retarr._itemsize,): retarr._contiguous = 0 else: for i in range(len(retarr._strides[:-1])): if retarr._strides[i] != \ retarr._strides[i+1]*retarr._shape[i+1]: retarr._contiguous = 0 break else: retarr._contiguous = 1 if value is None: if retarr._shape != (): return retarr else: retarr._shape = (1,) retarr._strides = (0,) return retarr[0] else: if retarr._shape == (): retarr._shape = (1,) retarr._strides = (0,) retarr._copyFrom(value) def _arrayIndexing(self, tkey, value): for item in tkey: if type(item) in [_types.SliceType, _types.EllipsisType]: raise IndexError("Cannot mix arrays and slices as indices") if value is None: return _take(self, tkey) else: return _put(self, tkey, value) def _isSlice(self, key): for item in key: if isinstance(item, _numarray.NumArray) or \ type(item) is _types.ListType: # array index return 0 if (item is not None) and \ type(item) not in [_types.IntType, _types.LongType, _types.SliceType, _types.EllipsisType]: print "Index: ", key, " contains illegal elements" raise IndexError("index is not of legal form") return 1 def _broadcast(self, arr, raiseOnFail=0): """Return broadcast view of arr, else return None.""" return _broadcast(arr, self._shape, raiseOnFail=raiseOnFail) def _dualbroadcast(self, arr, raiseOnFail=0): """Return broadcast views both self and arr, else return (None,None).""" sshape = self._shape ashape = arr._shape # Just return both if they have the same shape if sshape == ashape: return (self, arr) srank = len(sshape) arank = len(ashape) # do a special comparison of all dims with size>1 if srank > arank: newrank = srank sstrides = list(self._strides) ashape = sshape[:newrank-arank] + ashape astrides = [0]*(newrank-arank) + list(arr._strides) else: newrank = arank astrides = list(arr._strides) sshape = ashape[:newrank-srank] + sshape sstrides = [0]*(newrank-srank) + list(self._strides) newshape = list(sshape) for i in range(newrank): if sshape[i] != ashape[i]: if sshape[i] == 1: newshape[i] = ashape[i] sstrides[i] = 0 elif ashape[i] == 1: newshape[i] = sshape[i] astrides[i] = 0 else: if raiseOnFail: raise ValueError("Arrays have incompatible shapes"); else: return None, None newshape = tuple(newshape) tself, tarr = self, arr if self._shape != newshape: tself = self.view() tself._shape = newshape tself._contiguous = 0 tself._strides = tuple(sstrides) if arr._shape != newshape: tarr = arr.view() tarr._shape = newshape tarr._contiguous = 0 tarr._strides = tuple(astrides) return tself, tarr # THE FOLLOWING HAS NOT BEEN TESTED! def _copyFrom(self, arr, subShape=None, offset=None): """Copy elements from another array. This is the generic version. Subclasses (such as numarray) may override this method """ if subShape == None: subShape = self._shape if offset == None: offset = self._byteoffset # Arrays must be shape congruent and have the same itemsize. # xxx Don't handle broadcasting yet. if arr._shape != subShape: raise ValueError("Arrays have inconsistent sizes") if arr._itemsize != self._itemsize: raise ValueError("Arrays must have the same itemsize") cfunc = _bytes.functionDict['copyNbytes'] cfunc(subShape, arr._data, arr._byteoffset, arr._strides, self._data, offset, self._strides, self._itemsize) #XXX These slice methods are not needed after Python 2.1 if sys.version_info < (2,1,0): def __getslice__(self, i, j): return self.__getitem__((slice(i, j, None),)) def __setslice__(self, i, j, s): self.__setitem__((slice(i, j, None),), s) def __delslice__(self, i, j): raise TypeError("Can't delete array elements.") def setshape(self, shape, *args): """Change array shape in place. Call as setshape(i,j,k) or setshape((i,j,k)).""" if not self._contiguous: raise TypeError("Can't reshape non-contiguous arrays") shape = list(getShape(shape, *args)) # look for index = -1, which indicates an expandable dimension nelements = self.nelements() negcount = shape.count(-1) if negcount > 1: raise ValueError("no more than one dimension can have value -1") elif negcount == 1: tnelements = abs(product(shape)) shape[shape.index(-1)] = nelements/tnelements newnelements = product(shape) if newnelements == nelements: self._shape = tuple(shape) self._strides = self._stridesFromShape() else: raise ValueError("New shape is not consistent with the old shape") def getshape(self): return self._shape shape = property(getshape, setshape, doc="tuple of array dimensions") def getrank(self): """Returns the number of dimensions, ie 'rank' of an array.""" if self._shape != (0,): return len(self._shape) else: return 0 rank = property(getrank, doc="number of dimensions of an array.") def getflat(self): a = self.view() a.ravel() return a def setflat(self, flat): a = self.view() a.ravel() a[:] = flat flat = property(getflat, setflat, doc="access to array as 1D") def view(self): """Return a new array object, with the same reference to the data buffer""" # about 5x faster than call to copy.copy(self) arr = _NDArray() arr.__class__ = self.__class__ arr.__dict__.update(self.__dict__) return arr def copy(self): """Return a new array with the same shape and type, but a copy of the data""" arr = self.view() arr._data = memory.memory_buffer(arr._itemsize * arr.nelements()) arr._contiguous = 1 arr._byteoffset = 0 arr._bytestride = arr._itemsize arr._strides = arr._stridesFromShape() arr._aligned = 1 # now copy data, if possible using larger units fname = "copy"+str(self._itemsize)+"bytes" copyfunction = ((self._aligned and _bytes.functionDict.get(fname)) or _bytes.functionDict["copyNbytes"]) copyfunction(arr._shape, self._data, self._byteoffset, self._strides, arr._data, 0, arr._strides, arr._itemsize) return arr def tostring(self): """Return a string with a binary copy of the array Copies are always contiguous, but no conversions are implied """ return _bytes.copyToString(self._shape, self._data, self._byteoffset, self._strides, self._itemsize) def tofile(self, file): """Write the array to a file If file is a string, it attempts to open a file with that name, otherwise it assumes file is a file object. At the moment if special positioning is needed in the file one must do that with the file object beforehand. More options may be added to this method to allow positioning or appends. """ name = 0 if type(file) == type(""): name = 1 file = open(file, 'wb') niter = _IOBLOCKSIZE/self._itemsize indexlevel, blockingparameters = \ ufunc._getBlockingParameters(self._shape, niter) self._tofileByBlocks(file, [], indexlevel, blockingparameters) if name: file.close() def _tofileByBlocks(self, file, dims, indexlevel, blockingparameters): """Write the array to a file repeatedly in blocks This is done similarly to ufunc._doOverDimensions """ level = len(dims) if level == indexlevel: nregShapeIters, shape, leftover, leftoverShape, = blockingparameters currentIndex = 0 tshape = shape[:] for i in xrange(nregShapeIters + leftover): if i==nregShapeIters: tshape = leftoverShape tdims = dims + [currentIndex,] file.write(_bytes.copyToString( tshape, self._data, self._getByteOffset(tdims), self._strides, self._itemsize)) currentIndex += shape[0] else: # recurse for i in xrange(self._shape[level]): tdims = dims + [i] self._tofileByBlocks(file, tdims, indexlevel, blockingparameters) def iscontiguous(self): """iscontiguous returns 1 iff the array is contiguous.""" for i in range(len(self._shape)-1): if self._strides[i] != self._strides[i+1]*self._shape[i+1]: return 0 return ((self._strides[-1] == self._itemsize) and (0 not in self._strides)) def transpose(self, axes=None): """transpose(self, axes=None) re-shapes the array by permuting it's dimensions as specified by 'axes'. If 'axes' is none, transpose returns the array with it's dimensions reversed. """ slen = len(self._shape) if axes == None: axes = range(slen) axes.reverse() if len(axes) != slen: raise ValueError("Wrong number of axes in tranpose") tax = list(axes[:]) tax.sort() if tax != range(slen): raise ValueError("Duplicate or missing transpose axes") nshape, nstrides = [],[] for i in axes: nshape.append(self._shape[i]) nstrides.append(self._strides[i]) self._shape = tuple(nshape) self._strides = tuple(nstrides) self._contiguous = self.iscontiguous() def swapaxes(self, axis1, axis2): """swapaxes(self, axis1, axis2) interchanges axis1 and axis2. """ n = len(self._shape) if n <= 1: return # can't really swapaxes on a 0D or 1D array... if axis1 < 0: axis1 += n if axis2 < 0: axis2 += n if axis1 not in range(n) or axis2 not in range(n): raise ValueError("Specified dimension does not exist") if axis1 == axis2: return # Make sure that axes are strictly ordered so that the code below makes sense... if axis1 > axis2: axis1, axis2 = axis2, axis1 # Just swap the shape and stride elements for the swapped dimensions... self._shape = (self._shape[0:axis1] + (self._shape[axis2],) + self._shape[axis1+1:axis2] + (self._shape[axis1],) + self._shape[axis2+1:]) self._strides = (self._strides[0:axis1] + (self._strides[axis2],) + self._strides[axis1+1:axis2] + (self._strides[axis1],) + self._strides[axis2+1:]) self._contiguous = self.iscontiguous() def take(self, *indices, **keywords): return take(self, *indices, **keywords) def put(self, *indices, **keywords): return put(self, *indices, **keywords) def nonzero(self): return ufunc.nonzero(self) def resize(self, shape, *args): shape = getShape(shape, *args) nlen = product(shape) if nlen < 0: raise ValueError("Negative shape dims don't work with resize") olen = self.nelements() if isinstance(self._data, _types.BufferType): if self._contiguous: oself = self.view() else: oself = self.copy() self._data = memory.memory_buffer(nlen*self._itemsize) self._bytestride = self._itemsize self._contiguous = 1 self._byteoffset = 0 self._aligned = 1 blen = min(nlen, olen) self.ravel() oself.ravel() self[:blen] = oself[:blen] else: # Memmap self._data.resize(nlen*self._itemsize) self._shape = (nlen,) self._strides = self._stridesFromShape() if olen: # use any existing data as a pattern to be repeated if nlen > olen: offset = olen while offset+olen <= nlen: self[offset:offset+olen] = self[0:olen] offset += olen self[offset:nlen] = self[0:nlen-offset] else: # zero fill resized zero-length arrays self[:] = 0 self._shape = shape self._strides = self._stridesFromShape() def repeat(self, repeats, axis=0): return repeat(self, repeats, axis) def diagonal(self, k=0): """diagonal(self, k=0) returns the 'k'th diagonal of 'self' """ if k < 0: return transpose(self).diagonal(-k) maxi = min(self._shape[:2])-abs(k) ax = _numarray.arange(maxi) return _take(self, (ax, ax+k)) def trace(self, k=0): """trace(self, k=0) returns the sum of the elements of the 'k'th diagnoal of 'self' """ return ufunc.add.reduce(self.diagonal(k)) def ravel(self): """ravel(self) setshapes 'self' into an equivalent 1D array. """ self.setshape((self.nelements(),)) def array(self, *args, **keys): """array(...) calls the array(...) factory function defined in the source module where the class of 'self' was defined. """ module = sys.modules[self.__class__.__module__] return module.array(*args, **keys) def astype(self, type=None): #default so numarray.array works for NDArrays return self.copy() def reshape(arr, shape, *args): """Returns a reshaped *view* of array if possible, else a *copy*. Call either as reshape(i,j,k) or reshape((i,j,k)). """ v = _numarray.inputarray(arr) if v._contiguous: v = v.view() else: v = v.copy() v.setshape(shape, *args) return v def ravel(array): """Returns a *view* of array reshaped as 1D.""" return reshape(array, (array.nelements(),)) def resize(array, shape): """Returns a *copy* of array, replicated or truncated to match new shape.""" array = _numarray.array(array) array.resize(shape) # Assume array.resize() resizes in place. return array def fromstring(datastring): pass def transpose(array, axes=None): """Returns the transpose of a *view* of array""" v = _numarray.inputarray(array).view() v.transpose(axes) return v def diagonal(array, k=0): """Returns array[i,i] for all valid i.""" array = _numarray.inputarray(array) return array.diagonal(k) def trace(array, k=0): """Returns sum(diagonal(array, k)).""" array = _numarray.inputarray(array) return array.trace(k) def sort(array, axis=-1): """Returns a sorted *copy* of array.""" array = _numarray.array(array) array.sort(axis) return array def argsort(array, axis=-1): """Returns an array of indices which, if "taken" from 'array', would sort 'array'. """ array = _numarray.inputarray(array) return array.argsort(axis) def swapaxes(array, axis1, axis2): """Returns a *view* of array with axis1 and axis2 interchanged.""" v = _numarray.inputarray(array).view() v.swapaxes(axis1, axis2) return v def where(condition, x, y, out=None): """where(condition,x,y) returns an array shaped like 'condition' with elements selected from 'x' or 'y' by the 1 or 0 value of each condition element, respectively. """ return choose(ufunc.not_equal(condition, 0), (y,x), out) def clip(m, m_min, m_max): """clip(m, m_min, m_max) = every entry in m that is less than m_min is replaced by m_min, and every entry greater than m_max is replaced by m_max. """ selector = ufunc.less(m, m_min)+2*ufunc.greater(m, m_max) return choose(selector, (m, m_min, m_max)) def _concat(arrs): """_concat handles the simplest case of concatenating arrays along the zero-th axis. """ combinedLength = reduce(operator.add, [ a._shape[0] for a in arrs ]) rShape = arrs[0]._shape[1:] destShape = (combinedLength,) + tuple(rShape) convType = ufunc._maxPopType(arrs) dest = arrs[0].__class__(shape=destShape, type=convType) ix = 0 for a in arrs: if a._shape[1:] != rShape: raise ValueError("_concat array shapes must match except 1st dimension") if (convType is not None) and a._type != convType: a = a.astype(convType) dest._copyFrom( a, a._shape, dest._strides[0]*ix) ix += a._shape[0] return dest def concatenate(arrs, axis=0): """concatenate(arrs, axis=0) joins the sequence of arrays 'arrs' in a into a single array along the specified 'axis'. """ arrs = map(_numarray.inputarray, arrs) if axis == 0: return _concat(arrs) else: return swapaxes(_concat([swapaxes(m,axis,0) for m in arrs]), axis, 0) # import these last so module dependencies don't cause problems import numarray as _numarray import _ufunc, ufunc choose = ufunc._ChooseUFunc() choose.__doc__ ="""choose(selector, population) returns a new array shaped like 'selector' with elements chosen from members of sequence 'population' by the values of selector. The shape of each member of 'population' must be broadcastable to the shape of 'selector'. The value of each member of 'selector' must satisfy: 0 <= value < len(population).""" _take = ufunc._TakeUFunc("take") _take.__doc__ = ufunc._TakeUFunc.__doc__ def take(array, *indices, **keywords): """take(array, indices..., clipmode=(CLIP|WRAP), axis=0) picks elements of 'array' specified by index arrays 'indices'. parameters which must be specified by keyword: clipmode=CLIP if clipmode is CLIP, out of range indices are clipped at [0, shape[i]). if clipmode is WRAP, out of range indices are wrapped around from 0 to shape[i] (<0) or from shape[i] to 0 (>= shape[i]) axis=0 selects the axis along which the take should be performed. """ try: axis = keywords["axis"] del keywords["axis"] except KeyError: axis = 0 if axis == 0: return _take(array, indices, **keywords) else: return swapaxes(_take(swapaxes(array, 0, axis), indices, **keywords), 0, axis) _put = ufunc._PutUFunc("put") _put.__doc__ = ufunc._PutUFunc.__doc__ def put(array, *indices, **keywords): """put(array, indices..., values, clipmode=(CLIP|WRAP), axis=0) stores 'values' into 'array' at 'indices...'. parameters which must be specified by keyword: clipmode=CLIP if clipmode is CLIP, out of range indices are clipped at [0, shape[i]). if clipmode is WRAP, out of range indices are wrapped around from 0 to shape[i] (<0) or from shape[i] to 0 (>= shape[i]) axis=0 selects the axis along which the take should be performed. """ try: axis = keywords["axis"] del keywords["axis"] except KeyError: axis = 0 if axis == 0: return _put(array, indices[:-1], indices[-1], **keywords) else: return swapaxes( _put( swapaxes(array, 0, axis), indices[:-1], indices[-1], **keywords), 0, axis) def _compress(condition, a): """compress selects members of array 'a' along 'axis' which correspond to non-zero members of 'condition'. """ return _take(a, ufunc.nonzero(condition)) def compress(condition, a, axis=0): if axis == 0: return _compress(condition, a) else: return swapaxes( _compress(condition, swapaxes(a, 0, axis)), 0, axis) def _repeat(array, repeats): if ufunc._isScalar(repeats): repeats = (repeats,)*len(array) total = ufunc.add.reduce(_numarray.array(repeats)) newshape = (total,)+array._shape[1:] newarray = array.__class__(shape=newshape, type=array._type) newi = 0; for i in range(len(array)): for j in range(repeats[i]): newarray[newi] = array._simpleIndexing((i,), None) newi += 1 return newarray def repeat(array, repeats, axis=0): """repeat(a, r, axis=0) returns a new array with each element 'a[i]' repeated 'r[i]' times. """ if axis == 0: return _repeat(_numarray.array(array), repeats) else: return swapaxes( _repeat(swapaxes(array, 0, axis), repeats), 0, axis) def indices(shape, type=None): """indices(shape, type=None) returns an array representing a grid of indices with row-only, and column-only variation. """ a = concatenate(ufunc.nonzero(_numarray.ones(shape))) a.setshape((len(shape),)+shape) if type is not None: a = a.astype(type) return a def fromfunction(function, dimensions): # from Numeric """fromfunction(function, dimensions) returns an array constructed by calling function on a tuple of number grids. The function should accept as many arguments as there are dimensions which is a list of numbers indicating the length of the desired output for each axis. """ return apply(function, tuple(indices(dimensions)))