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`

and`cstride`

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), and`linewidth`

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.