装饰器+matplotlib实现快速生成线性相关热力图

装饰器部分：

from plt.Bar import init_plt
import matplotlib.pyplot as plt
import itertools
import numpy as np


def heat(func):
    """
    热力图
    func() return dict
    dict = {
    "cm": numpy.array,
    "classes": list,
    "title": str
    }
    """
    @init_plt
    def wrapper(**kwargs):
        d = func()
        cm = d.get("cm", None)
        classes = d.get("classes", None)
        title = d.get("title", None)
        # 设置配色方案
        cmap = plt.cm.Blues
        plt.imshow(cm, cmap=cmap,)
        plt.colorbar()  # 显示颜色条
        # 设置x，y轴刻度标签
        tick_marks = np.arange(len(classes))
        plt.xticks(tick_marks, classes)
        plt.yticks(tick_marks, classes)
        # 将数值显示在指定位置
        for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
            plt.text(j, i, cm[i, j], horizontalalignment="center")
        # plt.grid()
        plt.title(title)
        plt.show()

    return wrapper

        init_plt见上一章

实现部分：

@heat
def heat2():
    data = pd.read_excel(r'D:\pyproject\data_analysis_mining\01-数据和代码\chapter3\demo\data\catering_sale_all.xls')
    classes = data.columns[1:5]
    print(classes)
    corr = data.corr().iloc[1:5, 1:5]
    arr = []
    for index in range(len(classes)):
        list_c = corr.iloc[index,].tolist()
        arr.append(list_c)
    return {
        'cm': np.array(arr),
        'classes': classes,
        'title': '线性相关图'
    }


heat2()

        # 这里我没有做格式化