sanguo_quant_live/zhangfei-technical/research/task-20260417-atr-indicator/atr_indicator.py

"""
ATR (Average True Range) 指标计算工具
支持SMA和EMA两种计算方式

Author: 张飞 翼德
Date: 2026-04-17
"""

import numpy as np
import pandas as pd
from typing import Optional, Union


class ATRIndicator:
    """
    ATR (Average True Range) 指标计算类
    
    支持两种计算方式：
    - SMA: 简单移动平均
    - EMA: 指数移动平均
    """
    
    def __init__(self, period: int = 14, method: str = 'sma'):
        """
        初始化ATR指标
        
        Parameters:
            period: int, 默认14
                ATR计算周期
            method: str, 默认'sma'
                计算方法，可选 'sma' 或 'ema'
        """
        if period <= 0:
            raise ValueError(f"周期必须大于0，当前为 {period}")
        
        if method.lower() not in ['sma', 'ema']:
            raise ValueError(f"计算方法必须是 'sma' 或 'ema'，当前为 {method}")
        
        self.period = period
        self.method = method.lower()
        self._last_atr: Optional[float] = None
        
    @staticmethod
    def calculate_tr(
        high: Union[pd.Series, np.ndarray],
        low: Union[pd.Series, np.ndarray],
        close: Union[pd.Series, np.ndarray]
    ) -> Union[pd.Series, np.ndarray]:
        """
        计算True Range (TR)
        
        公式: TR = max[|high - low|, |high - prev_close|, |low - prev_close|]
        
        Parameters:
            high: 最高价序列
            low: 最低价序列
            close: 收盘价序列
            
        Returns:
            TR序列，第一个值为NaN
        """
        # 转换为numpy数组便于计算
        if isinstance(high, pd.Series):
            tr_values = ATRIndicator._calculate_tr_numpy(
                high.values, low.values, close.values
            )
            return pd.Series(tr_values, index=high.index, name='TR')
        else:
            return ATRIndicator._calculate_tr_numpy(high, low, close)
    
    @staticmethod
    def _calculate_tr_numpy(
        high: np.ndarray,
        low: np.ndarray,
        close: np.ndarray
    ) -> np.ndarray:
        """使用numpy计算TR"""
        if len(high) != len(low) or len(high) != len(close):
            raise ValueError("high, low, close序列长度必须一致")
        
        n = len(high)
        tr = np.full(n, np.nan)
        
        if n < 2:
            return tr
        
        # 计算三个成分
        high_low = high[1:] - low[1:]
        high_prev_close = np.abs(high[1:] - close[:-1])
        low_prev_close = np.abs(low[1:] - close[:-1])
        
        # 取最大值
        tr[1:] = np.maximum(np.maximum(high_low, high_prev_close), low_prev_close)
        
        return tr
    
    def calculate(
        self,
        high: Union[pd.Series, np.ndarray],
        low: Union[pd.Series, np.ndarray],
        close: Union[pd.Series, np.ndarray]
    ) -> Union[pd.Series, np.ndarray]:
        """
        计算ATR指标
        
        Parameters:
            high: 最高价序列
            low: 最低价序列
            close: 收盘价序列
            
        Returns:
            ATR序列
        """
        # 先计算TR
        tr = self.calculate_tr(high, low, close)
        
        if self.method == 'sma':
            return self._calculate_sma_atr(tr)
        else:
            return self._calculate_ema_atr(tr)
    
    def _calculate_sma_atr(self, tr: Union[pd.Series, np.ndarray]) -> Union[pd.Series, np.ndarray]:
        """使用SMA计算ATR"""
        if isinstance(tr, pd.Series):
            atr = tr.rolling(window=self.period, min_periods=self.period).mean()
            atr.name = f'ATR_{self.period}'
            if len(atr) > 0:
                self._last_atr = atr.iloc[-1]
            return atr
        else:
            n = len(tr)
            atr = np.full(n, np.nan)
            for i in range(self.period - 1, n):
                if i >= self.period:
                    atr[i] = np.mean(tr[i - self.period + 1:i + 1])
            if n > 0:
                self._last_atr = atr[-1]
            return atr
    
    def _calculate_ema_atr(self, tr: Union[pd.Series, np.ndarray]) -> Union[pd.Series, np.ndarray]:
        """使用EMA计算ATR，Wilder平滑方法"""
        if isinstance(tr, pd.Series):
            return self._calculate_ema_atr_pandas(tr)
        else:
            return self._calculate_ema_atr_numpy(tr)
    
    def _calculate_ema_atr_pandas(self, tr: pd.Series) -> pd.Series:
        """Pandas版本EMA-ATR计算，使用Wilder平滑"""
        n = len(tr)
        atr = pd.Series(np.full(n, np.nan), index=tr.index, name=f'ATR_{self.period}')
        
        if n < self.period:
            self._last_atr = None
            return atr
        
        # 第一个ATR用SMA计算
        first_atr = tr.iloc[1:self.period + 1].mean()
        atr.iloc[self.period] = first_atr
        self._last_atr = first_atr
        
        # Wilder平滑: ATR_t = (ATR_{t-1} * (period - 1) + TR_t) / period
        alpha = 1.0 / self.period
        
        for i in range(self.period + 1, n):
            if np.isnan(tr.iloc[i]):
                atr.iloc[i] = atr.iloc[i - 1]
            else:
                atr.iloc[i] = atr.iloc[i - 1] * (1 - alpha) + tr.iloc[i] * alpha
        
        if len(atr) > 0:
            self._last_atr = atr.iloc[-1]
        
        return atr
    
    def _calculate_ema_atr_numpy(self, tr: np.ndarray) -> np.ndarray:
        """Numpy版本EMA-ATR计算，使用Wilder平滑"""
        n = len(tr)
        atr = np.full(n, np.nan)
        
        if n < self.period:
            self._last_atr = None
            return atr
        
        # 第一个ATR用SMA计算
        first_atr = np.mean(tr[1:self.period + 1])
        atr[self.period] = first_atr
        self._last_atr = first_atr
        
        # Wilder平滑: ATR_t = (ATR_{t-1} * (period - 1) + TR_t) / period
        alpha = 1.0 / self.period
        
        for i in range(self.period + 1, n):
            if np.isnan(tr[i]):
                atr[i] = atr[i - 1]
            else:
                atr[i] = atr[i - 1] * (1 - alpha) + tr[i] * alpha
        
        self._last_atr = atr[-1]
        return atr
    
    def get_last_atr(self) -> Optional[float]:
        """获取最后一个ATR值"""
        return self._last_atr
    
    def update(self, high: float, low: float, prev_close: float) -> float:
        """
        更新单个ATR值（实盘时增量更新使用）
        
        Parameters:
            high: 当前K线最高价
            low: 当前K线最低价
            prev_close: 前一根K线收盘价
            
        Returns:
            新的ATR值
        """
        if self.method == 'sma':
            # SMA不适合增量更新，这里简化处理，用户应该重新计算完整序列
            raise ValueError("SMA方法不支持增量更新，请重新计算完整序列")
        
        # 计算当前TR
        tr1 = abs(high - low)
        tr2 = abs(high - prev_close)
        tr3 = abs(low - prev_close)
        tr = max(tr1, tr2, tr3)
        
        if self._last_atr is None:
            # 还没有足够数据，初始化ATR为TR
            self._last_atr = tr
            return tr
        
        # EMA支持增量更新
        alpha = 1.0 / self.period
        new_atr = self._last_atr * (1 - alpha) + tr * alpha
        self._last_atr = new_atr
        return new_atr


def calculate_atr(
    df: pd.DataFrame,
    period: int = 14,
    method: str = 'sma',
    high_col: str = 'high',
    low_col: str = 'low',
    close_col: str = 'close',
    drop_na: bool = False
) -> pd.DataFrame:
    """
    便捷函数：直接对DataFrame计算ATR并添加到原DataFrame
    
    Parameters:
        df: 输入DataFrame，必须包含high, low, close列
        period: ATR周期，默认14
        method: 计算方法 'sma' 或 'ema'，默认'sma'
        high_col: 最高价列名，默认'high'
        low_col: 最低价列名，默认'low'
        close_col: 收盘价列名，默认'close'
        drop_na: 是否删除NaN值，默认False
        
    Returns:
        添加了TR和ATR列的DataFrame
    """
    atr_indicator = ATRIndicator(period, method)
    
    result = df.copy()
    result['TR'] = atr_indicator.calculate_tr(
        df[high_col], df[low_col], df[close_col]
    )
    result[f'ATR_{period}'] = atr_indicator.calculate(
        df[high_col], df[low_col], df[close_col]
    )
    
    if drop_na:
        result = result.dropna()
        
    return result