python xyz data transformation for surface plot

Question

I want to plot a 3D surface in matplotlib from xy and z data. For this I need to do a seemingly simple data transformation but I am unsure on how to proceed.

My x and y data are uniform integers because it will be the surface of an image. My data is present in the form of a 512x512 numpy array. The values of the array are the z values, and the indexes are the x and y values respectively.

So if arr is my array, arr[x, y] would give my z . The form just looks like this:

z z z z ... z (512 columns)
z z z z ... z
z z z z ... z
z z z z ... z
. . . .
. . . .
z z z z (512 rows)

How do I get my data in the form of three columns x, y, z so I can do the surface plot? It should look like this after the transform:

  x | y | z
  ---------
  0 | 0 | z
  1 | 0 | z
  2 | 0 | z
  . | . | .
  . | . | .
511 | 0 | z
  0 | 1 | z
  1 | 1 | z
  2 | 1 | z
  . | . | .
  . | . | .

I tried to work with np.meshgrid and np.flatten but can't get it to work the way I want. Maybe theres an even easier pandas solution to this. Or maybe I can even plot it with the original form of the data?

Any suggestion is appreciated :)

Answer 1

You would use np.meshgrid like this:

# Make coordinate grids
x, y = np.meshgrid(np.arange(arr.shape[0]), np.arange(arr.shape[1]), indexing='ij')
# Flatten grid and data and stack them into a single array
data = np.stack([x.ravel(), y.ravel(), arr.ravel()], axis=1)

For example:

import numpy as np

arr = np.arange(12).reshape((3, 4))
x, y = np.meshgrid(np.arange(arr.shape[0]), np.arange(arr.shape[1]), indexing='ij')
data = np.stack([x.ravel(), y.ravel(), arr.ravel()], axis=1)
print(data)

Output:

[[ 0  0  0]
 [ 0  1  1]
 [ 0  2  2]
 [ 0  3  3]
 [ 1  0  4]
 [ 1  1  5]
 [ 1  2  6]
 [ 1  3  7]
 [ 2  0  8]
 [ 2  1  9]
 [ 2  2 10]
 [ 2  3 11]]

EDIT:

Actually, if you want to have your final array such that x values increase first (as in the example you gave), you could do it like this:

x, y = np.meshgrid(np.arange(arr.shape[0]), np.arange(arr.shape[1]), indexing='xy')
data = np.stack([x.ravel(), y.ravel(), arr.T.ravel()], axis=1)

In this case you would get:

[[ 0  0  0]
 [ 1  0  4]
 [ 2  0  8]
 [ 0  1  1]
 [ 1  1  5]
 [ 2  1  9]
 [ 0  2  2]
 [ 1  2  6]
 [ 2  2 10]
 [ 0  3  3]
 [ 1  3  7]
 [ 2  3 11]]