""" mice_fcs.py -- MULTIPLE IMPUTATION by FULLY CONDITIONAL SPECIFICATION / chained equations (scratch). Backs the notebooks in "Multiple Imputation by Chained Equations". Project 1 imputed by fitting ONE joint model (a multivariate normal) to all variables at once. That is elegant when everything is continuous, but real datasets mix continuous, binary and categorical columns, and no single joint distribution fits them all. MULTIPLE IMPUTATION BY CHAINED EQUATIONS (MICE, van Buuren; also "fully conditional specification", FCS) sidesteps the joint model: it imputes each incomplete variable from ITS OWN regression on all the others, cycling through the variables until the fills stabilise. A continuous column gets a linear regression, a binary column a probit/logistic one -- each variable is imputed by the model that suits it. Two ideas do the work. MULTIPLE IMPUTATION. A single filled-in value is a guess treated as if it were known, which makes the downstream analysis over-confident. Instead we create m > 1 completed datasets, each a different plausible draw, analyse every one, and COMBINE the results by Rubin's rules: qbar = mean of the m estimates T = Ubar + (1 + 1/m) B (within-imputation variance + between-imputation variance) The between-imputation term B is the extra uncertainty from not knowing the missing values, which single imputation simply omits. The fraction of missing information is (1+1/m)B / T. CHAINED EQUATIONS. Within each completed dataset we sweep the variables in turn; for the current one we draw a proper Bayesian regression on the others (conjugate linear regression for continuous targets, Albert-Chib probit augmentation for binary ones -- the same augmentation used throughout this portfolio) and replace its missing entries with a posterior-predictive draw. A dozen sweeps and the imputations are a sample from the implied joint distribution, never having specified one. """ import numpy as np from scipy.special import ndtr as Phi, ndtri as Phinv def _rtrunc_sign(mean, pos, rng): lo = np.where(pos, Phi(-mean), 0.0); hi = np.where(pos, 1.0, Phi(-mean)) u = lo + rng.random(mean.shape) * (hi - lo) return mean + Phinv(np.clip(u, 1e-12, 1 - 1e-12)) # --------------------------------------------------------------------------- # # per-variable Bayesian imputation models # # --------------------------------------------------------------------------- # def _blr_impute(y, X, obs, mis, rng): """Bayesian linear regression (conjugate reference prior): fit on observed rows, draw (beta, sigma^2) from the posterior, return a posterior-PREDICTIVE draw for the missing rows.""" Xo = X[obs]; yo = y[obs]; n, p = Xo.shape XtXi = np.linalg.inv(Xo.T @ Xo + 1e-8 * np.eye(p)) bhat = XtXi @ (Xo.T @ yo); resid = yo - Xo @ bhat sig2 = (resid @ resid) / rng.chisquare(max(n - p, 1)) # sigma^2 ~ RSS / chi^2_{n-p} beta = bhat + np.linalg.cholesky(sig2 * XtXi) @ rng.standard_normal(p) return X[mis] @ beta + np.sqrt(sig2) * rng.standard_normal(len(mis)) def _probit_impute(y, X, obs, mis, rng, aug=20): """Bayesian probit regression via Albert-Chib augmentation: fit on observed rows, then draw the missing entries as Bernoulli(Phi(x'beta)) -- the binary-variable imputer.""" Xo = X[obs]; yo = y[obs]; n, p = Xo.shape XtXi = np.linalg.inv(Xo.T @ Xo + 1e-6 * np.eye(p)); L0 = np.linalg.cholesky(XtXi) beta = np.zeros(p) for _ in range(aug): z = _rtrunc_sign(Xo @ beta, yo > 0.5, rng) beta = XtXi @ (Xo.T @ z) + L0 @ rng.standard_normal(p) return (rng.random(len(mis)) < Phi(X[mis] @ beta)).astype(float) # --------------------------------------------------------------------------- # # the chained-equations sampler # # --------------------------------------------------------------------------- # def fcs_impute(Y, types, rng, m=20, iters=12): """Fully conditional specification. Y:(N,p) with NaN; types: list of 'cont' / 'bin' / 'complete' per column ('complete' = never imputed, used only as a predictor). Returns a list of m completed (N,p) datasets -- the multiple imputations.""" Y = np.array(Y, float); N, p = Y.shape; Mask = np.isnan(Y) completed = [] for _ in range(m): Yc = Y.copy() for j in range(p): # initialise from random observed values if Mask[:, j].any(): Yc[Mask[:, j], j] = rng.choice(Yc[~Mask[:, j], j], Mask[:, j].sum()) for _ in range(iters): # cycle through variables for j in range(p): if types[j] == "complete" or not Mask[:, j].any(): continue mis = np.where(Mask[:, j])[0]; obs = np.where(~Mask[:, j])[0] others = [k for k in range(p) if k != j] X = np.column_stack([np.ones(N)] + [Yc[:, k] for k in others]) imp = _blr_impute(Yc[:, j], X, obs, mis, rng) if types[j] == "cont" \ else _probit_impute(Yc[:, j], X, obs, mis, rng) Yc[mis, j] = imp completed.append(Yc) return completed # --------------------------------------------------------------------------- # # analysis on completed data + Rubin's rules # # --------------------------------------------------------------------------- # def ols(y, X): """OLS fit -> (coefficients, coefficient covariance matrix).""" XtXi = np.linalg.inv(X.T @ X); b = XtXi @ (X.T @ y) resid = y - X @ b; s2 = (resid @ resid) / (len(y) - X.shape[1]) return b, s2 * XtXi def rubin_pool(betas, covs, dfcom=None): """Combine m complete-data fits by Rubin's rules. betas:(m,k), covs:(m,k,k). `dfcom` is the COMPLETE-DATA degrees of freedom (n - k). Supply it: Rubin's original (1987) df formula assumes dfcom is effectively infinite and can return more degrees of freedom than the complete data could ever provide -- on a 25-row dataset it happily reports 190. Barnard & Rubin (1999) correct that by combining it with the observed-data df; with dfcom given, `df` is the corrected value and `df_rubin` the uncorrected one for comparison.""" betas = np.asarray(betas); covs = np.asarray(covs); m, k = betas.shape qbar = betas.mean(0) Ubar = covs.mean(0).diagonal() # within-imputation variance B = betas.var(0, ddof=1) # between-imputation variance T = Ubar + (1 + 1 / m) * B # total variance fmi = (1 + 1 / m) * B / T # fraction of missing information with np.errstate(divide="ignore", invalid="ignore"): df_rubin = (m - 1) * (1 + Ubar / ((1 + 1 / m) * B)) ** 2 # Rubin (1987), large-sample df = df_rubin if dfcom is not None: # Barnard & Rubin (1999) nu_obs = (dfcom + 1) / (dfcom + 3) * dfcom * (1 - fmi) df = 1.0 / (1.0 / df_rubin + 1.0 / nu_obs) return dict(estimate=qbar, se=np.sqrt(T), within=Ubar, between=B, fmi=fmi, df=df, df_rubin=df_rubin) def complete_case_ols(Y, ycol, xcols): """listwise-deletion regression: keep only fully observed rows across the used columns.""" keep = ~np.isnan(Y[:, [ycol] + xcols]).any(axis=1) X = np.column_stack([np.ones(keep.sum())] + [Y[keep, c] for c in xcols]) b, cov = ols(Y[keep, ycol], X) return b, np.sqrt(np.diag(cov)), keep.sum()