Spaces:

pigracing
/

zg

Runtime error

zg / backtest /indicators.py

Upload 25 files

19a8231 verified 6 months ago

6.85 kB

# -*- coding: utf-8 -*-
"""
技术指标计算模块
包含：MACD, KDJ, RSI, 布林带, 均线, 量比, 换手率, 振幅, 波动率等
"""

import numpy as np
import pandas as pd
import config


def calc_all_indicators(df: pd.DataFrame) -> pd.DataFrame:
    """一次性计算全部技术指标，返回带新列的 DataFrame"""
    df = df.copy()
    
    # 均线
    for ma in config.MA_PERIODS:
        df[f"ma{ma}"] = df["close"].rolling(ma).mean()
    
    # 布林带
    df = calc_bollinger(df)
    
    # MACD
    df = calc_macd(df)
    
    # KDJ
    df = calc_kdj(df)
    
    # RSI
    df["rsi14"] = calc_rsi(df["close"], 14)
    
    # 成交量指标
    df["volume_ma5"] = df["volume"].rolling(5).mean()
    df["volume_ma20"] = df["volume"].rolling(20).mean()
    df["volume_ma60"] = df["volume"].rolling(60).mean()
    
    # 量比（当前量 / 20日均量）
    df["vol_ratio"] = df["volume"] / df["volume"].shift(1).rolling(20).mean()
    
    # 换手率（已有 turnover 列）
    if "turnover" in df.columns:
        df["turnover_ma5"] = df["turnover"].rolling(5).mean()
        df["turnover_ma20"] = df["turnover"].rolling(20).mean()
    
    # 振幅（使用前一日收盘价）
    df["pre_close"] = df["close"].shift(1)
    df["amplitude"] = (df["high"] - df["low"]) / df["pre_close"] * 100
    
    # 涨跌幅
    df["pct_chg"] = df["close"].pct_change() * 100
    
    # 波动率（近20日涨跌幅绝对值均值）
    df["volatility_20d"] = df["pct_chg"].abs().rolling(20).mean()
    
    # 最高价/最低价 滚动
    df["high_20"] = df["high"].rolling(20).max()
    df["low_20"] = df["low"].rolling(20).min()
    df["high_60"] = df["high"].rolling(60).max()
    df["low_60"] = df["low"].rolling(60).min()
    
    # 获利盘比例近似
    df["profit_ratio"] = calc_profit_ratio(df)
    
    # 筹码集中度近似（使用布林带宽度的逆指标和价格位置组合）
    df["chip_concentration"] = calc_chip_concentration(df)
    
    # V2.1.1: 两融指标（如有数据列）
    if "margin_net_buy" in df.columns:
        df["margin_ma5"] = df["margin_net_buy"].rolling(5).mean()
        df["margin_ma20"] = df["margin_net_buy"].rolling(20).mean()
    if "margin_short" in df.columns:
        df["margin_short_ma5"] = df["margin_short"].rolling(5).mean()
    
    return df


def calc_bollinger(df: pd.DataFrame, period: int = None, k: float = None) -> pd.DataFrame:
    """布林带指标"""
    period = period or config.BOLLINGER_PERIOD
    k = k or config.BOLLINGER_STD
    df = df.copy()
    df["boll_mid"] = df["close"].rolling(period).mean()
    df["boll_std"] = df["close"].rolling(period).std()
    df["boll_upper"] = df["boll_mid"] + k * df["boll_std"]
    df["boll_lower"] = df["boll_mid"] - k * df["boll_std"]
    df["boll_width"] = (df["boll_upper"] - df["boll_lower"]) / df["boll_mid"]
    df["boll_pct"] = (df["close"] - df["boll_lower"]) / (df["boll_upper"] - df["boll_lower"])
    df["boll_pct"] = df["boll_pct"].clip(0, 1)
    return df


