How to vectorize 3D Numpy arrays

Question

I have a 3D numpy array, for example a = np.zeros((100,100, 20)). I want to perform an operation on each position x,y, which includes all the elements on the axis z, and the result is stored in the type array b = np.zeros((100,100))in the same corresponding position x,y.

Now I do this with a for loop:

d_n = np.array([...]) # a parameter with the same shape as b
for (x,y), v in np.ndenumerate(b):
    C = a[x,y,:]

    ### calculate some_value using C
    minv = sys.maxint
    depth = -1
    C = a[x,y,:]
    for d in range(len(C)):
        e = 2.5 * float(math.pow(d_n[x,y] - d, 2)) + C[d] * 0.05
        if e < minv:
            minv = e
            depth = d

    some_value = depth
    if depth == -1:
        some_value = len(C) - 1
    ###

    b[x,y] = some_value

Now the problem is that this operation is much slower than others performed by Putin's method, for example. c = b * b(I actually profiled this function, and it is about 2 orders of magnitude slower than others using numpy built-in functions and vectorized functions compared to the same number of elements)

How to improve the performance of such functions that map a 3D array to a two-dimensional one?

Answer 1

, Z :

>>> a = a.transpose((2,0,1))
>>> a.shape
(20, 100, 100)

Z:

>>> for slice in a:
       do something

slice 100x100 3D-. , 2D- , . , a[10] 11- 2D- 100x100.

: , ( a = np.ascontiguousarray(a.transpose((2,0,1))), 2D- , .

Answer 2

, for, , , , C.

a = np.zeros((100,100, 20))
a[:,:] = np.linspace(1,20,20)    # example data: 1,2,3,.., 20 as "z" for every "x","y"

b = np.sum(a[:,:]**2, axis=2)

100 100 b "z" a, 1 + 4 + 9 +... + 400 = 2870.

Answer 3

,

for (x,y), v in np.ndenumerate(b):
    C = a[x,y,:]
    ...
    for d in range(len(C)):
        ... # complex, not vectorizable calc
    ...
    b[x,y] = some_value

1- , 2D- 1D, . (N,20) (N,). , .

- C some_value. , cumsum cumprod, . cython .

, , N . , , .

-. , some_value C d_n, , .

---

, e :

e = 2.5 * float(math.pow(d_n[x,y] - d, 2)) + C[d] * 0.05

E = 2.5 * (d_n[...,None] - np.arange(a.shape[-1]))**2 + a * 0.05  # (100,100,20)

E.min(axis=-1)  # smallest value along the last dimension
E.argmin(axis=-1)  # index of where that min occurs

: E.argmin - b, ( , ).

a d_n, , E.argmin(-1) b 66x.

Answer 4

, 3D- ?

Numpy "" ^* sum, any, std .. axis, None . argmin, e :

d = np.arange(a.shape[2])
e = 2.5 * (d_n[...,None] - d)**2 + a*0.05
b = np.argmin(e, axis=2)

[...,None] broadcasting. e , sys.maxint, :

I, J = np.indices(b.shape)
b[e[I,J,b] >= sys.maxint] = a.shape[2] - 1

_{* , reduce(operator, sequence) std argmin}