1import json
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 988
  6
  7def fit_joint(w, y, ridge=1e-6):
  8    mw, my = w.mean(0), y.mean(0)
  9    wc, yc = w - mw, y - my
 10    sww = wc.T @ wc / (len(w) - 1)
 11    syw = yc.T @ wc / (len(w) - 1)
 12    syy = yc.T @ yc / (len(w) - 1)
 13    gain = np.linalg.solve(sww + ridge * np.eye(w.shape[1]), syw.T).T
 14    cov = syy - gain @ syw.T
 15    return mw, my, gain, (cov + cov.T) / 2
 16
 17def condition(model, w):
 18    mw, my, gain, cov = model
 19    return my + (w - mw) @ gain.T, cov
 20
 21def nll(errors, cov, jitter=1e-8):
 22    c = (cov + cov.T) / 2 + jitter * np.eye(cov.shape[0])
 23    sign, ld = np.linalg.slogdet(c)
 24    if sign <= 0:
 25        return float("inf")
 26    q = np.mean(np.sum(errors * np.linalg.solve(c, errors.T).T, axis=1))
 27    return float(.5 * (q + ld + errors.shape[1] * np.log(2 * np.pi)))
 28
 29def coverage(errors, cov, alpha=.95):
 30    c = (cov + cov.T) / 2 + 1e-8 * np.eye(len(cov))
 31    q = np.sum(errors * np.linalg.solve(c, errors.T).T, axis=1)
 32    # chi-square(10) 95% quantile, avoiding a scipy dependency in the MVP
 33    threshold = 18.307038053275146 if len(cov) == 10 else np.quantile(q, alpha)
 34    return float(np.mean(q <= threshold)), float(np.mean(q))
 35
 36def ar(n, total, rho, sigma, rng):
 37    x = np.empty((n, total))
 38    x[:, 0] = rng.normal(size=n) * sigma / np.sqrt(1 - rho*rho)
 39    for t in range(1, total):
 40        x[:, t] = rho * x[:, t-1] + sigma * rng.normal(size=n)
 41    return x
 42
 43def main():
 44    rng = np.random.default_rng(SEED)
 45    out = {}
 46
 47    # Core algebra: Schur complement of a positive definite joint covariance.
 48    a = rng.normal(size=(10, 10))
 49    joint = a @ a.T + .2 * np.eye(10)
 50    sww, swy = joint[:4, :4], joint[:4, 4:]
 51    syw, syy = joint[4:, :4], joint[4:, 4:]
 52    schur = syy - syw @ np.linalg.solve(sww, swy)
 53    out["math_check"] = {
 54        "min_joint_eigenvalue": float(np.linalg.eigvalsh(joint).min()),
 55        "min_schur_eigenvalue": float(np.linalg.eigvalsh(schur).min()),
 56        "schur_trace": float(np.trace(schur)),
 57        "psd_verified": bool(np.linalg.eigvalsh(schur).min() > -1e-10),
 58    }
 59
 60    # Same process and windows for both predictors.
 61    rho, sigma, H = .85, 1.0, 10
 62    ntrain, ntest = 16000, 6000
 63    train = ar(ntrain, 20 + H, rho, sigma, rng)
 64    test = ar(ntest, 20 + H, rho, sigma, rng)
 65    wtr, ytr = train[:, 19:20], train[:, 20:]
 66    wte, yte = test[:, 19:20], test[:, 20:]
 67
 68    # Standard rollout: estimate one-step AR coefficient, recursively apply it.
 69    x0 = wtr[:, 0]
 70    rho_hat = float((x0 @ ytr[:, 0]) / (x0 @ x0))
 71    powers = rho_hat ** np.arange(1, H + 1)
 72    base_mean = wte * powers[None, :]
 73    # Innovation variance accumulated by recursive rollout; diagonal baseline.
 74    marginal_var = sigma*sigma * np.array([(1-rho_hat**(2*k))/(1-rho_hat**2) for k in range(1, H+1)])
 75    base_cov = np.diag(marginal_var)
 76
 77    # Idea: learn all future coordinates jointly and condition their covariance on w.
 78    joint_model = fit_joint(wtr, ytr, ridge=1e-6)
 79    idea_mean, idea_cov = condition(joint_model, wte)
 80    base_err, idea_err = yte - base_mean, yte - idea_mean
 81    out["rollout"] = {
 82        "rho_hat": rho_hat,
 83        "baseline_recursive": {
 84            "mse": float(np.mean(base_err**2)),
 85            "nll": nll(base_err, base_cov),
 86            "coverage95": coverage(base_err, base_cov)[0],
 87        },
 88        "covariance_conditioned": {
 89            "mse": float(np.mean(idea_err**2)),
 90            "nll": nll(idea_err, idea_cov),
 91            "coverage95": coverage(idea_err, idea_cov)[0],
 92            "conditional_trace": float(np.trace(idea_cov)),
 93            "min_cov_eigenvalue": float(np.linalg.eigvalsh(idea_cov).min()),
 94        },
 95    }
 96
 97    # Nested-prefix check at a fixed prediction origin. AR theory predicts p=1 helps;
 98    # older history is redundant, so this explicitly tests the limitation.
 99    prefix_rows = []
100    target = train[:, 20:30]
101    target_test = test[:, 20:30]
102    for p in [0, 1, 2, 4, 8, 12]:
103        if p == 0:
104            mean = target.mean(0)
105            cov = np.diag(np.var(target, axis=0, ddof=1))
106            pred = np.broadcast_to(mean, target_test.shape)
107        else:
108            model = fit_joint(train[:, 20-p:20], target, ridge=1e-6)
109            pred, cov = condition(model, test[:, 20-p:20])
110        prefix_rows.append({
111            "prefix": p,
112            "conditional_trace": float(np.trace(cov)),
113            "test_mse": float(np.mean((target_test-pred)**2)),
114            "min_eigenvalue": float(np.linalg.eigvalsh(cov).min()),
115        })
116    out["prefix_sweep"] = prefix_rows
117    Path("mvp_results.json").write_text(json.dumps(out, indent=2))
118    print(json.dumps(out, indent=2))
119
120if __name__ == "__main__":
121    main()