How to Create 3D Contour Plots with Matplotlib in Python

by

💡 Problem Formulation: Visualizing complex data can be challenging, especially when dealing with three-dimensional space. In this article, we solve the problem of representing three-dimensional surfaces by creating 3D contour plots using Python’s Matplotlib library. The input will be a set of data points in 3D, and the desired output is a visual representation that captures the contours of the surface formed by these points.

Method 1: Using the contourf() and scatter() Functions

This method involves plotting filled contours with the contourf() function and then using the scatter() function to overlay data points. The contourf() function creates a contour plot with filled areas between the lines, making surfaces distinguishable, while the scatter() function helps in marking the individual data points for reference.

Here’s an example:

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

X = np.linspace(-5, 5, 100)
Y = np.linspace(-5, 5, 100)
X, Y = np.meshgrid(X, Y)
Z = np.sin(np.sqrt(X**2 + Y**2))

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.contourf(X, Y, Z, 50)
ax.scatter(X, Y, Z, color='k')
plt.show()

The output of this code is a 3D contour plot with dark spots representing the data points.

This code snippet begins by generating a mesh grid of X and Y values, calculates Z as a function of X and Y, and then creates a 3D plot. The contourf() call adds filled contours to the plot while scatter() overlays the data points.

Method 2: Utilizing the plot_surface() Function with Contours

Here, the plot_surface() function is used in conjunction with contour plots to visualize the surface and its contours in two separate views. The plot_surface() function renders a 3D surface mesh, and by adding contours on the bottom plane, the relation between the two visualizations is clarified.

Here’s an example:

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

x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
x, y = np.meshgrid(x, y)
z = np.sin(np.sqrt(x**2 + y**2))

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(x, y, z, rstride=1, cstride=1, cmap='viridis', edgecolor='none')
ax.contourf(x, y, z, zdir='z', offset=-2, cmap='viridis')

plt.show()

The output is a 3-dimensional surface plot with an additional contour plot displayed on the xy-plane.

The code snippet above first defines the grid and function values for z. Then, it creates a 3D plot where plot_surface() is used to draw the surface and contourf() is used to add the contours to the xy-plane, giving a different view of the surface.

Method 3: Generating Contour Lines with contour()

In this method, contour() is used instead of contourf() to generate contour lines rather than filled contours. This can give a clearer view of the precise values of the contours. The contour() function is straightforward and allows for an easier interpretation of the data’s structure.

Here’s an example:

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

x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
x, y = np.meshgrid(x, y)
z = np.sin(np.sqrt(x**2 + y**2))

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.contour(x, y, z, 50, cmap='winter')
plt.show()

The output is a 3D plot featuring the contour lines of the function’s surface.

This concise snippet creates a mesh grid and calculates the function z. The contour() call then adds the contour lines to the 3D plot. By specifying the number of levels, the granularity of contour lines can be adjusted.

Method 4: Combining Surface and Contour Plots

This method combines the visual details of the surface plot with the contour plot by using plot_surface() and contour() together. The technique allows for a comprehensive view that includes both the smoothness of the surface and the precision of contour lines.

Here’s an example:

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

x = np.linspace(-5, 5, 100)
y = np.linspace(-5, 5, 100)
x, y = np.meshgrid(x, y)
z = np.sin(np.sqrt(x**2 + y**2))

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(x, y, z, cmap='coolwarm', alpha=0.5)
ax.contour(x, y, z, zdir='z', offset=np.min(z), cmap='coolwarm')
plt.show()

The output combines a translucent 3D surface plot with the contour lines on the bottom plane.

This technique initializes a 3D plot and employs both plot_surface() and contour(). The alpha parameter in plot_surface() makes the surface plot semi-transparent, which allows the contour lines beneath to be visible, offering a dual representation of the data.

Bonus One-Liner Method 5: Creating Inline 3D Contour Plot

This quick and compact method creates a 3D contour plot using a one-liner code with the combination of PyPlot’s features. Ideal for quick and simple visualizations without customization.

Here’s an example:

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

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.contour3D(*np.meshgrid(np.linspace(-5, 5, 100), np.linspace(-5, 5, 100)), np.sin(np.sqrt(X**2 + Y**2)), 50)
plt.show()

The output is a 3-dimensional plot showing the contours of the surface in question.

This one-liner is an effective method when speed and brevity are desired. It leverages Python’s unpacking with * to pass the meshgrid directly into the contour3D() function. This can be used to quickly draft a visualization when working with familiar datasets.

Summary/Discussion

  • Method 1: Using contourf() and scatter(). Strength: Provides filled contours with clear data point markers. Weakness: May not clearly represent precise contour levels.
  • Method 2: Utilizing plot_surface() with Contours. Strength: Offers a detailed 3D view of the surface with informative contours. Weakness: Can be computationally intensive for large datasets.
  • Method 3: Generating Contour Lines with contour(). Strength: Offers precise contour lines for clearer data structure. Weakness: Lacks the visual appeal of filled contours.
  • Method 4: Combining Surface and Contour Plots. Strength: Provides a comprehensive visualization combining both surface and contours. Weakness: Contour details may be obscured by the surface plot.
  • Bonus Method 5: Creating Inline 3D Contour Plot. Strength: Quick visualization for straightforward data. Weakness: Lacks flexibility for more complex or specific customization.