Creating 2D and 3D plots using Matplotlib

[box type="note" align="" class="" width=""]This article is an excerpt from a book written by L. Felipe Martins, Ruben Oliva Ramos and V Kishore Ayyadevara titled SciPy Recipes. This book provides data science recipes for users to effectively process, manipulate, and visualize massive datasets using SciPy.[/box]

In today’s tutorial, we will demonstrate how to create two-dimensional and three-dimensional plots for displaying graphical representation of data using a full-fledged scientific library -  Matplotlib.

Creating two-dimensional plots of functions and data

We will present the basic kind of plot generated by Matplotlib: a two-dimensional display, with axes, where datasets and functional relationships are represented by lines. Besides the data being displayed, a good graph will contain a title (caption), axes labels, and, perhaps, a legend identifying each line in the plot.

Getting ready

Start Jupyter and run the following commands in an execution cell:

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

How to do it…

Run the following code in a single Jupyter cell:

xvalues = np.linspace(-np.pi, np.pi)
yvalues1 = np.sin(xvalues)
yvalues2 = np.cos(xvalues)
plt.plot(xvalues, yvalues1, lw=2, color='red',
label='sin(x)')
plt.plot(xvalues, yvalues2, lw=2, color='blue',
label='cos(x)')
plt.title('Trigonometric Functions')
plt.xlabel('x')
plt.ylabel('sin(x), cos(x)')
plt.axhline(0, lw=0.5, color='black')
plt.axvline(0, lw=0.5, color='black')
plt.legend()
None

This code will insert the plot shown in the following screenshot into the Jupyter Notebook:

How it works…

We start by generating the data to be plotted, with the three following statements:

xvalues = np.linspace(-np.pi, np.pi, 300)
yvalues1 = np.sin(xvalues)
yvalues2 = np.cos(xvalues)

We first create an xvalues array, containing 300 equally spaced values between -π and π. We then compute the sine and cosine functions of the values in xvalues, storing the results in the yvalues1 and yvalues2 arrays. Next, we generate the first line plot with the following statement:

plt.plot(xvalues, yvalues1,
lw=2, color='red',
label='sin(x)')

The arguments to the plot() function are described as follows:

• xvalues and yvalues1 are arrays containing, respectively, the x and y coordinates of the points to be plotted. These arrays must have the same length.
• The remaining arguments are formatting options. lw specifies the line width and color the line color. The label argument is used by the legend() function, discussed as follows.

The next line of code generates the second line plot and is similar to the one explained previously. After the line plots are defined, we set the title for the plot and the legends for the axes with the following commands:

plt.title('Trigonometric Functions')
plt.xlabel('x')
plt.ylabel('sin(x), cos(x)')

We now generate axis lines with the following statements:

plt.axhline(0, lw=0.5, color='black')
plt.axvline(0, lw=0.5, color='black')

The first arguments in axhline() and axvline() are the locations of the axis lines and the options specify the line width and color.

We then add a legend for the plot with the following statement:

plt.legend()

Matplotlib tries to place the legend intelligently, so that it does not interfere with the plot. In the legend, one item is being generated by each call to the plot() function and the text for each legend is specified in the label option of the plot() function.

Generating multiple plots in a single figure

Wouldn't it be interesting to know how to generate multiple plots in a single figure? Well, let's get started with that.

Getting ready

Start Jupyter and run the following three commands in an execution cell:

%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

How to do it…

Run the following commands in a Jupyter cell:

