adaptive beta shrinkage

adaptive_beta.py

import numpy as np

def nearest_positive_definite(A: np.ndarray) -> np.ndarray:
    """
    Higham (2002) projection to the nearest symmetric positive definite matrix.
    """
    B = (A + A.T) / 2
    U, s, Vt = np.linalg.svd(B)
    H = (Vt.T * s) @ Vt
    A2 = (B + H) / 2
    A3 = (A2 + A2.T) / 2  # enforce symmetry

    # If still not PD (due to numerical issues), nudge the spectrum
    def is_pd(X):
        try:
            np.linalg.cholesky(X)
            return True
        except np.linalg.LinAlgError:
            return False

    if is_pd(A3):
        return A3

    spacing = np.spacing(np.linalg.norm(A))
    I = np.eye(A.shape[0])
    k = 1
    while not is_pd(A3):
        mineig = np.min(np.real(np.linalg.eigvals(A3)))
        A3 += I * (-mineig * (k**2) + spacing)
        k += 1
    return A3


def ABS(R: np.ndarray, lam: float | None = None, V: np.ndarray | None = None, B_prior: str = "0") -> np.ndarray:
    """
    Python port of the R ABS() function.

    Parameters
    ----------
    R : array-like, shape (t, n)
        Return matrix: t observations (rows) × n assets (cols).
    lam : float or None
        Ridge penalty λ. If None, set to max(log(n / t), 0).
    V : array-like or None
        Length-t weights. If None, uses all ones.
    B_prior : {"0", "MLE"}
        Prior for Vasicek shrinkage of pairwise betas.

    Returns
    -------
    P_hat : ndarray, shape (n, n)
        Shrunk, near-PD correlation matrix.
    """
    R = np.asarray(R, dtype=float)
    T, n = R.shape  # t = rows (time), n = columns (assets)

    # Set lambda
    if lam is None:
        lam = max(np.log(n / T), 0.0)

    # Set V
    if V is None:
        V = np.ones(T, dtype=float)
    else:
        V = np.asarray(V, dtype=float)
        if V.shape[0] != T:
            raise ValueError("the weight vector V must be of length t (number of rows of R)")

    # (Optional) sample correlation (not used later; kept for parity with R code)
    # P_sample = np.corrcoef(R, rowvar=False)

    # STEP (2): standardize returns (column-wise z-scores)
    mu = R.mean(axis=0, keepdims=True)
    std = R.std(axis=0, ddof=1, keepdims=True)
    # Avoid divide-by-zero if any std is zero
    std = np.where(std == 0, 1.0, std)
    S = (R - mu) / std

    # STEP (3): weighted ridge regression
    # M = S' W S, with W = diag(V)
    # Use vectorized form: W @ S == V[:, None] * S
    M = S.T @ (V[:, None] * S)  # shape (n, n)

    # Off-diagonal indices
    ii, jj = np.where(~np.eye(n, dtype=bool))  # (i_idx, j_idx)
    M_jj = M[jj, jj]
    M_ji = M[jj, ii]
    M_ii = M[ii, ii]

    # Ridge estimate of beta for off-diagonal entries
    B_hat_off = M_ji / (M_jj + lam)

    # STEP (4): Vasicek shrinkage
    if B_prior == "0":
        B_prior_off = np.zeros_like(B_hat_off)
    elif B_prior == "MLE":
        B_prior_off = np.full_like(B_hat_off, np.mean(B_hat_off))
    else:
        raise ValueError("B_prior must be '0' or 'MLE'")

    # Residual variance, sampling variance of beta
    ssr = M_ii - 2.0 * B_hat_off * M_ji + (B_hat_off**2) * M_jj
    regression_residual_variance = ssr / (T - 1)
    sigma2_hat_B_off = regression_residual_variance / (M_jj + lam)

    # Variance across all beta estimates
    sigma2_B_hat = np.var(B_hat_off, ddof=1)
    # Pairwise prior variance (clipped to avoid zeros/negatives)
    sigma2_prior_off = np.maximum(sigma2_B_hat - sigma2_hat_B_off, 1e-6)

    # Vasicek posterior mean
    B_post_off = ((B_prior_off / sigma2_prior_off) + (B_hat_off / sigma2_hat_B_off)) / \
                 ((1.0 / sigma2_prior_off) + (1.0 / sigma2_hat_B_off))

    # Populate the shrunk correlation matrix
    P_hat = np.eye(n)
    P_hat[ii, jj] = B_post_off

    # Average the two pairwise betas (symmetrize)
    P_hat = 0.5 * (P_hat + P_hat.T)

    # Ensure positive definiteness; if needed, project to nearest PD
    evals = np.linalg.eigvalsh(P_hat)
    if np.any(evals <= 0):
        P_hat = nearest_positive_definite(P_hat)
        # For correlation matrices, enforce unit diagonal post-projection
        d = np.sqrt(np.diag(P_hat))
        d = np.where(d == 0, 1.0, d)
        P_hat = P_hat / np.outer(d, d)
        np.fill_diagonal(P_hat, 1.0)

    return P_hat