python绘制热力图,数据来源pandas.dataframe

通过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()


python绘制热力图,数据来源pandas.dataframe_第1张图片

你可能感兴趣的:(python,python,pandas,开发语言)