plt.figure(figsize=(6,6))
xvalues = np.linspace(-2, 2, 100)
plt.subplot(2, 2, 1)
yvalues = xvalues
plt.plot(xvalues, yvalues, color='blue')
plt.xlabel('$x$')
plt.ylabel('$x$')
plt.subplot(2, 2, 2)
yvalues = xvalues ** 2
plt.plot(xvalues, yvalues, color='green')
plt.xlabel('$x$')
plt.ylabel('$x^2$')
plt.subplot(2, 2, 3)
yvalues = xvalues ** 3
plt.plot(xvalues, yvalues, color='red')
plt.xlabel('$x$')
plt.ylabel('$x^3$')
plt.subplot(2, 2, 4)
yvalues = xvalues ** 4
plt.plot(xvalues, yvalues, color='black')
plt.xlabel('$x$')
plt.ylabel('$x^3$')
plt.suptitle('Polynomial Functions')
plt.tight_layout()
plt.subplots_adjust(top=0.90)
None

Running this code will produce results like those in the following screenshot:

How it works…

To start the plotting constructions, we use the figure() function, as shown in the following line of code:

plt.figure(figsize=(6,6))

The main purpose of this call is to set the figure size, which needs adjustment, since we plan
to make several plots in the same figure. After creating the figure, we add four plots with
code, as demonstrated in the following segment:

plt.subplot(2, 2, 3)
yvalues = xvalues ** 3
plt.plot(xvalues, yvalues, color='red')
plt.xlabel('$x$')
plt.ylabel('$x^3$')

In the first line, the plt.subplot(2, 2, 3) call tells pyplot that we want to organize the plots in a two-by-two layout, that is, in two rows and two columns. The last argument specifies that all following plotting commands should apply to the third plot in the array. Individual plots are numbered, starting with the value 1 and counting across the rows and columns of the plot layout.

We then generate the line plot with the following statements:

yvalues = xvalues ** 3
plt.plot(xvalues, yvalues, color='red')

The first line of the preceding code computes the yvalues array, and the second draws the
corresponding graph. Notice that we must set options such as line color individually for each subplot.

After the line is plotted, we use the xlabel() and ylabel() functions to create labels for the axes. Notice that these have to be set up for each individual subplot too.

After creating the subplots, we explain the subplots:

• plt.suptitle('Polynomial Functions') sets a common title for all Subplots
• plt.tight_layout() adjusts the area taken by each subplot, so that axes' legends do not overlap
• plt.subplots_adjust(top=0.90) adjusts the overall area taken by the plots, so that the title displays correctly

Creating three-dimensional plots

Matplotlib offers several different ways to visualize three-dimensional data. In this recipe, we will demonstrate the following methods:

• Drawing surfaces plots
• Drawing two-dimensional contour plots
• Using color maps and color bars

Getting ready

Start Jupyter and run the following three commands in an execution cell:

%matplotlib inline

import numpy as np

import matplotlib.pyplot as plt

How to do it…

Run the following code in a Jupyter code cell:

from mpl_toolkits.mplot3d import Axes3D

from matplotlib import cm

f = lambda x,y: x**3 - 3*x*y**2

fig = plt.figure(figsize=(12,6))

ax = fig.add_subplot(1,2,1,projection='3d')

xvalues = np.linspace(-2,2,100)

yvalues = np.linspace(-2,2,100)

xgrid, ygrid = np.meshgrid(xvalues, yvalues)

zvalues = f(xgrid, ygrid)

surf = ax.plot_surface(xgrid, ygrid, zvalues,

rstride=5, cstride=5,

linewidth=0, cmap=cm.plasma)

ax = fig.add_subplot(1,2,2)

plt.contourf(xgrid, ygrid, zvalues, 30,

cmap=cm.plasma)

fig.colorbar(surf, aspect=18)

plt.tight_layout()

None

Running this code will produce a plot of the monkey saddle surface, which is a famous example of a surface with a non-standard critical point. The displayed graph is shown in the following screenshot:

How it works…

We start by importing the Axes3D class from the mpl_toolkits.mplot3d library, which is the Matplotlib object used for creating three-dimensional plots. We also import the cm class, which represents a color map. We then define a function to be plotted, with the following line of code:

f = lambda x,y: x**3 - 3*x*y**2

The next step is to define the Figure object and an Axes object with a 3D projection, as done in the following lines of code:

fig = plt.figure(figsize=(12,6))

