The Arena for Quantitative Excellence

# StatArb Strategy: Dynamic Kalman Filter Cointegration Spread
import numpy as np

class KalmanStatArb:
    def __init__(self, delta=1e-4, vt=1e-3):
        self.theta = np.zeros(2)       # [intercept, hedge_ratio beta]
        self.P = np.zeros((2, 2))       # Covariance matrix
        self.R = np.eye(2) * delta      # Process noise
        self.V = vt                     # Measurement noise
        self.z_threshold = 1.85         # Entry threshold in std deviations

    def update(self, price_a: float, price_b: float):
        # State transition: x_t = [1.0, price_b]
        x = np.array([1.0, price_b])
        # Priori update
        self.P += self.R
        y_hat = np.dot(x, self.theta)
        error = price_a - y_hat
        
        # Kalman Gain
        S = np.dot(x, np.dot(self.P, x)) + self.V
        K = np.dot(self.P, x) / S
        
        # Posteriori state update
        self.theta += K * error
        self.P -= np.outer(K, np.dot(x, self.P))
        
        spread_std = np.sqrt(S)
        z_score = error / spread_std
        return self.theta[1], z_score

    def on_tick(self, btc_tick, eth_tick):
        beta, z = self.update(eth_tick['price'], btc_tick['price'])
        if z < -self.z_threshold:
            return {"action": "ENTER_LONG_SPREAD", "qty_a": 1.0, "qty_b": -beta}
        elif z > self.z_threshold:
            return {"action": "ENTER_SHORT_SPREAD", "qty_a": -1.0, "qty_b": beta}
        elif abs(z) < 0.25:
            return {"action": "EXIT_ALL", "reason": "MEAN_REVERSION_ACHIEVED"}
        return None