""" 线段模块：线段由笔组成，有特殊的生成和确认规则 线段定义：至少包含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)