ax = fig.add_subplot(1,2,1,projection='3d')

Notice that the approach used here is somewhat different than the other recipes in this chapter. We are assigning the output of the figure() function call to the fig variable and then adding the subplot by calling the add_subplot() method from the fig object.

This is the recommended method of creating a three-dimensional plot in the most recent version of Matplotlib. Even in the case of a single plot, the add_subplot() method should be used, in which case the command would be ax = fig.add_subplot(1,1,1,projection='3d'). The next few lines of code, shown as follows, compute the data for the plot:

xvalues = np.linspace(-2,2,100)

yvalues = np.linspace(-2,2,100)

xgrid, ygrid = np.meshgrid(xvalues, yvalues)

zvalues = f(xgrid, ygrid)

The most important feature of this code is the call to meshgrid(). This is a NumPy convenience function that constructs grids suitable for three-dimensional surface plots. To understand how this function works, run the following code:

xvec = np.arange(0, 4)

yvec = np.arange(0, 3)

xgrid, ygrid = np.meshgrid(xvec, yvec)

After running this code, the xgrid array will contain the following values:

array([[0, 1, 2, 3],

[0, 1, 2, 3],

[0, 1, 2, 3]])

The ygrid array will contain the following values:

array([[0, 0, 0, 0],

[1, 1, 1, 1],

[2, 2, 2, 2]])

Notice that the two arrays have the same dimensions. Each grid point is represented by a pair of the (xgrid[i,j],ygrid[i,j]) type. This convention makes the computation of a vectorized function on a grid easy and efficient, with the f(xgrid, ygrid) expression. The next step is to generate the surface plot, which is done with the following function call:

surf = ax.plot_surface(xgrid, ygrid, zvalues,

rstride=5, cstride=5,

linewidth=0, cmap=cm.plasma)

The first three arguments, xgrid, ygrid, and zvalues, specify the data to be plotted. We then use the rstride and cstride options to select a subset of the grid points. Notice that the xvalues and yvalues arrays both have length 100, so that xgrid and ygrid will have 10,000 entries each. Using all grid points would be inefficient and produce a poor plot from the visualization point of view. Thus, we set rstride=5 and cstride=5, which results in a plot containing every fifth point across each row and column of the grid.

The next option, linewidth=0, sets the line width of the plot to zero, preventing the display of a wireframe. The final argument, cmap=cm.plasma, specifies the color map for the plot. We use the cm.plasma color map, which has the effect of plotting higher functional values with a hotter color. Matplotlib offer as large number of built-in color maps, listed at https:/​/​matplotlib.​org/​examples/​color/​colormaps_​reference.​html.​

Next, we add the filled contour plot with the following code:

ax = fig.add_subplot(1,2,2)

ax.contourf(xgrid, ygrid, zvalues, 30,

cmap=cm.plasma)

Notice that, when selecting the subplot, we do not specify the projection option, which is not necessary for two-dimensional plots. The contour plot is generated with the contourf() method. The first three arguments, xgrid, ygrid, zvalues, specify the data points, and the fourth argument, 30, sets the number of contours. Finally, we set the color map to be the same one used for the surface plot.

The final component of the plot is a color bar, which provides a visual representation of the value associated with each color in the plot, with the fig.colorbar(surf, aspect=18) method call.

Notice that we have to specify in the first argument which plot the color bar is associated to. The aspect=18 option is used to adjust the aspect ratio of the bar. Larger values will result in a narrower bar.

To finish the plot, we call the tight_layout() function. This adjusts the sizes of each plot, so that axis labels are displayed correctly.  

We generated 2D and 3D plots using Matplotlib and represented the results of technical computation in graphical manner.

If you want to explore other types of plots such as scatter plot or bar chart, you may read Visualizing 3D plots in Matplotlib 2.0.

Do check out the book SciPy Recipes to take advantage of other libraries of the SciPy stack and perform matrices, data wrangling and advanced computations with ease.