""" misscov.py -- MISSING COVARIATES in a regression, the joint-model way (from scratch). Backs the notebooks in "Missing Covariates in a Regression". The commonest missing-data situation is a fully observed OUTCOME y and PREDICTORS x with holes. The Bayesian remedy (Congdon; Ibrahim et al.) is to write a model for the covariates and augment the missing ones inside the regression -- a JOINT model of outcome and covariates: y_i = beta0 + beta' x_i + eps_i, eps_i ~ N(0, sigma^2) (the regression of interest) x_i ~ N(mu, Sigma) (a model for the covariates) Gibbs alternates three blocks: draw (beta, sigma^2) from the completed regression, draw (mu, Sigma) from the completed covariates (conjugate NIW), and IMPUTE each missing covariate from its full conditional given the observed covariates AND the outcome. That last "and the outcome" is the crux of missing-covariate imputation. A missing x_i is informed by two things -- the covariate model (x correlates with the observed covariates) and the regression (x must be consistent with the y actually seen). Combining a Gaussian covariate prior with the Gaussian y-likelihood gives a Gaussian conditional, and the y-term is what keeps the estimated beta unbiased. Impute the covariates WITHOUT the outcome (the common mistake -- filling predictors from the other predictors only) and the x-y association is diluted, ATTENUATING beta toward zero. The `use_y` switch below turns that term on and off so the attenuation can be shown directly. A second lesson specific to missing covariates: complete-case regression is unbiased for beta whenever missingness does not depend on y given x (a weaker condition than MCAR) -- so here deletion is more defensible than for missing outcomes, though still less efficient than imputation. """ import numpy as np from scipy.stats import invwishart def simulate_misscov(N, beta, sigma, mu_x, Sig_x, rng): """y = beta0 + beta[1:]'x + N(0,sigma^2); x ~ N(mu_x, Sig_x).""" X = rng.multivariate_normal(np.asarray(mu_x, float), np.asarray(Sig_x, float), N) y = beta[0] + X @ np.asarray(beta[1:], float) + sigma * rng.standard_normal(N) return y, X def _patterns(Mask): g = {} for i, row in enumerate(Mask): g.setdefault(tuple(np.where(row)[0]), []).append(i) return {k: np.array(v) for k, v in g.items()} def misscov_gibbs(y, X, rng, draws=2000, burn=1000, use_y=True): """Joint-model Gibbs for a regression with missing covariates. X:(N,p) with NaN, y:(N,) observed. If use_y=False the missing covariates are imputed from the covariate model ONLY (ignoring the outcome) -- which attenuates beta. Returns posterior draws of beta:(draws,p+1), sigma^2, the covariate moments, and the posterior-mean imputed X.""" y = np.asarray(y, float); X = np.array(X, float); N, p = X.shape Mask = np.isnan(X); pat = _patterns(Mask) Xc = X.copy() for j in range(p): Xc[Mask[:, j], j] = np.nanmean(X[:, j]) beta = np.zeros(p + 1); sig2 = np.var(y) B = np.empty((draws, p + 1)); S2 = np.empty(draws); MU = np.empty((draws, p)); imp_sum = np.zeros((N, p)); nimp = 0 for it in range(draws + burn): # ---- (beta, sigma^2) | completed data ---- D = np.column_stack([np.ones(N), Xc]); DtDi = np.linalg.inv(D.T @ D + 1e-8 * np.eye(p + 1)) bhat = DtDi @ (D.T @ y); resid = y - D @ bhat sig2 = (resid @ resid) / rng.chisquare(max(N - (p + 1), 1)) beta = bhat + np.linalg.cholesky(sig2 * DtDi) @ rng.standard_normal(p + 1) # ---- (mu, Sigma) | completed covariates (noninformative NIW) ---- xbar = Xc.mean(0); Sx = (Xc - xbar).T @ (Xc - xbar) Sigma = np.atleast_2d(invwishart.rvs(df=N - 1, scale=Sx, random_state=rng)) mu = xbar + np.linalg.cholesky(Sigma / N) @ rng.standard_normal(p) # ---- impute missing covariates | observed covariates (+ outcome) ---- b0 = beta[0]; bc = beta[1:] for mis, rows in pat.items(): if not mis: continue mis = list(mis); obs = [j for j in range(p) if j not in mis] if obs: Sooi = np.linalg.inv(Sigma[np.ix_(obs, obs)]); Bm = Sigma[np.ix_(mis, obs)] @ Sooi a = mu[mis] + (Xc[np.ix_(rows, obs)] - mu[obs]) @ Bm.T # covariate-model mean V = Sigma[np.ix_(mis, mis)] - Bm @ Sigma[np.ix_(obs, mis)] else: a = np.tile(mu[mis], (len(rows), 1)); V = Sigma[np.ix_(mis, mis)] Vi = np.linalg.inv(V); bm = bc[mis] if use_y: r = y[rows] - b0 - (Xc[np.ix_(rows, obs)] @ bc[obs] if obs else 0.0) # y residual P = Vi + np.outer(bm, bm) / sig2 # combine prior + y-likelihood Pi = np.linalg.inv(P) rhs = a @ Vi.T + np.outer(r, bm) / sig2 postmean = rhs @ Pi.T; Lp = np.linalg.cholesky(Pi) else: # ignore the outcome postmean = a; Lp = np.linalg.cholesky(V) Xc[np.ix_(rows, mis)] = postmean + rng.standard_normal((len(rows), len(mis))) @ Lp.T if it >= burn: i = it - burn; B[i] = beta; S2[i] = sig2; MU[i] = mu; imp_sum += Xc; nimp += 1 return dict(beta=B, sigma2=S2, mu=MU, Ximp=imp_sum / nimp, mask=Mask) def complete_case(y, X): """listwise-deletion OLS -- unbiased for beta if missingness is independent of y given x.""" y = np.asarray(y, float); X = np.asarray(X, float) keep = ~np.isnan(X).any(axis=1); D = np.column_stack([np.ones(keep.sum()), X[keep]]) DtDi = np.linalg.inv(D.T @ D); b = DtDi @ (D.T @ y[keep]) resid = y[keep] - D @ b; s2 = (resid @ resid) / (keep.sum() - D.shape[1]) return b, np.sqrt(np.diag(s2 * DtDi)), keep.sum() def mean_impute(y, X): """fill each missing covariate with its column mean, then OLS -- biased (attenuates beta).""" y = np.asarray(y, float); Xc = np.array(X, float) for j in range(Xc.shape[1]): Xc[np.isnan(Xc[:, j]), j] = np.nanmean(Xc[:, j]) D = np.column_stack([np.ones(len(y)), Xc]); DtDi = np.linalg.inv(D.T @ D); b = DtDi @ (D.T @ y) resid = y - D @ b; s2 = (resid @ resid) / (len(y) - D.shape[1]) return b, np.sqrt(np.diag(s2 * DtDi))