"""Numeric module defining a multi-dimensional array and useful procedures for Numerical computation. Functions ========= array NumPy Array construction zeros Return an array of all zeros fromstring Construct array from (byte) string take Select sub-arrays using sequence of indices put Set sub-arrays using sequence of 1-D indices putmask Set portion of arrays using a mask reshape Return array with new shape repeat Repeat elements of array choose Construct new array from indexed array tuple cross_correlate Correlate two 1-d arrays searchsorted Search for element in 1-d array sum Total sum over a specified dimension cumsum Cumulative sum over a specified dimension product Total product over a specified dimension cumproduct Cumulative product over a specified dimension alltrue Logical and over an entire axis sometrue Logical or over an entire axis allclose Tests if sequences are essentially equal arrayrange (arange) Return regularly spaced array asarray Guarantee NumPy array sarray Guarantee a NumPy array that keeps precision convolve Convolve two 1-d arrays swapaxes Exchange axes concatenate Join arrays together transpose Permute axes sort Sort elements of array argsort Indices of sorted array argmax Index of largest value argmin Index of smallest value innerproduct Innerproduct of two arrays dot Dot product (matrix multiplication) outerproduct Outerproduct of two arrays resize Return array with arbitrary new shape indices Tuple of indices fromfunction Construct array from universal function diagonal Return diagonal array trace Trace of array dump Dump array to file object (pickle) dumps Return pickled string representing data load Return array stored in file object loads Return array from pickled string ravel Return array as 1-D nonzero Indices of nonzero elements for 1-D array shape Shape of array where Construct array from binary result compress Elements of array where condition is true clip Clip array between two values zeros Array of all zeros ones Array of all ones identity 2-D identity array (matrix) (Universal) Math Functions ==================== add logical_or exp subtract logical_xor log multiply logical_not log10 divide maximum sin divide_safe minimum sinh conjugate bitwise_and sqrt power bitwise_or tan absolute bitwise_xor tanh negative invert ceil greater left_shift fabs greater_equal right_shift floor less arccos arctan2 less_equal arcsin fmod equal arctan hypot not_equal cos around logical_and cosh sign """ import numeric_version __version__ = numeric_version.version del numeric_version import multiarray from umath import * from Precision import * import _numpy # for freeze dependency resolution (at least on Mac) import string, types, math #Use this to add a new axis to an array NewAxis = None #The following functions are considered builtin, they all might be #in C some day def arrayrange(start, stop=None, step=1, typecode=None): """Just like range() except it returns a array whose type can be specfied by the keyword argument typecode. """ if (stop == None): stop = start start = 0 n = int(math.ceil(float(stop-start)/step)) if n <= 0: m = zeros( (0,) )+(step+start+stop) else: m = (add.accumulate(ones((n,), Int))-1)*step +(start+(stop-stop)) # the last bit is to deal with e.g. Longs -- 3L-3L==0L if typecode != None and m.typecode() != typecode: return m.astype(typecode) else: return m array = multiarray.array zeros = multiarray.zeros def asarray(a, typecode=None, savespace=0): """asarray(a,typecode=None) returns a as a NumPy array. Unlike array(), no copy is performed if a is already an array. """ return multiarray.array(a, typecode, copy=0, savespace=savespace) def sarray(a, typecode=None, copy=0): """sarray(a, typecode=None, copy=0) calls array with savespace=1.""" return multiarray.array(a, typecode, copy, savespace=1) fromstring = multiarray.fromstring take = multiarray.take reshape = multiarray.reshape repeat = multiarray.repeat choose = multiarray.choose cross_correlate = multiarray.cross_correlate def put (a, ind, v): """put(a, ind, v) results in a[n] = v[n] for all n in ind If v is shorter than mask it will be repeated as necessary. In particular v can be a scalar or length 1 array. """ multiarray.put (a, ind, array(v, copy=0).astype(a.typecode())) def putmask (a, mask, v): """putmask(a, mask, v) results in a = v for all places mask is true. If v is shorter than mask it will be repeated as necessary. In particular v can be a scalar or length 1 array. """ multiarray.putmask (a, mask, array(v, copy=0).astype(a.typecode())) def convolve(a,v,mode=0): """convolve(a,v,mode=0) returns the discrete, linear convolution of 1-D sequences a and v; mode can be 0 (full), 1 (same), or 2 (valid) to specify size of resulting sequence. """ if (len(v) > len(a)): temp = a a = v v = temp del temp return cross_correlate(a,asarray(v)[::-1],mode) ArrayType = multiarray.arraytype def swapaxes(a, axis1, axis2): """swapaxes(a, axis1, axis2) returns array a with axis1 and axis2 interchanged. """ n = len(shape(a)) if n <= 1: return a new_axes = arange(n) new_axes[axis1] = axis2 new_axes[axis2] = axis1 return multiarray.transpose(a, new_axes) arraytype = multiarray.arraytype #add extra intelligence to the basic C functions def concatenate(a, axis=0): """concatenate(a, axis=0) joins the tuple of sequences in a into a single NumPy array. """ if axis == 0: return multiarray.concatenate(a) else: new_list = [] for m in a: new_list.append(swapaxes(m, axis, 0)) return swapaxes(multiarray.concatenate(new_list), axis, 0) def transpose(a, axes=None): """transpose(a, axes=None) returns array with dimensions permuted according to axes. If axes is None (default) returns array with dimensions reversed. """ if axes == None: axes = arange(len(array(a).shape))[::-1] return multiarray.transpose(a, axes) def sort(a, axis=-1): """sort(a,axis=-1) returns array with elements sorted along given axis. """ if axis != -1: a = swapaxes(a, axis, -1) s = multiarray.sort(a) if axis != -1: s = swapaxes(s, axis, -1) return s def argsort(a, axis=-1): """argsort(a,axis=-1) return the indices into a of the sorted array along the given axis, so that take(a,result,axis) is the sorted array. """ if axis != -1: a = swapaxes(a, axis, -1) s = multiarray.argsort(a) if axis != -1: s = swapaxes(s, axis, -1) return s def argmax(a, axis=-1): """argmax(a,axis=-1) returns the indices to the maximum value of the 1-D arrays along the given axis. """ if axis != -1: a = swapaxes(a, axis, -1) s = multiarray.argmax(a) #probably need a swap here if > 2d if axis != -1: s = swapaxes(s, axis, -1) return s def argmin(x, axis=-1): """argmin(a,axis=-1) returns the indices to the minimum value of the 1-D arrays along the given axis. """ return argmax(negative(x), axis) searchsorted = multiarray.binarysearch def innerproduct(a,b): """innerproduct(a,b) returns the dot product of two arrays, which has shape a.shape[:-1] + b.shape[:-1] with elements computed by summing the product of the elements from the last dimensions of a and b. """ try: return multiarray.innerproduct(a,b) except TypeError,detail: if array(a).shape == () or array(b).shape == (): return a*b else: raise TypeError, detail or "invalid types for dot" def outerproduct(a,b): """outerproduct(a,b) returns the outer product of two vectors. result(i,j) = a(i)*b(j) when a and b are vectors Will accept any arguments that can be made into vectors. """ return asarray(a).flat[:,NewAxis]*asarray(b).flat[NewAxis,:] def dot(a, b): """dot(a,b) returns matrix-multiplication between a and b. The product-sum is over the last dimension of a and the second-to-last dimension of b. """ return innerproduct(a, swapaxes(b, -1, -2)) #This is obsolete, don't use in new code matrixmultiply = dot #Use Konrad's printing function (modified for both str and repr now) from ArrayPrinter import array2string def array_repr(a, max_line_width = None, precision = None, suppress_small = None): return array2string(a, max_line_width, precision, suppress_small, ', ', 1) def array_str(a, max_line_width = None, precision = None, suppress_small = None): return array2string(a, max_line_width, precision, suppress_small, ' ', 0) multiarray.set_string_function(array_str, 0) multiarray.set_string_function(array_repr, 1) #This is a nice value to have around #Maybe in sys some day LittleEndian = fromstring("\001"+"\000"*7, 'i')[0] == 1 def resize(a, new_shape): """resize(a,new_shape) returns a new array with the specified shape. The original array's total size can be any size. """ a = ravel(a) if not len(a): return zeros(new_shape, a.typecode()) total_size = multiply.reduce(new_shape) n_copies = total_size / len(a) extra = total_size % len(a) if extra != 0: n_copies = n_copies+1 extra = len(a)-extra a = concatenate( (a,)*n_copies) if extra > 0: a = a[:-extra] return reshape(a, new_shape) def indices(dimensions, typecode=None): """indices(dimensions,typecode=None) returns an array representing a grid of indices with row-only, and column-only variation. """ tmp = ones(dimensions, typecode) lst = [] for i in range(len(dimensions)): lst.append( add.accumulate(tmp, i, )-1 ) return array(lst) def fromfunction(function, dimensions): """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))) def diagonal(a, offset= 0, axis1=0, axis2=1): """diagonal(a, offset=0, axis1=0, axis2=1) returns the given diagonals defined by the last two dimensions of the array. """ a = array (a) if axis2 < axis1: axis1, axis2 = axis2, axis1 if axis2 > 1: new_axes = range (len (a.shape)) del new_axes [axis2]; del new_axes [axis1] new_axes [0:0] = [axis1, axis2] a = transpose (a, new_axes) s = a.shape if len (s) == 2: n1 = s [0] n2 = s [1] n = n1 * n2 s = (n,) a = reshape (a, s) if offset < 0: return take (a, range ( - n2 * offset, min(n2, n1+offset) * (n2+1) - n2 * offset, n2+1), 0) else: return take (a, range (offset, min(n1, n2-offset) * (n2+1) + offset, n2+1), 0) else : my_diagonal = [] for i in range (s [0]) : my_diagonal.append (diagonal (a [i], offset)) return array (my_diagonal) def trace(a, offset=0, axis1=0, axis2=1): """trace(a,offset=0, axis1=0, axis2=1) returns the sum along diagonals (defined by the last two dimenions) of the array. """ return add.reduce(diagonal(a, offset, axis1, axis2)) # These two functions are used in my modified pickle.py so that # matrices can be pickled. Notice that matrices are written in # binary format for efficiency, but that they pay attention to # byte-order issues for portability. def DumpArray(m, fp): if m.typecode() == 'O': raise TypeError, "Numeric Pickler can't pickle arrays of Objects" s = m.shape if LittleEndian: endian = "L" else: endian = "B" fp.write("A%s%s%d " % (m.typecode(), endian, m.itemsize())) for d in s: fp.write("%d "% d) fp.write('\n') fp.write(m.tostring()) def LoadArray(fp): ln = string.split(fp.readline()) if ln[0][0] == 'A': ln[0] = ln[0][1:] # Nasty hack showing my ignorance of pickle typecode = ln[0][0] endian = ln[0][1] shape = map(lambda x: string.atoi(x), ln[1:]) itemsize = string.atoi(ln[0][2:]) sz = reduce(multiply, shape)*itemsize data = fp.read(sz) m = fromstring(data, typecode) m = reshape(m, shape) if (LittleEndian and endian == 'B') or (not LittleEndian and endian == 'L'): return m.byteswapped() else: return m import pickle, copy class Unpickler(pickle.Unpickler): def load_array(self): self.stack.append(LoadArray(self)) dispatch = copy.copy(pickle.Unpickler.dispatch) dispatch['A'] = load_array class Pickler(pickle.Pickler): def save_array(self, object): DumpArray(object, self) dispatch = copy.copy(pickle.Pickler.dispatch) dispatch[ArrayType] = save_array #Convenience functions from StringIO import StringIO def dump(object, file): """dump(object, file) pickles (binary-writes) the object to an open file. """ Pickler(file).dump(object) def dumps(object): """dumps(object) pickles (binary-writes) the object and returns the byte stream. """ file = StringIO() Pickler(file).dump(object) return file.getvalue() def load(file): """load(file) returns an array from the open file pointing to pickled data. """ return Unpickler(file).load() def loads(str): """loads(str) returns an array from a byte stream containing its pickled representation. """ file = StringIO(str) return Unpickler(file).load() # slightly different format uses the copy_reg mechanism import copy_reg def array_constructor(shape, typecode, thestr, Endian=LittleEndian): x = fromstring(thestr, typecode) x.shape = shape if LittleEndian != Endian: return x.byteswapped() else: return x def pickle_array(a): return (array_constructor, (a.shape, a.typecode(), a.tostring(), LittleEndian)) copy_reg.pickle(ArrayType, pickle_array, array_constructor) # These are all essentially abbreviations # These might wind up in a special abbreviations module def ravel(m): """ravel(m) returns a 1d array corresponding to all the elements of it's argument. """ return reshape(m, (-1,)) def nonzero(a): """nonzero(a) returns the indices of the elements of a which are not zero, a must be 1d """ return repeat(arange(len(a)), not_equal(a, 0)) #Move this into C to do it right! def shape(a): """shape(a) returns the shape of a in functional form. """ return asarray(a).shape def where(condition, x, y): """where(condition,x,y) is shaped like condition and has elements of x and y where condition is respectively true or false. """ return choose(not_equal(condition, 0), (y, x)) def compress(condition, m, dimension=-1): """compress(condition, x, dimension=-1) = those elements of x corresponding to those elements of condition that are "true". condition must be the same size as the given dimension of x.""" return take(m, nonzero(condition), dimension) 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 = less(m, m_min)+2*greater(m, m_max) return choose(selector, (m, m_min, m_max)) def ones(shape, typecode='l', savespace=0): """ones(shape, typecode=Int, savespace=0) returns an array of the given dimensions which is initialized to all ones. """ a=zeros(shape, typecode, savespace) a[...]=1 return a def identity(n): """identity(n) returns the identity matrix of shape n x n. """ return resize([1]+n*[0], (n,n)) sum = add.reduce cumsum = add.accumulate product = multiply.reduce cumproduct = multiply.accumulate alltrue = logical_and.reduce sometrue = logical_or.reduce arange = arrayrange def around(m, decimals=0): """around(m, decimals=0) \ Round in the same way as standard python performs rounding. Returns always a float. """ m = asarray(m) s = sign(m) if decimals: m = absolute(m*10.**decimals) else: m = absolute(m) rem = m-asarray(m).astype(Int) m = where(less(rem,0.5), floor(m), ceil(m)) # convert back if decimals: m = m*s/(10.**decimals) else: m = m*s return m def sign(m): """sign(m) gives an array with shape of m with elements defined by sign function: where m is less than 0 return -1, where m greater than 0, a=1, elsewhere a=0. """ m = asarray(m) return zeros(shape(m))-less(m,0)+greater(m,0) def allclose (a, b, rtol=1.e-5, atol=1.e-8): """ allclose(a,b,rtol=1.e-5,atol=1.e-8) Returns true if all components of a and b are equal subject to given tolerances. The relative error rtol must be positive and << 1.0 The absolute error atol comes into play for those elements of y that are very small or zero; it says how small x must be also. """ x = array(a, copy=0) y = array(b, copy=0) d = less(absolute(x-y), atol + rtol * absolute(y)) return alltrue(ravel(d))