Correlation and Cointegration Analysis

import numpy as np

def kalman_hedge_ratio(
    y: pd.Series,
    x: pd.Series,
    delta: float = 1e-4,
    vt: float = 1.0,
) -> pd.DataFrame:
    """Estimate a dynamic hedge ratio with a Kalman filter.

    State equation:
        β_t = β_{t-1} + w_t,  w ~ N(0, Q)
    Observation equation:
        y_t = β_t · x_t + v_t,  v ~ N(0, R)

    Args:
        y: Price series of asset A
        x: Price series of asset B
        delta: State-noise intensity. Larger means faster hedge-ratio adaptation
        vt: Observation-noise variance

    Returns:
        DataFrame containing the dynamic hedge ratio and spread
    """
    n = len(y)
    # State: [β, α] = hedge ratio + intercept
    Wt = delta / (1 - delta) * np.eye(2)
    Vt = vt

    # Initialization
    theta = np.zeros((n, 2))
    P = np.zeros((n, 2, 2))
    P[0] = np.eye(2)

    spread = np.zeros(n)
    spread[0] = float("nan")

    for t in range(1, n):
        F = np.array([x.iloc[t], 1.0])

        # Predict
        theta_pred = theta[t - 1]
        P_pred = P[t - 1] + Wt

        # Innovation
        innovation = y.iloc[t] - F @ theta_pred
        S = F @ P_pred @ F.T + Vt

        # Kalman gain
        K = P_pred @ F.T / S

        # Update
        theta[t] = theta_pred + K * innovation
        P[t] = (np.eye(2) - np.outer(K, F)) @ P_pred

        spread[t] = y.iloc[t] - theta[t, 0] * x.iloc[t] - theta[t, 1]

    return pd.DataFrame({
        "hedge_ratio": theta[:, 0],
        "intercept": theta[:, 1],
        "spread": spread,
    }, index=y.index)

# Typical China A-share sector-correlation patterns
ASHARE_SECTOR_PATTERNS = {
    "strong_pairs_gt_0_7": [
        "Banks & insurance",
        "Baijiu & consumer staples",
        "New energy & solar",
        "Defense & aerospace",
    ],
    "medium_pairs_0_4_to_0_7": [
        "Pharma & consumer",
        "Technology & semiconductors",
        "Real estate & building materials",
    ],
    "low_or_negative_lt_0_3": [
        "Gold & technology",
        "Utilities & cyclicals",
        "Consumer & cyclicals",
    ],
}

def cross_market_correlation(
    markets: dict,  # {"China A-shares": series, "Hong Kong": series, "crypto": series, "US": series}
    rolling_window: int = 60,
    lag_days: list = [0, 1, 2, 3],
) -> dict:
    """Cross-market correlation plus lead-lag analysis.

    Args:
        markets: Daily return series for each market
        rolling_window: Rolling window
        lag_days: List of lags to test

    Returns:
        Correlation matrix, lead-lag analysis, and rolling correlation
    """
    df = pd.DataFrame(markets).dropna()

    # Static correlation matrix
    static_corr = df.corr()

    # Lead-lag correlation to detect cross-market transmission
    lead_lag = {}
    mkt_names = list(markets.keys())
    for i, m1 in enumerate(mkt_names):
        for m2 in mkt_names[i + 1:]:
            pair_key = f"{m1}_{m2}"
            lead_lag[pair_key] = {}
            for lag in lag_days:
                if lag == 0:
                    r, _ = pearsonr(df[m1], df[m2])
                else:
                    r, _ = pearsonr(df[m1].iloc[lag:], df[m2].iloc[:-lag])
                lead_lag[pair_key][f"lag_{lag}d"] = round(r, 4)

    # Rolling correlation
    rolling_corrs = {}
    for i, m1 in enumerate(mkt_names):
        for m2 in mkt_names[i + 1:]:
            key = f"{m1}_{m2}"
            rolling_corrs[key] = df[m1].rolling(rolling_window).corr(df[m2])

    return {
        "static_corr": static_corr,
        "lead_lag": pd.DataFrame(lead_lag).T,
        "rolling_corrs": pd.DataFrame(rolling_corrs),
    }

