Sherman-Morrison rolling regression

pandas_sherman.py

import pandas as pd

def pd_move_regress(X, Y, w, fit_intercept=True):
    assert (X.index == Y.index).all()
    B, A = move_regress(X.values, Y.values, w, fit_intercept=fit_intercept)
    
    idx = pd.MultiIndex.from_product([X.columns, Y.columns])
    B = pd.DataFrame(B.reshape(len(B), -1), index=X.index, columns=idx)
    A = pd.DataFrame(A, index=X.index, columns=Y.columns)
    return B, A
  

sherman.py

import numpy as np

def move_regress(X, Y, w, fit_intercept=True):
    """Moving window multi-regressions
    
    Solves the least-squares problems `Y[t-w+1:t+1] = X[t-w+1:t+1] @ B[t] + A[t]`
    for `B` and `A` for all `t` in `[w-1, len(Y)-1]`.
 
    Parameters
    ----------
    X : (T, N) ndarray
        Input array.
    Y : (T, M) ndarray
        Response array.
    w : int
        The number of elements in the moving window.
    fit_intercept : boolean, optional
        whether to include the intercept `A_t` in the equations.
        
    Returns
    -------
    B : (T, N, M) ndarray
        The moving coefficients.
    A : (T, M) ndarray
        The moving intercepts.
    """
    if fit_intercept:
        mX = X[:w].mean(0)
        mY = Y[:w].mean(0)

        H = (X[:w] - mX).T @ (Y[:w] - mY)
        C = (X[:w] - mX).T @ (X[:w] - mX)

        q = w ** -0.5
        mX /= q
        mY /= q
    else:
        H = X[:w].T @ Y[:w]
        C = X[:w].T @ X[:w]

    K = np.linalg.inv(C)

    B = np.full(X.shape + Y.shape[1:], np.nan)
    B[w-1] = K @ H
    
    A = np.full(Y.shape, np.nan)
    if fit_intercept:
        A[w-1] = mY - mX @ B[w-1]
    else:
        A[w-1:] = 0

    for i in range(w, len(B)):
        x_in, y_in, x_out, y_out = X[i], Y[i], X[i-w], Y[i-w]
        
        if fit_intercept:
            z = K @ mX
            H += mX[:, None] * mY
            K -= z[:, None] * z / (1.0 + mX @ z)

            mX += q * (x_in - x_out)
            mY += q * (y_in - y_out)

            z = K @ mX
            H -= mX[:, None] * mY
            K += z[:, None] * z / (1.0 - mX @ z)

        z = K @ x_in
        H += x_in[:, None] * y_in
        K -= z[:, None] * z / (1.0 + x_in @ z)
        
        z = K @ x_out
        H -= x_out[:, None] * y_out
        K += z[:, None] * z / (1.0 - x_out @ z)

        B[i] = K @ H
        if fit_intercept:
            A[i] = mY - mX @ B[i]

    if fit_intercept:
        A *= q
    
    return B, A