I was trying to get a good understanding of numpy apply along axis. Below is the code from the numpy documentation (https://numpy.org/doc/stable/reference/generated/numpy.apply_along_axis.html)
import numpy as np
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]])
print(np.apply_along_axis(my_func, 0, b))
#array([4., 5., 6.])
print(np.apply_along_axis(my_func, 1, b))
#array([2., 5., 8.])
According to webpage, the above code has a similar functionality to the code below which I took from the webpage and modified it (played around with it) to understand it:
arr = np.array([[1,2,3], [4,5,6], [7,8,9]])
axis = 0
def my_func(a):
"""Average first and last element of a 1-D array"""
print(a, a[0], a[-1])
return (a[0] + a[-1]) * 0.5
out = np.empty(arr.shape[axis+1:])
Ni, Nk = arr.shape[:axis], arr.shape[axis+1:]
print(Ni)
for ii in np.ndindex(Ni):
for kk in np.ndindex(Nk):
f = my_func(arr[ii + np.s_[:,] + kk])
Nj = f.shape
for jj in np.ndindex(Nj):
out[ii + jj + kk] = f[jj]
#The code below may help in understanding what I was trying to figure out.
#print(np.shape(np.asarray(1)))
#x = np.int32(1)
#print(x, type(x), x.shape)
I understand from the numpy documentation that scalars and arrays in numpy have the same attributes and methods. I am trying to understand the difference between ‘()’ and 0. I understand that () is a tuple. See below.
Example:
In the code below, the first for-loop does not iterate but the second for-loop iterates once. I am trying to understand why.
import numpy as np
for i in np.ndindex(0):
print(i) #does not run.
for i in np.ndindex(()):
print(i) #runs once
In summary: Given the above context, what is the difference between () and 0?
3
In summary: Given the above context, what is the difference between () and 0?
The first one represents a zero dimensional array with one element. The second one represents a one dimensional array with zero elements.
A zero dimensional array always has a single element.
Example:
>>> array = np.array(42)
>>> array
array(42)
Zero dimensional arrays have a shape of ()
.
>>> array.shape
()
Indexing into a zero dimensional array produces a scalar.
>>> array[()]
42
Zero dimensional arrays are kind of like scalars, in that both of them can only have a single element. However, they act differently in a few subtle ways. The differences between zero-dimensional arrays and scalars are out of scope of this post.
One dimensional arrays, unlike zero dimensional arrays, can contain any number of elements. For example, this one dimensional array contains zero elements:
>>> array = np.array([])
>>> array
array([], dtype=float64)
It has a shape of (0,)
.
>>> array.shape
(0,)
(When you provide a shape of 0
to np.nditer()
, this is implicitly converted to (0,)
. This is the same shape as your example.)
If you loop over each array with for i in np.nditer(array.shape):
, the loop over the array with one element will run once. The loop over the array with zero elements will run zero times.
1
ndindex
One returns an empty list, the other a list with one tuple:
In [39]: list(np.ndindex(0))
Out[39]: []
In [40]: list(np.ndindex(()))
Out[40]: [()]
According to ndindex
docs, it returns tuples, according to the shape of the input argument:
At each iteration a tuple
of indices is returned, the last dimension is
iterated over first.
This may be clearest with shapes like
In [55]: list(np.ndindex((2,3)))
Out[55]: [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]
In [56]: list(np.ndindex((2,)))
Out[56]: [(0,), (1,)]
It’s most useful when iterating over a n-d shape array. But usually we try to avoid iteration at all (even with apply_along_axis
).
If an array has 0 elements, its shape is (0,), and iterating over if does nothing. If the array is ‘scalar’, shape is (), and iterating over that occurs once.
In [89]: x=np.empty(0)
...: for i in np.ndindex(x.shape):
...: print(i, x[i])
...:
In [90]: x=np.array(2)
...: for i in np.ndindex(x.shape):
...: print(i, x[i])
...:
() 2
In [91]: x=np.array([[1,2]])
...: for i in np.ndindex(x.shape):
...: print(i, x[i])
...:
(0, 0) 1
(0, 1) 2
2d apply
Your apply example is the equivalent of
In [59]: (arr[0,:]+arr[-1,:])/2
Out[59]: array([4., 5., 6.])
In [60]: (arr[:,0]+arr[:,-1])/2
Out[60]: array([2., 5., 8.])
In the expanded apply...
, the use of ndindex
makes more sense when dealing with 3 or more dimensions. With only 2 (3,3) it isn’t really needed
In [61]: axis=0
In [62]: Ni, Nk = arr.shape[:axis], arr.shape[axis+1:]
In [63]: Ni,Nk
Out[63]: ((), (3,))
In [64]: arr.shape
Out[64]: (3, 3)
For this 2d array, the expanded code evaluates as:
In [67]: res=np.empty(3)
...: for i in range(3):
...: res[i] = (arr[0,i]+arr[-1,i])/2
...: res
Out[67]: array([4., 5., 6.])
That is it’s applying the first,last avg
to slices on the 2nd axis, or in other words, all dimensions except axis==0
.
3d array
It might be more interesting with a 3d array
In [77]: def myfunc(a1d):
...: print(a1d)
...: return (a1d[0]+a1d[-1])/2
...:
In [78]: arr = np.arange(24).reshape(2,3,4)
In [79]: arr
Out[79]:
array([[[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11]],
[[12, 13, 14, 15],
[16, 17, 18, 19],
[20, 21, 22, 23]]])
Iterating on all but the 0 axis does:
In [80]: np.apply_along_axis(myfunc,0,arr)
[ 0 12]
[ 1 13]
[ 2 14]
[ 3 15]
...
[11 23]
Out[80]:
array([[ 6., 7., 8., 9.],
[10., 11., 12., 13.],
[14., 15., 16., 17.]])
while iterating on all but the last axis does:
In [82]: np.apply_along_axis(myfunc,2,arr)
[0 1 2 3]
[4 5 6 7]
[ 8 9 10 11]
[12 13 14 15]
[16 17 18 19]
[20 21 22 23]
Out[82]:
array([[ 1.5, 5.5, 9.5],
[13.5, 17.5, 21.5]])
Here is sends size 4 arrays to the function, and does so for 6 times, a (2,3) shaped array.
That’s the same as indexing the last axis:
In [87]: (arr[:,:,0]+arr[:,:,-1])/2
Out[87]:
array([[ 1.5, 5.5, 9.5],
[13.5, 17.5, 21.5]])
and for axis=1, make a (2,4) result:
In [88]: (arr[:,0,:]+arr[:,-1,:])/2
Out[88]:
array([[ 4., 5., 6., 7.],
[16., 17., 18., 19.]])
Here (arr[:,0,:]+arr[:,-1,:])
is roughly equivalent to the use of ndindex
and f = my_func(arr[ii + np.s_[:,] + kk])
, a way of indexing ‘the-middle’ of multiple axes.
apply_along_axis
can be useful – but only if you can’t use other whole array operations. It is relatively slow, iterating in python for all but the designated axis.
apply...
code is python and readable via the [source] link. It simplifies the iteration on all by x-axis by transposing that axis to the end, making a simpler ind
(with ndindex
), and iterates with:
for ind in inds:
buff[ind] = asanyarray(
func1d(inarr_view[ind], *args, **kwargs))
https://github.com/numpy/numpy/blob/v2.1.0/numpy/lib/_shape_base_impl.py#L278-L419