def calc_macd(df: pd.DataFrame) -> pd.DataFrame:
    """MACD 指标"""
    df = df.copy()
    fast = config.MACD_FAST
    slow = config.MACD_SLOW
    signal = config.MACD_SIGNAL
    
    ema_fast = df["close"].ewm(span=fast, adjust=False).mean()
    ema_slow = df["close"].ewm(span=slow, adjust=False).mean()
    df["macd_dif"] = ema_fast - ema_slow
    df["macd_dea"] = df["macd_dif"].ewm(span=signal, adjust=False).mean()
    df["macd_bar"] = 2 * (df["macd_dif"] - df["macd_dea"])
    return df


def calc_kdj(df: pd.DataFrame) -> pd.DataFrame:
    """KDJ 指标"""
    df = df.copy()
    period = config.KDJ_PERIOD
    
    low_min = df["low"].rolling(period).min()
    high_max = df["high"].rolling(period).max()
    
    rsv = (df["close"] - low_min) / (high_max - low_min) * 100
    rsv = rsv.fillna(50)
    
    df["kdj_k"] = rsv.ewm(com=2, adjust=False).mean()
    df["kdj_d"] = df["kdj_k"].ewm(com=2, adjust=False).mean()
    df["kdj_j"] = 3 * df["kdj_k"] - 2 * df["kdj_d"]
    return df


def calc_rsi(close_series: pd.Series, period: int = 14) -> pd.Series:
    """RSI 指标"""
    delta = close_series.diff()
    gain = delta.where(delta > 0, 0.0)
    loss = (-delta).where(delta < 0, 0.0)
    
    avg_gain = gain.ewm(span=period, adjust=False).mean()
    avg_loss = loss.ewm(span=period, adjust=False).mean()
    
    rs = avg_gain / avg_loss.replace(0, np.nan)
    rsi = 100 - (100 / (1 + rs))
    return rsi


def calc_profit_ratio(df: pd.DataFrame) -> pd.Series:
    """
    获利盘比例近似计算
    基于：当前价格在60日区间内的位置
    profit_ratio = (close - low_60) / (high_60 - low_60)
    """
    low_60 = df["low"].rolling(60).min()
    high_60 = df["high"].rolling(60).max()
    ratio = (df["close"] - low_60) / (high_60 - low_60).replace(0, np.nan)
    return ratio.clip(0, 1)


def calc_chip_concentration(df: pd.DataFrame) -> pd.Series:
    """
    筹码集中度近似计算
    当价格在低位且振幅收窄、成交量萎缩时，视为筹码集中
    指标含义：值越高表示筹码越集中（0~1）
    """
    low_60 = df["low"].rolling(60).min()
    high_60 = df["high"].rolling(60).max()
    
    # 价格在低位（0~0.3区间为新低附近集中）
    price_pos = (df["close"] - low_60) / (high_60 - low_60).replace(0, np.nan)
    
    # 振幅收窄（近10日振幅均值与60日振幅均值的比值）
    amp_10 = (df["high"] - df["low"]).rolling(10).mean()
    amp_60 = (df["high"] - df["low"]).rolling(60).mean()
    amp_ratio = amp_10 / amp_60.replace(0, np.nan)
    amp_ratio = amp_ratio.clip(0, 1)
    
    # 成交量萎缩（近5日均量是60日均量的比例，越低越集中）
    vol_ratio = df["volume"].rolling(5).mean() / df["volume"].rolling(60).mean().replace(0, np.nan)
    vol_ratio = 1 - vol_ratio.clip(0, 1)  # 反转：量越小值越大
    
    # 加权综合
    concentration = (1 - price_pos.fillna(0.5)) * 0.4 + (1 - amp_ratio.fillna(0.5)) * 0.3 + vol_ratio.fillna(0.5) * 0.3
    return concentration.clip(0, 1)


