Python package to plot two heatmaps in one (split each square into two triangles)

plt.tripcolor colors a mesh of triangles similar to how plt.pcolormesh colors a rectangular mesh. Also similar to pcolormesh, care has to be taken that there is one row and one column of vertices less than there are triangles. Furthermore, the arrays need to be made 1D (np.ravel). All this renumbering to 1D can be a bit tricky.

As an example, the code below creates a coloring depending on x*y mod 10 and uses two different colormaps for the upper and the lower triangles.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.tri import Triangulation

M = 30
N = 20
x = np.arange(M + 1)
y = np.arange(N + 1)
xs, ys = np.meshgrid(x, y)
zs = (xs * ys) % 10
zs = zs[:-1, :-1].ravel()

triangles1 = [(i + j*(M+1), i+1 + j*(M+1), i + (j+1)*(M+1)) for j in range(N) for i in range(M)]
triangles2 = [(i+1 + j*(M+1), i+1 + (j+1)*(M+1), i + (j+1)*(M+1)) for j in range(N) for i in range(M)]
triang1 = Triangulation(xs.ravel(), ys.ravel(), triangles1)
triang2 = Triangulation(xs.ravel(), ys.ravel(), triangles2)
img1 = plt.tripcolor(triang1, zs, cmap=plt.get_cmap('inferno', 10), vmax=10)
img2 = plt.tripcolor(triang2, zs, cmap=plt.get_cmap('viridis', 10), vmax=10)

plt.colorbar(img2, ticks=range(10), pad=-0.05)
plt.colorbar(img1, ticks=range(10))
plt.xlim(x[0], x[-1])
plt.ylim(y[0], y[-1])
plt.xticks(x, rotation=90)
plt.yticks(y)
plt.show()

example plot

PS: to have the integer ticks nicely in the center of the cells (instead of at their borders), following changes would be needed:

triang1 = Triangulation(xs.ravel()-0.5, ys.ravel()-0.5, triangles1)
triang2 = Triangulation(xs.ravel()-0.5, ys.ravel()-0.5, triangles2)

# ...
plt.xlim(x[0]-0.5, x[-1]-0.5)
plt.ylim(y[0]-0.5, y[-1]-0.5)
plt.xticks(x[:-1], rotation=90)
plt.yticks(y[:-1]) 

Leave a Comment