diff options
author | Travis Oliphant <oliphant@enthought.com> | 2005-09-14 18:55:03 +0000 |
---|---|---|
committer | Travis Oliphant <oliphant@enthought.com> | 2005-09-14 18:55:03 +0000 |
commit | 951cdad1a9b0a15b673bf56aecc1fef5066061f3 (patch) | |
tree | f307f0f2ebf404526fcc671958e8ea836b53baf9 /base/function_base.py | |
parent | c261078e9c2d8fd958bf8b3758cdd07add5fc2e3 (diff) | |
download | numpy-951cdad1a9b0a15b673bf56aecc1fef5066061f3.tar.gz |
renamed directories
Diffstat (limited to 'base/function_base.py')
-rw-r--r-- | base/function_base.py | 536 |
1 files changed, 536 insertions, 0 deletions
diff --git a/base/function_base.py b/base/function_base.py new file mode 100644 index 000000000..369dd9f92 --- /dev/null +++ b/base/function_base.py @@ -0,0 +1,536 @@ +import types +import numerix as _nx +from numerix import ravel, nonzero, array, choose, ones, zeros, \ + sometrue, alltrue, reshape, alter_numeric, restore_numeric, arraymap, \ + pi, _insert, multiply, add, arctan2, maximum, minimum, any, all +from type_check import ScalarType, isscalar, asarray +from shape_base import squeeze, atleast_1d + +__all__ = ['round','logspace','linspace','fix','mod', + 'select','trim_zeros','amax','amin', 'alen', + 'ptp','cumsum','take', 'copy', + 'prod','cumprod','diff','gradient','angle','unwrap','sort_complex', + 'disp','unique','extract','insert','nansum','nanmax','nanargmax', + 'nanargmin','nanmin','sum','vectorize','asarray_chkfinite', + 'alter_numeric', 'restore_numeric','isaltered'] + +def isaltered(): + val = str(type(_nx.array([1]))) + return 'scipy' in val + +round = _nx.around + +def asarray_chkfinite(x): + """Like asarray except it checks to be sure no NaNs or Infs are present. + """ + x = asarray(x) + if not all(_nx.isfinite(x)): + raise ValueError, "Array must not contain infs or nans." + return x + +# Need this to change array type for low precision values +def sum(x,axis=0): # could change default axis here + x = asarray(x) + if x.typecode() in ['1','s','b','w']: + x = x.astype('l') + return _nx.sum(x,axis) + + +def logspace(start,stop,num=50,endpoint=1): + """ Evenly spaced samples on a logarithmic scale. + + Return num evenly spaced samples from 10**start to 10**stop. If + endpoint=1 then last sample is 10**stop. + """ + if num <= 0: return array([]) + if endpoint: + step = (stop-start)/float((num-1)) + y = _nx.arange(0,num) * step + start + else: + step = (stop-start)/float(num) + y = _nx.arange(0,num) * step + start + return _nx.power(10.0,y) + +def linspace(start,stop,num=50,endpoint=1,retstep=0): + """ Evenly spaced samples. + + Return num evenly spaced samples from start to stop. If endpoint=1 then + last sample is stop. If retstep is 1 then return the step value used. + """ + if num <= 0: return array([]) + if endpoint: + step = (stop-start)/float((num-1)) + y = _nx.arange(0,num) * step + start + else: + step = (stop-start)/float(num) + y = _nx.arange(0,num) * step + start + if retstep: + return y, step + else: + return y + +def fix(x): + """ Round x to nearest integer towards zero. + """ + x = asarray(x) + y = _nx.floor(x) + return _nx.where(x<0,y+1,y) + +def mod(x,y): + """ x - y*floor(x/y) + + For numeric arrays, x % y has the same sign as x while + mod(x,y) has the same sign as y. + """ + return x - y*_nx.floor(x*1.0/y) + +def select(condlist, choicelist, default=0): + """ Returns an array comprised from different elements of choicelist + depending on the list of conditions. + + condlist is a list of condition arrays containing ones or zeros + + choicelist is a list of choice matrices (of the "same" size as the + arrays in condlist). The result array has the "same" size as the + arrays in choicelist. If condlist is [c0,...,cN-1] then choicelist + must be of length N. The elements of the choicelist can then be + represented as [v0,...,vN-1]. The default choice if none of the + conditions are met is given as the default argument. + + The conditions are tested in order and the first one statisfied is + used to select the choice. In other words, the elements of the + output array are found from the following tree (notice the order of + the conditions matters): + + if c0: v0 + elif c1: v1 + elif c2: v2 + ... + elif cN-1: vN-1 + else: default + + Note, that one of the condition arrays must be large enough to handle + the largest array in the choice list. + """ + n = len(condlist) + n2 = len(choicelist) + if n2 != n: + raise ValueError, "List of cases, must be same length as the list of conditions." + choicelist.insert(0,default) + S = 0 + pfac = 1 + for k in range(1,n+1): + S += k * pfac * asarray(condlist[k-1]) + if k < n: + pfac *= (1-asarray(condlist[k-1])) + # handle special case of a 1-element condition but + # a multi-element choice + if type(S) in ScalarType or max(asarray(S).shape)==1: + pfac = asarray(1) + for k in range(n2+1): + pfac = pfac + asarray(choicelist[k]) + S = S*ones(asarray(pfac).shape) + return choose(S, tuple(choicelist)) + +def _asarray1d(arr): + """Ensure 1d array for one array. + """ + m = asarray(arr) + if len(m.shape)==0: + m = reshape(m,(1,)) + return m + +def copy(a): + """Return an array copy of the object. + """ + return array(a,copy=1) + +def take(a, indices, axis=0): + """Selects the elements in indices from array a along given axis. + """ + try: + a = _nx.take(a,indices,axis) + except ValueError: # a is scalar + pass + return a + +def _no_axis_is_all(function, m, axis): + if axis is None: + m = ravel(m) + axis = 0 + else: + m = _asarray1d(m) + if _nx.which[0] == "numeric": + r = function(m, axis) + else: + import numarray as _na + _na.Error.pushMode(overflow="raise") + try: + r = function(m, axis) + finally: + _na.Error.popMode() + return r + +# Basic operations +def amax(m,axis=-1): + """Returns the maximum of m along dimension axis. + """ + return _no_axis_is_all(maximum.reduce, m, axis) + +def amin(m,axis=-1): + """Returns the minimum of m along dimension axis. + """ + return _no_axis_is_all(minimum.reduce, m, axis) + +def alen(m): + """Returns the length of a Python object interpreted as an array + """ + return len(asarray(m)) + +# Actually from Basis, but it fits in so naturally here... + +def _amin_amax(m, axis): + return amax(m,axis)-amin(m,axis) + +def ptp(m,axis=-1): + """Returns the maximum - minimum along the the given dimension + """ + return _no_axis_is_all(_amin_amax, m, axis) + +def cumsum(m,axis=-1): + """Returns the cumulative sum of the elements along the given axis + """ + return _no_axis_is_all(add.accumulate, m, axis) + +def prod(m,axis=-1): + """Returns the product of the elements along the given axis + """ + return _no_axis_is_all(multiply.reduce, m, axis) + +def cumprod(m,axis=-1): + """Returns the cumulative product of the elments along the given axis + """ + return _no_axis_is_all(multiply.accumulate, m, axis) + +def gradient(f,*varargs): + """Calculate the gradient of an N-dimensional scalar function. + + Uses central differences on the interior and first differences on boundaries + to give the same shape. + + Inputs: + + f -- An N-dimensional array giving samples of a scalar function + + varargs -- 0, 1, or N scalars giving the sample distances in each direction + + Outputs: + + N arrays of the same shape as f giving the derivative of f with respect + to each dimension. + """ + N = len(f.shape) # number of dimensions + n = len(varargs) + if n==0: + dx = [1.0]*N + elif n==1: + dx = [varargs[0]]*N + elif n==N: + dx = list(varargs) + else: + raise SyntaxError, "Invalid number of arguments" + + # use central differences on interior and first differences on endpoints + + print dx + outvals = [] + + # create slice objects --- initially all are [:,:,...,:] + slice1 = [slice(None)]*N + slice2 = [slice(None)]*N + slice3 = [slice(None)]*N + + otype = f.typecode() + if otype not in ['f','d','F','D']: + otype = 'd' + + for axis in range(N): + # select out appropriate parts for this dimension + out = zeros(f.shape, f.typecode()) + slice1[axis] = slice(1,-1) + slice2[axis] = slice(2,None) + slice3[axis] = slice(None,-2) + # 1d equivalent -- out[1:-1] = (f[2:] - f[:-2])/2.0 + out[slice1] = (f[slice2] - f[slice3])/2.0 + slice1[axis] = 0 + slice2[axis] = 1 + slice3[axis] = 0 + # 1d equivalent -- out[0] = (f[1] - f[0]) + out[slice1] = (f[slice2] - f[slice3]) + slice1[axis] = -1 + slice2[axis] = -1 + slice3[axis] = -2 + # 1d equivalent -- out[-1] = (f[-1] - f[-2]) + out[slice1] = (f[slice2] - f[slice3]) + + # divide by step size + outvals.append(out / dx[axis]) + + # reset the slice object in this dimension to ":" + slice1[axis] = slice(None) + slice2[axis] = slice(None) + slice3[axis] = slice(None) + + if N == 1: + return outvals[0] + else: + return outvals + + +def diff(x, n=1,axis=-1): + """Calculates the nth order, discrete difference along given axis. + """ + if n==0: + return x + if n<0: + raise ValueError,'Order must be non-negative but got ' + `n` + x = _asarray1d(x) + nd = len(x.shape) + slice1 = [slice(None)]*nd + slice2 = [slice(None)]*nd + slice1[axis] = slice(1,None) + slice2[axis] = slice(None,-1) + if n > 1: + return diff(x[slice1]-x[slice2], n-1, axis=axis) + else: + return x[slice1]-x[slice2] + +def angle(z,deg=0): + """Return the angle of complex argument z.""" + if deg: + fact = 180/pi + else: + fact = 1.0 + z = asarray(z) + if z.typecode() in ['D','F']: + zimag = z.imag + zreal = z.real + else: + zimag = 0 + zreal = z + return arctan2(zimag,zreal) * fact + +def unwrap(p,discont=pi,axis=-1): + """unwraps radian phase p by changing absolute jumps greater than + discont to their 2*pi complement along the given axis. + """ + p = asarray(p) + nd = len(p.shape) + dd = diff(p,axis=axis) + slice1 = [slice(None,None)]*nd # full slices + slice1[axis] = slice(1,None) + ddmod = mod(dd+pi,2*pi)-pi + _nx.putmask(ddmod,(ddmod==-pi) & (dd > 0),pi) + ph_correct = ddmod - dd; + _nx.putmask(ph_correct,abs(dd)<discont,0) + up = array(p,copy=1,typecode='d') + up[slice1] = p[slice1] + cumsum(ph_correct,axis) + return up + +def sort_complex(a): + """ Doesn't currently work for integer arrays -- only float or complex. + """ + a = asarray(a,typecode=a.typecode().upper()) + def complex_cmp(x,y): + res = cmp(x.real,y.real) + if res == 0: + res = cmp(x.imag,y.imag) + return res + l = a.tolist() + l.sort(complex_cmp) + return array(l) + +def trim_zeros(filt,trim='fb'): + """ Trim the leading and trailing zeros from a 1D array. + + Example: + >>> import scipy + >>> a = array((0,0,0,1,2,3,2,1,0)) + >>> scipy.trim_zeros(a) + array([1, 2, 3, 2, 1]) + """ + first = 0 + if 'f' in trim or 'F' in trim: + for i in filt: + if i != 0.: break + else: first = first + 1 + last = len(filt) + if 'b' in trim or 'B' in trim: + for i in filt[::-1]: + if i != 0.: break + else: last = last - 1 + return filt[first:last] + +def unique(inseq): + """Returns unique items in 1-dimensional sequence. + """ + set = {} + for item in inseq: + set[item] = None + return asarray(set.keys()) + +def where(condition,x=None,y=None): + """If x and y are both None, then return the (1-d equivalent) indices + where condition is true. Otherwise, return an array shaped like + condition with elements of x and y in the places where condition is + true or false respectively. + """ + if (x is None) and (y is None): + # Needs work for multidimensional arrays + return nonzero(ravel(condition)) + else: + return choose(not_equal(condition, 0), (y,x)) + +def extract(condition, arr): + """Elements of ravel(condition) where ravel(condition) is true (1-d) + + Equivalent of compress(ravel(condition), ravel(arr)) + """ + return _nx.take(ravel(arr), nonzero(ravel(condition))) + +def insert(arr, mask, vals): + """Similar to putmask arr[mask] = vals but 1d array vals has the + same number of elements as the non-zero values of mask. Inverse of extract. + """ + return _nx._insert(arr, mask, vals) + +def nansum(x,axis=-1): + """Sum the array over the given axis treating nans as missing values. + """ + x = _asarray1d(x).copy() + _nx.putmask(x,isnan(x),0) + return _nx.sum(x,axis) + +def nanmin(x,axis=-1): + """Find the minimium over the given axis ignoring nans. + """ + x = _asarray1d(x).copy() + _nx.putmask(x,isnan(x),inf) + return amin(x,axis) + +def nanargmin(x,axis=-1): + """Find the indices of the minimium over the given axis ignoring nans. + """ + x = _asarray1d(x).copy() + _nx.putmask(x,isnan(x),inf) + return argmin(x,axis) + + +def nanmax(x,axis=-1): + """Find the maximum over the given axis ignoring nans. + """ + x = _asarray1d(x).copy() + _nx.putmask(x,isnan(x),-inf) + return amax(x,axis) + +def nanargmax(x,axis=-1): + """Find the maximum over the given axis ignoring nans. + """ + x = _asarray1d(x).copy() + _nx.putmask(x,isnan(x),-inf) + return argmax(x,axis) + +def disp(mesg, device=None, linefeed=1): + """Display a message to device (default is sys.stdout) with(out) linefeed. + """ + if device is None: + import sys + device = sys.stdout + if linefeed: + device.write('%s\n' % mesg) + else: + device.write('%s' % mesg) + device.flush() + return + +class vectorize: + """ + vectorize(somefunction) Generalized Function class. + + Description: + + Define a vectorized function which takes nested sequence + objects or numerix arrays as inputs and returns a + numerix array as output, evaluating the function over successive + tuples of the input arrays like the python map function except it uses + the broadcasting rules of numerix Python. + + Input: + + somefunction -- a Python function or method + + Example: + + def myfunc(a,b): + if a > b: + return a-b + else + return a+b + + vfunc = vectorize(myfunc) + + >>> vfunc([1,2,3,4],2) + array([3,4,1,2]) + + """ + def __init__(self,pyfunc,otypes=None,doc=None): + if not callable(pyfunc) or type(pyfunc) is types.ClassType: + raise TypeError, "Object is not a callable Python object." + self.thefunc = pyfunc + if doc is None: + self.__doc__ = pyfunc.__doc__ + else: + self.__doc__ = doc + if otypes is None: + self.otypes='' + else: + if isinstance(otypes,types.StringType): + self.otypes=otypes + else: + raise ValueError, "Output types must be a string." + + def __call__(self,*args): + try: + return squeeze(arraymap(self.thefunc,args,self.otypes)) + except IndexError: + return self.zerocall(*args) + + def zerocall(self,*args): + # one of the args was a zeros array + # return zeros for each output + # first --- find number of outputs + # get it from self.otypes if possible + # otherwise evaluate function at 0.9 + N = len(self.otypes) + if N==1: + return zeros((0,),'d') + elif N !=0: + return (zeros((0,),'d'),)*N + newargs = [] + args = atleast_1d(args) + for arg in args: + if arg.typecode() != 'O': + newargs.append(0.9) + else: + newargs.append(arg[0]) + newargs = tuple(newargs) + try: + res = self.thefunc(*newargs) + except: + raise ValueError, "Zerocall is failing. "\ + "Try using otypes in vectorize." + if isscalar(res): + return zeros((0,),'d') + else: + return (zeros((0,),'d'),)*len(res) + |