def fx_adjusted_correlation(
    foreign_price: pd.Series,   # foreign-market price, denominated in foreign currency
    domestic_price: pd.Series,  # domestic-market price
    fx_rate: pd.Series,         # foreign currency / domestic currency, e.g. USD/CNY
) -> dict:
    """Cross-market correlation adjusted for FX effects.

    Args:
        foreign_price: Foreign-market price series in foreign currency
        domestic_price: Domestic-market price series in domestic currency
        fx_rate: FX series expressed as foreign / domestic

    Returns:
        Raw correlation vs FX-adjusted correlation
    """
    # Domestic-currency foreign return = foreign return + FX return
    foreign_ret = foreign_price.pct_change(fill_method=None)
    fx_ret = fx_rate.pct_change(fill_method=None)
    foreign_ret_cny = (1 + foreign_ret) * (1 + fx_ret) - 1

    domestic_ret = domestic_price.pct_change(fill_method=None)

    df = pd.concat([foreign_ret.rename("foreign_raw"),
                    foreign_ret_cny.rename("foreign_domestic"),
                    domestic_ret.rename("domestic"),
                    fx_ret.rename("fx")], axis=1).dropna()

    raw_corr, _ = pearsonr(df["foreign_raw"], df["domestic"])
    adj_corr, _ = pearsonr(df["foreign_domestic"], df["domestic"])
    fx_corr, _ = pearsonr(df["fx"], df["domestic"])

    return {
        "raw_corr_foreign_domestic": round(raw_corr, 4),
        "fx_adjusted_corr": round(adj_corr, 4),
        "fx_domestic_corr": round(fx_corr, 4),
        "fx_contribution": round(adj_corr - raw_corr, 4),
        "note": "fx_contribution > 0 means FX amplified cross-market correlation",
    }

def correlation_breakdown_test(
    returns: pd.DataFrame,
    crisis_threshold: float = -0.02,  # one-day benchmark drop threshold for crisis days
    benchmark_col: str = None,
    window: int = 20,
) -> dict:
    """Detect jumps in correlation during crisis periods.

    Args:
        returns: Multi-asset daily return matrix
        crisis_threshold: Benchmark return below this level defines a crisis day
        benchmark_col: Benchmark column name. If None, use cross-sectional mean return
        window: Window for rolling average correlation

    Returns:
        Comparison of correlation in normal periods vs crisis periods
    """
    if benchmark_col:
        bm = returns[benchmark_col]
    else:
        bm = returns.mean(axis=1)

    crisis_mask = bm < crisis_threshold
    normal_mask = ~crisis_mask

    # Average pairwise correlation for each period
    def avg_corr(df_subset: pd.DataFrame) -> float:
        if len(df_subset) < 5:
            return float("nan")
        c = df_subset.corr()
        upper = c.where(np.triu(np.ones(c.shape), k=1).astype(bool))
        return float(upper.stack().mean())

    crisis_corr = avg_corr(returns[crisis_mask])
    normal_corr = avg_corr(returns[normal_mask])

    # Rolling average correlation to detect structural change
    rolling_avg_corr = pd.Series(dtype=float, index=returns.index)
    for i in range(window, len(returns)):
        sub = returns.iloc[i - window:i]
        rolling_avg_corr.iloc[i] = avg_corr(sub)

    return {
        "normal_avg_corr": round(normal_corr, 4),
        "crisis_avg_corr": round(crisis_corr, 4),
        "corr_jump": round(crisis_corr - normal_corr, 4),
        "crisis_days": int(crisis_mask.sum()),
        "normal_days": int(normal_mask.sum()),
        "rolling_avg_corr": rolling_avg_corr,
    }

Step 1: Asset screening
  - Run scan_correlated_assets and keep candidate pairs with Pearson > 0.6
  - Run engle_granger_coint and keep pairs with p < 0.05

Step 2: Spread quality assessment
  - Run compute_half_life and keep pairs with half-life between 5 and 60 days
  - Test spread stationarity with ADF and require p < 0.05
  - Measure how often the absolute Z-Score exceeds 1.5 over the last 12 months
    to estimate trading frequency

Step 3: Hedge-ratio selection
  - Use static OLS for stable pairs
  - Use Kalman Filter for long-lived or drifting pairs

Step 4: Signal generation
  - Compute rolling Z-Score with a lookback of 2-3× half-life
  - Generate long / short / exit signals based on thresholds

Step 5: Signal monitoring
  - Recompute Z-Score daily
  - Re-run cointegration monthly to avoid broken relationships
  - Warn if half-life exceeds 2× the original value

