Reshape 1D numpy array to 3D with x,y,z ordering

Question

Say I have a 1D array of values corresponding to x, y, and z values like this:

x  y  z  arr_1D
0  0  0  0
1  0  0  1
0  1  0  2
1  1  0  3
0  2  0  4
1  2  0  5
0  0  1  6
...
0  2  3  22
1  2  3  23

I want to get arr_1D into a 3D array arr_3D with shape (nx,ny,nz) (in this case (2,3,4)). I'd like to the values to be referenceable using arr_3D[x_index, y_index, z_index], so that, for example, arr_3D[1,2,0]=5. Using numpy.reshape(arr_1D, (2,3,4)) gives me a 3D matrix of the right dimensions, but not ordered the way I want. I know I can use the following code, but I'm wondering if there's a way to avoid the clunky nested for loops.

arr_1d = np.arange(24)
nx = 2
ny = 3
nz = 4
arr_3d = np.empty((nx,ny,nz))
count = 0
for k in range(nz):
    for j in range(ny):
        for i in range(nx):
            arr_3d[i,j,k] = arr_1d[count]
            count += 1

print arr_3d[1,2,0]

output: 5

What would be the most pythonic and/or fast way to do this? I'll typically want to do this for arrays of length on the order of 100,000.

Solution

You where really close, but since you want the x axis to be the one that is iterated trhough the fastest, you need to use something like

arr_3d = arr_1d.reshape((4,3,2)).transpose()

So you create an array with the right order of elements but the dimensions in the wrong order and then you correct the order of the dimensions.

Comments

[Ben Lindsay]

Worked like a charm, thanks! I tried some other transpose thing that didn't work, not sure why I didn't try that one.