356 lines
12 KiB
Python
356 lines
12 KiB
Python
"""
|
|
线段模块：线段由笔组成，有特殊的生成和确认规则
|
|
线段定义：至少包含3笔，且满足特定的破坏条件
|
|
"""
|
|
|
|
import pandas as pd
|
|
import numpy as np
|
|
from typing import List, Tuple, Optional, Dict
|
|
from dataclasses import dataclass
|
|
import logging
|
|
from .stroke import StrokeElement
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class SegmentElement:
|
|
"""线段元素数据类"""
|
|
strokes: List[StrokeElement] # 组成线段的笔列表
|
|
direction: int # 线段方向：1为上升，-1为下降
|
|
start_price: float # 起始价格
|
|
end_price: float # 结束价格
|
|
start_time: pd.Timestamp # 开始时间
|
|
end_time: pd.Timestamp # 结束时间
|
|
length: float # 线段长度
|
|
strength: float # 线段强度
|
|
confirmed: bool = False # 是否已确认
|
|
|
|
|
|
class Segment:
|
|
"""线段识别和处理类"""
|
|
|
|
def __init__(self, strokes: List[StrokeElement]):
|
|
"""
|
|
初始化线段处理器
|
|
|
|
Args:
|
|
strokes: 笔列表
|
|
"""
|
|
self.strokes = sorted(strokes, key=lambda x: x.start_fractal.index)
|
|
self.segments = []
|
|
self.up_segments = [] # 上升线段
|
|
self.down_segments = [] # 下降线段
|
|
|
|
def can_form_segment(self, stroke_group: List[StrokeElement]) -> bool:
|
|
"""
|
|
判断笔组是否可以形成线段
|
|
|
|
Args:
|
|
stroke_group: 笔组
|
|
|
|
Returns:
|
|
是否可以形成线段
|
|
"""
|
|
if len(stroke_group) < 3:
|
|
return False
|
|
|
|
# 检查是否有主导方向
|
|
up_count = sum(1 for s in stroke_group if s.direction == 1)
|
|
down_count = sum(1 for s in stroke_group if s.direction == -1)
|
|
|
|
# 主导方向的笔至少比反向笔多1个
|
|
return abs(up_count - down_count) >= 1
|
|
|
|
def determine_segment_direction(self, stroke_group: List[StrokeElement]) -> int:
|
|
"""
|
|
确定线段方向
|
|
|
|
Args:
|
|
stroke_group: 笔组
|
|
|
|
Returns:
|
|
线段方向
|
|
"""
|
|
if not stroke_group:
|
|
return 0
|
|
|
|
# 比较起点和终点价格
|
|
start_price = stroke_group[0].start_fractal.price
|
|
end_price = stroke_group[-1].end_fractal.price
|
|
|
|
if end_price > start_price:
|
|
return 1 # 上升线段
|
|
elif end_price < start_price:
|
|
return -1 # 下降线段
|
|
else:
|
|
# 价格相等时，根据笔的数量决定
|
|
up_count = sum(1 for s in stroke_group if s.direction == 1)
|
|
down_count = sum(1 for s in stroke_group if s.direction == -1)
|
|
return 1 if up_count >= down_count else -1
|
|
|
|
def create_segment(self, stroke_group: List[StrokeElement]) -> SegmentElement:
|
|
"""
|
|
创建线段元素
|
|
|
|
Args:
|
|
stroke_group: 组成线段的笔组
|
|
|
|
Returns:
|
|
线段元素
|
|
"""
|
|
if not stroke_group:
|
|
raise ValueError("笔组不能为空")
|
|
|
|
direction = self.determine_segment_direction(stroke_group)
|
|
|
|
start_price = stroke_group[0].start_fractal.price
|
|
end_price = stroke_group[-1].end_fractal.price
|
|
start_time = stroke_group[0].start_fractal.timestamp
|
|
end_time = stroke_group[-1].end_fractal.timestamp
|
|
|
|
length = abs(end_price - start_price)
|
|
strength = self._calculate_segment_strength(stroke_group, length)
|
|
|
|
return SegmentElement(
|
|
strokes=stroke_group,
|
|
direction=direction,
|
|
start_price=start_price,
|
|
end_price=end_price,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
length=length,
|
|
strength=strength
|
|
)
|
|
|
|
def _calculate_segment_strength(self, stroke_group: List[StrokeElement],
|
|
length: float) -> float:
|
|
"""
|
|
计算线段强度
|
|
|
|
Args:
|
|
stroke_group: 笔组
|
|
length: 线段长度
|
|
|
|
Returns:
|
|
线段强度
|
|
"""
|
|
# 基础强度：价格变化幅度
|
|
base_strength = length
|
|
|
|
# 笔的数量因子
|
|
stroke_count_factor = len(stroke_group) / 3 # 3笔为基准
|
|
|
|
# 笔的平均强度
|
|
avg_stroke_strength = np.mean([s.strength for s in stroke_group])
|
|
|
|
# 方向一致性：主导方向笔的比例
|
|
main_direction = self.determine_segment_direction(stroke_group)
|
|
main_direction_count = sum(1 for s in stroke_group if s.direction == main_direction)
|
|
direction_consistency = main_direction_count / len(stroke_group)
|
|
|
|
# 综合强度
|
|
strength = (base_strength *
|
|
(1 + stroke_count_factor * 0.1) *
|
|
(1 + avg_stroke_strength * 0.01) *
|
|
direction_consistency)
|
|
|
|
return strength
|
|
|
|
def detect_segment_break(self, segment: SegmentElement,
|
|
new_stroke: StrokeElement) -> bool:
|
|
"""
|
|
检测线段是否被破坏
|
|
|
|
Args:
|
|
segment: 当前线段
|
|
new_stroke: 新的笔
|
|
|
|
Returns:
|
|
是否被破坏
|
|
"""
|
|
if segment.direction == 1: # 上升线段
|
|
# 如果新笔是下降笔且跌破线段起点
|
|
if (new_stroke.direction == -1 and
|
|
new_stroke.end_fractal.price < segment.start_price):
|
|
return True
|
|
else: # 下降线段
|
|
# 如果新笔是上升笔且涨破线段起点
|
|
if (new_stroke.direction == 1 and
|
|
new_stroke.end_fractal.price > segment.start_price):
|
|
return True
|
|
|
|
return False
|
|
|
|
def detect_segments(self) -> List[SegmentElement]:
|
|
"""
|
|
检测所有线段
|
|
|
|
Returns:
|
|
线段列表
|
|
"""
|
|
if len(self.strokes) < 3:
|
|
logger.warning("笔数量不足，无法生成线段")
|
|
return []
|
|
|
|
segments = []
|
|
current_stroke_group = []
|
|
|
|
for stroke in self.strokes:
|
|
current_stroke_group.append(stroke)
|
|
|
|
# 当有足够笔时，尝试形成线段
|
|
if len(current_stroke_group) >= 3:
|
|
if self.can_form_segment(current_stroke_group):
|
|
# 检查是否有现有线段被破坏
|
|
if segments:
|
|
last_segment = segments[-1]
|
|
if self.detect_segment_break(last_segment, stroke):
|
|
# 线段被破坏，确认上一个线段
|
|
last_segment.confirmed = True
|
|
|
|
# 开始新的线段
|
|
current_stroke_group = [stroke]
|
|
continue
|
|
|
|
# 尝试扩展或创建新线段
|
|
if len(current_stroke_group) >= 5: # 限制线段长度
|
|
segment = self.create_segment(current_stroke_group[:-2])
|
|
segments.append(segment)
|
|
current_stroke_group = current_stroke_group[-2:] # 保留最后两笔
|
|
|
|
# 处理最后一组笔
|
|
if len(current_stroke_group) >= 3 and self.can_form_segment(current_stroke_group):
|
|
segment = self.create_segment(current_stroke_group)
|
|
segments.append(segment)
|
|
|
|
# 分类存储
|
|
for segment in segments:
|
|
if segment.direction == 1:
|
|
self.up_segments.append(segment)
|
|
else:
|
|
self.down_segments.append(segment)
|
|
|
|
self.segments = segments
|
|
logger.info(f"检测到 {len(segments)} 个线段：上升线段 {len(self.up_segments)} 个，下降线段 {len(self.down_segments)} 个")
|
|
|
|
return segments
|
|
|
|
def get_segment_sequence(self) -> List[SegmentElement]:
|
|
"""
|
|
获取连续的线段序列
|
|
|
|
Returns:
|
|
连续线段序列
|
|
"""
|
|
return sorted(self.segments, key=lambda x: x.start_time)
|
|
|
|
def find_segment_overlaps(self) -> List[Tuple[SegmentElement, SegmentElement]]:
|
|
"""
|
|
寻找线段重叠区域（可能的中枢）
|
|
|
|
Returns:
|
|
重叠线段对列表
|
|
"""
|
|
overlaps = []
|
|
sequence = self.get_segment_sequence()
|
|
|
|
for i in range(len(sequence) - 1):
|
|
for j in range(i + 1, len(sequence)):
|
|
seg1, seg2 = sequence[i], sequence[j]
|
|
|
|
# 检查价格区间是否重叠
|
|
if (min(seg1.start_price, seg1.end_price) <= max(seg2.start_price, seg2.end_price) and
|
|
max(seg1.start_price, seg1.end_price) >= min(seg2.start_price, seg2.end_price)):
|
|
overlaps.append((seg1, seg2))
|
|
|
|
return overlaps
|
|
|
|
def analyze_segment_patterns(self) -> Dict:
|
|
"""
|
|
分析线段模式
|
|
|
|
Returns:
|
|
模式分析结果
|
|
"""
|
|
if not self.segments:
|
|
return {}
|
|
|
|
sequence = self.get_segment_sequence()
|
|
|
|
# 分析趋势
|
|
trend_changes = 0
|
|
for i in range(1, len(sequence)):
|
|
if sequence[i].direction != sequence[i-1].direction:
|
|
trend_changes += 1
|
|
|
|
# 统计信息
|
|
avg_length = np.mean([s.length for s in self.segments])
|
|
avg_stroke_count = np.mean([len(s.strokes) for s in self.segments])
|
|
max_length = max([s.length for s in self.segments]) if self.segments else 0
|
|
|
|
# 寻找重叠区域
|
|
overlaps = self.find_segment_overlaps()
|
|
|
|
return {
|
|
'total_segments': len(self.segments),
|
|
'up_segments': len(self.up_segments),
|
|
'down_segments': len(self.down_segments),
|
|
'trend_changes': trend_changes,
|
|
'avg_length': avg_length,
|
|
'max_length': max_length,
|
|
'avg_stroke_count': avg_stroke_count,
|
|
'overlaps': len(overlaps),
|
|
'confirmed_segments': sum(1 for s in self.segments if s.confirmed)
|
|
}
|
|
|
|
def filter_segments_by_strength(self, min_strength: float) -> List[SegmentElement]:
|
|
"""
|
|
按强度过滤线段
|
|
|
|
Args:
|
|
min_strength: 最小强度要求
|
|
|
|
Returns:
|
|
过滤后的线段列表
|
|
"""
|
|
return [s for s in self.segments if s.strength >= min_strength]
|
|
|
|
def filter_segments_by_length(self, min_length: float) -> List[SegmentElement]:
|
|
"""
|
|
按长度过滤线段
|
|
|
|
Args:
|
|
min_length: 最小长度要求
|
|
|
|
Returns:
|
|
过滤后的线段列表
|
|
"""
|
|
return [s for s in self.segments if s.length >= min_length]
|
|
|
|
def to_dataframe(self) -> pd.DataFrame:
|
|
"""
|
|
将线段转换为DataFrame
|
|
|
|
Returns:
|
|
包含线段信息的DataFrame
|
|
"""
|
|
if not self.segments:
|
|
return pd.DataFrame()
|
|
|
|
data = []
|
|
for i, segment in enumerate(self.segments):
|
|
data.append({
|
|
'segment_id': i,
|
|
'start_time': segment.start_time,
|
|
'end_time': segment.end_time,
|
|
'start_price': segment.start_price,
|
|
'end_price': segment.end_price,
|
|
'direction': segment.direction,
|
|
'length': segment.length,
|
|
'strength': segment.strength,
|
|
'stroke_count': len(segment.strokes),
|
|
'confirmed': segment.confirmed
|
|
})
|
|
|
|
return pd.DataFrame(data) |