Question

I have a 3d matrix grid_z0, whose dimension are (let’s say) 50x25x36. Each point of this matrix represents a cell. I would like to convert this matrix into a 1D array of size 50x25x36. I would also like to create 3 array of the same size showing the coordinate of the cell center.
The array cx,cy and cz stores the coordinate of the cell center in one direction.

This example work but it rather slow, especially for large data set. Is there a way to make it faster?

data={"x":[],"y":[],"z":[],"rho":[]}

for i in arange(0,50):
  for j in arange(0,25):
    for k in arange(0,36):
      data["x"].append(cx[i])
      data["y"].append(cy[j])
      data["z"].append(cz[k])
      data["f"].append(grid_z0[i][j][k])

Answer 1

You should consider using NumPy.

>>> import numpy as np
>>> a = np.random.rand(50,25,36) # Create a fake array
>>> print a.shape
(50, 25, 36)
>>> a.shape = -1 # Flatten the array in place
>>> print a.shape
(45000,)

When flattening your array, you’re doing the equivalent of :

>>> b = []
>>> for i in range(a.shape[0]):
...     for j in range(a.shape[1]):
...         for k in range(a.shape[2]):
...             b.append(a[i,j,k])

That is, the last axis is parsed first, then the second, then the first.

Given your three 1D lists cx, cy, cz of length N, you can construct a 2D array with:

>>> centers = np.array([cx,cy,cz])
>>> print centers.shape
(3, N)