import os, sys, time

_conda_lib = r"C:\Users\user\anaconda3\envs\pymc-env\Library\bin"
if _conda_lib not in os.environ.get("PATH", ""):
    os.environ["PATH"] = _conda_lib + os.pathsep + os.environ.get("PATH", "")

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.special import ndtr
import importlib, mvprobit_gibbs
importlib.reload(mvprobit_gibbs)
from mvprobit_gibbs import mvprobit_gibbs as gibbs_sampler, make_prior, make_init

print("Imports OK")

Imports OK

# Long format: one row per child-year (537 x 4 = 2148 rows)
# Ohio Six Cities wheeze data (geepack::ohio) — shipped alongside this notebook
long = pd.read_csv("ohio.csv")
print("Long format:", long.shape)
print(long.head(8))
print()

# Reshape to wide: one row per child (537 rows)
wide = long.pivot(index="id", columns="age", values="resp")
wide.columns = [f"wheeze_{7 + i}" for i in range(4)]   # age 7,8,9,10
wide = wide.reset_index()

# Attach maternal smoking (time-invariant: take first obs per child)
smoke = long.groupby("id")["smoke"].first().reset_index()
wide  = wide.merge(smoke, on="id")

OUTCOMES   = ["wheeze_7", "wheeze_8", "wheeze_9", "wheeze_10"]
Y = wide[OUTCOMES].values.astype(int)
X = np.column_stack([np.ones(len(wide)), wide["smoke"].values])

N, J = Y.shape
K    = X.shape[1]

print(f"Wide format: N={N}, J={J}, K={K}")
print(f"Maternal smoking: {wide['smoke'].mean():.1%} of mothers smoke")
print()
summary = pd.DataFrame({
    "n_wheeze": Y.sum(axis=0),
    "rate"    : Y.mean(axis=0).round(3),
    "smoke=1" : [Y[wide["smoke"]==1, j].mean().round(3) for j in range(J)],
    "smoke=0" : [Y[wide["smoke"]==0, j].mean().round(3) for j in range(J)],
}, index=OUTCOMES)
print(summary)

Long format: (2148, 5)
   rownames  resp  id  age  smoke
0         1     0   0   -2      0
1         2     0   0   -1      0
2         3     0   0    0      0
3         4     0   0    1      0
4         5     0   1   -2      0
5         6     0   1   -1      0
6         7     0   1    0      0
7         8     0   1    1      0

Wide format: N=537, J=4, K=2
Maternal smoking: 34.8% of mothers smoke

           n_wheeze   rate  smoke=1  smoke=0
wheeze_7         87  0.162    0.166    0.160
wheeze_8         91  0.169    0.209    0.149
wheeze_9         85  0.158    0.187    0.143
wheeze_10        63  0.117    0.139    0.106

ages    = [7, 8, 9, 10]
rate_s1 = [Y[wide["smoke"]==1, j].mean() for j in range(J)]
rate_s0 = [Y[wide["smoke"]==0, j].mean() for j in range(J)]

fig, ax = plt.subplots(figsize=(7, 4))
ax.plot(ages, rate_s1, "o-", color="tomato",    label="Maternal smoke = Yes")
ax.plot(ages, rate_s0, "o-", color="steelblue", label="Maternal smoke = No")
ax.set_xlabel("Child age")
ax.set_ylabel("Wheeze rate")
ax.set_xticks(ages)
ax.set_title("Wheeze prevalence by age and maternal smoking\n(Ohio Six Cities Study, N=537)")
ax.legend()
plt.tight_layout()
plt.show()

# A_J: age design matrix (J x 2), age centred at 9
A_J = np.array([[1., -2.],   # age 7
                [1., -1.],   # age 8
                [1.,  0.],   # age 9
                [1.,  1.]])  # age 10

prior = make_prior(J, K, v=J+1, V=np.eye(J), beta_var=100.0)
prior['beta_var'] = 100.0    # passed through to constrained update
b0    = make_init(Y, X, J, K)

