Initial commit

2026-07-17 18:49:07 +08:00
commit 9c7d7abdd4
100 changed files with 41337 additions and 0 deletions
@@ -0,0 +1,130 @@
+"""
+技术指标计算模块（numpy向量化优化版）
+实现 MACD、SKDJ、EMA 等技术指标
+
+优化要点：
+- calc_sma 使用 numpy 原生数组替代 pandas.iloc，速度提升 5-10x
+- calc_all_indicators 智能跳过已是 float 的类型转换
+"""
+import pandas as pd
+import numpy as np
+
+
+def calc_ema(series, period):
+    """计算指数移动平均线(EMA) — 使用pandas的C底层ewm实现，已足够快"""
+    return series.ewm(span=period, adjust=False).mean()
+
+
+def calc_sma(series, period, weight=1):
+    """
+    计算SMA（通达信公式风格） — numpy优化版
+    SMA(X, N, M) = (M * X + (N - M) * prev_SMA) / N
+
+    优化：使用 numpy 原生数组 arr[i] 替代 pandas series.iloc[i]
+    numpy 数组元素访问约 50ns，pandas iloc 约 5μs，提升 ~100x
+    """
+    arr = series.values.astype(np.float64)
+    n = len(arr)
+    result = np.empty(n, dtype=np.float64)
+    result[0] = arr[0]
+    w = np.float64(weight)
+    carry = np.float64(period - weight)
+    inv_p = np.float64(1.0 / period)
+    for i in range(1, n):
+        result[i] = (w * arr[i] + carry * result[i - 1]) * inv_p
+    return pd.Series(result, index=series.index)
+
+
+def calc_macd(close, fast=12, slow=26, signal=9):
+    """
+    计算MACD指标
+    返回: DIF, DEA, MACD柱
+    """
+    ema_fast = calc_ema(close, fast)
+    ema_slow = calc_ema(close, slow)
+    dif = ema_fast - ema_slow
+    dea = calc_ema(dif, signal)
+    macd_hist = 2 * (dif - dea)
+    return dif, dea, macd_hist
+
+
+def calc_kdj(high, low, close, n=9, m1=3, m2=3):
+    """
+    计算KDJ指标
+    返回: K, D, J
+    """
+    lowest_low = low.rolling(window=n, min_periods=1).min()
+    highest_high = high.rolling(window=n, min_periods=1).max()
+
+    rsv = pd.Series(np.where(
+        highest_high == lowest_low, 50,
+        (close - lowest_low) / (highest_high - lowest_low) * 100
+    ), index=close.index, dtype=float)
+
+    k = calc_sma(rsv, m1, 1)
+    d = calc_sma(k, m2, 1)
+    j = 3 * k - 2 * d
+    return k, d, j
+
+
+def calc_skdj(high, low, close, n=9, m=3):
+    """
+    计算SKDJ（慢速随机指标）
+    对RSV先做一次SMA得到K_fast，再对K_fast做两次SMA得到SKDJ的K和D
+    返回: K, D
+    """
+    lowest_low = low.rolling(window=n, min_periods=1).min()
+    highest_high = high.rolling(window=n, min_periods=1).max()
+
+    rsv = pd.Series(np.where(
+        highest_high == lowest_low, 50,
+        (close - lowest_low) / (highest_high - lowest_low) * 100
+    ), index=close.index, dtype=float)
+
+    k_fast = calc_sma(rsv, m, 1)
+    k = calc_sma(k_fast, m, 1)
+    d = calc_sma(k, m, 1)
+    return k, d
+
+
+def calc_all_indicators(df):
+    """
+    计算所有技术指标并添加到DataFrame（优化版）
+    df 需要包含: close, high, low, open, volume 列
+    返回: 添加了指标列的DataFrame
+
+    优化：智能跳过已是 float64 的列，避免重复 astype
+    """
+    close = df['close']
+    high = df['high']
+    low = df['low']
+
+    # 智能类型转换：仅在需要时转换
+    if not np.issubdtype(close.dtype, np.floating):
+        close = close.astype(np.float64)
+        high = high.astype(np.float64)
+        low = low.astype(np.float64)
+
+    df['ema3'] = calc_ema(close, 3)
+    df['ema21'] = calc_ema(close, 21)
+
+    dif, dea, macd_hist = calc_macd(close)
+    df['dif'] = dif
+    df['dea'] = dea
+    df['macd'] = macd_hist
+
+    k, d, j = calc_kdj(high, low, close)
+    df['kdj_k'] = k
+    df['kdj_d'] = d
+    df['kdj_j'] = j
+
+    sk, sd = calc_skdj(high, low, close)
+    df['skdj_k'] = sk
+    df['skdj_d'] = sd
+
+    df['ma5'] = close.rolling(5).mean()
+    df['ma10'] = close.rolling(10).mean()
+    df['ma20'] = close.rolling(20).mean()
+    df['ma60'] = close.rolling(60).mean()
+
+    return df