Cointegration & Pairs Trading

import numpy as np
import pandas as pd
from statsmodels.tsa.stattools import adfuller, coint

# --- Download price data ---
# prices_df columns: ['HDFCBANK.NS', 'ICICIBANK.NS']

Y = prices_df['HDFCBANK.NS']
X = prices_df['ICICIBANK.NS']

# Step 1: OLS for hedge ratio
beta = np.polyfit(X, Y, 1)[0]   # slope = hedge ratio
spread = Y - beta * X             # the "leash"

# Step 2: ADF test on spread
result = adfuller(spread.dropna())
adf_stat, p_value = result[0], result[1]

print(f"β (hedge ratio): {beta:.4f}")
print(f"ADF statistic:   {adf_stat:.4f}")
print(f"p-value:         {p_value:.4f}")

if p_value < 0.05:
    print("✓ COINTEGRATED — pairs trade is valid")
else:
    print("✗ Not cointegrated — don't trade this pair")

# Shortcut: statsmodels coint() does both steps
score, pval, _ = coint(Y, X)
print(f"Coint p-value: {pval:.4f}")

# Estimate half-life via OLS on lagged spread
from statsmodels.regression.linear_model import OLS
from statsmodels.tools import add_constant

delta_spread = spread.diff().dropna()
lag_spread   = spread.shift(1).dropna()

# Align
delta_spread, lag_spread = delta_spread.align(lag_spread, join='inner')

# Regress: ΔS_t = a + κ·S_{t-1} + ε_t
X_reg = add_constant(lag_spread)
res   = OLS(delta_spread, X_reg).fit()

kappa     = -res.params.iloc[1]           # mean reversion speed
half_life = np.log(2) / kappa

print(f"Mean reversion speed (κ): {kappa:.4f}")
print(f"Half-life: {half_life:.1f} days")

# Rule of thumb:
# half-life < 5 days   → too fast, probably noise
# half-life 5–30 days  → sweet spot for daily trading
# half-life > 60 days  → too slow, capital tied up too long

# Full pairs trading signal generator
LOOKBACK   = 20    # rolling window for mean/std
ENTRY_Z    = 2.0   # enter at ±2 sigma
EXIT_Z     = 0.5   # exit when spread normalises
STOP_Z     = 3.5   # stop loss at ±3.5 sigma

spread_mean = spread.rolling(LOOKBACK).mean()
spread_std  = spread.rolling(LOOKBACK).std()
z_score     = (spread - spread_mean) / spread_std

positions = pd.Series(0, index=z_score.index)

for i in range(1, len(z_score)):
    z = z_score.iloc[i]
    prev_pos = positions.iloc[i - 1]

    if prev_pos == 0:               # flat — look for entry
        if z < -ENTRY_Z:  positions.iloc[i] =  1   # long spread
        elif z > ENTRY_Z: positions.iloc[i] = -1   # short spread
    elif prev_pos == 1:             # long spread — look for exit/stop
        if z >= -EXIT_Z:  positions.iloc[i] =  0   # take profit
        elif z < -STOP_Z: positions.iloc[i] =  0   # stop loss
        else:             positions.iloc[i] =  1   # hold
    elif prev_pos == -1:            # short spread — mirror logic
        if z <= EXIT_Z:   positions.iloc[i] =  0
        elif z > STOP_Z:  positions.iloc[i] =  0
        else:             positions.iloc[i] = -1

from statsmodels.tsa.vector_ar.vecm import coint_johansen

# price_matrix: DataFrame with N asset columns
result = coint_johansen(price_matrix, det_order=0, k_ar_diff=1)

# Trace statistic test
trace_stat   = result.lr1          # test statistics
crit_values  = result.cvt          # [90%, 95%, 99%] critical values
eigen_vectors = result.evec        # each column = one cointegrating vector

# Count how many cointegrating vectors exist
n_cointegrations = sum(trace_stat > crit_values[:, 1])  # at 95%
print(f"Number of cointegrating vectors: {n_cointegrations}")

# First cointegrating vector = the hedge weights
hedge_weights = eigen_vectors[:, 0]
spread_basket = price_matrix @ hedge_weights  # dot product