""" missdata.py -- FOUNDATIONS OF BAYESIAN MISSING-DATA IMPUTATION (from scratch). Backs the notebooks in "Missing Data -- Foundations". The Bayesian view of missing data is disarmingly simple: a missing value is just another unknown, so give it a distribution and sample it alongside the parameters. Under a multivariate normal model that is Schafer's DATA AUGMENTATION (Little & Rubin; Tanner & Wong) -- the exact "treat the unknown as a parameter and draw it" move already used for the latent z in Albert-Chib probit, the censored draws in Tobit, and the latent class in LCA. Missing data is the general case; those were special cases. Rubin's taxonomy sets the ground rules. Missingness is MCAR if it is independent of all data, MAR if it depends only on the OBSERVED data, and MNAR if it depends on the values that are missing. Under MCAR/MAR the missingness mechanism is IGNORABLE: we may impute from the data model alone and need not model why data are missing. Complete-case (listwise) deletion is unbiased only under MCAR and generally wastes information; Bayesian imputation uses every observed entry. The sampler is a two-step Gibbs on an incomplete N x p multivariate normal Y ~ N(mu, Sigma): I-step (impute): for each row, draw the missing entries from their conditional normal given the observed entries, y_mis | y_obs ~ N( mu_mis + B (y_obs - mu_obs), S_mis.obs ), with B = Sigma_mis,obs Sigma_obs,obs^{-1} and Schur-complement covariance. P-step (params): given the completed data, draw (mu, Sigma) from the conjugate normal-inverse- Wishart posterior (here the standard noninformative limit): Sigma ~ InvWishart(N-1, S), mu | Sigma ~ N(ybar, Sigma / N). Rows are grouped by missingness PATTERN so each conditional regression B is formed once per pattern (Schafer's efficiency trick). Storing spaced completed datasets yields the multiple imputations that feed Rubin's rules in the next project. """ import numpy as np from scipy.stats import invwishart # --------------------------------------------------------------------------- # # induce missingness under a chosen mechanism # # --------------------------------------------------------------------------- # def induce_missing(Y, rng, mechanism="MCAR", frac=0.3, target=0, driver=1): """Return a copy of Y with NaNs in column `target` under MCAR / MAR / MNAR. MCAR: missing completely at random (rate `frac`). MAR : probability of missing rises with the OBSERVED `driver` column. MNAR: probability of missing rises with the target column's OWN (would-be) value.""" Y = np.array(Y, float); N = Y.shape[0] if mechanism == "MCAR": p = np.full(N, frac) elif mechanism == "MAR": z = (Y[:, driver] - Y[:, driver].mean()) / Y[:, driver].std() p = 1 / (1 + np.exp(-(z * 2))) # higher driver -> more missing else: # MNAR z = (Y[:, target] - Y[:, target].mean()) / Y[:, target].std() p = 1 / (1 + np.exp(-(z * 2))) # higher own value -> more missing p = p * (frac / p.mean()) # scale to the requested average rate Y[rng.random(N) < np.clip(p, 0, 1), target] = np.nan return Y # --------------------------------------------------------------------------- # # simulation # # --------------------------------------------------------------------------- # def simulate_mvn(N, mu, Sigma, rng): return rng.multivariate_normal(np.asarray(mu, float), np.asarray(Sigma, float), size=N) # --------------------------------------------------------------------------- # # data-augmentation Gibbs sampler # # --------------------------------------------------------------------------- # def _patterns(Mask): """group row indices by identical missingness pattern (tuple of missing columns).""" groups = {} for i, row in enumerate(Mask): groups.setdefault(tuple(np.where(row)[0]), []).append(i) return {k: np.array(v) for k, v in groups.items()} def da_gibbs(Y, rng, draws=2000, burn=1000, m_keep=20): """Schafer data-augmentation Gibbs for an incomplete multivariate normal. Y : (N,p) with np.nan at missing entries. Returns posterior draws of mu (draws,p) and Sigma (draws,p,p), `m_keep` completed datasets (multiple imputations) and the posterior-mean imputation.""" Y = np.array(Y, float); N, p = Y.shape Mask = np.isnan(Y); pat = _patterns(Mask) Yc = Y.copy() col_mean = np.nanmean(Y, axis=0) for j in range(p): Yc[Mask[:, j], j] = col_mean[j] # start from mean imputation MU = np.empty((draws, p)); SIG = np.empty((draws, p, p)) keep_idx = set(np.linspace(0, draws - 1, m_keep).astype(int)) completed = []; imp_sum = np.zeros((N, p)); n_imp = 0 for it in range(draws + burn): # ---- P-step: (mu, Sigma) | completed data (noninformative NIW limit) ---- ybar = Yc.mean(0); D = Yc - ybar; S = D.T @ D Sigma = invwishart.rvs(df=N - 1, scale=S, random_state=rng) Sigma = np.atleast_2d(Sigma) mu = ybar + np.linalg.cholesky(Sigma / N) @ rng.standard_normal(p) # ---- I-step: impute missing entries per pattern ---- 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: So_o = Sigma[np.ix_(obs, obs)] B = Sigma[np.ix_(mis, obs)] @ np.linalg.inv(So_o) cmean = mu[mis] + (Yc[np.ix_(rows, obs)] - mu[obs]) @ B.T # (n_rows, n_mis) ccov = Sigma[np.ix_(mis, mis)] - B @ Sigma[np.ix_(obs, mis)] else: # whole row missing -> marginal cmean = np.tile(mu[mis], (len(rows), 1)); ccov = Sigma[np.ix_(mis, mis)] Lc = np.linalg.cholesky(ccov + 1e-10 * np.eye(len(mis))) draw = cmean + rng.standard_normal((len(rows), len(mis))) @ Lc.T Yc[np.ix_(rows, mis)] = draw if it >= burn: i = it - burn; MU[i] = mu; SIG[i] = Sigma imp_sum += Yc; n_imp += 1 if i in keep_idx: completed.append(Yc.copy()) return dict(mu=MU, Sigma=SIG, completed=completed, imputed_mean=imp_sum / n_imp, mask=Mask) # --------------------------------------------------------------------------- # # EM (frequentist counterpart) and complete-case moments # # --------------------------------------------------------------------------- # def em_mvn(Y, tol=1e-6, max_iter=500): """EM for the incomplete-multivariate-normal MLE of (mu, Sigma) -- the frequentist twin of the Bayesian data augmentation (expectation replaces the imputation draw, maximisation the parameter draw).""" Y = np.array(Y, float); N, p = Y.shape; Mask = np.isnan(Y); pat = _patterns(Mask) cc = ~Mask.any(axis=1) # init from complete-case moments: if cc.sum() > p: # a DIAGONAL Sigma is a degenerate fixed mu = Y[cc].mean(0); Sigma = np.cov(Y[cc], rowvar=False) # point (B=0), so start full-rank else: mu = np.nanmean(Y, axis=0); Sigma = np.diag(np.nanvar(Y, axis=0)) for _ in range(max_iter): T1 = np.zeros(p); T2 = np.zeros((p, p)) for mis, rows in pat.items(): obs = [j for j in range(p) if j not in mis]; mis = list(mis) Yr = Y[rows] if mis and obs: So_o = Sigma[np.ix_(obs, obs)] B = Sigma[np.ix_(mis, obs)] @ np.linalg.inv(So_o) yhat = Yr.copy() yhat[:, mis] = mu[mis] + (Yr[:, obs] - mu[obs]) @ B.T ccov = Sigma[np.ix_(mis, mis)] - B @ Sigma[np.ix_(obs, mis)] elif mis: yhat = np.tile(mu, (len(rows), 1)); ccov = Sigma[np.ix_(mis, mis)] else: yhat = Yr; ccov = None T1 += yhat.sum(0); T2 += yhat.T @ yhat if mis and ccov is not None: add = np.zeros((p, p)); add[np.ix_(mis, mis)] = ccov * len(rows); T2 += add mu_new = T1 / N; Sigma_new = T2 / N - np.outer(mu_new, mu_new) if np.max(np.abs(mu_new - mu)) < tol: mu, Sigma = mu_new, Sigma_new; break mu, Sigma = mu_new, Sigma_new return mu, Sigma def complete_case(Y): """listwise-deletion moments -- unbiased only under MCAR.""" Y = np.array(Y, float); keep = ~np.isnan(Y).any(axis=1); Yc = Y[keep] return Yc.mean(0), np.cov(Yc, rowvar=False), keep.sum()