Python量化因子测算与绘图超详细流程代码

# 测算用
import numpy as np
import pandas as pd
from copy import deepcopy
from tqdm import tqdm
from datetime import datetime
import talib
# 绘图用
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
%matplotlib inline
# 绘图现实中文
sns.set()
plt.rcParams["figure.figsize"] = (20,10)
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # 当前字体支持中文
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
# 其他
import warnings
warnings.filterwarnings("ignore")

import os
def readfile(path, limit=None):
    files = os.listdir(path)
    file_list = []
    for file in files:  # 遍历文件夹
        if not os.path.isdir(file):
            file_list.append(path + '/' + file)
    if limit:
        return file_list[:limit]
    return file_list
stock_dict = {}
for _file in tqdm(readfile("../data/stock_data")):
    if not _file.endswith(".pkl"):
        continue
    # TODO 这里可以添加筛选，是否需要将当前的股票添加到测算的股票池中
    file_df = pd.read_pickle(_file)
    file_df.set_index(["日期"], inplace=True)
    file_df.index.name = ""
    file_df.index = pd.to_datetime(file_df.index)
    file_df.rename(columns={'开盘':'open',"收盘":"close","最高":"high","最低":"low","成交量":"volume"},inplace=True)
    stock_code = _file.split("/")[-1].replace(".pkl", '')
    # TODO 这里可以添加日期，用来截取一部分数据
    stock_dict[stock_code] = file_df

for _index,_stock_df in tqdm(stock_dict.items()):
    measure_df = deepcopy(_stock_df)

# 开始测算
trade_record_list = []
this_trade:dict = None
for _mea_i, _mea_series in measure_df.iterrows(): # 循环每一天
    if 发出买入信号:
        if this_trade is None:  # 当前没有持仓，则买入
            this_trade = {
                "buy_date": _mea_i,
                "close_record": [_mea_series['close']],
            }
    elif 发出卖出信号:
        if this_trade is not None:  # 要执行卖出
            this_trade['sell_date'] = _mea_i
            this_trade['close_record'].append(_mea_series['close'])
            trade_record_list.append(this_trade)
            this_trade = None
    else:
        if this_trade is not None:  # 当前有持仓
            this_trade['close_record'].append(_mea_series['close'])

{
    'buy_date': Timestamp('2015-08-31 00:00:00'), # 买入时间
    'close_record': [41.1,42.0,40.15,40.65,36.6,32.97], # 收盘价的记录
    'sell_date': Timestamp('2015-10-12 00:00:00')} # 卖出时间
    # TODO 也可以添加自定义记录的指标
}

trade_record_df = pd.DataFrame(trade_record_list)
for _,_trade_series in trade_record_df.iterrows():
    trade_record_df.loc[_i,'年化收益率'] = (_trade_series['close_record'][-1] - _trade_series['close_record'][0])/_trade_series['close_record'][0]/(_trade_series['sell_date'] - _trade_series['buy_date']).days * 365 # 年化收益
    # TODO 这里根据自己想要的结果添加更多的测算指标

# 清理绘图缓存
plt.cla()
plt.clf()
# 开始绘图
plt.figure(figsize=(10, 14), dpi=100)
# 使用seaborn绘制胜率图
fig = sns.heatmap(pd.DataFrame(total_measure_record).T.round(2), annot=True, cmap="RdBu_r",center=0.5)
plt.title("胜率图")
scatter_fig = fig.get_figure()
# 保存到本地
scatter_fig.savefig("胜率图")
scatter_fig.show() # 最后显示