def generate_pair_signals(
    y_price: pd.Series,
    x_price: pd.Series,
    lookback: int = 60,
    entry_z: float = 2.0,
    exit_z: float = 0.5,
    stop_z: float = 3.5,
    use_kalman: bool = False,
) -> pd.DataFrame:
    """Generate pair-trading signals.

    Args:
        y_price, x_price: Two price series
        lookback: Rolling Z-Score lookback, usually 2-3× half-life
        entry_z: Entry threshold
        exit_z: Exit threshold, usually near mean reversion
        stop_z: Stop threshold. Crossing it suggests cointegration may have broken
        use_kalman: Whether to use a Kalman dynamic hedge ratio

    Returns:
        DataFrame containing signals, Z-Score, and positions
    """
    if use_kalman:
        kf = kalman_hedge_ratio(y_price, x_price)
        spread = kf["spread"]
    else:
        y_ret = y_price.pct_change(fill_method=None)
        x_ret = x_price.pct_change(fill_method=None)
        res = bivariate_correlation_analysis(y_ret, x_ret, lookback)
        hedge_ratio = abs(res["beta"])
        spread = np.log(y_price) - hedge_ratio * np.log(x_price)

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

    # Signal state machine to avoid repeated re-entry.
    signal_y = pd.Series(0.0, index=y_price.index)
    signal_x = pd.Series(0.0, index=x_price.index)
    position = 0  # 0=flat, 1=long spread, -1=short spread

    for i in range(lookback, len(z_score)):
        z = z_score.iloc[i]
        if np.isnan(z):
            continue

        if position == 0:
            if z < -entry_z:
                position = 1   # spread is too low: buy y, sell x
            elif z > entry_z:
                position = -1  # spread is too high: sell y, buy x
        elif position == 1:
            if z > -exit_z or z > stop_z:
                position = 0
        elif position == -1:
            if z < exit_z or z < -stop_z:
                position = 0

        signal_y.iloc[i] = 0.5 * position
        signal_x.iloc[i] = -0.5 * position

    return pd.DataFrame({
        "spread": spread,
        "z_score": z_score,
        "spread_mean": spread_mean,
        "spread_std": spread_std,
        "signal_y": signal_y,
        "signal_x": signal_x,
        "position": signal_y * 2,  # 1=long spread, -1=short spread, 0=flat
    })

def monitor_spread_health(
    spread: pd.Series,
    original_half_life: float,
    original_corr: float,
    warning_hl_multiple: float = 2.0,
    warning_corr_drop: float = 0.2,
) -> dict:
    """Monitor spread stability and judge whether cointegration still holds.

    Args:
        spread: Live spread series
        original_half_life: Half-life at entry, in days
        original_corr: Correlation at entry
        warning_hl_multiple: Warn if half-life exceeds this multiple of the original
        warning_corr_drop: Warn if correlation drops by more than this amount

    Returns:
        Health-status report
    """
    recent = spread.iloc[-60:] if len(spread) > 60 else spread

    current_hl = compute_half_life(recent)
    current_adf = adfuller(recent.dropna())[1]

    hl_ratio = current_hl / original_half_life if original_half_life > 0 else float("inf")

    # Cointegration health score
    health_score = 100
    warnings = []

    if current_adf > 0.10:
        health_score -= 40
        warnings.append(f"Spread ADF p={current_adf:.3f} > 0.10, stationarity has weakened")

    if hl_ratio > warning_hl_multiple:
        health_score -= 30
        warnings.append(
            f"Half-life {current_hl:.1f}d is {hl_ratio:.1f}x the original {original_half_life:.1f}d"
        )

    if current_adf > 0.20:
        health_score -= 20
        warnings.append("Spread may no longer be stationary. Re-test cointegration immediately")

    status = "healthy" if health_score >= 70 else "warning" if health_score >= 40 else "danger"

    return {
        "health_score": health_score,
        "status": status,
        "current_half_life": round(current_hl, 1),
        "hl_ratio": round(hl_ratio, 2),
        "spread_adf_p": round(current_adf, 4),
        "warnings": warnings,
        "action": "hold" if status == "healthy" else "reduce" if status == "warning" else "exit_now",
    }

import matplotlib.pyplot as plt

def plot_rolling_correlation(
    rolling_corrs: pd.DataFrame,
    title: str = "Rolling Correlation",
    figsize: tuple = (14, 5),
) -> plt.Figure:
    """Plot rolling correlation time series across multiple windows."""
    fig, ax = plt.subplots(figsize=figsize)
    colors = ["#2196F3", "#FF9800", "#4CAF50"]
    for i, col in enumerate(rolling_corrs.columns):
        ax.plot(rolling_corrs.index, rolling_corrs[col],
                label=col, color=colors[i % len(colors)], alpha=0.8)
    ax.axhline(0, color="black", linestyle="--", linewidth=0.8)
    ax.axhline(0.6, color="green", linestyle=":", linewidth=0.8, label="high-correlation threshold (0.6)")
    ax.axhline(-0.6, color="red", linestyle=":", linewidth=0.8)
    ax.set_title(title)
    ax.set_ylabel("Correlation")
    ax.legend(loc="best")
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    return fig

