Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- def apply_along_axis(func1d,axis,arr,*args,**moreargs):
- """
- Apply a function to 1-D slices along the given axis.
- Execute `func1d(a, *args)` where `func1d` operates on 1-D arrays and `a`
- is a 1-D slice of `arr` along `axis`.
- Parameters
- ----------
- func1d : function
- This function should accept 1-D arrays. It is applied to 1-D
- slices of `arr` along the specified axis.
- axis : integer
- Axis along which `arr` is sliced.
- arr : ndarray
- Input array.
- args : any
- Additional arguments to `func1d`.
- Returns
- -------
- apply_along_axis : ndarray
- The output array. The shape of `outarr` is identical to the shape of
- `arr`, except along the `axis` dimension, where the length of `outarr`
- is equal to the size of the return value of `func1d`. If `func1d`
- returns a scalar `outarr` will have one fewer dimensions than `arr`.
- See Also
- --------
- apply_over_axes : Apply a function repeatedly over multiple axes.
- Examples
- --------
- >>> def my_func(a):
- ... \"\"\"Average first and last element of a 1-D array\"\"\"
- ... return (a[0] + a[-1]) * 0.5
- >>> b = np.array([[1,2,3], [4,5,6], [7,8,9]])
- >>> np.apply_along_axis(my_func, 0, b)
- array([ 4., 5., 6.])
- >>> np.apply_along_axis(my_func, 1, b)
- array([ 2., 5., 8.])
- For a function that doesn't return a scalar, the number of dimensions in
- `outarr` is the same as `arr`.
- >>> b = np.array([[8,1,7], [4,3,9], [5,2,6]])
- >>> np.apply_along_axis(sorted, 1, b)
- array([[1, 7, 8],
- [3, 4, 9],
- [2, 5, 6]])
- """
- 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 = list(range(nd))
- indlist.remove(axis)
- i[axis] = slice(None, None)
- outshape = asarray(arr.shape).take(indlist)
- i.put(indlist, ind)
- print
- res = func1d(arr[tuple(i.tolist())],*args,**moreargs)
- # 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,**moreargs)
- 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,**moreargs)
- outarr[tuple(i.tolist())] = res
- k += 1
- return outarr
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement