Hedge ratios: OLS, total least squares and the Kalman filter

"""M13A · Hedge ratios: OLS, total least squares and the Kalman filter. Read notes §13A.9 first.

State-space model (eq. 13A.31):
    y_t = x_t' beta_t + eps_t,        eps_t ~ N(0, R)
    beta_t = beta_{t-1} + eta_t,      eta_t ~ N(0, Q)
Filter (Prop. 13A.13, eq. 13A.32–13A.34):
    predict:  b_{t|t-1} = b_{t-1|t-1},  P_{t|t-1} = P_{t-1|t-1} + Q
    innovate: e_t = y_t - x_t' b_{t|t-1},  F_t = x_t' P_{t|t-1} x_t + R
    update:   K_t = P_{t|t-1} x_t / F_t,  b_{t|t} = b_{t|t-1} + K_t e_t,
              P_{t|t} = P_{t|t-1} - K_t x_t' P_{t|t-1}

Run:  python -m pytest modules/M13A-cointegration-stat-arb/tests/test_m13a_kalman.py --import-mode=importlib
"""

from __future__ import annotations

import numpy as np


def ols_slope(y: np.ndarray, x: np.ndarray) -> float:
    """Slope of the OLS regression of y on a constant and x: s_xy / s_xx (centred sums)."""
    raise NotImplementedError


def tls_slope(y: np.ndarray, x: np.ndarray) -> float:
    """Total-least-squares (orthogonal regression) slope of y on x, with intercept (eq. 13A.30):
        beta = [ (s_yy - s_xx) + sqrt((s_yy - s_xx)^2 + 4 s_xy^2) ] / (2 s_xy),
    with centred sums s_xx, s_yy, s_xy.
    """
    raise NotImplementedError


def simulate_tv_beta(n: int, q: float, r: float, beta0: float = 1.0, seed: int = 0):
    """Simulate (y, x, beta) with x a positive random-walk 'price', beta a Gaussian random walk
    with variance q per step, and y = x beta + eps, Var(eps) = r. (Given.)"""
    rng = np.random.default_rng(seed)
    x = 100.0 * np.exp(np.cumsum(0.01 * rng.standard_normal(n)))
    beta = beta0 + np.cumsum(np.sqrt(q) * rng.standard_normal(n))
    y = x * beta + np.sqrt(r) * rng.standard_normal(n)
    return y, x, beta


def kalman_regression(y: np.ndarray, X: np.ndarray, Q, R: float, b0, P0):
    """Kalman filter for the time-varying regression (eq. 13A.31–13A.34).

    y: (n,), X: (n, k) regressors (for the scalar model pass x[:, None]).
    Q: (k, k) state-noise covariance (a scalar is accepted when k == 1; use np.atleast_2d).
    R: observation-noise variance. b0 (k,), P0 (k, k): prior for beta_0 (before the first predict).
    Returns dict with 'b' (n, k) filtered means b_{t|t}, 'P' (n, k, k) filtered covariances,
    'e' (n,) innovations, 'F' (n,) innovation variances, 'b_pred' (n, k) predictions b_{t|t-1}.
    """
    raise NotImplementedError


def kalman_loglik(y, X, Q, R, b0, P0, burn: int = 10) -> float:
    """Gaussian log-likelihood by prediction-error decomposition (eq. 13A.35), skipping the first
    `burn` observations:
        -1/2 sum_{t >= burn} [ ln(2 pi F_t) + e_t^2 / F_t ].
    """
    raise NotImplementedError


def steady_state(q: float, r: float, x: float = 1.0):
    """Steady state of the scalar filter with constant regressor x (Prop. 13A.14, eq. 13A.36).

    Returns (P_pred, K) with P_pred = [q + sqrt(q^2 + 4 q r / x^2)] / 2 the limit of P_{t|t-1}
    and K = P_pred x / (x^2 P_pred + r) the limiting gain.
    """
    raise NotImplementedError


def rolling_ols_beta(y, x, window: int) -> np.ndarray:
    """Rolling no-intercept OLS beta_t = sum y x / sum x^2 over the last `window` points (NaN before). (Given.)"""
    y, x = np.asarray(y, float), np.asarray(x, float)
    num = np.convolve(y * x, np.ones(window), "full")[: y.size]
    den = np.convolve(x * x, np.ones(window), "full")[: y.size]
    out = num / den
    out[: window - 1] = np.nan
    return out