diff options
Diffstat (limited to 'shape_base.py')
| -rw-r--r-- | shape_base.py | 633 |
1 files changed, 633 insertions, 0 deletions
diff --git a/shape_base.py b/shape_base.py new file mode 100644 index 000000000..9ac77666f --- /dev/null +++ b/shape_base.py @@ -0,0 +1,633 @@ +__all__ = ['atleast_1d','atleast_2d','atleast_3d','vstack','hstack', + 'column_stack','row_stack', 'dstack','array_split','split','hsplit', + 'vsplit','dsplit','apply_over_axes','expand_dims', + 'apply_along_axis', 'kron', 'tile', 'get_array_wrap'] + +import numpy.core.numeric as _nx +from numpy.core.numeric import asarray, zeros, newaxis, outer, \ + concatenate, isscalar, array, asanyarray +from numpy.core.fromnumeric import product, reshape + +def apply_along_axis(func1d,axis,arr,*args): + """ Execute func1d(arr[i],*args) where func1d takes 1-D arrays + and arr is an N-d array. i varies so as to apply the function + along the given axis for each 1-d subarray in arr. + """ + arr = asarray(arr) + nd = arr.ndim + if axis < 0: + axis += nd + if (axis >= nd): + raise ValueError("axis must be less than arr.ndim; axis=%d, rank=%d." + % (axis,nd)) + ind = [0]*(nd-1) + i = zeros(nd,'O') + indlist = range(nd) + indlist.remove(axis) + i[axis] = slice(None,None) + outshape = asarray(arr.shape).take(indlist) + i.put(indlist, ind) + res = func1d(arr[tuple(i.tolist())],*args) + # if res is a number, then we have a smaller output array + if isscalar(res): + outarr = zeros(outshape,asarray(res).dtype) + outarr[tuple(ind)] = res + Ntot = product(outshape) + k = 1 + while k < Ntot: + # increment the index + ind[-1] += 1 + n = -1 + while (ind[n] >= outshape[n]) and (n > (1-nd)): + ind[n-1] += 1 + ind[n] = 0 + n -= 1 + i.put(indlist,ind) + res = func1d(arr[tuple(i.tolist())],*args) + outarr[tuple(ind)] = res + k += 1 + return outarr + else: + Ntot = product(outshape) + holdshape = outshape + outshape = list(arr.shape) + outshape[axis] = len(res) + outarr = zeros(outshape,asarray(res).dtype) + outarr[tuple(i.tolist())] = res + k = 1 + while k < Ntot: + # increment the index + ind[-1] += 1 + n = -1 + while (ind[n] >= holdshape[n]) and (n > (1-nd)): + ind[n-1] += 1 + ind[n] = 0 + n -= 1 + i.put(indlist, ind) + res = func1d(arr[tuple(i.tolist())],*args) + outarr[tuple(i.tolist())] = res + k += 1 + return outarr + + +def apply_over_axes(func, a, axes): + """Apply a function repeatedly over multiple axes, keeping the same shape + for the resulting array. + + func is called as res = func(a, axis). The result is assumed + to be either the same shape as a or have one less dimension. + This call is repeated for each axis in the axes sequence. + """ + val = asarray(a) + N = a.ndim + if array(axes).ndim == 0: + axes = (axes,) + for axis in axes: + if axis < 0: axis = N + axis + args = (val, axis) + res = func(*args) + if res.ndim == val.ndim: + val = res + else: + res = expand_dims(res,axis) + if res.ndim == val.ndim: + val = res + else: + raise ValueError, "function is not returning"\ + " an array of correct shape" + return val + +def expand_dims(a, axis): + """Expand the shape of a by including newaxis before given axis. + """ + a = asarray(a) + shape = a.shape + if axis < 0: + axis = axis + len(shape) + 1 + return a.reshape(shape[:axis] + (1,) + shape[axis:]) + + +def atleast_1d(*arys): + """ Force a sequence of arrays to each be at least 1D. + + Description: + Force an array to be at least 1D. If an array is 0D, the + array is converted to a single row of values. Otherwise, + the array is unaltered. + Arguments: + *arys -- arrays to be converted to 1 or more dimensional array. + Returns: + input array converted to at least 1D array. + """ + res = [] + for ary in arys: + res.append(array(ary,copy=False,subok=True,ndmin=1)) + if len(res) == 1: + return res[0] + else: + return res + +def atleast_2d(*arys): + """ Force a sequence of arrays to each be at least 2D. + + Description: + Force an array to each be at least 2D. If the array + is 0D or 1D, the array is converted to a single + row of values. Otherwise, the array is unaltered. + Arguments: + arys -- arrays to be converted to 2 or more dimensional array. + Returns: + input array converted to at least 2D array. + """ + res = [] + for ary in arys: + res.append(array(ary,copy=False,subok=True,ndmin=2)) + if len(res) == 1: + return res[0] + else: + return res + +def atleast_3d(*arys): + """ Force a sequence of arrays to each be at least 3D. + + Description: + Force an array each be at least 3D. If the array is 0D or 1D, + the array is converted to a single 1xNx1 array of values where + N is the orginal length of the array. If the array is 2D, the + array is converted to a single MxNx1 array of values where MxN + is the orginal shape of the array. Otherwise, the array is + unaltered. + Arguments: + arys -- arrays to be converted to 3 or more dimensional array. + Returns: + input array converted to at least 3D array. + """ + res = [] + for ary in arys: + ary = asarray(ary) + if len(ary.shape) == 0: + result = ary.reshape(1,1,1) + elif len(ary.shape) == 1: + result = ary[newaxis,:,newaxis] + elif len(ary.shape) == 2: + result = ary[:,:,newaxis] + else: + result = ary + res.append(result) + if len(res) == 1: + return res[0] + else: + return res + + +def vstack(tup): + """ Stack arrays in sequence vertically (row wise) + + Description: + Take a sequence of arrays and stack them vertically + to make a single array. All arrays in the sequence + must have the same shape along all but the first axis. + vstack will rebuild arrays divided by vsplit. + Arguments: + tup -- sequence of arrays. All arrays must have the same + shape. + Examples: + >>> import numpy + >>> a = array((1,2,3)) + >>> b = array((2,3,4)) + >>> numpy.vstack((a,b)) + array([[1, 2, 3], + [2, 3, 4]]) + >>> a = array([[1],[2],[3]]) + >>> b = array([[2],[3],[4]]) + >>> numpy.vstack((a,b)) + array([[1], + [2], + [3], + [2], + [3], + [4]]) + + """ + return _nx.concatenate(map(atleast_2d,tup),0) + +def hstack(tup): + """ Stack arrays in sequence horizontally (column wise) + + Description: + Take a sequence of arrays and stack them horizontally + to make a single array. All arrays in the sequence + must have the same shape along all but the second axis. + hstack will rebuild arrays divided by hsplit. + Arguments: + tup -- sequence of arrays. All arrays must have the same + shape. + Examples: + >>> import numpy + >>> a = array((1,2,3)) + >>> b = array((2,3,4)) + >>> numpy.hstack((a,b)) + array([1, 2, 3, 2, 3, 4]) + >>> a = array([[1],[2],[3]]) + >>> b = array([[2],[3],[4]]) + >>> numpy.hstack((a,b)) + array([[1, 2], + [2, 3], + [3, 4]]) + + """ + return _nx.concatenate(map(atleast_1d,tup),1) + +row_stack = vstack + +def column_stack(tup): + """ Stack 1D arrays as columns into a 2D array + + Description: + Take a sequence of 1D arrays and stack them as columns + to make a single 2D array. All arrays in the sequence + must have the same first dimension. 2D arrays are + stacked as-is, just like with hstack. 1D arrays are turned + into 2D columns first. + + Arguments: + tup -- sequence of 1D or 2D arrays. All arrays must have the same + first dimension. + Examples: + >>> import numpy + >>> a = array((1,2,3)) + >>> b = array((2,3,4)) + >>> numpy.column_stack((a,b)) + array([[1, 2], + [2, 3], + [3, 4]]) + + """ + arrays = [] + for v in tup: + arr = array(v,copy=False,subok=True) + if arr.ndim < 2: + arr = array(arr,copy=False,subok=True,ndmin=2).T + arrays.append(arr) + return _nx.concatenate(arrays,1) + +def dstack(tup): + """ Stack arrays in sequence depth wise (along third dimension) + + Description: + Take a sequence of arrays and stack them along the third axis. + All arrays in the sequence must have the same shape along all + but the third axis. This is a simple way to stack 2D arrays + (images) into a single 3D array for processing. + dstack will rebuild arrays divided by dsplit. + Arguments: + tup -- sequence of arrays. All arrays must have the same + shape. + Examples: + >>> import numpy + >>> a = array((1,2,3)) + >>> b = array((2,3,4)) + >>> numpy.dstack((a,b)) + array([[[1, 2], + [2, 3], + [3, 4]]]) + >>> a = array([[1],[2],[3]]) + >>> b = array([[2],[3],[4]]) + >>> numpy.dstack((a,b)) + array([[[1, 2]], + <BLANKLINE> + [[2, 3]], + <BLANKLINE> + [[3, 4]]]) + + """ + return _nx.concatenate(map(atleast_3d,tup),2) + +def _replace_zero_by_x_arrays(sub_arys): + for i in range(len(sub_arys)): + if len(_nx.shape(sub_arys[i])) == 0: + sub_arys[i] = _nx.array([]) + elif _nx.sometrue(_nx.equal(_nx.shape(sub_arys[i]),0)): + sub_arys[i] = _nx.array([]) + return sub_arys + +def array_split(ary,indices_or_sections,axis = 0): + """ Divide an array into a list of sub-arrays. + + Description: + Divide ary into a list of sub-arrays along the + specified axis. If indices_or_sections is an integer, + ary is divided into that many equally sized arrays. + If it is impossible to make an equal split, each of the + leading arrays in the list have one additional member. If + indices_or_sections is a list of sorted integers, its + entries define the indexes where ary is split. + + Arguments: + ary -- N-D array. + Array to be divided into sub-arrays. + indices_or_sections -- integer or 1D array. + If integer, defines the number of (close to) equal sized + sub-arrays. If it is a 1D array of sorted indices, it + defines the indexes at which ary is divided. Any empty + list results in a single sub-array equal to the original + array. + axis -- integer. default=0. + Specifies the axis along which to split ary. + Caveats: + Currently, the default for axis is 0. This + means a 2D array is divided into multiple groups + of rows. This seems like the appropriate default, + """ + try: + Ntotal = ary.shape[axis] + except AttributeError: + Ntotal = len(ary) + try: # handle scalar case. + Nsections = len(indices_or_sections) + 1 + div_points = [0] + list(indices_or_sections) + [Ntotal] + except TypeError: #indices_or_sections is a scalar, not an array. + Nsections = int(indices_or_sections) + if Nsections <= 0: + raise ValueError, 'number sections must be larger than 0.' + Neach_section,extras = divmod(Ntotal,Nsections) + section_sizes = [0] + \ + extras * [Neach_section+1] + \ + (Nsections-extras) * [Neach_section] + div_points = _nx.array(section_sizes).cumsum() + + sub_arys = [] + sary = _nx.swapaxes(ary,axis,0) + for i in range(Nsections): + st = div_points[i]; end = div_points[i+1] + sub_arys.append(_nx.swapaxes(sary[st:end],axis,0)) + + # there is a wierd issue with array slicing that allows + # 0x10 arrays and other such things. The following cluge is needed + # to get around this issue. + sub_arys = _replace_zero_by_x_arrays(sub_arys) + # end cluge. + + return sub_arys + +def split(ary,indices_or_sections,axis=0): + """ Divide an array into a list of sub-arrays. + + Description: + Divide ary into a list of sub-arrays along the + specified axis. If indices_or_sections is an integer, + ary is divided into that many equally sized arrays. + If it is impossible to make an equal split, an error is + raised. This is the only way this function differs from + the array_split() function. If indices_or_sections is a + list of sorted integers, its entries define the indexes + where ary is split. + + Arguments: + ary -- N-D array. + Array to be divided into sub-arrays. + indices_or_sections -- integer or 1D array. + If integer, defines the number of (close to) equal sized + sub-arrays. If it is a 1D array of sorted indices, it + defines the indexes at which ary is divided. Any empty + list results in a single sub-array equal to the original + array. + axis -- integer. default=0. + Specifies the axis along which to split ary. + Caveats: + Currently, the default for axis is 0. This + means a 2D array is divided into multiple groups + of rows. This seems like the appropriate default + """ + try: len(indices_or_sections) + except TypeError: + sections = indices_or_sections + N = ary.shape[axis] + if N % sections: + raise ValueError, 'array split does not result in an equal division' + res = array_split(ary,indices_or_sections,axis) + return res + +def hsplit(ary,indices_or_sections): + """ Split ary into multiple columns of sub-arrays + + Description: + Split a single array into multiple sub arrays. The array is + divided into groups of columns. If indices_or_sections is + an integer, ary is divided into that many equally sized sub arrays. + If it is impossible to make the sub-arrays equally sized, the + operation throws a ValueError exception. See array_split and + split for other options on indices_or_sections. + Arguments: + ary -- N-D array. + Array to be divided into sub-arrays. + indices_or_sections -- integer or 1D array. + If integer, defines the number of (close to) equal sized + sub-arrays. If it is a 1D array of sorted indices, it + defines the indexes at which ary is divided. Any empty + list results in a single sub-array equal to the original + array. + Returns: + sequence of sub-arrays. The returned arrays have the same + number of dimensions as the input array. + Related: + hstack, split, array_split, vsplit, dsplit. + Examples: + >>> import numpy + >>> a= array((1,2,3,4)) + >>> numpy.hsplit(a,2) + [array([1, 2]), array([3, 4])] + >>> a = array([[1,2,3,4],[1,2,3,4]]) + >>> hsplit(a,2) + [array([[1, 2], + [1, 2]]), array([[3, 4], + [3, 4]])] + + """ + if len(_nx.shape(ary)) == 0: + raise ValueError, 'hsplit only works on arrays of 1 or more dimensions' + if len(ary.shape) > 1: + return split(ary,indices_or_sections,1) + else: + return split(ary,indices_or_sections,0) + +def vsplit(ary,indices_or_sections): + """ Split ary into multiple rows of sub-arrays + + Description: + Split a single array into multiple sub arrays. The array is + divided into groups of rows. If indices_or_sections is + an integer, ary is divided into that many equally sized sub arrays. + If it is impossible to make the sub-arrays equally sized, the + operation throws a ValueError exception. See array_split and + split for other options on indices_or_sections. + Arguments: + ary -- N-D array. + Array to be divided into sub-arrays. + indices_or_sections -- integer or 1D array. + If integer, defines the number of (close to) equal sized + sub-arrays. If it is a 1D array of sorted indices, it + defines the indexes at which ary is divided. Any empty + list results in a single sub-array equal to the original + array. + Returns: + sequence of sub-arrays. The returned arrays have the same + number of dimensions as the input array. + Caveats: + How should we handle 1D arrays here? I am currently raising + an error when I encounter them. Any better approach? + + Should we reduce the returned array to their minium dimensions + by getting rid of any dimensions that are 1? + Related: + vstack, split, array_split, hsplit, dsplit. + Examples: + import numpy + >>> a = array([[1,2,3,4], + ... [1,2,3,4]]) + >>> numpy.vsplit(a,2) + [array([[1, 2, 3, 4]]), array([[1, 2, 3, 4]])] + + """ + if len(_nx.shape(ary)) < 2: + raise ValueError, 'vsplit only works on arrays of 2 or more dimensions' + return split(ary,indices_or_sections,0) + +def dsplit(ary,indices_or_sections): + """ Split ary into multiple sub-arrays along the 3rd axis (depth) + + Description: + Split a single array into multiple sub arrays. The array is + divided into groups along the 3rd axis. If indices_or_sections is + an integer, ary is divided into that many equally sized sub arrays. + If it is impossible to make the sub-arrays equally sized, the + operation throws a ValueError exception. See array_split and + split for other options on indices_or_sections. + Arguments: + ary -- N-D array. + Array to be divided into sub-arrays. + indices_or_sections -- integer or 1D array. + If integer, defines the number of (close to) equal sized + sub-arrays. If it is a 1D array of sorted indices, it + defines the indexes at which ary is divided. Any empty + list results in a single sub-array equal to the original + array. + Returns: + sequence of sub-arrays. The returned arrays have the same + number of dimensions as the input array. + Caveats: + See vsplit caveats. + Related: + dstack, split, array_split, hsplit, vsplit. + Examples: + >>> a = array([[[1,2,3,4],[1,2,3,4]]]) + >>> dsplit(a,2) + [array([[[1, 2], + [1, 2]]]), array([[[3, 4], + [3, 4]]])] + + """ + if len(_nx.shape(ary)) < 3: + raise ValueError, 'vsplit only works on arrays of 3 or more dimensions' + return split(ary,indices_or_sections,2) + +def get_array_wrap(*args): + """Find the wrapper for the array with the highest priority. + + In case of ties, leftmost wins. If no wrapper is found, return None + """ + wrappers = [(getattr(x, '__array_priority__', 0), -i, + x.__array_wrap__) for i, x in enumerate(args) + if hasattr(x, '__array_wrap__')] + wrappers.sort() + if wrappers: + return wrappers[-1][-1] + return None + +def kron(a,b): + """kronecker product of a and b + + Kronecker product of two arrays is block array + [[ a[ 0 ,0]*b, a[ 0 ,1]*b, ... , a[ 0 ,n-1]*b ], + [ ... ... ], + [ a[m-1,0]*b, a[m-1,1]*b, ... , a[m-1,n-1]*b ]] + """ + wrapper = get_array_wrap(a, b) + b = asanyarray(b) + a = array(a,copy=False,subok=True,ndmin=b.ndim) + ndb, nda = b.ndim, a.ndim + if (nda == 0 or ndb == 0): + return _nx.multiply(a,b) + as_ = a.shape + bs = b.shape + if not a.flags.contiguous: + a = reshape(a, as_) + if not b.flags.contiguous: + b = reshape(b, bs) + nd = ndb + if (ndb != nda): + if (ndb > nda): + as_ = (1,)*(ndb-nda) + as_ + else: + bs = (1,)*(nda-ndb) + bs + nd = nda + result = outer(a,b).reshape(as_+bs) + axis = nd-1 + for _ in xrange(nd): + result = concatenate(result, axis=axis) + if wrapper is not None: + result = wrapper(result) + return result + + +def tile(A, reps): + """Repeat an array the number of times given in the integer tuple, reps. + + If reps has length d, the result will have dimension of max(d, A.ndim). + If reps is scalar it is treated as a 1-tuple. + + If A.ndim < d, A is promoted to be d-dimensional by prepending new axes. + So a shape (3,) array is promoted to (1,3) for 2-D replication, + or shape (1,1,3) for 3-D replication. + If this is not the desired behavior, promote A to d-dimensions manually + before calling this function. + + If d < A.ndim, tup is promoted to A.ndim by pre-pending 1's to it. Thus + for an A.shape of (2,3,4,5), a tup of (2,2) is treated as (1,1,2,2) + + + Examples: + >>> a = array([0,1,2]) + >>> tile(a,2) + array([0, 1, 2, 0, 1, 2]) + >>> tile(a,(1,2)) + array([[0, 1, 2, 0, 1, 2]]) + >>> tile(a,(2,2)) + array([[0, 1, 2, 0, 1, 2], + [0, 1, 2, 0, 1, 2]]) + >>> tile(a,(2,1,2)) + array([[[0, 1, 2, 0, 1, 2]], + <BLANKLINE> + [[0, 1, 2, 0, 1, 2]]]) + + See Also: + repeat + """ + try: + tup = tuple(reps) + except TypeError: + tup = (reps,) + d = len(tup) + c = _nx.array(A,copy=False,subok=True,ndmin=d) + shape = list(c.shape) + n = max(c.size,1) + if (d < c.ndim): + tup = (1,)*(c.ndim-d) + tup + for i, nrep in enumerate(tup): + if nrep!=1: + c = c.reshape(-1,n).repeat(nrep,0) + dim_in = shape[i] + dim_out = dim_in*nrep + shape[i] = dim_out + n /= max(dim_in,1) + return c.reshape(shape) |