量化python：使用热力图heatmap绘制胜率图方法及工具函数

胜率图是分析策略的一种图形，对于胜率的分析需要三种数据：

1. 策略/参数集：表现在胜率图的横纵轴含义

2. 对比的场景编号：表现在对比两个策略时使用的场景；通常见于不同的时间、不同的参数、不同的周期

3. 场景对应的值：场景对应的值，可能代表收益率、模型的分数等

使用seaborn可以进行如下绘图：

def plot_victory_rate(dataframe, hue: str, compare_index: str, compare_value: str):
    """绘制胜率图
    :param dataframe: 数据源
    :param hue: 合并/展示的列名
    :param compare_index: 场景编号对应在dataframe中的列名称
    :param compare_value: 场景值对应在dataframe中的列名称
    """
    total_mark_index = dataframe[hue].value_counts().index
    one_total_victory_matrix = pd.DataFrame(0, index=total_mark_index, columns=total_mark_index)
    for _first_mark in tqdm(total_mark_index):
        for _second_mark in total_mark_index:
            first_df = dataframe[dataframe[hue] == _first_mark].sort_values(compare_index)
            second_df = dataframe[dataframe[hue] == _second_mark].sort_values(compare_index)
            one_total_victory_matrix.loc[_first_mark, _second_mark] = sum(
                (first_df[compare_value].values - second_df[compare_value].values) >= 0) / first_df.shape[0]
    # 开始绘图
    plt.cla()
    plt.clf()
    plt.figure(figsize=(10, 10), dpi=100)
    fig = sns.heatmap(one_total_victory_matrix.round(2), annot=True, cmap="RdBu_r")
    plt.title("胜率图")
    scatter_fig = fig.get_figure()
    scatter_fig.show()

示例代码

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from tqdm import tqdm

sns.set_style("whitegrid")  # 需要先写这个，否则plt的配置会被覆盖掉

plt.rcParams["font.sans-serif"] = ["Arial Unicode MS"]  # drop_duplicates持中文
plt.rcParams["figure.figsize"] = 20, 10
plt.rcParams["axes.unicode_minus"] = False  # 解决保存图像是负号'-'显示为方块的问题


def plot_victory_rate(dataframe, hue: str, compare_index: str, compare_value: str):
    """绘制胜率图
    :param dataframe: 数据源
    :param hue: 合并/展示的列名
    :param compare_index: 场景编号对应在dataframe中的列名称
    :param compare_value: 场景值对应在dataframe中的列名称
    """
    total_mark_index = dataframe[hue].value_counts().index
    one_total_victory_matrix = pd.DataFrame(0, index=total_mark_index, columns=total_mark_index)
    for _first_mark in tqdm(total_mark_index):
        for _second_mark in total_mark_index:
            first_df = dataframe[dataframe[hue] == _first_mark].sort_values(compare_index)
            second_df = dataframe[dataframe[hue] == _second_mark].sort_values(compare_index)
            one_total_victory_matrix.loc[_first_mark, _second_mark] = sum(
                (first_df[compare_value].values - second_df[compare_value].values) >= 0) / first_df.shape[0]
    # 开始绘图
    plt.cla()
    plt.clf()
    plt.figure(figsize=(10, 10), dpi=100)
    fig = sns.heatmap(one_total_victory_matrix.round(2), annot=True, cmap="RdBu_r")
    plt.title("胜率图")
    scatter_fig = fig.get_figure()
    scatter_fig.show()


def main():
    one_total_df = pd.DataFrame({
        "策略": [i for i in "abcde"] * 10,  # 策略
        "对比的场景编号": [i for i in range(10) for _ in range(5)],
        "对比的场景值": np.random.random(50),  # 具体值
    })
    plot_victory_rate(one_total_df, "策略", "对比的场景编号", "对比的场景值")


if __name__ == '__main__':
    main()

得到图像如下所示：