myquant/djapi/api/stock/smoothBrush.py

import pandas as pd
import numpy as np

def smooth_series_brush(series: pd.Series, window_size: int = 7, threshold_factor: float = 0.5, max_brush_length: int = 5) -> pd.Series:
    """
    平滑处理pandas序列中的连续毛刺数据，使用前值或后值填充。

    参数:
        series (pd.Series): 输入的pandas序列。
        window_size (int): 用于检测毛刺的滑动窗口大小。必须为奇数。默认为7。
        threshold_factor (float): 判断毛刺的阈值因子。如果 abs(value - median) / median > threshold_factor，则认为是毛刺。默认为0.5 (50%)。
        max_brush_length (int): 允许的最大连续毛刺长度。超过此长度的连续点将不被处理。默认为5。

    返回:
        pd.Series: 处理后的平滑序列。
    """
    if not isinstance(series, pd.Series):
        raise TypeError("输入必须是 pandas Series 对象。")

    if window_size % 2 == 0:
        raise ValueError("window_size 必须是奇数。")

    # 创建副本以避免修改原始数据
    smoothed_series = series.copy()
    
    # 用于标记是否为毛刺的布尔序列
    is_brush = pd.Series([False] * len(series), index=series.index)
    
    half_window = window_size // 2

    # --- 第一步：检测毛刺 ---
    for i in range(len(series)):
        start_idx = max(0, i - half_window)
        end_idx = min(len(series), i + half_window + 1)
        
        # 获取当前窗口数据
        window_data = series.iloc[start_idx:end_idx]
        
        if len(window_data) < 2:
            continue

        # 计算窗口中位数
        window_median = window_data.median()
        
        # 避免除以零
        if window_median == 0:
            continue 
            
        current_value = series.iloc[i]
        
        # 计算偏差比例
        deviation_ratio = abs(current_value - window_median) / abs(window_median)
        
        # 如果偏差超过阈值，则标记为毛刺
        if deviation_ratio > threshold_factor:
            is_brush.iloc[i] = True

    # --- 第二步：处理连续的毛刺段 ---
    # 使用 cumsum 技巧识别连续毛刺段
    brush_groups = (is_brush != is_brush.shift()).cumsum() * is_brush
    
    # 遍历每个被标记为毛刺的组
    for group_id in brush_groups[brush_groups != 0].unique():
        if pd.isna(group_id):
            continue
            
        brush_indices = brush_groups[brush_groups == group_id].index
        
        # 检查连续毛刺长度
        if len(brush_indices) > max_brush_length:
            print(f"警告: 发现长度为 {len(brush_indices)} 的连续毛刺段 (超过 max_brush_length={max_brush_length})，将不进行平滑处理。")
            continue

        # --- 平滑处理：使用前值或后值填充 ---
        # 查找前一个非毛刺点
        prev_valid_val = None
        start_loc = series.index.get_loc(brush_indices[0])
        for j in range(start_loc - 1, -1, -1):
            if not is_brush.iloc[j]:
                prev_valid_val = series.iloc[j]
                break
        
        # 查找后一个非毛刺点
        next_valid_val = None
        end_loc = series.index.get_loc(brush_indices[-1])
        for j in range(end_loc + 1, len(series)):
            if not is_brush.iloc[j]:
                next_valid_val = series.iloc[j]
                break

        # 决定使用哪个值填充
        if prev_valid_val is not None:
            fill_value = prev_valid_val
        elif next_valid_val is not None:
            fill_value = next_valid_val
        else:
            print(f"警告: 毛刺段 {brush_indices} 无有效邻居，使用全局中位数填充。")
            fill_value = series.median()

        # 用 fill_value 填充整个毛刺段
        for idx in brush_indices:
            smoothed_series.loc[idx] = fill_value
                
    return smoothed_series


def smooth_dataframe_brush(df: pd.DataFrame, target_columns: list, **kwargs) -> pd.DataFrame:
    """
    对DataFrame中的指定列进行毛刺平滑处理。

    参数:
        df (pd.DataFrame): 输入的pandas DataFrame。
        target_columns (list): 需要去毛刺处理的列名列表。
        **kwargs: 传递给 smooth_series_brush 函数的参数 (如 window_size, threshold_factor, max_brush_length)。

    返回:
        pd.DataFrame: 处理后的DataFrame，指定列已平滑，其余列不变。
    """
    if not isinstance(df, pd.DataFrame):
        raise TypeError("输入必须是 pandas DataFrame 对象。")
    
    # 创建副本以避免修改原始数据
    result_df = df.copy()
    
    # 检查目标列是否都存在于DataFrame中
    missing_cols = [col for col in target_columns if col not in df.columns]
    if missing_cols:
        raise ValueError(f"以下列不在DataFrame中: {missing_cols}")
    
    # 对每个目标列应用平滑函数
    for col in target_columns:
        print(f"正在处理列: {col}")
        try:
            # 应用去毛刺函数
            result_df[col] = smooth_series_brush(df[col], **kwargs)
        except Exception as e:
            print(f"处理列 {col} 时出错: {e}")
            # 可以选择保留原始数据或抛出异常
            # 这里选择保留原始数据
            continue
    
    return result_df


# --- 示例 ---
if __name__ == "__main__":
    # 1. 创建示例 DataFrame
    dates = pd.date_range('2023-01-01', periods=20, freq='D')
    # 需要处理的列
    values_to_smooth = [10, 11, 10.5, 12, 11.8, 50, 12.1, 11.9, 10, 10.2, 
                        9.8, 100, 105, 99, 10.1, 9.9, 10.3, 5, 10.2, 10.1]
    # 不需要处理的列 (例如，另一个传感器数据)
    other_data = np.random.randn(20).cumsum() + 100  # 累积和，模拟趋势
    
    # 构建 DataFrame
    df_original = pd.DataFrame({
        'Date': dates,
        'Sensor_A': values_to_smooth,     # 需要去毛刺
        'Sensor_B': other_data,           # 不需要处理
        'Other_Info': range(20)           # 其他信息，不需要处理
    })
    # 设置日期为索引 (常见做法)
    df_original.set_index('Date', inplace=True)
    
    print("原始 DataFrame:")
    print(df_original.head(10))
    print("\n" + "="*50 + "\n")

    # 2. 应用平滑函数
    # 指定需要处理的列
    columns_to_smooth = ['Sensor_A']
    # 调用新函数
    df_smoothed = smooth_dataframe_brush(
        df_original, 
        target_columns=columns_to_smooth, 
        window_size=5, 
        threshold_factor=0.3, 
        max_brush_length=5
    )
    
    print("平滑后的 DataFrame:")
    print(df_smoothed.head(10))
    print("\n" + "="*50 + "\n")
    
    # 3. 比较
    comparison_df = df_original.copy()
    comparison_df['Sensor_A_Smoothed'] = df_smoothed['Sensor_A']
    comparison_df['Difference'] = comparison_df['Sensor_A'] - comparison_df['Sensor_A_Smoothed']
    print("对比 (原始 Sensor_A vs 平滑后 vs 差异):")
    print(comparison_df[['Sensor_A', 'Sensor_A_Smoothed', 'Difference']].head(10))