print(f"Constrained model: B = A_J @ Gamma")
print(f"  A_J shape : {A_J.shape}  (J x n_a)")
print(f"  X shape   : {X.shape}   (N x K)")
print(f"  Gamma     : ({A_J.shape[1]} x {K}) = {A_J.shape[1]*K} free parameters")
print(f"  R         : {J*(J+1)//2} parameters  ({J*(J-1)//2} off-diagonal)")

Constrained model: B = A_J @ Gamma
  A_J shape : (4, 2)  (J x n_a)
  X shape   : (537, 2)   (N x K)
  Gamma     : (2 x 2) = 4 free parameters
  R         : 10 parameters  (6 off-diagonal)

np.random.seed(603)          # reproducible Gibbs draws
NDRAWS = 30000
BURN   = 3000

print(f"Running Gibbs sampler (C&G constrained): {NDRAWS:,} draws, {BURN:,} burn-in ...")
t0 = time.time()

results = gibbs_sampler(Y, X, prior, ndraws=NDRAWS, b0=b0, burn=BURN,
                        verbose=500, A_age=A_J)

elapsed = time.time() - t0
print(f"Done in {elapsed:.1f} seconds.")

Running Gibbs sampler (C&G constrained): 30,000 draws, 3,000 burn-in ...

  Iteration   500 / 30000

  Iteration  1000 / 30000

  Iteration  1500 / 30000

  Iteration  2000 / 30000

  Iteration  2500 / 30000

  Iteration  3000 / 30000

  Iteration  3500 / 30000

  Iteration  4000 / 30000

  Iteration  4500 / 30000

B_draws     = results["B_draws"]       # (keep, K*J)  — normalized B
Sigma_draws = results["Sigma_draws"]   # (keep, J*J)
Gamma_draws = results["Gamma_draws"]   # (keep, n_a*K) — normalized Gamma
keep        = B_draws.shape[0]

# B_mean: J x K (equation x regressor)
B_mean = B_draws.mean(axis=0).reshape(J, K, order="F")
B_std  = B_draws.std(axis=0).reshape(J, K, order="F")

# Gamma_mean: n_a x K
n_a = A_J.shape[1]
Gamma_mean = Gamma_draws.mean(axis=0).reshape(K, n_a, order="F").T   # (n_a, K)
Gamma_std  = Gamma_draws.std(axis=0).reshape(K, n_a, order="F").T

# R_mean: J x J
R_lower = Sigma_draws.reshape(keep, J, J, order="F")
R_draws = R_lower + np.triu(R_lower.transpose(0, 2, 1), k=1)
R_mean  = R_draws.mean(axis=0)
R_std   = R_draws.std(axis=0)

print(f"Stored draws  : {keep:,}")
print(f"B_draws shape : {B_draws.shape}  -> B_mean: {B_mean.shape}")
print(f"Gamma_draws   : {Gamma_draws.shape}  -> Gamma_mean: {Gamma_mean.shape}")
print(f"R_mean shape  : {R_mean.shape}")

Stored draws  : 27,000
B_draws shape : (27000, 8)  -> B_mean: (4, 2)
Gamma_draws   : (27000, 4)  -> Gamma_mean: (2, 2)
R_mean shape  : (4, 4)

from statsmodels.tsa.stattools import acf as compute_acf

N_MONITOR = 3

