通过pandas.dataframe绘制热力图,并标出颜色最深即z轴数据最大的点。
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
class Heatmap:
def __init__(self, data, marker='*', marker_color='white', marker_size=10):
self.data = data
self.size = len(data)
self.marker = marker
self.marker_color = marker_color
self.marker_size = marker_size
self.title = None
self.xlabel = None
self.ylabel = None
self.cbar_label = None
self.cbar_location = "right"
def plot_heatmap(self):
# 创建一个新的图像
fig, ax = plt.subplots()
# 使用imshow创建热力图
cax = ax.imshow(self.data, cmap='turbo', interpolation='nearest')
# 找到最大值的坐标并绘制标记
max_pos = self._get_max_pos()
self._plot_marker(ax, max_pos)
# 设置刻度
self.set_ticks(ax)
# 设置标题和x,y轴标签
ax.set_title(self.title)
ax.set_xlabel(self.xlabel)
ax.set_ylabel(self.ylabel)
# 设置colorbar
cbar = plt.colorbar(cax, ax=ax, label=self.cbar_label, pad=0.01, fraction=0.046)
cbar.ax.yaxis.set_label_position(self.cbar_location) # 设置colorbar的标签位置
# 显示图像
plt.show()
def _get_max_pos(self):
# 找到最大值的坐标
max_value = self.data.values.max()
max_pos = np.unravel_index(self.data.values.argmax(), self.data.values.shape)
return max_pos
def _plot_marker(self, ax, position):
# 在指定位置绘制标记
ax.plot(position[1], position[0], marker=self.marker, color=self.marker_color, markersize=self.marker_size)
def set_labels(self, title=None, xlabel=None, ylabel=None):
# 设置标题和x,y轴标签
self.title = title
self.xlabel = xlabel
self.ylabel = ylabel
def set_ticks(self, ax):
# 动态设置x轴和y轴的刻度
if self.size <= 10:
interval = 1
else:
interval = self.size // 10
ticks = np.arange(0, len(self.data.columns), interval)
ax.set_xticks(ticks)
ax.set_yticks(ticks)
ax.set_xticklabels(self.data.columns[ticks])
ax.set_yticklabels(self.data.index[ticks])
def set_colorbar(self, label=None, location="right"):
# 设置colorbar的标签和位置
self.cbar_label = label
self.cbar_location = location
# 创建一个随机数据 DataFrame,使用你所要求的行名和列名
index = range(-20, 21)
columns = range(-20, 21)
data = pd.DataFrame(np.random.rand(41, 41), index=index, columns=columns)
# 创建 Heatmap 对象
heatmap = Heatmap(data)
# 设置热力图标题和 x, y 轴标签
heatmap.set_labels('Heatmap Title', 'X-axis Label', 'Y-axis Label')
# 设置 colorbar 的标签和位置
heatmap.set_colorbar('Colorbar Label', 'right')
# 创建并显示热力图
heatmap.plot_heatmap()