Выравнивание дендрограммы с тепловой картой с помощью matplotlib

Я хотел бы создать рисунок, который сочетает в себе дендрограмму и тепловую карту. Пакет в R делает это очень хорошо, но я не могу найти способ сделать это в python. Вот пример из блога:

Прямо сейчас у меня есть скрипт на Python, который выглядит так:

%matplotlib inline
from sklearn.metrics.pairwise import pairwise_distances
from scipy.spatial.distance import squareform
from scipy.cluster.hierarchy import linkage
from scipy.cluster.hierarchy import dendrogram
import numpy as np
import matplotlib.pyplot as plt

X = np.random.rand(4, 8)

Z = linkage(squareform(np.around(pairwise_distances(X), 4)))

fig = plt.figure()
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
dendrogram(Z, ax=ax1)
ax2.pcolor(X, cmap=plt.cm.Reds)

И полученная на основе этого цифра показана следующим образом:

Я хотел бы выровнять ось y тепловой карты с осью x дендрограммы. Какие-либо предложения?