How do I create a surface plot in matplotlib when Z is already calculated?

Question

I have seen the following example for surface plot here, in the Surface Plot section:

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np


fig = plt.figure()
ax = fig.gca(projection='3d')

# Make data.
X = np.arange(-5, 5, 0.25)
Y = np.arange(-5, 5, 0.25)
X, Y = np.meshgrid(X, Y)
R = np.sqrt(X**2 + Y**2)
Z = np.sin(R)

# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False)

# Customize the z axis.
ax.set_zlim(-1.01, 1.01)
ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.02f'))

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

In the example, X and Y coordinates are provided as vectors, to later be reused as a XY coordinate mesh with the mesh command.

In the example, Z is calculated directly from X,Y. However in my case I am using a previously calculated value for Z (so somehow I already did the mesh function in previous steps):

array([[1.00000000e+00, 1.00000000e+01, 1.59135708e+06],
       [1.00000000e+00, 1.10000000e+01, 2.44938222e+06],
       [1.00000000e+00, 1.20000000e+01, 3.96364528e+06],
       [1.00000000e+00, 1.30000000e+01, 6.13414627e+06],
       [1.00000000e+00, 1.40000000e+01, 8.96088517e+06],
       [2.00000000e+00, 1.00000000e+01, 1.59386643e+06],
       [2.00000000e+00, 1.10000000e+01, 2.45962860e+06],
       [2.00000000e+00, 1.20000000e+01, 3.98162868e+06],
       [2.00000000e+00, 1.30000000e+01, 6.15986668e+06],
       [2.00000000e+00, 1.40000000e+01, 8.99434261e+06],
       [3.00000000e+00, 1.00000000e+01, 1.59647579e+06],
       [3.00000000e+00, 1.10000000e+01, 2.46997497e+06],
       [3.00000000e+00, 1.20000000e+01, 3.99971208e+06],
       [3.00000000e+00, 1.30000000e+01, 6.18568710e+06],
       [3.00000000e+00, 1.40000000e+01, 9.02790005e+06],
       [4.00000000e+00, 1.00000000e+01, 1.59918515e+06],
       [4.00000000e+00, 1.10000000e+01, 2.48042135e+06],
       [4.00000000e+00, 1.20000000e+01, 4.01789548e+06],
       [4.00000000e+00, 1.30000000e+01, 6.21160752e+06],
       [4.00000000e+00, 1.40000000e+01, 9.06155749e+06]])

So the first column is the X, the second the Y and the third one the Z.

Z comes from R which is calculated using np.sqrt(X**2 + Y**2) and X, Y are already variables created with mesh.

How shall those values be assigned to variables X, Y and Z?

Sergey Bushmanov · Accepted Answer · 2020-10-24 19:33:16Z

1

Suppose you have an aray arr. Then you can cast it to the format expected by matplotlib as:

# Make data.
X = arr[:,0].reshape((4,5)).T
Y = arr[:,1].reshape((4,5)).T
Z = arr[:,2].reshape((4,5)).T

and plot:

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator, FormatStrFormatter
import numpy as np


fig = plt.figure()
ax = fig.gca(projection='3d')
# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False)

# Customize the z axis.
# ax.set_zlim(-1.01, 1.01)
# ax.zaxis.set_major_locator(LinearLocator(10))
ax.zaxis.set_major_formatter(FormatStrFormatter('%.1E'))

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=.7, aspect=5, pad=.1);

You also may wish to generelaize like:

from itertools import groupby
total_length = arr.shape[0]
group_length = [len(list(g)) for _,g in groupby(arr[:,0])]
print(f"Reshape 1st: {int(total_length/group_length[0])}",f"Reshape 2nd: {int(group_length[0])}",sep="\n")
Reshape 1st: 4 
Reshape 2nd: 5

edited Oct 24, 2020 at 19:33

answered Oct 24, 2020 at 19:10

Sergey Bushmanov

25.5k8 gold badges63 silver badges84 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

M.E. Over a year ago

Thanks Sergey, why are you using specifically 4 and 5 (.reshape(4,5))?

Sergey Bushmanov Over a year ago

This is because how your array is structured, 4 groups of 5 points

Collectives™ on Stack Overflow

How do I create a surface plot in matplotlib when Z is already calculated?

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related