使用 meshgrid 将 X,Y,Z 三元组转换为三个 2D arrays for surface plot in matplotlib

Question

I'm new to Python so please be patient. I appreciate any help!

What I have: three 1D lists ( xr, yr, zr ), one containing x-values, the other two y- and z-values
What I want to do: create a 3D contour plot in matplotlib

I realized that I need to convert the three 1D lists into three 2D lists, by using the meshgrid function.

Here's what I have so far:

xr = np.asarray(xr) 
yr = np.asarray(yr)
zr = np.asarray(zr)

X, Y = np.meshgrid(xr,yr)
znew = np.array([zr for x,y in zip(np.ravel(X), np.ravel(Y))])
Z = znew.reshape(X.shape)

Running this gives me the following error (for the last line I entered above):

 total size of new array must be unchanged

I went digging around stackoverflow, and tried using suggestions from people having similar problems. Here are the errors I get from each of those suggestions:

Changing the last line to:

Z = znew.reshape(X.shape[0])

Gives the same error.

Changing the last line to:

Z = znew.reshape(X.shape[0], len(znew))

Gives the error:

Shape of x does not match that of z: found (294, 294) instead of (294, 86436).

Changing it to:

Z = znew.reshape(X.shape, len(znew))

Gives the error:

an integer is required

Any ideas?

Answer 1

Well,sample code below works for me

import numpy as np
import matplotlib.pyplot as plt

xr = np.linspace(-20, 20, 100)
yr = np.linspace(-25, 25, 110)
X, Y = np.meshgrid(xr, yr)

#Z = 4*X**2 + Y**2

zr = []
for i in range(0, 110):
    y = -25.0 + (50./110.)*float(i)
    for k in range(0, 100):
        x = -20.0 + (40./100.)*float(k)

        v = 4.0*x*x + y*y

        zr.append(v)

Z = np.reshape(zr, X.shape)

print(X.shape)
print(Y.shape)
print(Z.shape)

plt.contour(X, Y, Z)
plt.show()

Answer 2

TL;DR

import matplotlib.pyplot as plt
import numpy as np

def get_data_for_mpl(X, Y, Z):
    result_x = np.unique(X)
    result_y = np.unique(Y)
    result_z = np.zeros((len(result_x), len(result_y)))
    # result_z[:] = np.nan

    for x, y, z in zip(X, Y, Z):
        i = np.searchsorted(result_x, x)
        j = np.searchsorted(result_y, y)
        result_z[i, j] = z

    return result_x, result_y, result_z


xr, yr, zr = np.genfromtxt('data.txt', unpack=True)

plt.contourf(*get_data_for_mpl(xr, yr, zr), 100)
plt.show()

Detailed answer

At the beginning, you need to find out for which values of x and y the graph is being plotted. This can be done using the numpy.unique function:

result_x = numpy.unique(X)
result_y = numpy.unique(Y)

Next, you need to create a numpy.ndarray with function values for each point (x, y) from zip(X, Y) :

result_z = numpy.zeros((len(result_x), len(result_y)))

for x, y, z in zip(X, Y, Z):
    i = search(result_x, x)
    j = search(result_y, y)
    result_z[i, j] = z

If the array is sorted, then the search in it can be performed not in linear time, but in logarithmic time, so it is enough to use the numpy.searchsorted function to search. but to use it, the arrays result_x and result_y must be sorted. Fortunately, sorting is part of the numpy.unique method and there are no additional actions to do. It is enough to replace the search (this method is not implemented anywhere and is given simply as an intermediate step) method with np.searchsorted .

Finally, to get the desired image, it is enough to call the matplotlib.pyplot.contour or matplotlib.pyplot.contourf method.

If the function value does not exist for (x, y) for all x from result_x and all y from result_y , and you just want to not draw anything, then it is enough to replace the missing values with NaN. Or, more simply, create result_z as numpy.ndarray` from NaN and then fill it in:

result_z = numpy.zeros((len(result_x), len(result_y)))
result_z[:] = numpy.nan