Matplotlib: Тепловые карты

import numpy as np

import matplotlib.pyplot as plt

from matplotlib.colors import TwoSlopeNorm # позволяет задать "середину" шкалы цветов

# 1. Генерируем случайные данные (100 строк, 6 переменных)

rng = np.random.RandomState(1)

data = rng.randn(100, 6) # нормальное распределение

# 2. Вычисляем корреляционную матрицу (6x6)

corr = np.corrcoef(data, rowvar=False)  

labels = [f'V{i}' for i in range(corr.shape[0])]

fig, ax = plt.subplots()

# 3. Нормализация шкалы так, чтобы 0 был белым цветом, -1 и 1 — крайние тона

norm = TwoSlopeNorm(vmin=-1, vcenter=0, vmax=1)

# 4. Рисуем тепловую карту

im = ax.imshow(corr, cmap='RdBu_r', norm=norm)

# 5. Настраиваем подписи осей

ax.set_xticks(np.arange(len(labels)))

ax.set_yticks(np.arange(len(labels)))

ax.set_xticklabels(labels, rotation=45)

ax.set_yticklabels(labels)

# 6. Добавляем текстовые значения корреляций в каждую ячейку

for i in range(corr.shape[0]):

  for j in range(corr.shape[1]):

    ax.text(j, i, f"{corr[i, j]:.2f}", ha='center', va='center', fontsize=8)

# 7. Добавляем цветовую шкалу

fig.colorbar(im, ax=ax)

plt.show()