Populate numpy array based on cell indices

Question

I am trying to create a 3d-array, whose cell entries are to be calculated from the cell indices. Specifically, I want that cell (i,j,k) = sqrt(i+j+k).

This is very easily done with the following for loops:

N=10
A=np.zeros((N,N,N))

for i in range(N):
    for j in range(N):
        for k in range(N):
            A[i][j][k] = np.sqrt(i+j+k)

I was wondering if numpy has inbuilt functions that make these nested for loops redundant.

Answer 1

Simplest and performant one would be with open grids using np.ogrid and then performing the relevant operation(s) -

I,J,K = np.ogrid[:N,:N,:N]
A = np.sqrt(I+J+K)

Or with np.sum for the broadcasted summations of open grids for a one-liner -

A = np.sqrt(np.sum(np.ogrid[:N,:N,:N]))

Relevant : General workflow on vectorizing loops involving range iterators.