def plot_zscore_signals(
    signal_df: pd.DataFrame,
    entry_z: float = 2.0,
    stop_z: float = 3.5,
    figsize: tuple = (14, 8),
) -> plt.Figure:
    """Plot spread Z-Score and pair-trading signals."""
    fig, axes = plt.subplots(2, 1, figsize=figsize, sharex=True)

    # Top chart: spread
    axes[0].plot(signal_df["spread"], label="Spread", color="#1565C0")
    axes[0].plot(signal_df["spread_mean"], label="Mean", color="orange", linestyle="--")
    axes[0].fill_between(signal_df.index,
                         signal_df["spread_mean"] - signal_df["spread_std"],
                         signal_df["spread_mean"] + signal_df["spread_std"],
                         alpha=0.2, color="orange", label="±1σ")
    axes[0].set_title("Spread and Mean")
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)

    # Bottom chart: Z-Score + signals
    axes[1].plot(signal_df["z_score"], label="Z-Score", color="#1565C0")
    axes[1].axhline(entry_z, color="red", linestyle="--", label=f"entry threshold (±{entry_z})")
    axes[1].axhline(-entry_z, color="red", linestyle="--")
    axes[1].axhline(stop_z, color="darkred", linestyle=":", label=f"stop threshold (±{stop_z})")
    axes[1].axhline(-stop_z, color="darkred", linestyle=":")
    axes[1].axhline(0, color="black", linestyle="-", linewidth=0.8)

    # Annotate entry / exit points
    long_entry = signal_df["position"].diff() > 0
    short_entry = signal_df["position"].diff() < 0
    exit_pos = (signal_df["position"] == 0) & (signal_df["position"].shift(1) != 0)

    axes[1].scatter(signal_df.index[long_entry], signal_df["z_score"][long_entry],
                    color="green", marker="^", s=80, label="long spread", zorder=5)
    axes[1].scatter(signal_df.index[short_entry], signal_df["z_score"][short_entry],
                    color="red", marker="v", s=80, label="short spread", zorder=5)
    axes[1].scatter(signal_df.index[exit_pos], signal_df["z_score"][exit_pos],
                    color="gray", marker="o", s=40, label="exit", zorder=5)

    axes[1].set_title("Z-Score and Trading Signals")
    axes[1].legend(loc="best")
    axes[1].grid(True, alpha=0.3)

    plt.tight_layout()
    return fig

pip install pandas numpy scipy statsmodels matplotlib seaborn

## Correlation and Cointegration Analysis Report

### Pair: [Asset A] vs [Asset B] ([Start Date] - [End Date])

#### Correlation Statistics
| Metric | Value | Interpretation |
|------|------|------|
| Pearson r | 0.82 | Strong linear positive correlation |
| Spearman ρ | 0.80 | Consistent monotonic relationship |
| Beta (A/B) | 1.15 | Sensitivity of A to B |
| R² | 0.67 | 67% of return variance in A is explained by B |

#### Cointegration Tests
| Method | Statistic | p-value | Conclusion |
|------|--------|------|------|
| Engle-Granger | -4.12 | 0.008 | Cointegrated ** |
| Johansen trace test | 28.3 | — | 1 cointegrating vector |
| Spread ADF | -3.95 | 0.002 | Spread is stationary ** |

#### Mean-Reversion Characteristics
| Metric | Value |
|------|------|
| OLS hedge ratio | 1.23 |
| Half-life | 18.5 days |
| Suggested holding window | 10-30 days |
| Suggested lookback window | 40-60 days |

#### Conditional Correlation (Regime Analysis)
| Regime | Correlation | Sample Size |
|------|---------|--------|
| Bull | 0.76 | 312 days |
| Bear | 0.88 | 198 days |
| High volatility | 0.91 | 87 days |
| Sideways | 0.71 | 645 days |

#### Recommended Pair-Trading Signal Parameters
| Parameter | Value |
|------|-----|
| entry_z | 2.0 |
| exit_z | 0.5 |
| stop_z | 3.5 |
| lookback | 60 days |

#### Current Spread Status
| Metric | Value | Alert |
|------|-----|------|
| Current Z-Score | -2.3 | Near entry zone |
| Health score | 85/100 | Healthy |
| Half-life (last 60 days) | 21.2 days | Normal |