"""NumArray Universal Function Module $Id: ufunc.py,v 1.98 2004/03/12 18:14:19 jaytmiller Exp $ """ import types import memory import _sort, _bytes from operator import and_ import _ufuncall import generic as _gen import _ufunc import _converter import numerictypes as _nt import numarraycore as _nc import safethread import warnings as _warnings # note more imports are at end of file #import pywin.debugger """DEBUGGING NOTE: A reasonable amount of care went into maintaining the Python prototype for _callOverDimensions, _UnaryUFunc, _BinaryUFunc, _Converter, and _Operator. Thus, problems in the C implementation can sometimes be isolated by substituting the original Python code back in, either in whole or part. Substituting in the whole is easy: set _PROTOTYPE to 1 below. """ _PROTOTYPE = 0 _FPE_DIVIDE_BY_ZERO = 1 _FPE_OVERFLOW = 2 _FPE_UNDERFLOW = 4 _FPE_INVALID = 8 class _NumErrorMode: _errorset = ("ignore", "warn", "raise") def _num_mode_test(self, val, all, default): if val not in self._errorset and val is not None: raise ValueError("Error mode value must be in " + repr(self._errorset) + " or None.") if all is not None: if val is not default: return 1 else: return 0 else: return 1 def __init__(self, all=None, overflow="warn", underflow="ignore", dividebyzero="warn", invalid="warn"): if all in self._errorset: self.overflow = all self.underflow = all self.dividebyzero = all self.invalid = all if self._num_mode_test(overflow, all, "warn"): self.overflow = overflow if self._num_mode_test(underflow, all, "ignore"): self.underflow = underflow if self._num_mode_test(dividebyzero, all, "warn"): self.dividebyzero = dividebyzero if self._num_mode_test(invalid, all, "warn"): self.invalid = invalid def __repr__(self): return ("_NumErrorMode(overflow='%s', underflow='%s', dividebyzero='%s', invalid='%s')" % (self.overflow, self.underflow, self.dividebyzero, self.invalid)) class NumError: """Defines how numeric errors should be handled. Error configuration is managed as a mapping (by thread) of stacks. Thus it is possible for each thread to have an independent error configuration, and also for multiple routines within the same thread to save and restore error configurations as needed. """ def __init__(self, all=None, overflow="warn", underflow="ignore", dividebyzero="warn", invalid="warn"): self._modestack = {} # map of stacks indexed by thread id self._defaultmode = _NumErrorMode() self.setMode(all=all, underflow=underflow, overflow=overflow, dividebyzero=dividebyzero, invalid=invalid) def setDefaultMode(self, all=None, overflow="warn", underflow="ignore", dividebyzero="warn", invalid="warn"): self._defaultmode = _NumErrorMode(all, overflow, underflow, dividebyzero, invalid) def setMode(self, all=None, overflow="warn", underflow="ignore", dividebyzero="warn", invalid="warn"): self._setmodes( _NumErrorMode(all, overflow, underflow, dividebyzero, invalid) ) def pushMode(self, all=None, overflow="warn", underflow="ignore", dividebyzero="warn", invalid="warn"): self._pushmodes( _NumErrorMode(all, overflow, underflow, dividebyzero, invalid) ) def _getmodestack(self): id = safethread.get_ident() try: l = self._modestack[id] except KeyError: l = [ self._defaultmode ] self._modestack[id] = l if l == []: l = [ self._defaultmode ] return l def _setmodes(self, modes): l = self._getmodestack() l[-1] = modes def getMode(self): l = self._getmodestack() return l[-1] def _pushmodes(self, modes): l = self._getmodestack() l.append(modes) def popMode(self): l = self._getmodestack() return l.pop() Error = NumError() class MathDomainError(ArithmeticError): pass class UnderflowError(ArithmeticError): pass def handleError(errorStatus, sourcemsg): """Take error status and use error mode to handle it.""" modes = Error.getMode() if errorStatus & _FPE_INVALID: if modes.invalid == "warn": print "Warning: Encountered invalid numeric result(s)", sourcemsg if modes.invalid == "raise": raise MathDomainError(sourcemsg) if errorStatus & _FPE_DIVIDE_BY_ZERO: if modes.dividebyzero == "warn": print "Warning: Encountered divide by zero(s)", sourcemsg if modes.dividebyzero == "raise": raise ZeroDivisionError(sourcemsg) if errorStatus & _FPE_OVERFLOW: if modes.overflow == "warn": print "Warning: Encountered overflow(s)", sourcemsg if modes.overflow == "raise": raise OverflowError(sourcemsg) if errorStatus & _FPE_UNDERFLOW: if modes.underflow == "warn": print "Warning: Encountered underflow(s)", sourcemsg if modes.underflow == "raise": raise UnderflowError(sourcemsg) def ufuncFactory(operator, ucfuncs, identity=None): """Create UFunc instance based on signature of functions""" ninputs, noutputs = _nIOArgs(ucfuncs) if ninputs==1 and noutputs==1: return _UnaryUFunc(operator, ucfuncs, identity) elif ninputs==2 and noutputs==1: return _BinaryUFunc(operator, ucfuncs, identity) elif noutputs==1: raise ValueError("ufunc with more than 2 input arguments") else: raise ValueError("ufunc with more than 1 output argument") def _nIOArgs(ucfuncs): """Determine number of arguments from signature Returns tuple of number of input and output arguments. Also check for consistency in numbers of arguments. """ # The number of items in the optype key implies the number of inputs, # hence the number of inputs = # of underscores+1 noptypeinputs = [] ninputs = [] noutputs = [] for key in ucfuncs.keys(): if key in ['reduce', 'accumulate']: continue noptypeinputs.append(key.count('_')) for insigkey in ucfuncs[key].keys(): ninputs.append(len(insigkey)) noutputs.append(len(ucfuncs[key][insigkey][0])) if max(noptypeinputs) != min(noptypeinputs): raise StandardError( "inconsistency in operator types for operator: "+operator) if max(ninputs) != min(ninputs): raise StandardError( "not all input signatures match in number of arguments\n" + "for operator: "+operator) if max(noutputs) != min(noutputs): raise StandardError( "not all output signatures match in number of arguments\n" + "for operator: "+operator) return max(ninputs), max(noutputs) def _getBlockingMiss(shape, niter, overlap=0): """Figure out how to split the input array into subnumarray close (but less than) the size of the block buffers. overlap is used in cases where it is necessary to overlap the last dimension in computations (for example for reduce and accumulate.) returns: level: index at which the c functions should be called blockingparameters: tuple consisting of number of full subnumarray at the given level subarray shape a flag indicating if there is a leftover part of the array which is smaller than the regular subarray (leftover) the shape of the leftover subarray """ # Now in C: # retval = _blockingParametersCache.get((shape,niter,overlap)) # if retval is not None: # return retval if shape == (): retval = 0, (1, (), 0, None) _blockingParametersCache[(shape,niter,overlap)] = retval return retval ndim = len(shape) csize = long(shape[-1]) # cumulative shape if csize > niter: # last dimension is too big for buffer # special overlap handling in this case nblocks, leftover = divmod(csize, niter-overlap) nblocks, leftover = int(nblocks), int(leftover) if leftover: retval = ndim-1, (nblocks, (niter,), 1, (leftover,)) else: retval = ndim-1, (nblocks, (niter,), 0, None) else: for i in xrange(2,ndim+1): csize *= shape[-i] if csize >= niter: # need to break up this dimension tdim = long(shape[-i]) blockdim = niter*tdim/csize nblocks, leftover = divmod(tdim, blockdim) nblocks, leftover = int(nblocks), int(leftover) stail = shape[-i+1:] if leftover: retval = ndim-i, (nblocks, (blockdim,)+stail, 1, (leftover,) + stail) else: retval = ndim-i, (nblocks, (blockdim,)+stail, 0, None) break else: # don't need to break up retval = 0, (1, shape, 0, None) _blockingParametersCache[(shape,niter,overlap)] = retval return retval # _doOverDimensions has been re-implemented in C. def _doOverDimensions(objects, outshape, dims, indexlevel, blockingparameters, overlap=0, level=0): """This version handles alignment, byteswapping and n-dimensions This is a recursive method dims starts out as an empty list, and accumulates the indices for each dimension. Eg. if the array has shape (7, 3, 5) on the initial call, the next level call will specify the the 'first' index (corresponding to the 7), the next will have the first two indices (and that would be the last level since the lowest dimension is taken care of by the cfuncs.) """ if level == indexlevel: nregShapeIters, shape, leftover, leftoverShape, = blockingparameters if shape != (): dimval = shape[0]-overlap for i in xrange(nregShapeIters): for o in objects: o.compute(dims+[i*dimval], shape) if leftover: for o in objects: o.compute(dims+[(i+1)*dimval], leftoverShape) else: for o in objects: o.compute([0], shape) else: # recurse for i in xrange(outshape[level]): _doOverDimensions(objects, outshape, dims+[i], indexlevel, blockingparameters, overlap, level+1) def _callOverDimensions(objects, outshape, index, blocking, overlap=0, level=0): return _doOverDimensions(objects, outshape, [], index, blocking, overlap, level) # Replace the Python version of _callOverDimensions with one implemented in C. # Although _doOverDimensions is not replaced here, it has been re-implemented # in C as well if not _PROTOTYPE: from _ufunc import _callOverDimensions def _copyFromAndConvertMiss(inarr, outarr): key = (_digest(inarr), _digest(outarr), safethread.get_ident()) # Type conversion done using inputconverter input = _InputConverter(inarr, outarr._type.name) # Ensure output step uses output buffer of input step output = _OutputConverter(outarr, outarr._type.name, inbuffer=input.result_buff, forcestride=1) maxitemsize = max(inarr._type.bytes, outarr._type.bytes) niter = _ufunc.getBufferSize()/maxitemsize if (isinstance(inarr, _nc.NumArray) and isinstance(outarr, _nc.NumArray)): _copyCache[ key ] = input, output, niter return input, output, niter if _PROTOTYPE: def _copyFromAndConvert(inarr, outarr): """Copy from one array to another handling strides, form & conv""" if inarr._shape != outarr._shape: raise ValueError("Arrays must have the same shape") key = (_digest(inarr), _digest(outarr), safethread.get_ident()) try: input, output, niter = _copyCache[ key ] except KeyError: input, output, niter = _copyFromAndConvertMiss(inarr, outarr) input.rebuffer(inarr) output.rebuffer(outarr, input.result_buff) indexlevel, blockingparameters = \ _ufunc._getBlockingParameters(outarr._shape, niter) objects = (input, output) _ufunc.CheckFPErrors() _callOverDimensions(objects, outarr._shape, indexlevel, blockingparameters) errorstatus = _ufunc.CheckFPErrors() input.rebuffer(None) output.rebuffer(None) if errorstatus: handleError(errorstatus, " during type conversion") else: from _ufunc import _copyFromAndConvert def _noConversionNeeded(x, xt): """returns 1 iff x is a c_array of type xt.""" return x.is_c_array() and (x._type.name == xt or x._type == xt) def _inputcheck(*inargs): """Check input args for type, convert sequences to numnumarray Return list of 1) input arg tuple with scalars and numnumarray only 2) signature string indicating which args are scalars and which are numnumarray 3) flag indicating if only scalar arguments are present """ retargs = [] retsig = [] scalar = 1 for inarg in inargs: if isinstance(inarg, _nc.NumArray): scalar = 0 retsig.append('vector') elif _nc.PyNUMERIC_TYPES.has_key(type(inarg)): retsig.append('scalar') else: scalar = 0 retsig.append('vector') # see if it can be made into an array try: inarg = _nc.array(inarg) except TypeError: raise TypeError( "UFunc arguments must be numarray, scalars or numeric sequences") retargs.append(inarg) return retargs, '_'.join(retsig), scalar def _pseudoArrayFromScalars(scalarvalues, type): """Wrap a scalar in a buffer so it can be used as an array""" arr = _bufferPool.getBuffer() arr._check_overflow = 1 newtype = _numtypedict[type] arr._strides = (newtype.bytes,) arr._shape = (len(scalarvalues),) arr._type = newtype arr._itemsize = newtype.bytes for i in xrange(len(scalarvalues)): arr[i] = scalarvalues[i]; arr._strides = None # Modify block buffer attributes to look like vector/vector setup. return arr # Portions of the UFunc classes have been re-implemented in C. class _UFunc(_ufunc._ufunc): """Class to handle all element-by-element mathematical functions This is the base class -- subclasses must implement __call__ and (if allowed) the outer and _cumulative methods. """ def __init__(self, operator, ufuncs, identity): self.operator = operator self._ucfuncs = ufuncs self._identity = identity self._coercions = _nt.genericCoercions self._promotionExclusions = _nt.genericPromotionExclusions self._typerank = _nt.genericTypeRank self._cache = {} self.arity = 0 def _getidentity(self): return _nc.array([self._identity]) if _PROTOTYPE: def __call__(self, *args): raise NotImplementedError("__call__ is not implemented by base UFunc class") def __repr__(self): return "" % self.operator def _typematch(self, inputSignature, bintypeDict, scalar=0): """Determine which C function is most appropriate for the given types returns the input types the inputs must be converted to, the output type, and the cfunc object """ # First see if there is an exact match if bintypeDict.has_key(inputSignature): return inputSignature + bintypeDict[inputSignature] else: # Look for 'closest' diagonal element. # First find 'lowest' common type. intype = self._coercions[inputSignature] keys = bintypeDict.keys() atypelist = [] for key in keys: if key[0] == key[1]: atypelist.append(key[0]) newInputSignature = (self._typePromoter(intype, atypelist),)*2 return newInputSignature + bintypeDict[newInputSignature] def _typePromoter(self, intype, typelist): """Type a given type should be promoted to given a list of choices""" # type as string for now # Not an efficient version for now # Generate the type specific rank list tranks = self._typerank[:] # copy since we will be deleting elements exclusions = self._promotionExclusions[intype] for exclusion in exclusions: del tranks[exclusion] index = tranks.index(intype) tranks = tranks[index:] # Which one has lowest rank? minimumRank = 100 for atype in typelist: try: rank = tranks.index(atype) if rank < minimumRank: minimumRank = rank except ValueError: pass if minimumRank < 100: return tranks[minimumRank] else: raise TypeError("unable to find type to promote to") def _typeDemoter(self, intype, typelist): """The opposite of what typePromoter does--not implemented yet""" pass def outer(self, inarr1, inarr2, outarr=None): raise ValueError("Outer only available for binary functions") _blockingParametersCache = {} _copyCache = {} def flush_caches(): """flush_cache elminates all ufunc cache entries.""" import sys module = sys.modules[__name__] for n in module.__dict__.keys(): f = module.__dict__[n] if isinstance(f, _UFunc): f._cache_flush() _blockingParametersCache = {} _copyCache = {} class _BufferPool: """A Pool of available buffers available for block compuations Buffers can be gotten. When they are deleted they are automatically returned to the buffer pool so they can be reused. """ def __init__(self, buffersize=100000): # buffers are initially allocated as needed self.buffers = [] self.setBufferSize(buffersize) def getBuffer(self): try: return _nc._UBuffer(self.buffers.pop()) except IndexError: # out of buffers, create a new one return _nc._UBuffer(memory.new_memory(_ufunc.getBufferSize())) def getBufferSize(self): return _ufunc.getBufferSize() def setBufferSize(self, buffersize): """Change buffer size (useful mainly for testing)""" oldsize = _ufunc.getBufferSize() if buffersize != oldsize: _ufunc.setBufferSize(buffersize) # delete existing (wrong-sized) buffers self.buffers = [] flush_caches() return oldsize _bufferPool = _BufferPool() def _sequence(x): return isinstance(x, list) or isinstance(x, tuple) if _PROTOTYPE: def _digest(x): if isinstance(x, _nc.NumArray): return ("array", x.iscontiguous(), x.isaligned(), x.isbyteswapped(), x.type()) elif _isScalar(x): return type(x), x # for demo only elif x is None: return None else: raise KeyError, "_digest force cache miss" def _restuff_pseudo(p, x): if _isScalar(x): if p is None: raise RuntimeError, "scalar value with no pseudo array" else: # restuff p with x # p[0] = x return p else: if p is not None: return p else: return x else: _digest = _ufunc.digest _restuff_pseudo = _ufunc.restuff_pseudo def _cache(original_input, pseudo_array): if _isScalar(original_input): return pseudo_array else: return None def _rank(seq, r=0): if isinstance(seq, _nc.NumArray): return r + len(seq._shape) try: x = len(seq) except: return r else: if x == 0: return r else: return _rank(seq[0], r+1) class _UnaryUFunc(_UFunc): """Class for ufuncs with 1 input and 1 output argument""" def __init__(self, *args, **keys): _UFunc.__init__(self, *args, **keys) self.arity = 1 if _PROTOTYPE: self._cache = {} if _PROTOTYPE: def __call__(self, inarr1, outarr=None): """The standard calling interface for UFuncs""" return self._cached_dispatch1(inarr1, outarr) def _fast_exec1(self, in1, out, cached): mode, outtype, cfunc, ufargs_junk, insig, cin1 = cached ufargs = (out.nelements(), 1, 1, ((in1._data, in1._byteoffset), (out._data, out._byteoffset))) apply(cfunc, ufargs) def _slow_exec1(self, in1, out, cached): mode, outtype, cfunc, ufargs, insig, cin1 = cached inputs, outputs, maxitemsize = ufargs niter = _ufunc.getBufferSize()/maxitemsize indexlevel, blockingparameters = \ _ufunc._getBlockingParameters(out._shape, niter) operator = _Operator(cfunc, [inputs[0].rebuffer(in1)], [outputs[0].rebuffer(out)]) objects = (inputs[0],operator,outputs[0]) _callOverDimensions(objects, out._shape, indexlevel, blockingparameters) for o in inputs+outputs: o.rebuffer(None) def _cache_exec1(self, in1, out, cached): if out.nelements(): # skip 0-element arrays _ufunc.CheckFPErrors() if cached[0] == "fast": self._fast_exec1(in1, out, cached) else: self._slow_exec1(in1, out, cached) error = _ufunc.CheckFPErrors() if error: handleError(error, " in "+self.operator) if out._shape == (): return out[()] else: return out def _cache_lookup1(self, in1, out): """_cache_lookup2 checks the ufunc cache for an entry corresponding to (in1, in2, out) and returns it if it is found. If the cache misses, _cache_lookup2 performs a full ufunc setup and returns it. In either case, vector_vector inputs are dual broadcast to match shapes. """ win1 = in1 try: key = (_digest(win1), _digest(out), safethread.get_ident()) cached = self._cache[ key ] except KeyError: return self._cache_miss1(win1, out) else: mode, otype, cfunc, ufargs, insig, cin1 = cached cin1 = _restuff_pseudo(cin1, in1) if out is None: if insig == "vector": wout = win1.new(cached[1]) else: # scalar wout = _nc.zeros((), type=otype) else: if insig == "vector": shape = win1._shape else: # scalar shape = () if shape != out._shape: raise ValueError("Supplied output array does not have appropriate shape") wout = out return cin1, wout, cached def _cached_dispatch1(self, inarr1, outarr): params = self._cache_lookup1(inarr1, outarr) result = self._cache_exec1(*params) if outarr is None: return result else: return None def _cache_miss1(self, n1, out): (in1,), insig, scalar = _inputcheck(n1) mode, win1, wout, cfunc, ufargs = \ self._setup(in1, insig, out) if mode == "fast": cached = mode, wout._type, cfunc, None, insig, _cache(n1, win1) else: cached = mode, wout._type, cfunc, ufargs, insig, _cache(n1, win1) try: key = (_digest(n1), _digest(out), safethread.get_ident()) except KeyError: pass else: if _PROTOTYPE: self._cache[ key ] = cached else: self._cache_insert( cached, n1, None, out) return win1, wout, cached def _setup(self, inarr, insig, outarr): """Setup for unary ufunc""" ufdict = self._ucfuncs["vector"] if insig == "scalar": t = _nc.array(inarr) # create rank-0 t._check_overflow = 1 # enable overflow checking for t t[()] = inarr # force any overflow inarr = t shape = inarr._shape typekeys = [key[0] for key in ufdict.keys()] intype = self._typePromoter(inarr._type.name, typekeys) outtype, ucfunc = ufdict[(intype,)] if outarr is None: outarr = inarr.new(_numtypedict[outtype[0]]) if shape != outarr._shape: raise ValueError("Supplied output array does not have appropriate shape") # If fast case, return the arguments for an immediate call to the # C function if (_noConversionNeeded(inarr, intype) and _noConversionNeeded(outarr, outtype[0])): ufargs = (outarr.nelements(), 1, 1, ((inarr._data, inarr._byteoffset), (outarr._data, outarr._byteoffset))) return ("fast", inarr, outarr, ucfunc, ufargs) # slower, general case inputs = (_InputConverter(inarr, intype), ) outputs = (_OutputConverter(arr=outarr, type=outtype[0]),) # find largest itemsize involved, it will determine how many # iterations can be done on the temporary buffers (i.e., niter) maxitemsize = max(_numtypedict[inarr._type.name].bytes, _numtypedict[intype].bytes, _numtypedict[outtype[0]].bytes, _numtypedict[outarr._type.name].bytes) preprocessing_outbuffers = (inputs[0].result_buff,) postprocessing_inbuffers = (outputs[0].result_buff,) return ("slow", inarr, outarr, ucfunc, (inputs, outputs, maxitemsize)) def _firstcol(arr): rval = arr.view() rval._shape = rval._shape[:-1] rval._strides = rval._strides[:-1] return rval def _moveToLast(dim, iseq): oseq = [] for i in range(len(iseq)): if i != dim: oseq.append(iseq[i]) else: s = iseq[i] oseq.append(s) return oseq def _fixdim(axis, keys): """support the deprecation of the 'dim' keyword in favor of 'axis'""" if "dim" in keys.keys(): if axis != 0: raise RuntimeError("Specify 'axis' or 'dim', but not both. 'dim' is deprecated.") _warnings.warn("The 'dim' keyword is deprecated. Specify 'axis' instead.", DeprecationWarning, stacklevel=3) return keys["dim"] return axis def _fixdimaux(axis, dim): """provide glue so C-code can easily call _fixdim. Note that C-code cannot detect the default case: dim=0.""" if dim != 0: return _fixdim(axis, {"dim":dim}) else: return axis class _BinaryUFunc(_UFunc): """Class for ufuncs with 2 input and 1 output arguments""" def __init__(self, *args, **keys): _UFunc.__init__(self, *args, **keys) self.arity = 2 if _PROTOTYPE: def _fast_exec2(self, in1, in2, out, cached): mode, outtype, cfunc, ufargs_junk, insig, cin1, cin2 = cached ufargs = (out.nelements(), 2, 1, ((in1._data, in1._byteoffset), (in2._data, in2._byteoffset), (out._data, out._byteoffset))) apply(cfunc, ufargs) def _slow_exec2(self, in1, in2, out, cached): mode, outtype, cfunc, ufargs, insig, cin1, cin2 = cached inputs, outputs, maxitemsize = ufargs niter = _ufunc.getBufferSize()/maxitemsize indexlevel, blockingparameters = \ _ufunc._getBlockingParameters(out._shape, niter) operator = _Operator(cfunc,[inputs[0].rebuffer(in1), inputs[1].rebuffer(in2)], [outputs[0].rebuffer(out)]) objects = inputs + (operator,) + outputs _callOverDimensions(objects, out._shape, indexlevel, blockingparameters) for o in inputs+outputs: o.rebuffer(None) def _cache_exec2(self, in1, in2, out, cached): if out.nelements(): # skip 0-element arrays _ufunc.CheckFPErrors() if cached[0] == "fast": self._fast_exec2(in1, in2, out, cached) else: self._slow_exec2(in1, in2, out, cached) error = _ufunc.CheckFPErrors() if error: handleError(error, " in "+self.operator) if out._shape == (): return out[()] else: return out def _cache_lookup2(self, in1, in2, out): """_cache_lookup2 checks the ufunc cache for an entry corresponding to (in1, in2, out) and returns it if it is found. If the cache misses, _cache_lookup2 performs a full ufunc setup and returns it. In either case, vector_vector inputs are dual broadcast to match shapes. """ if (isinstance(in1, _nc.NumArray) and isinstance(in2, _nc.NumArray)): win1, win2 = in1._dualbroadcast(in2) else: win1, win2 = in1, in2 try: key = (_digest(win1), _digest(win2), _digest(out), safethread.get_ident()) cached = self._cache[ key ] except KeyError: return self._cache_miss2(win1, win2, out) else: mode, otype, cfunc, ufargs, insig, cin1, cin2 = cached cin1 = _restuff_pseudo(cin1, win1) cin2 = _restuff_pseudo(cin2, win2) if out is None: if insig in ["vector_scalar", "vector_vector"]: wout = win1.new(cached[1]) elif insig == "scalar_vector": wout = win2.new(cached[1]) else: # scalar_scalar wout = _nc.zeros((), type=otype) else: if insig in ["vector_scalar", "vector_vector"]: shape = win1._shape elif insig == "scalar_vector": shape = win2._shape else: # scalar_scalar shape = () if shape != out._shape: raise ValueError("Supplied output array does not have appropriate shape") wout = out return cin1, cin2, wout, cached def _cached_dispatch2(self, inarr1, inarr2, outarr): params = self._cache_lookup2(inarr1, inarr2, outarr) result = self._cache_exec2(*params) if isinstance(result, _nc.NumArray) and result.rank == 0: result = result[()] if outarr is None: return result else: return None def __call__(self, inarr1, inarr2, outarr=None): """The standard calling interface for UFuncs""" return self._cached_dispatch2(inarr1, inarr2, outarr) def _cache_miss2(self, n1, n2, out): (in1, in2), insig, scalar = _inputcheck(n1, n2) mode, win1, win2, wout, cfunc, ufargs = \ self._setup(in1, in2, insig, out) if mode == "fast": cached = mode, wout._type, cfunc, None, insig, \ _cache(n1, win1), _cache(n2, win2) else: cached = mode, wout._type, cfunc, ufargs, insig, \ _cache(n1, win1), _cache(n2, win2) try: key = (_digest(n1), _digest(n2), _digest(out), safethread.get_ident()) except KeyError: pass else: if _PROTOTYPE: self._cache[ key ] = cached else: self._cache_insert(cached, n1, n2, out) return win1, win2, wout, cached def _setup(self, in1, in2, insig, out): """Setup for binary ufunc""" wout = out if insig == "scalar_scalar": # do it as two 0-d numarray. in1, in2 = _nc.array(in1), _nc.array(in2) insig = "vector_vector" if out is None: wout = in1.copy() elif insig == "vector_scalar": intypes = (in1._type.name, type(in2)) convtype1, convtype2, outtype, ucfunc \ = self._typematch(intypes, self._ucfuncs[insig]) inarr1, inarr2 = in1, _pseudoArrayFromScalars([in2], convtype2) if out is None: wout = in1.new(outtype[0]) elif insig == "scalar_vector": intypes = (type(in1), in2._type.name) convtype1, convtype2, outtype, ucfunc \ = self._typematch(intypes, self._ucfuncs[insig]) inarr1, inarr2 = _pseudoArrayFromScalars([in1], convtype1), in2 if out is None: wout = in2.new(outtype[0]) if insig == "vector_vector": intypes = (in1._type.name, in2._type.name) inarr1, inarr2 = in1._dualbroadcast(in2) convtype1, convtype2, outtype, ucfunc \ = self._typematch(intypes, self._ucfuncs[insig]) if out is None: wout = inarr1.new(outtype[0]) #If fast case, return the arguments for an immediate call to the # C function if (_noConversionNeeded(inarr1, convtype1) and _noConversionNeeded(inarr2, convtype2) and _noConversionNeeded(wout, _numtypedict[outtype[0]])): ufargs = (wout.nelements(), 2, 1, ((inarr1._data, inarr1._byteoffset), (inarr2._data, inarr2._byteoffset), (wout._data, wout._byteoffset))) return ("fast", inarr1, inarr2, wout, ucfunc, ufargs) # slower, general case inputs = (_InputConverter(inarr1, convtype1), _InputConverter(inarr2, convtype2)) outputs = (_OutputConverter(arr=wout, type=outtype[0]),) # find largest itemsize involved, it will determine how many # iterations can be done on the temporary buffers (i.e., niter) maxitemsize = max(_numtypedict[inarr1._type.name].bytes, _numtypedict[inarr2._type.name].bytes, _numtypedict[convtype1].bytes, _numtypedict[convtype2].bytes, _numtypedict[outtype[0]].bytes, _numtypedict[wout._type.name].bytes) return ("slow", inarr1, inarr2, wout, ucfunc, (inputs, outputs, maxitemsize)) def outer(self, inarr1, inarr2, outarr=None): """Return outer product of 2 numarray""" (inarr1, inarr2), insig, scalar = _inputcheck(inarr1, inarr2) ndim2 = len(inarr2._shape) indexarg = (slice(None, None, None),)*len(inarr1._shape) + (None,)*ndim2 if outarr is not None: self(inarr1.__getitem__(indexarg), inarr2, outarr) else: return self(inarr1.__getitem__(indexarg), inarr2) def zreduce(self, array, axis=0, out=None, type=None, **keys): """zreduce returns a rank-0 array as the result of the reduction of array <= rank-1 array, similar to Numeric. """ axis = _fixdim(axis, keys) in1 = _nc.asarray(array) r = self._cum_swapped(in1, axis, out, "reduce", type) if in1.rank <= 1: r.shape = () r.strides = () return r if _PROTOTYPE: def accumulate(self, array, axis=0, out=None, type=None, **keys): """accumulate applies the binary operator 'self' at successive pairs of elements of 'array' along dimension 'dim', storing the result in 'out'. If no 'out' is provided, the result of the accumulation will be returned. The result of accumulating 'array' has the same shape as 'array'. """ axis = _fixdim(axis, keys) in1 = _nc.asarray(array) return self._cum_swapped(in1, axis, out, "accumulate", type) def areduce(self, array, axis=0, out=None, type=None, **keys): """areduce applies the operator of ufunc 'self' to reduce 'array' along its 'dim' axis. If no 'out' is provided, the result of the reduction is returned. The result of reducing an N-dimensional array is an N-1 dimensional array. The result of reducing a rank-1 array is a rank-0 array. see also 'reduce'. """ axis = _fixdim(axis, keys) in1 = _nc.asarray(array) return self._cum_swapped(in1, axis, out, "reduce", type) def reduce(self, array, axis=0, out=None, type=None, **keys): """areduce applies the operator of ufunc 'self' to reduce 'array' along its 'dim' axis. If no 'out' is provided, the result of the reduction is returned. The result of reducing an N-dimensional array is an N-1 dimensional array. The result of reducing a rank-1 array is a scalar. see also 'areduce'. """ axis = _fixdim(axis, keys) in1 = _nc.asarray(array) r = self._cum_swapped(in1, axis, out, "reduce", type) if r is not None: if r._shape == (1,) and in1.rank <= 1: r = r[0] return r def _cum_cached(self, cumop, in1, out, type): if (out is not None): if not isinstance(out, _nc.NumArray): raise TypeError("output array must be a NumArray") if (not out.isaligned() or out.isbyteswapped()): raise ValueError("Reduce/Accumulate: no support for misaligned / byteswapped output numarray") if cumop not in ["reduce", "accumulate"]: raise ValueError("Unknown cumulative option") if in1.rank == 0: if self._identity is not None: return self._getidentity() # 0D array form else: raise ValueError("Empty array in a ufunc that has no identity value") params = self._cum_lookup(cumop, in1, out, type) wout = self._cum_exec(*params) if cumop=="reduce": wout._strides = wout._strides[:-1] if wout._shape == (): wout._shape = (1,) wout._strides = (wout._itemsize,) if out is None: return wout i, o, cached = params mode, otype, cfunc, ufargs = cached if out._type != otype: _copyFromAndConvert(wout, out) return out def _cum_lookup(self, cumop, in1, out, type): try: key = (_digest(in1), _digest(out), cumop, safethread.get_ident(), type) cached = self._cache[ key ] except KeyError: # lists always miss return self._cum_cache_miss(cumop, in1, out, type) else: mode, otype, cfunc, ufargs = cached if cumop == "reduce": wout = self._reduce_out(in1, out, otype) else: wout = self._accumulate_out(in1, out, otype) return in1, wout, cached def _cum_exec(self, in1, out, cached): mode, otype, cfunc, ufargs = cached if otype == _nt.Bool: if in1._type != _nt.Bool: in1 = in1.astype(_nt.Bool) if in1.nelements(): # skip 0-element arrays _ufunc.CheckFPErrors() if mode == "fast": self._cum_fast_exec(in1, out, cached) else: self._cum_slow_exec(in1, out, cached) errorstatus = _ufunc.CheckFPErrors() if errorstatus: handleError(errorstatus, sourcemsg) return out def _cum_fast_exec(self, in1, out, cached): mode, otype, cfunc, ufargs = cached cfunc(in1._shape, in1._data, in1._byteoffset, in1._strides, out._data, out._byteoffset, out._strides) def _cum_slow_exec(self, in1, out, cached): mode, otype, cfunc, ufargs = cached input, output, maxitemsize = ufargs niter = _ufunc.getBufferSize()/maxitemsize if in1._shape[-1] > niter: overlap = 1 else: overlap = 0 indexlevel, blockingparameters = \ _ufunc._getBlockingParameters(in1._shape, niter, overlap) operator = _Operator(cfunc, [input.rebuffer(in1)], [output.rebuffer(out)], otype.bytes) objects = (input, operator, output) _callOverDimensions(objects, out._shape, indexlevel, blockingparameters, overlap) input.rebuffer(None) output.rebuffer(None) def _accumulate_out(self, inarr, outarr, outtype): """_accumulate_out creates output array for accumulate""" # Create output array if not supplied. if outarr is None: toutarr = inarr.new( outtype ) else: if outarr._shape != inarr._shape: raise ValueError( "Supplied output array does not have the appropriate shape") toutarr = outarr # It is necessary to initialize the first subarray of the output # to the first subarray of the input (C functions depend on it). # toutarr[...,0] = inarr[...,0] if inarr.nelements(): _firstcol(toutarr)._copyFrom(_firstcol(inarr)) return toutarr def _reduce_out(self, inarr, outarr, outtype): """_reduce_out creates output array for reduce""" firstcol = _firstcol(inarr) # inarr[..., 0] # Create output array if not supplied or wrong type if outarr is None or outarr._type != outtype: toutarr = firstcol.new(outtype) else: toutarr = outarr.view() if len(inarr._shape) == 1 and toutarr._shape == (1,): toutarr._shape, toutarr._strides = (), () if firstcol._shape != toutarr._shape: raise ValueError("output array shape mismatch") if inarr.nelements(): # skip 0-element arrays toutarr._copyFrom(firstcol) # Assign first subarray of output else: toutarr._copyFrom(self._identity) toutarr._strides += (0,) return toutarr def _cum_swapped(self, in1, dim, out, cumop, type=None): if dim == in1.rank-1: dim = -1 if dim != -1: _in1 = _gen.swapaxes(in1, -1, dim) _out = _gen.swapaxes(out, -1, dim) else: _in1, _out = in1, out _out1 = self._cum_cached(cumop, _in1, _out, type) if cumop == "accumulate": if dim != -1: _out1.swapaxes(-1, dim) # swap axes to "fix" the result else: if dim != -1: _out1._shape = _moveToLast(dim, _out1._shape) _out1._strides = _moveToLast(dim, _out1._strides) if out is None: return _out1 def _cum_cache_miss(self, cumop, in1, out, type): mode, win1, wout, cfunc, ufargs = \ self._cum_setup(cumop, in1, out, type) cached = mode, wout._type, cfunc, ufargs try: key = (_digest(in1), _digest(out), cumop, safethread.get_ident(), type) except KeyError: pass else: if _PROTOTYPE: self._cache[ key ] = cached else: self._cache_insert(cached, in1, None, out, cumop, type); return win1, wout, cached def _cum_setup(self, cumop, in1, out, type=None): """Used by both reduce and accumulate to compute accumulations. Assumes that the dimension of accumulation/reduction is 0. in1 -- The input array to accumulate/reduce out -- Optional output array The output array must have a consistent shape with the input Array (the same for accumulate and minus the accumulated dimension for reduce). The output type must match that produced by default and cannot be nonaligned or byteswapped. (It is a real mess to remove this restriction. (The original numeric doesn't even support output numarray.) """ (win1,), insig, scalar = _inputcheck(in1) # win1 substitute for lists if scalar: raise ValueError("First argument must be an array or sequence") # The original numeric returns the operations identity for empty numarray. sourcemsg = "in %s.%s" % (self.operator, cumop) # For error message # Get appropriate function, figure out any needed input type conversions # and what the output type will be ufdict = self._ucfuncs[cumop] typekeys = [key[0] for key in ufdict.keys()] try: intype = _numtypedict[self._typePromoter(win1._type.name, typekeys)] except TypeError: if typekeys == ['Bool']: # Hack for logical operators return self._cum_setup( cumop, win1.astype('Bool'), out) if out is None: if type is None: optype = intype else: optype = _nt.getType(type) else: if type is not None: raise ValueError("Can't specify both 'type' and 'out'") else: optype = out.type() otype, cfunc = ufdict[(optype,)] otype = _numtypedict[otype[0]] # A few words on shape/stride manipulations: # The C function assumes the last dimension is being accumulated. # Create the accumulate/reduce specific output array, tout # type of reduce tout is always otype to avoid output conversion # with mismatched shapes. if cumop == "reduce": tout = self._reduce_out(win1, out, otype) else: tout = self._accumulate_out(win1, out, otype) # check for fast case (Unlike Ufuncs, contiguous numarray not required) if (not win1.isbyteswapped() and win1.isaligned() and optype == win1._type and not tout.isbyteswapped() and tout.isaligned()): return "fast", win1, tout, cfunc, None else: # slow case # And now a few words about overlap. It is used to allow the # accumulate C function to work across blocks. Overlap in effect # makes the first subarray of the next block the last subarray # of the previous block (or last array value for 1-d numarray) so # the accumulate can carry on. The computation of the block size # must take it into account as does computeOverDimensions in its # computation of the offsets for the blocks. This is only needed # where the accumulated dimension is larger than a block. It is # not needed at all for reduce since the output array contains # the net accumulation. # check special case where last dimension is bigger than blocksize maxitemsize = max(intype.bytes, optype.bytes, otype.bytes) niter = _ufunc.getBufferSize()/maxitemsize # This is a bit tricky. If the accumulated dimension is bigger # than the block size, it is necessary to overlap the block # computations by one element so that the last value of the # previous block can continue the count into the next block. if win1._shape[-1] > niter: overlap = 1 else: overlap = 0 indexlevel, blockingparameters = \ _ufunc._getBlockingParameters(win1._shape, niter, overlap) # The usual compute stuff input = _InputConverter(win1, type=optype.name) # Since cfunc is selected by output type, no type conversion # is required. Since the cfuncs are striding, no striding # is required. Byteswapping and aligning are disallowed. # Thus, the output converter is a NULL converter. output = _OutputConverter(tout, type=otype.name, nonstriding=1) return "slow", win1, tout, cfunc, (input, output, maxitemsize) # Portions of the following class are implemented in C. See _convertermodule.c. # In particular, compute, stride, convert, and rebuffer. class _Converter(_converter._converter): def __init__(self, arr, type, inbuffer=None, forcestride=0, nonstriding=0): self.stridefunction = None self.convfunction = None self.type = type if arr.isbyteswapped(): if isinstance(arr._type, _nt.ComplexType): fname = "byteswap" + arr._type.name else: fname = "byteswap"+`arr._itemsize`+"bytes" self.stridefunction = _bytes.functionDict[fname] elif not arr.isaligned(): fname = "align"+`arr._itemsize`+"bytes" self.stridefunction = (_bytes.functionDict.get(fname) or _bytes.functionDict["copyNbytes"]) elif not nonstriding and not arr.iscontiguous() or forcestride: fname = "copy"+`arr._itemsize`+"bytes" self.stridefunction = (_bytes.functionDict.get(fname) or _bytes.functionDict["copyNbytes"]) if type != arr._type.name: self.convfunction = \ _numtypedict[type]._conv.astype[arr._type.name] if self.stridefunction: # Since the output array is iterated over dimensions, even # the input array, always contiguous, needs strides defined. # But first see if input buffer is a real array. if (inbuffer is not None) and inbuffer._strides: if not inbuffer.is_c_array(): raise ValueError("input buffer is not a c_array") self.bytestrides = [ inbuffer._strides, list(arr._strides) ] self.generated = 0 else: self.bytestrides = [ _stridesFromShape(arr._shape, arr._itemsize), list(arr._strides) ] self.generated = 1 else: self.bytestrides = [ None, None ] # If inbuffer not specified, allocate one if (inbuffer is None) and (self.stridefunction or self.convfunction): inbuffer = _bufferPool.getBuffer() self.conversion_required = 1 if self.convfunction: if self.stridefunction: buff = _bufferPool.getBuffer() buffers = [inbuffer, buff, buff, arr] self.arr_position = 3 self.inb_position = 0 else: buffers = [inbuffer, arr, None, None] self.arr_position = 1 self.inb_position = 0 else: if self.stridefunction: buffers = [None, None,inbuffer, arr] self.arr_position = 3 self.inb_position = 2 else: if inbuffer is not None: raise ValueError( "Specified both input and output but no conversion implied") buffers = [None, None, None, None] inbuffer = arr self.conversion_required = 0 self.buffers = buffers self.result_buff = inbuffer if _PROTOTYPE: # The following method has been re-implemented in C. def rebuffer(self, arr, inbuffer=None): if not self.conversion_required: self.result_buff = arr return arr l = self.buffers l[self.arr_position] = arr if not self.generated and inbuffer is not None: self.result_buff = inbuffer l[self.inb_position] = inbuffer self.buffers = l if arr is not None and self.bytestrides != [None, None]: if self.generated: ustrides = _stridesFromShape(arr._shape, arr._itemsize) else: if inbuffer is not None: ustrides = inbuffer._strides else: ustrides = self.bytestrides[ not self.direction ] if self.direction: self.bytestrides = [ ustrides, arr._strides ] else: self.bytestrides = [ arr._strides, ustrides ] return self.result_buff def convert(self, buf, indices, shape): """Perform type conversions.""" if self.convfunction: niter = _gen.product(shape) inoffset = self.buffers[buf*2+0]._getByteOffset(indices) offset = self.buffers[buf*2+1]._getByteOffset(indices) self.convfunction(niter, 1, 1, ((self.buffers[buf*2+0]._data, inoffset), (self.buffers[buf*2+1]._data, offset))) def stride(self, buf, indices, shape): """Perform copies and re-alignments.""" if self.stridefunction: inoffset = self.buffers[buf*2+0]._getByteOffset(indices) offset = self.buffers[buf*2+1]._getByteOffset(indices) self.stridefunction(shape, self.buffers[buf*2+0]._data, inoffset, self.bytestrides[0][-len(shape):], self.buffers[buf*2+1]._data, offset, self.bytestrides[1][-len(shape):]) # The following code has been implemented in C. # Renaming "compute" below affects compute, stride, and convert. def compute(self, indices, shape): """ Perform the necessary transformations on the array. """ if self.direction: self.convert(0, indices, shape) self.stride(1, indices, shape) else: self.stride(0, indices, shape) self.convert(1, indices, shape) class _InputConverter(_Converter): """This class handles alignment, byteswaping, copying and type conversions Stride here is a bit of a misnomer, it refers to one of functions (three currently) that deal with strides. These are: byteswap, align, copy operations. The conv functions are type conversion operations. They never deal with strides. At most one of each is necessary; neither is required. Strides and offsets are *byte*! """ def __init__(self, arr, type, inbuffer=None, forcestride=0): _Converter.__init__(self, arr, type, inbuffer, forcestride) self.direction = 0 # input if type != arr._type.name: self.convfunction = arr._type._conv.astype[type] l = self.bytestrides l.reverse() self.bytestrides = l l = self.buffers l.reverse() self.buffers = l if self.conversion_required: self.arr_position = {1:2, 3:0}[self.arr_position] self.inb_position = {0:3, 2:1}[self.inb_position] class _OutputConverter(_Converter): """This class handles alignment, byteswaping, copying and type conversions processing the output of a ufunc. """ def __init__(self, arr, type, inbuffer=None, forcestride=0, nonstriding=0): _Converter.__init__(self, arr, type, inbuffer, forcestride, nonstriding) self.direction = 1 # output # The following class has been re-implemented in C. See _operatormodule.c. class _Operator: """performs the operation""" def __init__(self, cfunction, inputs, outputs, striding=0): self.cfunction = cfunction self.inputs = inputs self.outputs = outputs self.striding = striding def compute(self, indices, shape): niter = _gen.product(shape) args = [] for item in self.inputs+self.outputs: if item._strides: offset = item._getByteOffset(indices) else: offset = 0 args.append((item._data, offset)) if self.striding: if item._strides: args.append(item._strides) else: args.append(_stridesFromShape(shape, self.striding)) if not self.striding: self.cfunction(niter, len(self.inputs), len(self.outputs), tuple(args)) else: self.cfunction(shape, args[0][0], args[0][1], args[1], args[2][0], args[2][1], args[3]) # Replace the Python version of "_Operator" with a C version. if not _PROTOTYPE: from _operator import _operator as _Operator def _stridesFromShape(shape, bytestride): """Compute the strides from shape for a contiguous array, sort of""" if shape != (): ndim = len(shape) strides = [0]*ndim strides[-1] = bytestride for i in xrange(ndim-2, -1, -1): strides[i] = strides[i+1] * shape[i+1] else: strides = () return strides def _makeCUFuncDict(functionDict): """Organize C functions by function name and signature""" dict = {} for keystr in functionDict.keys(): operator, optype, signature = eval(keystr) if not dict.has_key(operator): dict[operator] = {} if not dict[operator].has_key(optype): dict[operator][optype] = {} dict[operator][optype][signature[0]] = (signature[1],functionDict[keystr]) return dict def _isScalar(x): return _nc.PyNUMERIC_TYPES.has_key(type(x)) def _maxPopType(xs): """_maxPopType determines the maximum type of a sequence of numarray and scalars. see _nt.genericCoercions. """ if isinstance(xs[0], types.InstanceType) and \ not isinstance(xs[0], _nc.NumArray): return None # If they're not all NumArrays, quit now. maxT = None for x in xs: if isinstance(x, _nc.NumArray): t = x._type elif _nc.PyNUMERIC_TYPES.has_key(type(x)): t = type(x) else: t = _nc._maxtype(x) if maxT is None: maxT = t else: maxT = _nt.genericCoercions[(maxT, t)] return maxT CLIP = 0 # Peg indices > N-1 at N-1, indices < 0 at 0 WRAP = 1 # Index = Index % N RAISE = 2 # Raise an exception for out of range indices class _ChooseUFunc(_UFunc): """Class for building the "choose" ufunc. Class for ufuncs with 2 input (S, [~S...]) and 1 output (S) arguments""" def __init__(self, operator = "choose"): self.operator = operator def _doit(self, computation_mode, outarr, cfunc, ufargs, scalar): _ufunc.CheckFPErrors() if computation_mode == "fast": apply(cfunc, ufargs) else: inputs, outputs, preprocessing_outbuffers, \ postprocessing_inbuffers, maxitemsize = ufargs niter = _ufunc.getBufferSize()/maxitemsize outshape = outarr._shape indexlevel, blockingparameters = \ _ufunc._getBlockingParameters(outshape, niter) operator = _Operator(cfunc, preprocessing_outbuffers, postprocessing_inbuffers) objects = inputs + (operator,)+ outputs _callOverDimensions(objects, outshape, indexlevel, blockingparameters) errorstatus = _ufunc.CheckFPErrors() if errorstatus: handleError(errorstatus, " in "+self.operator) if scalar: return outarr[()] else: return outarr def __call__(self, inarr1, inarr2, outarr=None, clipmode=CLIP): """The standard calling interface for UFuncs""" computation_mode, woutarr, cfunc, ufargs = \ self._setup(inarr1, inarr2, outarr, clipmode) result = self._doit(computation_mode, woutarr, cfunc, ufargs, 0) if outarr is None: return result else: return None def _setup(self, in1, in2, outarr=None, clipmode=CLIP): """Setup for Nary ufunc""" in1 = _nc.asarray(in1, type=MaybeLong) in2 = list(in2) if outarr is None: convType = _maxPopType(in2) else: convType = outarr._type for i in range(len(in2)): if _isScalar(in2[i]): in2[i] = _nc.array([in2[i]], type=convType) elif isinstance(in2[i], _gen.NDArray): pass else: in2[i] = _nc.array(in2[i], type=convType) result = _gen._nWayBroadcast( [in1] + in2 ) in1, in2 = result[0], result[1:] if outarr is None: outarr = in1.new(convType) if in1._shape != outarr._shape: raise ValueError("Supplied output array does not have appropriate shape") N = outarr._itemsize if N in [1,2,4,8,16]: ucfname = "choose" + `N` + "bytes" ucfunc = _bytes.functionDict[ucfname] else: ucfunc = _bytes.functionDict["chooseNbytes"] # pars = clipmode, population count, itemsize pars = _pseudoArrayFromScalars([clipmode, len(in2),N], _nt.MaybeLong) args = [pars, in1] + in2 + [outarr] fastargs = reduce(and_, [ _noConversionNeeded(x, convType) for x in in2 + [outarr]]) fastargs = fastargs and _noConversionNeeded(in1, _nt.MaybeLong) if fastargs: # If fast case, return the arguments for an immediate call to the # C function assert outarr._byteoffset == 0 fastparms = tuple([ (x._data, x._byteoffset) for x in args]) return ("fast", outarr, ucfunc, (outarr.nelements(), len(args)-1, 1, fastparms)) # slower, general case cpars = _InputConverter(pars, _nt.MaybeLong) selector = _InputConverter(in1, _nt.MaybeLong) population = [_InputConverter(x, convType) for x in in2 ] inputs = (cpars, selector,) + tuple(population) outputs = (_OutputConverter(arr=outarr, type=convType),) alltypes = [_numtypedict[convType], _nt.MaybeLong] + \ [a.type() for a in in2] maxitemsize = max([t.bytes for t in alltypes]) preprocessing_outbuffers = tuple([x.result_buff for x in inputs]) postprocessing_inbuffers = tuple([x.result_buff for x in outputs]) return ("slow", outarr, ucfunc, (inputs, outputs, preprocessing_outbuffers, postprocessing_inbuffers, maxitemsize)) _choose = _ChooseUFunc() def choose(selector, population, outarr=None, clipmode=CLIP): """ choose() 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). clipmode=CLIP if clipmode is CLIP, out of range selector values are clipped at [0, shape[i]). if clipmode is WRAP, out of range selector values are wrapped around from 0 to shape[i] (<0) or from shape[i] to 0 (>= shape[i]). """ return _choose(selector, population, outarr, clipmode) def _scatteredPseudos( scattered ): scatteredStrides = _pseudoArrayFromScalars(scattered._strides, type=MaybeLong) scatteredShape = _pseudoArrayFromScalars(scattered._shape, type=MaybeLong) return scatteredStrides, scatteredShape class _TakeUFunc(_ChooseUFunc): """take(scattered, indexArrays, gathered=None) scattered: input, array from which elements are "gathered" indexArrays: input, tuple of index numarray or scalars gathered: input/output result array take plucks the elements of 'scattered' specified by the tuple of index numarray, 'indexArrays', and stores the result in the output array 'gathered'. 'indexArrays' may be partial, i.e. incompletely specified, hence the shape of 'gathered' is derived from both the shape of 'indexArrays' (primarily) and the shape of the trailing unspecified dimensions of 'scattered'. """ def _setup(self, scattered, indexArrays, gathered=None, clipmode=RAISE): """Setup for Nary ufunc """ scattered = _nc.asarray(scattered) indexArrays = list(indexArrays) gatheredWasSpecified = (gathered is not None) if gatheredWasSpecified: ctype = gathered._type else: ctype = scattered._type # *All* of the scattered array must be simultaneously available. # So blocked conversions won't work. if not _noConversionNeeded(scattered, ctype): scattered = scattered.astype(ctype) impliedShape, N = _gen._takeShape(scattered, indexArrays) scatteredStrides, scatteredShape = _scatteredPseudos(scattered) if gatheredWasSpecified: if gathered._shape != impliedShape: raise ValueError("Inconsistent array shapes...") else: gathered = scattered.__class__(shape=impliedShape, type=scattered._type) # scattered and gathered should now be identically typed. # Since N can be arbitrarily large depending on the trailing # dimensions of scattered, may have to make gathered contiguous # to ensure a buffer large enough to store at least N. # if (N > _ufunc.getBufferSize() if not _noConversionNeeded(gathered, scattered._type): raise ValueError("take() destination array must be contiguous, aligned, not byteswapped, and of the same type as the source array") niter = (gathered.nelements()*gathered._itemsize)/N ucfunc = _bytes.functionDict[self.operator + "Nbytes"] pars = _pseudoArrayFromScalars([clipmode, N], type=MaybeLong) # Index arrays must be aligned, not byteswapped, and correctly typed. # Throw in contiguous for now. for i in range(len(indexArrays)): x = indexArrays[i] if not _noConversionNeeded(x, MaybeLong): indexArrays[i] = x.astype(MaybeLong) args = [pars, scattered, scatteredStrides, scatteredShape] + \ indexArrays + [gathered] fastparms = tuple([ (x._data, x._byteoffset) for x in args]) return ("fast", gathered, ucfunc, (niter, len(args)-1, 1, fastparms)) _take = _TakeUFunc("take") def take(array, indices, axis=0, outarr=None, clipmode=RAISE): """take() picks elements of 'array' specified by sequence of numerical sequences 'indices'. parameters which must be specified by keyword: array data to be indexed & collected (taken from). indices An integer sequence, or tuple of integer sequences specifying the array coordinates from which data is to be taken. Partial indices result in entire inner blocks being taken. axis=0 selects the axis (or axes) along which the take should be performed. 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]). """ if axis == 0: array = _nc.asarray(array) return array._take((indices,), outarr=outarr, clipmode=clipmode) elif isinstance(axis, (int, long)): if isinstance(indices, (int,long,float)): raise ValueError("indices must be a sequence") work = _gen.swapaxes(array, 0, axis) work = work._take((indices,), outarr=outarr, clipmode=clipmode) return _gen.swapaxes(work, 0, axis) else: def_axes = range(array.rank) for x in axis: def_axes.remove(x) axis = list(axis) + def_axes work = _gen.transpose(array, axis) return work._take(indices, outarr=outarr, clipmode=clipmode) def _nonIteratedArray(a): ni = a.view() ni._shape=(0,) ni._strides = (0,) ni._contiguous = 1 return ni class _PutUFunc(_TakeUFunc): """put(scattered, indexArrays, gathered) is the inverse function of "take", and scatters the elements of array 'gathered' across array 'scattered' as specified by tuple of index numarray, 'indexArrays'. scattered: input/output Array where "stuff" gets put indexArrays: input Arrays of indexArrays of scattered where corresponding pieces of gathered go. gathered: input Source array gathered can be an array which is broadcastable to the array shape implied by scattered and indexArrays -or- gathered can be a scalar. """ def _setup(self, scattered, indexArrays, gathered, clipmode=RAISE): """Setup for Nary ufunc """ if not isinstance(scattered, _gen.NDArray): raise TypeError("put only works on NDArray and its subclasses") indexArrays = list(indexArrays) impliedShape, N = _gen._takeShape(scattered, indexArrays) scatteredStrides, scatteredShape = _scatteredPseudos(scattered) scattered = _nonIteratedArray(scattered) pars = _pseudoArrayFromScalars([clipmode, N], type=MaybeLong) gathered = _nc.asarray(gathered, type=scattered._type) gathered = _gen._broadcast(gathered, impliedShape) if not gathered.iscontiguous(): gathered = gathered.copy() # gathered and scattered should now be identically typed. # gathered need not be contiguous since it is "bufferable". # Since N can be arbitrarily large depending on the trailing # dimensions of scattered, may have to pre-process gathered. # if N > _ufunc.getBufferSize() and \ if not _noConversionNeeded(gathered, scattered._type): gathered = gathered.astype( scattered._type ) niter = (gathered.nelements()*gathered._itemsize)/N ucfunc = _bytes.functionDict[self.operator + "Nbytes"] for i in range(len(indexArrays)): x = indexArrays[i] if not _noConversionNeeded(x, MaybeLong): indexArrays[i] = x.astype(MaybeLong) args = [pars, gathered, scatteredStrides, scatteredShape] + \ indexArrays + [scattered] fastparms = tuple([ (x._data, x._byteoffset) for x in args]) return ("fast", scattered, ucfunc, (niter, len(args)-1, 1, fastparms)) _put = _PutUFunc("put") def put(array, indices, values, axis=0, clipmode=RAISE): """put(array, indices, values, clipmode=(CLIP|WRAP), axis=0) stores 'values' into 'array' at 'indices...'. parameters which must be specified by keyword: array data to be indexed & stuffed (put to). indices An integer sequence, or tuple of integer sequences specifying the array coordinates to hich data is to be put. Partial indices result in entire inner blocks being overwritten. values A sequence of values to be written to the specified indices of array. axis=0 selects the axis along which the put should be performed. 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]) """ work = _nc.asarray(array) if axis == 0: work._put((indices,), values, clipmode=clipmode) elif isinstance(axis, (int, long)): if isinstance(indices, (int,long,float)): raise ValueError("indices must be a sequence") work = _gen.swapaxes(work, 0, axis) work._put((indices,), values, clipmode=clipmode) work = _gen.swapaxes(work, 0, axis) else: def_axes = range(work.rank) for x in axis: def_axes.remove(x) axis = list(axis) + def_axes work = _gen.transpose(work, axis) work._put(indices, values, clipmode=clipmode) work = _gen.transpose(work, axis) if work is not array: if isinstance(array, _gen.NDArray): array[:] = work class _NonzeroUFunc: """nonzero(array) array: input, array scanned for non-zero elements Nonzero scans "array" for nonzero elements, and returns a tuple of coordinate numarray corresponding to the nonzero elements. """ def __call__(self, array): """The standard calling interface for UFuncs""" computation_mode, outarr, cfunc, ufargs = self._setup(array) apply(cfunc, ufargs) return tuple(outarr) def _setup(self, inarr1): """Setup for Nonzero """ nz = _nc.asarray(not_equal(inarr1, 0)) if nz.rank < 1: raise ValueError("Can't compute nonzero of a rank-0 array") nonzeroCount = add.reduce(_gen.ravel(nz).astype(MaybeLong)) outarr = [ _nc.NumArray(shape=(nonzeroCount,), type=Long) for i in range(nz.rank) ] sstrides = _nc.array(nz._strides, type=MaybeLong) sstrides /= nz._bytestride niter = nz.nelements() ucfunc = _sort.functionDict[repr(("Bool", "nonzeroCoords"))] args = [ nz, sstrides] + outarr fastparms = tuple([ (x._data, x._byteoffset) for x in args]) return ("fast", outarr, ucfunc, (niter, 2, len(args)-2, fastparms)) def _sort1(inarr1): """1D in-place Sort""" conversionNeeded = not _noConversionNeeded(inarr1, inarr1._type) contigarr1 = _nc.array(inarr1, copy=conversionNeeded) cfunc = _sort.functionDict[repr((contigarr1._type.name, 'sort'))] cfunc(contigarr1.nelements(), 0, 1, ((contigarr1._data, contigarr1._byteoffset),)) if conversionNeeded: inarr1._copyFrom(contigarr1) def searchsorted(bins, values): """searchsort(bins, values) returns the array A[j] of greatest indices 'i' such that each values[j] <= bins[i]. """ bins = _nc.asarray(bins) values = _nc.asarray(values) if len(bins._shape) > 1 or len(values._shape) > 1: raise ValueError("Multi-dimensional searchsort not supported.") outarr = _nc.NumArray(shape=values.shape, type=Long) if not _noConversionNeeded(bins, bins._type): bins = bins.copy() if not _noConversionNeeded(values, bins._type): values = values.astype(bins._type) cfunc = _sort.functionDict[repr((bins._type.name, 'searchsorted'))] nbins = _pseudoArrayFromScalars([len(bins)], MaybeLong) args = [(x._data, x._byteoffset) for x in [nbins, bins, values, outarr]] cfunc(values.nelements(), 3, 1, tuple(args)) if outarr.rank == 0: return outarr[()] else: return outarr def _argsort1(inarr1, witness=None): """1D ArgSort""" niter = inarr1.nelements() inarr1 = _nc.array(inarr1) # make a discardable copy if witness is None: witness =_nc.arange(niter,type=Long) elif not _noConversionNeeded(witness, Long): raise ValueError("No support for converting the witness array") cfunc = _sort.functionDict[repr((inarr1._type.name,'asort'))] cfunc(niter, 0, 2, tuple([ (x._data, x._byteoffset) for x in [inarr1, witness]])) return witness # ===================================================================== nonzero = _NonzeroUFunc() def _fbroadcast(f, N, shape, args, params=()): """_fbroadcast(f, N, args, shape, params=()) calls 'f' for each of the 'N'-dimensional inner subnumarray of 'args'. Each subarray has .shape == 'shape'[-N:]. There are a total of product(shape[:-N]) calls to 'f'. """ if len(shape) == N: apply(f, tuple(args)+params) else: for i in range(shape[0]): _fbroadcast(f, N, shape[1:], [x[i] for x in args], params) def _sortN(a): """_sortN implements N-D in-place sort in terms of 1D sort""" _fbroadcast(_sort1, 1, a._shape, (a,)) def _argsortN(a, w): """_argsortN implements N-D in-place argsort in terms of 1D argsort""" _fbroadcast(_argsort1, 1, a._shape, (a,w)) def _chktype(a): if isinstance(a._type, _nt.IntegralType): artype = _nt.Long elif isinstance(a._type, _nt.FloatingType): artype = _nt.Float elif isinstance(a._type, _nt.ComplexType): artype = _nt.Complex else: raise ValueError("Only integer, float, and complex matrixmultiply supported.") return artype def _matrixtype(a,b): artype, brtype = _chktype(a), _chktype(b) return max(artype, brtype) def innerproduct(a,b): """innerproduct(a,b) computes dot products of the last dimensions of MxL matrix 'a' and NxL matrix 'b' resulting in an MxN matrix. innerproduct coerces its inputs to the maximum type of the inputs' typeclass, either Long, Float, or Complex """ a = _nc.asarray(a) b = _nc.asarray(b) if a.rank == 0: a._shape = (1,) a._strides = (a._itemsize,) if b.rank == 0: b._shape = (1,) b._strides = (b._itemsize,) if a._shape[-1] != b._shape[-1]: raise ValueError("incompatible array dimensions.") rtype = _matrixtype(a,b) if not _noConversionNeeded(a, rtype): a = a.astype(rtype) if not _noConversionNeeded(b, rtype): b = b.astype(rtype) adots, bdots = a._shape[:-1], b._shape[:-1] r = a.__class__(shape=adots+bdots, type=rtype) maxijk = _pseudoArrayFromScalars([_gen.product(adots), _gen.product(bdots), a._shape[-1]], MaybeLong) _ufunc.CheckFPErrors() cfunc = _sort.functionDict[repr((rtype.name, "mxmultiply"))] cfunc(1, 3, 1, tuple([(x._data, x._byteoffset) for x in \ [ a, b, maxijk, r]])) errorstatus = _ufunc.CheckFPErrors() if errorstatus: handleError(errorstatus, " in innerproduct") if len(adots+bdots): return r else: return r[()] def divide_remainder(a,b): """returns (a/b, a%b).""" a, b = _nc.asarray(a), _nc.asarray(b) return (a/b,a%b) Long = _nt.Long MaybeLong = _nt.MaybeLong # module-level variables _numtypedict = _nt.typeDict # Short term solution -- identity needs to be defined in _ufuncall module XXX _identities = {"add":0, "subtract":0, "multiply":1, "divide":1, "bitwise_and":1, "bitwise_or":0, "bitwise_xor":0, "logical_and":1, "logical_or":0, "logical_xor":0, } _cufuncs = _makeCUFuncDict(_ufuncall.functionDict) _UFuncs = {} for operator in _cufuncs.keys(): _UFuncs[operator] = ufuncFactory(operator, _cufuncs[operator], identity=_identities.get(operator)) globals().update(_UFuncs) del operator, _UFuncs