def get_support_resistance(df: pd.DataFrame) -> dict:
    """
    识别支撑位和压力位
    基于前期高点和低点
    """
    if len(df) < 60:
        return {"support": None, "resistance": None}
    
    recent = df.tail(60)
    # 支撑位：60日最低点
    support = recent["low"].min()
    # 压力位：60日最高点
    resistance = recent["high"].max()
    
    # 均线支撑：20日线位置
    ma20 = recent["close"].rolling(20).mean().iloc[-1]
    
    return {
        "support": round(support, 2),
        "resistance": round(resistance, 2),
        "ma20": round(ma20, 2),
    }

	# -- coding: utf-8 --
	"""
	技术指标计算模块
	包含：MACD, KDJ, RSI, 布林带, 均线, 量比, 换手率, 振幅, 波动率等
	"""

	import numpy as np
	import pandas as pd
	import config


	def calc_all_indicators(df: pd.DataFrame) -> pd.DataFrame:
	"""一次性计算全部技术指标，返回带新列的 DataFrame"""
	df = df.copy()

	# 均线
	for ma in config.MA_PERIODS:
	df[f"ma{ma}"] = df["close"].rolling(ma).mean()

	# 布林带
	df = calc_bollinger(df)

	# MACD
	df = calc_macd(df)

	# KDJ
	df = calc_kdj(df)

	# RSI
	df["rsi14"] = calc_rsi(df["close"], 14)

	# 成交量指标
	df["volume_ma5"] = df["volume"].rolling(5).mean()
	df["volume_ma20"] = df["volume"].rolling(20).mean()
	df["volume_ma60"] = df["volume"].rolling(60).mean()

	# 量比（当前量 / 20日均量）
	df["vol_ratio"] = df["volume"] / df["volume"].shift(1).rolling(20).mean()

	# 换手率（已有 turnover 列）
	if "turnover" in df.columns:
	df["turnover_ma5"] = df["turnover"].rolling(5).mean()
	df["turnover_ma20"] = df["turnover"].rolling(20).mean()

	# 振幅（使用前一日收盘价）
	df["pre_close"] = df["close"].shift(1)
	df["amplitude"] = (df["high"] - df["low"]) / df["pre_close"] * 100

	# 涨跌幅
	df["pct_chg"] = df["close"].pct_change() * 100

	# 波动率（近20日涨跌幅绝对值均值）
	df["volatility_20d"] = df["pct_chg"].abs().rolling(20).mean()

	# 最高价/最低价滚动
	df["high_20"] = df["high"].rolling(20).max()
	df["low_20"] = df["low"].rolling(20).min()
	df["high_60"] = df["high"].rolling(60).max()
	df["low_60"] = df["low"].rolling(60).min()

	# 获利盘比例近似
	df["profit_ratio"] = calc_profit_ratio(df)

	# 筹码集中度近似（使用布林带宽度的逆指标和价格位置组合）
	df["chip_concentration"] = calc_chip_concentration(df)

	# V2.1.1: 两融指标（如有数据列）
	if "margin_net_buy" in df.columns:
	df["margin_ma5"] = df["margin_net_buy"].rolling(5).mean()
	df["margin_ma20"] = df["margin_net_buy"].rolling(20).mean()
	if "margin_short" in df.columns:
	df["margin_short_ma5"] = df["margin_short"].rolling(5).mean()

	return df


	def calc_bollinger(df: pd.DataFrame, period: int = None, k: float = None) -> pd.DataFrame:
	"""布林带指标"""
	period = period or config.BOLLINGER_PERIOD
	k = k or config.BOLLINGER_STD
	df = df.copy()
	df["boll_mid"] = df["close"].rolling(period).mean()
	df["boll_std"] = df["close"].rolling(period).std()
	df["boll_upper"] = df["boll_mid"] + k * df["boll_std"]
	df["boll_lower"] = df["boll_mid"] - k * df["boll_std"]
	df["boll_width"] = (df["boll_upper"] - df["boll_lower"]) / df["boll_mid"]
	df["boll_pct"] = (df["close"] - df["boll_lower"]) / (df["boll_upper"] - df["boll_lower"])
	df["boll_pct"] = df["boll_pct"].clip(0, 1)
	return df


	def calc_macd(df: pd.DataFrame) -> pd.DataFrame:
	"""MACD 指标"""
	df = df.copy()
	fast = config.MACD_FAST
	slow = config.MACD_SLOW
	signal = config.MACD_SIGNAL

	ema_fast = df["close"].ewm(span=fast, adjust=False).mean()
	ema_slow = df["close"].ewm(span=slow, adjust=False).mean()
	df["macd_dif"] = ema_fast - ema_slow
	df["macd_dea"] = df["macd_dif"].ewm(span=signal, adjust=False).mean()
	df["macd_bar"] = 2 * (df["macd_dif"] - df["macd_dea"])
	return df


	def calc_kdj(df: pd.DataFrame) -> pd.DataFrame:
	"""KDJ 指标"""
	df = df.copy()
	period = config.KDJ_PERIOD

	low_min = df["low"].rolling(period).min()
	high_max = df["high"].rolling(period).max()

	rsv = (df["close"] - low_min) / (high_max - low_min) * 100
	rsv = rsv.fillna(50)

	df["kdj_k"] = rsv.ewm(com=2, adjust=False).mean()
	df["kdj_d"] = df["kdj_k"].ewm(com=2, adjust=False).mean()
	df["kdj_j"] = 3 * df["kdj_k"] - 2 * df["kdj_d"]
	return df


	def calc_rsi(close_series: pd.Series, period: int = 14) -> pd.Series:
	"""RSI 指标"""
	delta = close_series.diff()
	gain = delta.where(delta > 0, 0.0)
	loss = (-delta).where(delta < 0, 0.0)

	avg_gain = gain.ewm(span=period, adjust=False).mean()
	avg_loss = loss.ewm(span=period, adjust=False).mean()

	rs = avg_gain / avg_loss.replace(0, np.nan)
	rsi = 100 - (100 / (1 + rs))
	return rsi


	def calc_profit_ratio(df: pd.DataFrame) -> pd.Series:
	"""
	获利盘比例近似计算
	基于：当前价格在60日区间内的位置
	profit_ratio = (close - low_60) / (high_60 - low_60)
	"""
	low_60 = df["low"].rolling(60).min()
	high_60 = df["high"].rolling(60).max()
	ratio = (df["close"] - low_60) / (high_60 - low_60).replace(0, np.nan)
	return ratio.clip(0, 1)


	def calc_chip_concentration(df: pd.DataFrame) -> pd.Series:
	"""
	筹码集中度近似计算
	当价格在低位且振幅收窄、成交量萎缩时，视为筹码集中
	指标含义：值越高表示筹码越集中（0~1）
	"""
	low_60 = df["low"].rolling(60).min()
	high_60 = df["high"].rolling(60).max()

	# 价格在低位（0~0.3区间为新低附近集中）
	price_pos = (df["close"] - low_60) / (high_60 - low_60).replace(0, np.nan)

	# 振幅收窄（近10日振幅均值与60日振幅均值的比值）
	amp_10 = (df["high"] - df["low"]).rolling(10).mean()
	amp_60 = (df["high"] - df["low"]).rolling(60).mean()
	amp_ratio = amp_10 / amp_60.replace(0, np.nan)
	amp_ratio = amp_ratio.clip(0, 1)

	# 成交量萎缩（近5日均量是60日均量的比例，越低越集中）
	vol_ratio = df["volume"].rolling(5).mean() / df["volume"].rolling(60).mean().replace(0, np.nan)
	vol_ratio = 1 - vol_ratio.clip(0, 1) # 反转：量越小值越大

	# 加权综合
	concentration = (1 - price_pos.fillna(0.5)) * 0.4 + (1 - amp_ratio.fillna(0.5)) * 0.3 + vol_ratio.fillna(0.5) * 0.3
	return concentration.clip(0, 1)


	def get_support_resistance(df: pd.DataFrame) -> dict:
	"""
	识别支撑位和压力位
	基于前期高点和低点
	"""
	if len(df) < 60:
	return {"support": None, "resistance": None}

	recent = df.tail(60)
	# 支撑位：60日最低点
	support = recent["low"].min()
	# 压力位：60日最高点
	resistance = recent["high"].max()

	# 均线支撑：20日线位置
	ma20 = recent["close"].rolling(20).mean().iloc[-1]

	return {
	"support": round(support, 2),
	"resistance": round(resistance, 2),
	"ma20": round(ma20, 2),
	}