scripts/claim5_boston.py

"""Claim 5 -- Table 3: Boston-housing linear regression under strong model
misspecification.  Reproduce the covariance errors of CT / LR+WS / DQ+exact
(and, additionally, DQ+const, which the claim wording mentions but Table 3 does
not actually contain) at B = 16 and B = floor(0.1 N).

Protocol (matched to the paper and to the authors' released script):
  * 12 covariates (all Boston columns except CHAS), standardised; response MEDV
    untransformed; no intercept.  N = 506, D = 12.
  * working Gaussian model with sigma = 1  ->  strong misspecification.
  * target S* = Jhat^{-1} Ihat Jhat^{-1} / N  (sandwich covariance at theta_hat).
  * 500 epochs of preconditioned SGD from theta_hat, 50% burn-in, 30 seeds.
  * covariance error = ||Shat - S*||_F / ||S*||_F.
Data: the original Harrison & Rubinfeld table from lib.stat.cmu.edu/datasets/boston
(byte-identical to the authors' boston.csv, which is checked here).
"""

import json
import sys
import time
import zlib
import numpy as np

sys.path.insert(0, "/home/ubuntu/samuel/sgmcmc-uq-repro/scripts")
from common import NoiseModel, relf
from linreg import LinReg, sandwich, tune_Lambda

OUT = "/home/ubuntu/samuel/sgmcmc-uq-repro/outputs"
ROOT = "/home/ubuntu/samuel/sgmcmc-uq-repro"
SEED = 20260725
COLS = [
    "CRIM",
    "ZN",
    "INDUS",
    "NOX",
    "RM",
    "AGE",
    "DIS",
    "RAD",
    "TAX",
    "PTRATIO",
    "B",
    "LSTAT",
]
PAPER = {  # Table 3 of the paper
    ("log", 16): dict(
        Posterior=0.358, CT=0.247, LRWS=9.23e8, DQexact=0.337, NUTS=2.528
    ),
    ("log", 50): dict(
        Posterior=0.358, CT=0.589, LRWS=1.40e7, DQexact=0.352, NUTS=2.528
    ),
    ("beta", 16): dict(CT=2.054, LRWS=float("inf"), DQexact=2.782),
    ("beta", 50): dict(CT=3.126, LRWS=float("inf"), DQexact=1.398),
}


def load_boston():
    """Parse the raw CMU file (each record spans two lines)."""
    raw = open(f"{ROOT}/data_boston_raw.txt").read().splitlines()
    vals = []
    for ln in raw[22:]:
        vals += [float(v) for v in ln.split()]
    arr = np.array(vals).reshape(-1, 14)
    names = [
        "CRIM",
        "ZN",
        "INDUS",
        "CHAS",
        "NOX",
        "RM",
        "AGE",
        "DIS",
        "RAD",
        "TAX",
        "PTRATIO",
        "B",
        "LSTAT",
        "MEDV",
    ]
    X = np.column_stack([arr[:, names.index(c)] for c in COLS])
    X = (X - X.mean(0)) / X.std(0)
    return X, arr[:, names.index("MEDV")]


t0 = time.time()
X, y = load_boston()
N, D = X.shape
_ols = np.linalg.lstsq(X, y, rcond=None)[0]
SIGMA = float(np.std(y - X @ _ols))  # working scale, as in the authors' script
print("sigma_working =", SIGMA)
print(f"Boston: N={N} D={D}")
# integrity check against the authors' csv
import csv

with open(f"{ROOT}/refcode/boston.csv") as f:
    rd = list(csv.DictReader(f))
Xa = np.array([[float(r[c]) for c in COLS] for r in rd])
Xa = (Xa - Xa.mean(0)) / Xa.std(0)
ya = np.array([float(r["MEDV"]) for r in rd])
data_match = dict(
    max_abs_X_diff=float(np.abs(Xa - X).max()),
    max_abs_y_diff=float(np.abs(ya - y).max()),
)
print("data check vs authors' boston.csv:", data_match)

