Upload stat_arb_features.py

stat_arb_features.py

@@ -0,0 +1,49 @@
+"""Statistical Arbitrage Features - Cointegration, spread, relative value"""
+import numpy as np
+import pandas as pd
+from statsmodels.tsa.stattools import coint
+
+
+class StatArbFeatures:
+    """Pairs trading and cointegration features"""
+    
+    @staticmethod
+    def cointegration_spread(price_a, price_b, window=126):
+        features = pd.DataFrame(index=price_a.index)
+        ret_a = price_a.pct_change()
+        ret_b = price_b.pct_change()
+        rolling_cov = ret_a.rolling(window).cov(ret_b)
+        rolling_var = ret_b.rolling(window).var()
+        hedge = rolling_cov / rolling_var.replace(0, np.nan)
+        spread = np.log(price_a) - hedge * np.log(price_b)
+        spread_mean = spread.rolling(window).mean()
+        spread_std = spread.rolling(window).std().replace(0, 1)
+        features['coint_hedge'] = hedge
+        features['coint_spread'] = spread
+        features['coint_zscore'] = (spread - spread_mean) / spread_std
+        return features
+    
+    @staticmethod
+    def relative_value(close, sector_close, window=21):
+        features = pd.DataFrame(index=close.index)
+        sector_avg = sector_close.mean(axis=1)
+        features['rv_ratio'] = close / sector_avg
+        features['rv_zscore'] = (features['rv_ratio'] - features['rv_ratio'].rolling(window).mean()) / \
+                                  features['rv_ratio'].rolling(window).std().replace(0, 1)
+        return features
+    
+    @staticmethod
+    def half_life(series, window=60):
+        """Compute rolling half-life of mean reversion"""
+        result = pd.Series(index=series.index, dtype=float)
+        for i in range(window, len(series)):
+            y = series.iloc[i-window:i].diff().dropna()
+            x = series.iloc[i-window:i-1].values.reshape(-1, 1)
+            if len(y) < 10:
+                result.iloc[i] = 21.0
+                continue
+            from sklearn.linear_model import LinearRegression
+            lr = LinearRegression()
+            lr.fit(x[:len(y)], y.values)
+            result.iloc[i] = -np.log(2) / lr.coef_[0] if lr.coef_[0] != 0 else 21.0
+        return result.clip(1, 252)