NewStock/main/factor/operator_base.py

"""
因子算子基础框架 - 简化版本
提供股票截面和日期截面两个基础函数
"""

import polars as pl
from typing import Callable, Any, Optional, Union
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def apply_stockwise(
    df: pl.DataFrame,
    operator_func: Callable[[pl.DataFrame, Any], pl.DataFrame],
    *args,
    **kwargs
) -> pl.DataFrame:
    """
    在股票截面上应用算子函数
    
    Args:
        df: 输入的polars DataFrame，必须包含ts_code和trade_date列
        operator_func: 算子函数，接收单个股票的数据和参数，返回处理后的DataFrame
        *args, **kwargs: 传递给算子函数的额外参数
    
    Returns:
        处理后的完整DataFrame
    """
    # 验证必需列
    required_cols = ['ts_code', 'trade_date']
    missing_cols = [col for col in required_cols if col not in df.columns]
    if missing_cols:
        raise ValueError(f"缺少必需列: {missing_cols}")
    
    # 获取股票列表
    stock_list = df['ts_code'].unique().to_list()
    results = []
    
    # 按股票分组处理
    for ts_code in stock_list:
        try:
            # 获取单个股票的数据并按日期排序
            stock_df = df.filter(pl.col('ts_code') == ts_code).sort('trade_date')
            
            # 应用算子函数
            result_df = operator_func(stock_df, *args, **kwargs)
            results.append(result_df)
            
        except Exception as e:
            logger.error(f"股票 {ts_code} 处理失败: {e}")
            # 失败时返回原始数据
            stock_df = df.filter(pl.col('ts_code') == ts_code).sort('trade_date')
            results.append(stock_df)
    
    # 合并结果并排序
    if results:
        return pl.concat(results).sort(['ts_code', 'trade_date'])
    else:
        return df


def apply_datewise(
    df: pl.DataFrame,
    operator_func: Callable[[pl.DataFrame, Any], pl.DataFrame],
    *args,
    **kwargs
) -> pl.DataFrame:
    """
    在日期截面上应用算子函数
    
    Args:
        df: 输入的polars DataFrame，必须包含ts_code和trade_date列
        operator_func: 算子函数，接收单个日期的数据和参数，返回处理后的DataFrame
        *args, **kwargs: 传递给算子函数的额外参数
    
    Returns:
        处理后的完整DataFrame
    """
    # 验证必需列
    required_cols = ['ts_code', 'trade_date']
    missing_cols = [col for col in required_cols if col not in df.columns]
    if missing_cols:
        raise ValueError(f"缺少必需列: {missing_cols}")
    
    # 获取日期列表
    date_list = df['trade_date'].unique().to_list()
    results = []
    
    # 按日期分组处理
    for trade_date in date_list:
        try:
            # 获取单个日期的数据
            date_df = df.filter(pl.col('trade_date') == trade_date)
            
            # 应用算子函数
            result_df = operator_func(date_df, *args, **kwargs)
            results.append(result_df)
            
        except Exception as e:
            logger.error(f"日期 {trade_date} 处理失败: {e}")
            # 失败时返回原始数据
            date_df = df.filter(pl.col('trade_date') == trade_date)
            results.append(date_df)
    
    # 合并结果并排序
    if results:
        return pl.concat(results).sort(['ts_code', 'trade_date'])
    else:
        return df


# 常用算子函数示例
def rolling_mean_operator(df: pl.DataFrame, column: str, window: int, output_col: str = None) -> pl.DataFrame:
    """
    滚动均值算子 - 股票截面
    
    Args:
        df: 单个股票的数据
        column: 要计算均值的列
        window: 窗口大小
        output_col: 输出列名，默认为f'{column}_mean_{window}'
    
    Returns:
        添加均值列的DataFrame
    """
    if output_col is None:
        output_col = f'{column}_mean_{window}'
    
    return df.with_columns(
        pl.col(column).rolling_mean(window_size=window).alias(output_col)
    )


def rolling_std_operator(df: pl.DataFrame, column: str, window: int, output_col: str = None) -> pl.DataFrame:
    """
    滚动标准差算子 - 股票截面
    
    Args:
        df: 单个股票的数据
        column: 要计算标准差的列
        window: 窗口大小
        output_col: 输出列名，默认为f'{column}_std_{window}'
    
    Returns:
        添加标准差列的DataFrame
    """
    if output_col is None:
        output_col = f'{column}_std_{window}'
    
    return df.with_columns(
        pl.col(column).rolling_std(window_size=window).alias(output_col)
    )


def rank_operator(df: pl.DataFrame, column: str, ascending: bool = True, output_col: str = None) -> pl.DataFrame:
    """
    排名算子 - 日期截面
    
    Args:
        df: 单个日期的数据
        column: 要排名的列
        ascending: 是否升序
        output_col: 输出列名，默认为f'{column}_rank'
    
    Returns:
        添加排名列的DataFrame
    """
    if output_col is None:
        output_col = f'{column}_rank'
    
    return df.with_columns(
        pl.col(column).rank(method='dense', descending=not ascending).alias(output_col)
    )


def pct_change_operator(df: pl.DataFrame, column: str, periods: int = 1, output_col: str = None) -> pl.DataFrame:
    """
    百分比变化算子 - 股票截面
    
    Args:
        df: 单个股票的数据
        column: 要计算变化的列
        periods: 期数
        output_col: 输出列名，默认为f'{column}_pct_change_{periods}'
    
    Returns:
        添加变化率列的DataFrame
    """
    if output_col is None:
        output_col = f'{column}_pct_change_{periods}'
    
    return df.with_columns(
        ((pl.col(column) / pl.col(column).shift(periods)) - 1).alias(output_col)
    )