How to plot a smooth 2D color plot for z = f(x, y) in Matplotlib?

To plot a smooth 2D color plot for z = f(x, y) in Matplotlib, we create a function that maps two variables to a color-coded surface. This visualization is useful for displaying mathematical functions, heat maps, and scientific data.

Basic Steps

Follow these steps to create a smooth 2D color plot:

• Set the figure size and adjust the padding between and around the subplots
• Create x and y data points using numpy
• Get z data points using f(x, y)
• Display the data as an image on a 2D regular raster with z data points
• Use interpolation to create smooth color transitions
• To display the figure, use show() method

Example

Here's how to create a smooth 2D color plot for the function z = x² + y² :

import numpy as np
from matplotlib import pyplot as plt

plt.rcParams["figure.figsize"] = [7.50, 3.50]
plt.rcParams["figure.autolayout"] = True

def f(x, y):
    return np.array([i * i + j * j for j in y for i in x]).reshape(5, 5)

x = y = np.linspace(-1, 1, 5)
z = f(x, y)

plt.imshow(z, interpolation='bilinear')
plt.colorbar()
plt.title('2D Color Plot: z = x² + y²')

plt.show()

Enhanced Version with Higher Resolution

For smoother plots, increase the grid resolution and use meshgrid for better coordinate handling :

import numpy as np
from matplotlib import pyplot as plt

# Create high-resolution grid
x = np.linspace(-2, 2, 100)
y = np.linspace(-2, 2, 100)
X, Y = np.meshgrid(x, y)

# Define function
Z = X**2 + Y**2

# Create smooth color plot
plt.figure(figsize=(8, 6))
plt.imshow(Z, extent=[-2, 2, -2, 2], interpolation='bilinear', 
           cmap='viridis', origin='lower')
plt.colorbar(label='z value')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Smooth 2D Color Plot: z = x² + y²')

plt.show()

Alternative Functions

You can plot different mathematical functions by changing the Z calculation :

import numpy as np
from matplotlib import pyplot as plt

x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
X, Y = np.meshgrid(x, y)

# Gaussian function
Z = np.exp(-(X**2 + Y**2))

plt.figure(figsize=(8, 6))
plt.imshow(Z, extent=[-3, 3, -3, 3], interpolation='bilinear', 
           cmap='hot', origin='lower')
plt.colorbar(label='z value')
plt.xlabel('x')
plt.ylabel('y')
plt.title('2D Gaussian Function')

plt.show()

Key Parameters

Conclusion

Use imshow() with bilinear interpolation to create smooth 2D color plots. Higher grid resolution and proper extent settings produce better visualizations. The meshgrid() function provides cleaner coordinate handling for complex functions.