def ess(chain, max_lag=500):
    n     = len(chain)
    acf_v = compute_acf(chain, nlags=min(max_lag, n // 2), fft=True)
    ci    = 1.96 / np.sqrt(n)
    rho   = 0.0
    for k in range(1, len(acf_v)):
        if abs(acf_v[k]) < ci:
            break
        rho += acf_v[k]
    return n / max(1.0, 1 + 2 * rho)

print("Computing ESS ...", end=" ", flush=True)
b_ess    = np.array([ess(B_draws[:, p]) for p in range(K * J)])
rows_tri, cols_tri = np.tril_indices(J, k=-1)
corr_ess = np.array([ess(Sigma_draws[:, r + c * J])
                     for r, c in zip(rows_tri, cols_tri)])
print("done.")

REGRESSOR_NAMES = ["intercept", "smoke"]
def b_label(p):
    j, k = p % J, p // J
    return f"B[{OUTCOMES[j]},{REGRESSOR_NAMES[k]}]"

worst_b    = np.argsort(b_ess)[:N_MONITOR]
worst_corr = np.argsort(corr_ess)[:N_MONITOR]
corr_pairs = [(int(rows_tri[i]), int(cols_tri[i])) for i in worst_corr]

monitor_chains = np.column_stack(
    [B_draws[:, p] for p in worst_b] +
    [Sigma_draws[:, r + c * J] for r, c in corr_pairs]
)
monitor_labels = (
    [b_label(p) for p in worst_b] +
    [f"R[age{7+r},age{7+c}]" for r, c in corr_pairs]
)
n_mon = monitor_chains.shape[1]
iters = np.arange(1, keep + 1)

print(f"\nWorst-mixing B parameters:")
for p in worst_b:
    print(f"  {b_label(p):<30s}  ESS = {b_ess[p]:6.0f}")
print(f"\nWorst-mixing correlation pairs:")
for idx, (r, c) in zip(worst_corr, corr_pairs):
    print(f"  R[age{7+r}, age{7+c}]  ESS = {corr_ess[idx]:6.0f}")

ess_vals = [ess(monitor_chains[:, i]) for i in range(n_mon)]
ess_df   = pd.DataFrame({
    "draws"   : keep,
    "ESS"     : [round(e) for e in ess_vals],
    "ESS/draw": [f"{e/keep:.3f}" for e in ess_vals],
}, index=monitor_labels)
print()
print(ess_df.to_string())

Computing ESS ...

done.

Worst-mixing B parameters:
  B[wheeze_10,smoke]              ESS =   4935
  B[wheeze_9,intercept]           ESS =   5253
  B[wheeze_8,intercept]           ESS =   5258

Worst-mixing correlation pairs:
  R[age10, age7]  ESS =   1569
  R[age10, age9]  ESS =   1743
  R[age8, age7]  ESS =   1805

                       draws   ESS ESS/draw
B[wheeze_10,smoke]     27000  4935    0.183
B[wheeze_9,intercept]  27000  5253    0.195
B[wheeze_8,intercept]  27000  5258    0.195
R[age10,age7]          27000  1569    0.058
R[age10,age9]          27000  1743    0.065
R[age8,age7]           27000  1805    0.067

fig, axes = plt.subplots(n_mon, 1, figsize=(12, 2.2 * n_mon), sharex=True)
for ax, chain, lbl in zip(axes, monitor_chains.T, monitor_labels):
    ax.plot(iters, chain, lw=0.3, color="steelblue", alpha=0.8)
    ax.axhline(chain.mean(), color="red", lw=1, linestyle="--")
    ax.set_ylabel(lbl, fontsize=8, rotation=0, ha="right", labelpad=90)
axes[-1].set_xlabel("Post-burn-in draw")
fig.suptitle("Trace plots — worst-mixing parameters", fontsize=11)
plt.tight_layout()
plt.show()

REGRESSOR_NAMES = ["intercept", "smoke"]
GAMMA_NAMES     = ["beta_0 (intercept at age9)", "beta_1 (age slope)",
                   "beta_2 (smoke at age9)",     "beta_3 (age x smoke)"]

print("=" * 60)
print("Gamma — C&G shared parameters (n_a=2 x K=2 = 4 parameters)")
print("=" * 60)
gamma_df = pd.DataFrame({
    "post_mean": Gamma_mean.ravel(order="C").round(3),
    "post_std" : Gamma_std.ravel(order="C").round(3),
}, index=GAMMA_NAMES)
print(gamma_df.to_string())
print()
print("=" * 60)
print("B = A_J @ Gamma — equation-specific coefficients")
print("=" * 60)
mean_df = pd.DataFrame(B_mean.round(3), index=OUTCOMES, columns=REGRESSOR_NAMES)
std_df  = pd.DataFrame(B_std.round(3),  index=OUTCOMES, columns=REGRESSOR_NAMES)
print("Posterior means:")
print(mean_df.to_string())
print()
print("Posterior std devs:")
print(std_df.to_string())
print()
print("Model-implied wheeze rates:")
for j, out in enumerate(OUTCOMES):
    rate_ns = ndtr(B_mean[j, 0])
    rate_s  = ndtr(B_mean[j, 0] + B_mean[j, 1])
    obs_ns  = Y[wide["smoke"]==0, j].mean()
    obs_s   = Y[wide["smoke"]==1, j].mean()
    print(f"  {out}  smoke=0: {rate_ns:.3f} (obs {obs_ns:.3f})  "
          f"smoke=1: {rate_s:.3f} (obs {obs_s:.3f})")

============================================================
Gamma — C&G shared parameters (n_a=2 x K=2 = 4 parameters)
============================================================
                            post_mean  post_std
beta_0 (intercept at age9)     -1.114     0.063
beta_1 (age slope)             -0.079     0.032
beta_2 (smoke at age9)          0.150     0.101
beta_3 (age x smoke)            0.037     0.052

============================================================
B = A_J @ Gamma — equation-specific coefficients
============================================================
Posterior means:
           intercept  smoke
wheeze_7      -1.007  0.080
wheeze_8      -0.991  0.110
wheeze_9      -1.046  0.140
wheeze_10     -1.253  0.194

Posterior std devs:
           intercept  smoke
wheeze_7       0.078  0.124
wheeze_8       0.076  0.095
wheeze_9       0.076  0.096
wheeze_10      0.088  0.134

Model-implied wheeze rates:
  wheeze_7  smoke=0: 0.157 (obs 0.160)  smoke=1: 0.177 (obs 0.166)
  wheeze_8  smoke=0: 0.161 (obs 0.149)  smoke=1: 0.189 (obs 0.209)
  wheeze_9  smoke=0: 0.148 (obs 0.143)  smoke=1: 0.183 (obs 0.187)
  wheeze_10  smoke=0: 0.105 (obs 0.106)  smoke=1: 0.145 (obs 0.139)

fig, ax = plt.subplots(figsize=(6, 5))
sns.heatmap(
    R_mean, annot=True, fmt=".3f", center=0,
    cmap="RdBu_r", vmin=-0.5, vmax=1.0,
    xticklabels=["age 7","age 8","age 9","age 10"],
    yticklabels=["age 7","age 8","age 9","age 10"],
    ax=ax, linewidths=0.5, annot_kws={"size": 11},
)
ax.set_title(
    f"Posterior mean correlation matrix R\n"
    f"Gibbs  |  {keep:,} draws  |  N={N}, J={J}, K={K}"
)
plt.tight_layout()
plt.show()

# Full table with posterior SDs and 95% CIs
rows_r, cols_r = np.tril_indices(J, k=-1)
corr_tbl = pd.DataFrame({
    "pair"     : [f"R[age{7+r}, age{7+c}]" for r, c in zip(rows_r, cols_r)],
    "post_mean": R_mean[rows_r, cols_r].round(3),
    "post_std" : R_std[rows_r, cols_r].round(3),
    "CI_2.5"   : np.percentile(R_draws[:, rows_r, cols_r], 2.5, axis=0).round(3),
    "CI_97.5"  : np.percentile(R_draws[:, rows_r, cols_r], 97.5, axis=0).round(3),
})
print(corr_tbl.to_string(index=False))

          pair  post_mean  post_std  CI_2.5  CI_97.5
 R[age8, age7]      0.585     0.067   0.447    0.707
 R[age9, age7]      0.529     0.071   0.382    0.659
 R[age9, age8]      0.688     0.055   0.571    0.785
R[age10, age7]      0.566     0.071   0.416    0.695
R[age10, age8]      0.565     0.072   0.414    0.694
R[age10, age9]      0.632     0.066   0.492    0.751

# ── C&G Table 4 M1 MLE values (column-major order: (2,1),(3,1),(4,1),(3,2),(4,2),(4,3)) ──
cg_mle = {
    "R[age8, age7]" : 0.584,
    "R[age9, age7]" : 0.521,
    "R[age10, age7]": 0.586,
    "R[age9, age8]" : 0.688,
    "R[age10, age8]": 0.562,
    "R[age10, age9]": 0.631,
}

rows_r, cols_r = np.tril_indices(J, k=-1)

print("Correlation comparison — Gibbs posterior mean vs C&G Table 4 M1 MLE")
print(f"{'Pair':<18}  {'Ours':>7}  {'SD':>6}  {'95% CI':<16}  {'C&G MLE':>8}  {'Diff':>6}")
print("-" * 68)
for r, c in zip(rows_r, cols_r):
    key  = f"R[age{7+r}, age{7+c}]"
    ours = R_mean[r, c]
    std  = R_std[r, c]
    lo   = np.percentile(R_draws[:, r, c], 2.5)
    hi   = np.percentile(R_draws[:, r, c], 97.5)
    cg   = cg_mle.get(key, None)
    diff = f"{ours - cg:+.3f}" if cg is not None else "   —"
    cg_s = f"{cg:.3f}" if cg is not None else "    —"
    print(f"{key:<18}  {ours:7.3f}  {std:6.3f}  [{lo:.3f}, {hi:.3f}]  {cg_s:>8}  {diff:>6}")

print()
print("Gamma comparison — Gibbs posterior mean vs C&G Table 4 M1 MLE")
cg_gamma = {"beta_0": -1.118, "beta_1": -0.079, "beta_2": None, "beta_3": None}
# use the SAME ravel order as the Gamma table in section 6 (order='C'), so the
# age-slope / smoke rows line up with their labels (previously [1,0] and [0,1] were swapped)
gflat = Gamma_mean.ravel(order="C"); gsd = Gamma_std.ravel(order="C")
gamma_flat = [(gflat[0], gsd[0], "beta_0 (intercept at age9)", -1.118),
              (gflat[1], gsd[1], "beta_1 (age slope)",         -0.079),
              (gflat[2], gsd[2], "beta_2 (smoke at age9)",      None),
              (gflat[3], gsd[3], "beta_3 (age x smoke)",        None)]
print(f"{'Parameter':<36}  {'Ours':>7}  {'SD':>6}  {'C&G MLE':>8}  {'Diff':>6}")
print("-" * 70)
for mean, std, name, cg in gamma_flat:
    cg_s = f"{cg:.3f}" if cg is not None else "    —"
    diff = f"{mean - cg:+.3f}" if cg is not None else "   —"
    print(f"{name:<36}  {mean:+7.3f}  {std:6.3f}  {cg_s:>8}  {diff:>6}")

Correlation comparison — Gibbs posterior mean vs C&G Table 4 M1 MLE
Pair                   Ours      SD  95% CI             C&G MLE    Diff
--------------------------------------------------------------------
R[age8, age7]         0.585   0.067  [0.447, 0.707]     0.584  +0.001
R[age9, age7]         0.529   0.071  [0.382, 0.659]     0.521  +0.008
R[age9, age8]         0.688   0.055  [0.571, 0.785]     0.688  -0.000
R[age10, age7]        0.566   0.071  [0.416, 0.695]     0.586  -0.020
R[age10, age8]        0.565   0.072  [0.414, 0.694]     0.562  +0.003
R[age10, age9]        0.632   0.066  [0.492, 0.751]     0.631  +0.001

Gamma comparison — Gibbs posterior mean vs C&G Table 4 M1 MLE
Parameter                                Ours      SD   C&G MLE    Diff
----------------------------------------------------------------------
beta_0 (intercept at age9)             -1.114   0.063    -1.118  +0.004
beta_1 (age slope)                     -0.079   0.032    -0.079  -0.000
beta_2 (smoke at age9)                 +0.150   0.101         —       —
beta_3 (age x smoke)                   +0.037   0.052         —       —

Dimension	Value	Detail
$N$	537	children
$J$	4	wheeze at ages 7, 8, 9, 10
$K$	2	intercept + maternal smoking

Ohio Children's Wheeze Study — Multivariate Probit via Gibbs Sampler¶

Model¶

1. Load and reshape data¶

2. Set up model¶

3. Run Gibbs sampler¶

4. Extract posterior draws¶

5. MCMC Diagnostics¶

6. Regression coefficients¶

7. Posterior correlation matrix R¶

8. Comparison with Chib & Greenberg (1998)¶