数据获取/choose.py

import math
import pandas as pd
import talib


def technology(df):  # 定义计算技术指标的函数
    slippage = 0.001  # 设置滑点千分之一
    commission = 0.0013  # 设置滑点千分之一
    try:
        df=df.sort_values(by="日期")  # 以日期列为索引,避免计算错误
        if "换手率" in df.columns:
            df['涨跌幅（开盘原值）'] = df['涨跌幅（开盘）'].shift(-1)
            df['涨跌幅（原值）'] = df['涨跌幅']+1
            for n in range(1, 30):
                if n == 1:
                    df[f"{n}日后总涨跌幅（未来函数）"] = df['涨跌幅（原值）'].shift(-1)-1
                if n > 1:
                    df[f"{n}日后总涨跌幅（未来函数）"] = (
                        (df[f"{n-1}日后总涨跌幅（未来函数）"]+1)*df['涨跌幅（原值）'].shift(-n))-1
            # for n in range(1, 81):
            #     # 计算加滑点之后的收益，A股扣一分钱版本
            #     df["买入（未来函数）"] = df["开盘"].apply(
            #         lambda x: math.ceil(x * (1 + slippage) * 100) / 100)
            #     df["卖出（未来函数）"] = df["开盘"].apply(lambda x: math.floor(
            #         x * (1 - slippage) * (1 - commission) * 100) / 100)
            #     df[f"{n}日后总涨跌幅（未来函数）"] = (
            #         df["卖出（未来函数）"].copy().shift(-n) / df["买入（未来函数）"]) - 1
            df[f"{5}日量比"] = df["成交量"] / \
                talib.MA(df["成交量"], timeperiod=5, matype=0)
            df[f"{40}日成交量低点"] = df['成交量'].rolling(40).min()
            df[f"{40}日成交量高点"] = df['成交量'].rolling(40).max()
            df[f"{5}日收盘高点"] = df['收盘'].rolling(40).min()
            # df = df[df[f"收盘"] > 4].copy()
        else:
            for n in range(1, 50):
                # 计算不加滑点的收益
                df[f"{n}日后总涨跌幅（未来函数）"] = (
                    df["收盘"].copy().shift(-n) / df["收盘"]) - 1
            df[f"{5}日量比"] = df["成交量"] / \
                talib.MA(df["成交量"], timeperiod=5, matype=0)
            df[f"{40}日成交量低点"] = df['成交量'].rolling(40).min()
            df[f"{40}日成交量高点"] = df['成交量'].rolling(40).max()
            df[f"{5}日收盘高点"] = df['收盘'].rolling(40).min()
            # df = df[df[f"收盘"] > 4].copy()
    except Exception as e:
        print(f"发生bug: {e}")
    return df


def rank(df):  # 计算每个标的的各个指标在当日的排名，并将排名映射到 [0, 1] 的区间中
    # 计算每个指标的排名
    for column in df.columns:  # 从大到小排序
        if (("未来函数" or "日期" or "代码") not in str(column)):
            df = pd.concat([df, (df[str(column)].rank(method="max", ascending=False) / len(df)).rename(f"{str(column)}_rank")], axis=1)
    return df


def choose(name, df):
    df = df.sort_values(by="日期")  # 以日期列为索引,避免计算错误
    code = df[df["日期"] == df["日期"].min()]["代码"]  # 获取首日标的数量，杜绝未来函数
    print(name, "首日标的数量", len(code))
    if ("股票" in name):  # 数据截取
        df = df[(df["开盘"] >= 4)].copy()  # 过滤垃圾股
        df = df[(df["涨跌幅（开盘原值）"] <= 0.09)].copy()  # 过滤垃圾股
        m = 0.01  # 设置手续费
        n = 25  # 设置持仓周期
    if ("COIN" in name):
        m = 0.001  # 设置手续费
        n = 25  # 设置持仓周期
    print(name, n)
    return df, m, n