238 lines
7.6 KiB
Python
238 lines
7.6 KiB
Python
"""
|
|
K线包含关系处理模块
|
|
"""
|
|
|
|
import pandas as pd
|
|
import numpy as np
|
|
from typing import List, Tuple, Optional
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class KLineElement:
|
|
"""单个K线元素类"""
|
|
|
|
def __init__(self, high: float, low: float, open_price: float,
|
|
close: float, volume: float, timestamp: pd.Timestamp):
|
|
"""
|
|
初始化K线元素
|
|
|
|
Args:
|
|
high: 最高价
|
|
low: 最低价
|
|
open_price: 开盘价
|
|
close: 收盘价
|
|
volume: 成交量
|
|
timestamp: 时间戳
|
|
"""
|
|
self.high = high
|
|
self.low = low
|
|
self.open = open_price
|
|
self.close = close
|
|
self.volume = volume
|
|
self.timestamp = timestamp
|
|
self.direction = 1 if close >= open_price else -1
|
|
|
|
def __repr__(self):
|
|
return f"KLineElement(H:{self.high}, L:{self.low}, O:{self.open}, C:{self.close})"
|
|
|
|
|
|
class KLine:
|
|
"""K线包含关系处理类"""
|
|
|
|
def __init__(self, data: pd.DataFrame):
|
|
"""
|
|
初始化K线处理器
|
|
|
|
Args:
|
|
data: 包含OHLCV数据的DataFrame
|
|
"""
|
|
self.original_data = data.copy()
|
|
self.processed_data = None
|
|
self.kline_elements = []
|
|
self._process_data()
|
|
|
|
def _process_data(self):
|
|
"""处理原始数据为K线元素"""
|
|
self.kline_elements = []
|
|
|
|
for idx, row in self.original_data.iterrows():
|
|
element = KLineElement(
|
|
high=row['high'],
|
|
low=row['low'],
|
|
open_price=row['open'],
|
|
close=row['close'],
|
|
volume=row['volume'],
|
|
timestamp=idx
|
|
)
|
|
self.kline_elements.append(element)
|
|
|
|
def has_containment(self, k1: KLineElement, k2: KLineElement) -> bool:
|
|
"""
|
|
判断两根K线是否存在包含关系
|
|
|
|
Args:
|
|
k1: 第一根K线
|
|
k2: 第二根K线
|
|
|
|
Returns:
|
|
是否存在包含关系
|
|
"""
|
|
# K1包含K2：K1的高低点完全包含K2
|
|
k1_contains_k2 = (k1.high >= k2.high and k1.low <= k2.low)
|
|
|
|
# K2包含K1：K2的高低点完全包含K1
|
|
k2_contains_k1 = (k2.high >= k1.high and k2.low <= k1.low)
|
|
|
|
return k1_contains_k2 or k2_contains_k1
|
|
|
|
def merge_contained_klines(self, k1: KLineElement, k2: KLineElement,
|
|
direction: int) -> KLineElement:
|
|
"""
|
|
合并包含关系的K线
|
|
|
|
Args:
|
|
k1: 第一根K线
|
|
k2: 第二根K线
|
|
direction: 当前趋势方向（1为上升，-1为下降）
|
|
|
|
Returns:
|
|
合并后的K线
|
|
"""
|
|
if direction > 0: # 上升趋势中
|
|
# 取两根K线的最高点的最大值，最低点的最大值
|
|
merged_high = max(k1.high, k2.high)
|
|
merged_low = max(k1.low, k2.low)
|
|
else: # 下降趋势中
|
|
# 取两根K线的最高点的最小值，最低点的最小值
|
|
merged_high = min(k1.high, k2.high)
|
|
merged_low = min(k1.low, k2.low)
|
|
|
|
# 开盘价和收盘价使用第一根K线的值
|
|
merged_open = k1.open
|
|
merged_close = k2.close
|
|
merged_volume = k1.volume + k2.volume
|
|
merged_timestamp = k2.timestamp # 使用最后一根K线的时间
|
|
|
|
return KLineElement(
|
|
high=merged_high,
|
|
low=merged_low,
|
|
open_price=merged_open,
|
|
close=merged_close,
|
|
volume=merged_volume,
|
|
timestamp=merged_timestamp
|
|
)
|
|
|
|
def handle_containment(self) -> List[KLineElement]:
|
|
"""
|
|
处理所有K线的包含关系
|
|
|
|
Returns:
|
|
处理包含关系后的K线列表
|
|
"""
|
|
if len(self.kline_elements) < 2:
|
|
return self.kline_elements.copy()
|
|
|
|
processed_klines = [self.kline_elements[0]] # 第一根K线
|
|
|
|
# 初始方向：根据前两根K线确定
|
|
if len(self.kline_elements) >= 2:
|
|
k1, k2 = self.kline_elements[0], self.kline_elements[1]
|
|
if k2.high > k1.high:
|
|
current_direction = 1 # 上升
|
|
elif k2.high < k1.high:
|
|
current_direction = -1 # 下降
|
|
else:
|
|
current_direction = 1 if k2.low >= k1.low else -1
|
|
else:
|
|
current_direction = 1
|
|
|
|
i = 1
|
|
while i < len(self.kline_elements):
|
|
current_k = self.kline_elements[i]
|
|
last_processed = processed_klines[-1]
|
|
|
|
if self.has_containment(last_processed, current_k):
|
|
# 存在包含关系，进行合并
|
|
merged = self.merge_contained_klines(
|
|
last_processed, current_k, current_direction
|
|
)
|
|
processed_klines[-1] = merged # 替换最后一个
|
|
else:
|
|
# 不存在包含关系，直接添加
|
|
processed_klines.append(current_k)
|
|
|
|
# 更新方向
|
|
if current_k.high > last_processed.high:
|
|
current_direction = 1
|
|
elif current_k.high < last_processed.high:
|
|
current_direction = -1
|
|
# 如果high相等，保持原方向
|
|
|
|
i += 1
|
|
|
|
logger.info(f"包含关系处理完成：{len(self.kline_elements)} -> {len(processed_klines)}")
|
|
return processed_klines
|
|
|
|
def to_dataframe(self, processed_klines: Optional[List[KLineElement]] = None) -> pd.DataFrame:
|
|
"""
|
|
将处理后的K线转换为DataFrame
|
|
|
|
Args:
|
|
processed_klines: 处理后的K线列表，如果为None则使用默认处理结果
|
|
|
|
Returns:
|
|
包含处理后K线的DataFrame
|
|
"""
|
|
if processed_klines is None:
|
|
processed_klines = self.handle_containment()
|
|
|
|
data = []
|
|
for kline in processed_klines:
|
|
data.append({
|
|
'timestamp': kline.timestamp,
|
|
'open': kline.open,
|
|
'high': kline.high,
|
|
'low': kline.low,
|
|
'close': kline.close,
|
|
'volume': kline.volume,
|
|
'direction': kline.direction
|
|
})
|
|
|
|
df = pd.DataFrame(data)
|
|
if not df.empty:
|
|
df.set_index('timestamp', inplace=True)
|
|
|
|
return df
|
|
|
|
def get_processed_data(self) -> pd.DataFrame:
|
|
"""
|
|
获取处理包含关系后的数据
|
|
|
|
Returns:
|
|
处理后的DataFrame
|
|
"""
|
|
if self.processed_data is None:
|
|
processed_klines = self.handle_containment()
|
|
self.processed_data = self.to_dataframe(processed_klines)
|
|
|
|
return self.processed_data
|
|
|
|
def visualize_containment(self) -> dict:
|
|
"""
|
|
生成包含关系可视化信息
|
|
|
|
Returns:
|
|
包含可视化信息的字典
|
|
"""
|
|
original_count = len(self.kline_elements)
|
|
processed_klines = self.handle_containment()
|
|
processed_count = len(processed_klines)
|
|
|
|
return {
|
|
'original_count': original_count,
|
|
'processed_count': processed_count,
|
|
'merged_count': original_count - processed_count,
|
|
'merge_ratio': (original_count - processed_count) / original_count if original_count > 0 else 0
|
|
} |