tkowt - 8 months ago 56

Python Question

I want to apply arbitrary function to 3d-ndarray as element, which use (3rd-dimensional) array for its arguments and return scalar.As a result, we should get 2d-Matrix.

e.g) pseudo code

`A = [[[1,2,3],[4,5,6]],`

[[7,8,9],[10,11,12]]]

A.apply_3d_array(sum) ## or apply_3d_array(A,sum) is Okey.

>> [[6,15],[24,33]]

I understand it's possible with loop using ndarray.shape function,but direct index access is inefficient as official document says.

Is there more effective way than using loop?

Answer

`apply_along_axis`

is designed to make this task easy:

```
In [683]: A=np.arange(1,13).reshape(2,2,3)
In [684]: A
Out[684]:
array([[[ 1, 2, 3],
[ 4, 5, 6]],
[[ 7, 8, 9],
[10, 11, 12]]])
In [685]: np.apply_along_axis(np.sum, 2, A)
Out[685]:
array([[ 6, 15],
[24, 33]])
```

It, in effect, does

```
for all i,j:
out[i,j] = func( A[i,j,:])
```

taking care of the details. It's not faster than doing that iteration yourself, but it makes it easier.

Another trick is to reshape your input to 2d, perform the simpler 1d iteration, and the reshape the result

```
A1 = A.reshape(-1, A.shape[-1])
for i in range(A1.shape[0]):
out[i] = func(A1[i,:])
out.reshape(A.shape[:2])
```

To do things faster, you need to dig into the guts of the function, and figure out how to use compile numpy operations on more than one dimension. In the simple case of `sum`

, that function already can work on selected axes.