How to vectorize 3D Numpy arrays

Question

I have a 3D numpy array like a = np.zeros((100,100, 20)). I want to perform an operation over every x,y position that involves all the elements over the z axis and the result is stored in an array like b = np.zeros((100,100)) on the same corresponding x,y position.

Now i'm doing it using 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

The problem now is that this operation is much slower than others done the pythonic way, e.g. c = b * b (I actually profiled this function and it's around 2 orders of magnitude slower than others using numpy built in functions and vectorized functions, over a similar number of elements)

How can I improve the performance of such kind of functions mapping a 3D array to a 2D one?

Answer 1

What is usually done in 3D images is to swap the Z axis to the first index:

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

And now you can easily iterate over the Z axis:

>>> for slice in a:
       do something

The slice here will be each of your 100x100 fractions of your 3D matrix. Additionally, by transpossing allows you to access each of the 2D slices directly by indexing the first axis. For example a[10] will give you the 11th 2D 100x100 slice.

Bonus: If you store the data contiguosly, without transposing (or converting to a contiguous array using a = np.ascontiguousarray(a.transpose((2,0,1))) the access to you 2D slices will be faster since they are mapped contiguosly in memory.