To generate a 3D waterfall plot with colored heights create a 2D sine wave using the NumPy meshgrid()
function, then apply a colormap to the heights using Matplotlib’s Normalize
function. The plot_surface()
function generates the 3D plot, while the color gradient is added using a ScalarMappable
object.
Here’s a code example for copy and paste: ๐
import numpy as np import matplotlib.pyplot as plt from matplotlib import cm # Create the X, Y, and Z coordinate arrays. x = np.linspace(-5, 5, 101) y = np.linspace(-5, 5, 101) X, Y = np.meshgrid(x, y) Z = np.sin(np.sqrt(X**2 + Y**2)) # Create a surface plot and projected filled contour plot under it. fig = plt.figure(figsize=(8, 6)) ax = fig.add_subplot(111, projection='3d') # Select Colormap cmap = cm.viridis # Norm for color mapping norm = plt.Normalize(Z.min(), Z.max()) # Plot surface with color mapping surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, facecolors=cmap(norm(Z)), alpha=0.9, linewidth=0) # Add a color bar which maps values to colors m = cm.ScalarMappable(cmap=cmap, norm=norm) m.set_array(Z) fig.colorbar(m) ax.set_xlabel('X') ax.set_ylabel('Y') ax.set_zlabel('Z') plt.show()
If you run this code in your Python shell, you’ll get a beautiful interactive 3D waterfall plot:

In your case you’ll probably need to replace the data with your own. Feel free to check out our guide on linspace()
with video if you need some background there.
Let’s dive into the steps in the code to get this plot done next. ๐
Step 1: Import Libraries
import numpy as np import matplotlib.pyplot as plt from matplotlib import cm
Here we are importing the necessary libraries. Numpy is for numerical operations, Matplotlib’s pyplot
is for plotting, and cm
from Matplotlib is for working with colormaps.
Step 2: Generate Data
x = np.linspace(-5, 5, 101) y = np.linspace(-5, 5, 101) X, Y = np.meshgrid(x, y) Z = np.sin(np.sqrt(X**2 + Y**2))
Here, linspace
generates 101 evenly spaced points between -5 and 5 for both x and y.
You can watch our explainer video on the function here: ๐
๐ก Recommended: How to Use np.linspace()
in Python? A Helpful Illustrated Guide
The NumPy function meshgrid
takes two 1D arrays representing the Cartesian coordinates in the x and y axis and produces two 2D arrays.
The Z array is a 2D array that represents our “heights” and is calculated by applying the sine function to the square root of the sum of the squares of X and Y, essentially creating a 2D sine wave.
Step 3: Prepare the Figure
fig = plt.figure(figsize=(8, 6)) ax = fig.add_subplot(111, projection='3d')
We first create a figure object, and then add a subplot to it. The '111'
argument means that we want to create a grid with 1 row and 1 column and place the subplot in the first (and only) cell of this grid. The projection='3d'
argument means that we want this subplot to be a 3D plot.
Do you need a quick refresher on subplots? Check out our video on the Finxter blog: ๐
๐ก Recommended: Matplotlib Subplot โ A Helpful Illustrated Guide
Step 4: Prepare Colormap
cmap = cm.viridis norm = plt.Normalize(Z.min(), Z.max())
Here we’re choosing a colormap (cm.viridis
), and then creating a normalization object (plt.Normalize
) using the minimum and maximum values of Z. This normalization object will later be used to map the heights in Z to colors in the colormap.
Step 5: Plot the Surface
surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, facecolors=cmap(norm(Z)), alpha=0.9, linewidth=0)
The method ax.plot_surface()
plots the 3D surface.
- The arguments
X, Y, Z
are the coordinates for the plot. - The
rstride
andcstride
parameters determine the stride (step size) for row and column data respectively, facecolors
parameter takes the colormap applied on the normalized Z,alpha
is used for blending value, between 0 (transparent) and 1 (opaque), andlinewidth
determines the line width of the surface plot.
Step 6: Add a Color Bar
m = cm.ScalarMappable(cmap=cmap, norm=norm) m.set_array(Z) fig.colorbar(m)
A ScalarMappable
object is created with the same colormap and normalization as our surface plot. Then we associate this object with our Z array using the set_array
function. Finally, we add a color bar to the figure that represents how the colors correspond to the Z values.
Step 7: Set Labels and Show the Plot
ax.set_xlabel('X') ax.set_ylabel('Y') ax.set_zlabel('Z') plt.show()
Here, we’re setting the labels for each axis (X, Y, Z). Finally, plt.show()
is called to display the plot. The plot will remain visible until all figures are closed.
Instead of using the plot.show()
function to see the output, you can use the plt.savefig('output.jpeg')
statement to save it in a file 'output.jpeg'
.
๐ก Recommended: Matplotlib โ A Simple Guide with Videos

While working as a researcher in distributed systems, Dr. Christian Mayer found his love for teaching computer science students.
To help students reach higher levels of Python success, he founded the programming education website Finxter.com that has taught exponential skills to millions of coders worldwide. He’s the author of the best-selling programming books Python One-Liners (NoStarch 2020), The Art of Clean Code (NoStarch 2022), and The Book of Dash (NoStarch 2022). Chris also coauthored the Coffee Break Python series of self-published books. He’s a computer science enthusiast, freelancer, and owner of one of the top 10 largest Python blogs worldwide.
His passions are writing, reading, and coding. But his greatest passion is to serve aspiring coders through Finxter and help them to boost their skills. You can join his free email academy here.