Step1:写独立模块utils
首先,写一个独立的python文件,汇总一些常用的函数,实现数据分组、仓位/权重计算等功能。
废话不说,上代码。
# 导入必要的模块
import pandas as pd
import numpy as np
from scipy.stats import rankdata
然后,写一个函数,根据某条件(如:市盈率)对传入的数据进行分组。
# 定义一个函数,根据市盈率PE Ratio对原始数据进行分组
def group_1d(arr, nog):
"""
根据某条件对一维数据进行分组
:param arr: 分组前的原始数据
:param nog: 组数,要把原始数据分成几组
:return: 以Numpy数组形式返回分组的结果
"""
# 创建一个全零数组,用于存放最终的分组结果
results = np.zeros_like(arr, dtype=int)
# 接受一个数组,返回该数组中元素对应的顺序编号,如果元素相同,则返回该元素对应顺序的平均值
ranks = rankdata(arr)
total_num = len(arr) # 数据总量
nan_num = np.isnan(arr).sum() # nan数据量
no_nan_num = total_num - nan_num # 非空数据量
# 下面计算每个分组中有多少个元素,考虑到实际意义,在此四舍五入
d = round(no_nan_num / nog)
# 下面开始分组
# 除了最后一组,其余分组中包含的全部元素个数是相同的
for i in range(1, nog):
rank1 = 1 + (i - 1) * d # 当前分组开始元素对应的顺序编号值
rank2 = rank1 + d # 当前分组结束元素对应的顺序编号值
results[(ranks >= rank1) & (ranks < rank2)] = i # 写入对应元素的分组结果
# 最后一组的分组情况
rank1 = 1 + (nog - 1) * d
rank2 = no_nan_num
results[(ranks >= rank1) & (ranks < rank2)] = nog # 剩下的全部是最后一组
return results # 返回分组结果
再写一个函数,对二维数据分组。
def grouping(arr2d, nog):
"""
对二维数据进行分组
:arr2d:传入需要分组的二维数据集
:nog:组数
:return:返回分组的结果
"""
results = np.zeros_like(arr2d, dtype=int) # 创建一个全零数组,用于保存分组的结果
# 分组
for i in range(len(results)): # 对第i行数据进行分组
results[i] = group_1d(arr2d[i], nog) # 调用函数,对某一行数据进行分组
return results # 返回二维数据的分组结果
再写一个函数,根据分组结果计算权重。
def getting_weight(group