Chan/chanlun/pipeline/builders/indicators.py

"""TF_DF builder mixin — 由 split_tfdf_builders 自动生成，逻辑与原 TF_DF 一致。"""
from __future__ import annotations

from datetime import timedelta
from decimal import Decimal

import numpy as np
import pandas as pd
from chanlun.indicators import ta
from pandas import DataFrame

from chanlun.core.ChanBI import ChanBI
from chanlun.core.ChanBIZS import ChanBIZS
from chanlun.core.ChanBSP import ChanBSP
from chanlun.core.ChanEnum import (
    Chan_BI_DIR,
    Chan_BSP_DIR,
    Chan_BSP_TYPE,
    Chan_FX_TYPE,
    Chan_K_DIR,
    Chan_KLC_FX,
    Chan_KLC_STATE,
    Chan_KLINE_DIR,
    Chan_KLU_PATTERN,
    Chan_PRICE_TREND,
    Chan_SEG_DIR,
    Chan_ZS_DIR,
)
from chanlun.core.ChanKLC import ChanKLC
from chanlun.core.ChanKLU import ChanKLU
from chanlun.core.ChanSBI import ChanSBI
from chanlun.core.ChanSEG import ChanSEG
from chanlun.core.ChanZS import ChanZS, ChanZS_Big
from chanlun.indicators.ChanMACD import ChanMACD


class IndicatorsBuilderMixin:
	def get_ema52(self, index=-1):
		if self.klu_list:
			ema52_value = self.klu_list[index].ema52
			# 处理NaN值
			if pd.isna(ema52_value) or ema52_value is None:
				return None
			return float(ema52_value)
		return None

	def get_ema24(self, index=-1):
		if self.klu_list:
			ema24_value = self.klu_list[index].ema24
			# 处理NaN值
			if pd.isna(ema24_value) or ema24_value is None:
				return None
			return float(ema24_value)
		return None

	def add_indicators(self, df):
		"""算指标并一次性挂到 df 上。

		这里不逐列 `df['x'] = ...`：那样每一列都触发一次 BlockManager 插入，
		30 多列的开销比全部 TA 计算本身还大（2001 行实测 TA 合计 2.5ms，
		逐列赋值 3.6ms）。增量路径每根都要走一遍，这笔开销是白付的。
		"""
		fast = 26
		slow = 52
		period = 9
		macd = ta.MACD(df, fastperiod=fast, slowperiod=slow, signalperiod=period)
		bb365 = ta.BBANDS(df, timeperiod=365, nbdevup=3.0, nbdevdn=3.0, matype=0)
		bb120 = ta.BBANDS(df, timeperiod=120, nbdevup=3.0, nbdevdn=3.0, matype=0)
		bb30 = ta.BBANDS(df, timeperiod=41, nbdevup=2.3, nbdevdn=2.3, matype=0)
		bb302 = ta.BBANDS(df, timeperiod=41, nbdevup=2.0, nbdevdn=2.0, matype=0)
		bb30 = ta.BBANDS(df, timeperiod=20, nbdevup=2.0, nbdevdn=2.0, matype=0)
		bb302 = ta.BBANDS(df, timeperiod=20, nbdevup=2.0, nbdevdn=2.0, matype=0)
		bb2633 = ta.BBANDS(df, timeperiod=26, nbdevup=3.0, nbdevdn=3.0, matype=0)
		# 计算布林带中轨（移动平均线）
		bb30_middle = ta.SMA(df, timeperiod=90)
		
		# 手动计算布林带 %B 指标 (BBP)
		# %B = (Price - Lower Band) / (Upper Band - Lower Band)
		bbp365 = (df['close'] - bb365['lowerband']) / (bb365['upperband'] - bb365['lowerband'])
		bbp120 = (df['close'] - bb120['lowerband']) / (bb120['upperband'] - bb120['lowerband'])
		bbp30 = (df['close'] - bb30['lowerband']) / (bb30['upperband'] - bb30['lowerband'])
		bbp302 = (df['close'] - bb302['lowerband']) / (bb302['upperband'] - bb302['lowerband'])
		bbp2633 = (df['close'] - bb2633['lowerband']) / (bb2633['upperband'] - bb2633['lowerband'])
		cols = {
			'bb2633upper': bb2633['upperband'],
			'bb2633lower': bb2633['lowerband'],
			'bbp2633': bbp2633,
			'bb2633middle': bb2633['middleband'],
			'atr': ta.ATR(df, timeperiod=14),
			'bbup365': bb365['upperband'],
			'bblow365': bb365['lowerband'],
			'bbp365': bbp365,
			'bbup120': bb120['upperband'],
			'bblow120': bb120['lowerband'],
			'bbp120': bbp120,
			'bbup30': bb30['upperband'],
			'bblow30': bb30['lowerband'],
			'bbmiddle30': bb30_middle,
			'bbp30': bbp30,
			'bbup302': bb302['upperband'],
			'bblow302': bb302['lowerband'],
			'bbp302': bbp302,
			'macd': macd['macd'],
			'macdsignal': macd['macdsignal'],
			'macdhist': macd['macdhist'],
		}
		for _p, _n in ((5, 'ema5'), (10, 'ema10'), (24, 'ema24'), (52, 'ema52'),
		               (104, 'ema104'), (156, 'ema156'), (208, 'ema208'),
		               (26, 'ema26'), (13, 'ema13'), (7, 'ema7')):
			cols[_n] = ta.EMA(df, timeperiod=_p)
		cols['rsi'] = ta.RSI(df, timeperiod=14)
		cols['volume_ratio'] = self.cal_volume_ratio(df)

		# 重复调用（增量路径每根都会）时先摘掉旧列，否则 concat 出重名列。
		# 摘掉再接回末尾，列序与逐列覆盖的结果一致。
		new = pd.DataFrame(cols, index=df.index)
		dup = [c for c in new.columns if c in df.columns]
		if dup:
			df = df.drop(columns=dup)
		return pd.concat([df, new], axis=1)

	def get_ema_state(self, dataframe):
		klu_list = self.get_klu_list(dataframe)
		klc_list = self.get_klc_list(klu_list)
		bi_list = self.cal_bi_list(klc_list)
		klu_state_list = []
		for klu in klu_list:
			if klu.near0_return == 1:
				klu_state_list.append("1")
			elif klu.near0_return == 9:
				klu_state_list.append("-1")
			elif klu.candle_dir == Chan_K_DIR.BULL:
				klu_state_list.append("2")
			elif klu.candle_dir == Chan_K_DIR.BEAR:
				klu_state_list.append("-2")
			else:
				klu_state_list.append("0")
		return klu_state_list

	def get_decimal(self, value):
		return Decimal("{:.2f}".format(value))