res = {
    "claim": "Table 3 (Boston housing)",
    "sigma_working": SIGMA,
    "seed": SEED,
    "N": N,
    "D": D,
    "columns": COLS,
    "data_check": data_match,
    "paper_table3": {f"{k[0]}_B{k[1]}": v for k, v in PAPER.items()},
}
rows = []
BATCHES = [16, int(0.1 * N)]
ZS = np.random.default_rng(999).normal(0.0, SIGMA * 5.0, size=50)  # authors' n_mc=50
for loss in ["log", "beta", "beta_mc"]:
    mdl = LinReg(X, y, sigma=SIGMA, loss=loss, beta=1.5, z_samples=ZS)
    that = mdl.fit()
    g = mdl.grad_n(that)
    Jn = mdl.hess_n(that)
    Ihat = g.T @ g / N
    Jhat = Jn.mean(0)
    Sstar = sandwich(Jhat, Ihat, N)
    H = Jhat  # no prior -> Hhat = Jbar
    nm = NoiseModel(g, Jn, that, None, N)
    A_cov = X.T @ X / N
    sig2 = float(np.var(y - X @ that))
    print(
        f"\n--- {loss} loss --- ||grad L(that)||={np.linalg.norm(g.mean(0)):.2e} "
        f"cond(Jhat)={np.linalg.cond(Jhat):.2f}  minEig(Jhat)={np.linalg.eigvalsh(Jhat)[0]:.3e} ||S*||_F={np.linalg.norm(Sstar):.3e}"
    )
    if loss == "log":
        Spost = SIGMA**2 * np.linalg.inv(X.T @ X)  # flat-prior Gaussian posterior
        e_post = relf(Spost, Sstar)
        print(
            f"    exact Gaussian posterior covariance error = {e_post:.4f} "
            f"(paper: {PAPER[('log',16)]['Posterior']})"
        )
        res["posterior_cov_error"] = e_post
        res["posterior_cov_error_paper"] = PAPER[("log", 16)]["Posterior"]
    for B in BATCHES:
        n_iters = int(500 * N / B)
        for method in ["CT", "LR+WS", "DQ+const", "DQ+exact",
                       "LR+WS-cf", "DQ+const-cf", "DQ+exact-cf"]:
            # "-cf" = the closed-form Lambda0 = (S H + H S)(C + H S H)^{-1} that the
            # authors' released code uses, instead of the scipy.optimize.root
            # (Powell hybrid) solve of Eq. (11) that Section 6 of the paper describes.
            base = method.replace("-cf", "")
            tl0 = time.time()
            try:
                Lam, rr, Lam0 = tune_Lambda(
                    base, H, Sstar, Ihat, Jhat, B, nm=nm, A_cov=A_cov,
                    sig2=sig2, sigma=SIGMA
                )
                if method.endswith("-cf"):
                    Lam = Lam0
                tune_t = time.time() - tl0
            except Exception as e:
                rows.append(dict(loss=loss, B=B, method=method, error=str(e)))
                print(f"  {method}: tuning failed: {e}")
                continue
            rng = np.random.default_rng(
                SEED + 1000 * B + zlib.crc32(method.encode()) % 100000
            )
            ms0 = time.time()
            covs, div = mdl.sgd_paths(that, Lam, B, n_iters, R=30, rng=rng)
            mcmc_t = time.time() - ms0
            errs = np.array([relf(c, Sstar) for c in covs])
            errs[div] = np.inf
            fin = errs[np.isfinite(errs)]
            mean = float(np.mean(errs)) if len(fin) == len(errs) else float("inf")
            med = float(np.median(errs))
            lo = float(np.percentile(fin, 2.5)) if len(fin) else float("nan")
            hi = float(np.percentile(fin, 97.5)) if len(fin) else float("nan")
            rows.append(
                dict(
                    loss=loss,
                    B=B,
                    method=method,
                    cov_err_mean=mean,
                    cov_err_median=med,
                    ci95=[lo, hi],
                    n_diverged=int(div.sum()),
                    tuning_residual=float(rr),
                    lam_trace=float(np.trace(Lam)),
                    lam_max_eig=float(np.max(np.abs(np.linalg.eigvals(Lam)))),
                    tune_time_s=tune_t,
                    mcmc_time_s=mcmc_t,
                    tune_over_mcmc=tune_t / mcmc_t,
                    n_iters=n_iters,
                )
            )
            print(
                f"  B={B:<3} {method:<9} cov err mean={mean:<12.4g} "
                f"median={med:<12.4g} 95%CI=[{lo:.3g},{hi:.3g}] diverged={div.sum()}/30 "
                f"tune/mcmc={tune_t/mcmc_t:.2e}"
            )
res["results"] = rows
res["runtime_s"] = time.time() - t0
with open(f"{OUT}/claim5_boston.json", "w") as f:
    json.dump(res, f, indent=1)
print("\nwrote outputs/claim5_boston.json  runtime %.1f s" % res["runtime_s"])