Join us and get access to thousands of tutorials and a community of expert Pythonistas.
This lesson is for members only. Join us and get access to thousands of tutorials and a community of expert Pythonistas.
Creating a Colored Matrix
A heatmap is a type of visualization that looks really cool. matplotlib may not have any kind of heatmap method you can apply to an axes, but you can create a grid of colored squares from a two-dimensional ndarray.
To get started, you’ll need to import an obscure function:
from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable
00:00 One of my favorite types of visualizations is a heat map. I can’t show a picture of one here, but look one up if you’ve never seen it before. They’re really cool.
00:12
matplotlib doesn’t have any sort of heat map method we can apply to an Axes, but we can create a grid of colored squares from a two-dimensional ndarray.
00:23
Let’s do that. I’m here in a new file called plot4.py and I’ve already imported pyplot and numpy. We need one more import, this really obscure function called make_axes_locatable().
00:41
We can import that like this: from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable. We’ll plot two separate two-dimensional ndarrays.
01:05 You can think of these like matrices. I’ll show them onscreen when we’re coding to help make things more clear.
01:13
The first data set will be called x, and it’ll be obtained by calling np.diag(), passing in np.arange(2, 12). I’ll flip this vertically with [::-1].
01:32 This results in a 10 by 10 matrix of zeros with a line of non-zero numbers running through it diagonally.
01:43 This next line of code will add another diagonal line in the opposite direction.
01:50
Now, I’ll create one more data set called x2, which will come from np.arange(). Same size as x and I’ll reshape that ndarray so it matches the 10 by 10 grid we had before.
02:08 Next, we have to disable all the tick marks and numeric labels that appear on the sides. We’ll use some fancy dictionary comprehension for this. First, I’m going to make a tuple containing a string for each side.
02:25
Now I’ll create a new dictionary called nolabels,
02:31
which will contain a key for each side with values of False.
02:38
That takes care of the ticks, and fortunately, the labels work in a similar way. We just have to add four more sides to this dictionary, with the string 'label' appended to the front of each one. And that looks like this.
02:55
To actually plot this data, we’re going to use a context manager with no grid. That’s the rc_context() function, passing in this dictionary.
03:08
Inside this with block, I’ll create a Figure and two Axes. We can call the .matshow() method on ax1 to create a color-coded matrix out of our x array. We’ll do something similar for ax2, but this time we’ll use a reverse Red-Yellow-Green color map.
03:32 A color map simply maps a color to a numeric value, so if before low numbers were red and high numbers were green, now low numbers are green and high numbers are red.
03:47
That’s what reversing it does. This function actually returns an AxesImage object, so we’ll capture that to use it later on. For each Axes object, we need to use our nolabels dictionary to disable the tick marks and labels.
04:05
So I’ll write a for loop to iterate through the Axes. Now we can use the tick_params() function, specifying 'both' for the axis and 'both' for the which value.
04:21
And finally, our dictionary unpacked. At this point, we’ve got a grid of colored squares, but it still needs labels on our diagonal color strips. To do this, I’ll write for i, j in zip(*x.nonzero()):.
04:46
This will let us iterate through two lists that represent the x and y coordinates of each non-zero value in our x matrix.
04:56
Then we can say ax1.text(), passing in j and i for the coordinates, and the label x[i, j], which will fetch the value behind each color.
05:14
We’ll set the color to 'white' and set the horizontal and vertical alignment to 'center' so that the text label sits in the center of its colored square.
05:27
The last thing we want to do is create a color bar. To do that, we must first make the second Axes locatable, create a divider at the right border of that Axes, append a new Axes, and fill that with a color bar.
05:45
First, I’ll create a variable called divider and use the make_axes_locatable() function we imported at the top to make ax2 locatable.
05:58
Now I’ll set cax, short for child axes, equal to divider.append(), passing in 'right' for a side, a size of '5%', and no padding—a.k.a. spacing between the axes and the color bar.
06:19
To create the color bar, I’ll call the colorbar() function from pyplot, passing in the image object from earlier and the child axes it should occupy.
06:30
The last thing I’m going to do is create a subtitle for the entire figure, which is a centered title that sits at the top. I’ll make it say 'Heatmaps with `Axes.matshow`', and I’ll give it a fontsize of 16.
06:50 All that’s left is to show it. And would you look at that!
Desmond on April 21, 2020
There are a lot of concepts covered in this presentation. Whenever, I needed some clarification on the functionality of the Matplotlib methods that were used in this presentation, I went to matplotlib.org/users/index.html and searched through the links for Documentation, Examples, and Tutorials. The Examples link has a gallery of plot examples and the source code is available. The Tutorials link covers core ideas with code examples, as well. It’s a lot of material to review, but it looks very promising.
Satish on Nov. 4, 2021
I too found this lesson difficult to follow. It would have been helpful if the explanations were more insightful.
Bartosz Wilk on Aug. 18, 2022
I cannot run this code, I get an error:
Process finished with exit code 137 (interrupted by signal 9: SIGKILL)
I use python 3.10 and macos 12.5.1, numpy and matplotlib installed by pip.
Become a Member to join the conversation.
Qasim Albaqali on March 17, 2020
I don’t know. This course was a bit difficult to follow. The instructor felt like he was just talking about what he is doing, but not explaining the concepts behind it, and how the things actually work. It was difficult to understand everything that was being done. A lot of new concepts are being introduced that are difficult to grasp. Just my thoughts.