How to convert A[x,y] = z to [ [ x0…xN ], [ y0…yN], [ z0…zN] ]

Question

I have a 2D Numpy array that represents an image, and I want to create a surface plot of image intensity using matplotlib.surface_plot. For this function, I need to convert the 2D array A[x,y] => z into three arrays: [x0,...,xN] , [y0,...,yN] and [z0,...,zN] . I can see how to do this conversion element-by-element:

X = []
Y = []
Z = []
for x in range( A.shape[ 0 ] ):
    for y in range( A.shape[ 1 ] ):
        X.append( x )
        Y.append( y )
        Z.append( A[x,y] )

but I'm wondering whether there is a more Pythonic way to do this?

Answer 1

a very simple way to do this could be to basically use the code shown in the matplotlib example . assuming x and y representing the sizes of the two dims in your image array A , you could do

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

# generate some input data that looks nice on a color map:
A = np.mgrid[0:10:0.1,0:10:0.1][0]

X = np.arange(0, A.shape[0], 1)
Y = np.arange(0, A.shape[1], 1)
X, Y = np.meshgrid(X, Y)

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, A, cmap='viridis',
                       linewidth=0, antialiased=False)

gives

Answer 2

You possibly don't need to construct the actual grid, because some pyplot functions accept 1d arrays for x and y, implying that a grid is to be constructed. It seems that Axes3D.plot_surface (which I presume you meant) does need 2d arrays as input, though.

So to get your grid the easiest way is using np.indices to get the indices corresponding to your array:

>>> import numpy as np 
...  
... # dummy data 
... A = np.random.random((3,4))  # randoms of shape (3,4) 
...  
... # get indices 
... x,y = np.indices(A.shape)  # both arrays have shape (3,4) 
...  
... # prove that the indices correspond to the values of A 
... print(all(A[i,j] == A[x[i,j], y[i,j]] for i in x.ravel() for j in y.ravel()))
True

The resulting arrays all have the same shape as A , which should be correct for most use cases. If for any reason you really need a flattened 1d array, you should use x.ravel() etc. to get a flattened view of the same 2d array.

---

I should note though that the standard way to visualize images (due to the short-wavelength variation of the data) is pyplot.imshow or pyplot.pcolormesh which can give you pixel-perfect visualization, albeit in two dimensions.

Answer 3

We agree X, Y and Z have different sizes (N for X and Y and N^2 for Z) ? If yes:

1. X looks not correct (you add several times the same values)
2. something like:

    X = list(range(A.shape[0])
    Y = list(range(A.shape[1])
    Z = [A[x,y] for x in X for y in Y]