Gibbs vs NUTS Profiling: Cross-Sectional Models¶

import time

import arviz as az
import numpy as np
import pandas as pd
from libpysal.graph import Graph

from neighbayes.models import SAR, SDEM, SDM, SEM

# Try importing JAX to check availability
try:
    import jax  # noqa: F401

    HAS_JAX = True
except ImportError:
    HAS_JAX = False

print(f"JAX available: {HAS_JAX}")

SIDE = 7  # 7x7 grid → n=49
SEED = 42
DRAWS = 2000
TUNE = 1000
CHAINS = 4


def make_rook_W(side: int) -> np.ndarray:
    """Row-standardized rook-contiguity weights on a side x side grid."""
    n = side * side
    W = np.zeros((n, n))
    for r in range(side):
        for c in range(side):
            i = r * side + c
            if r > 0:
                W[i, (r - 1) * side + c] = 1
            if r < side - 1:
                W[i, (r + 1) * side + c] = 1
            if c > 0:
                W[i, r * side + (c - 1)] = 1
            if c < side - 1:
                W[i, r * side + (c + 1)] = 1
    row_sums = W.sum(axis=1, keepdims=True)
    return W / np.where(row_sums == 0, 1, row_sums)


def W_to_graph(W_dense: np.ndarray) -> Graph:
    """Convert a dense weight matrix to a libpysal Graph."""
    n = W_dense.shape[0]
    focal, neighbor, weight = [], [], []
    for i in range(n):
        for j in range(n):
            if W_dense[i, j] != 0:
                focal.append(i)
                neighbor.append(j)
                weight.append(W_dense[i, j])
    return Graph.from_arrays(
        np.array(focal),
        np.array(neighbor),
        np.array(weight, dtype=float),
    ).transform("r")


W_dense = make_rook_W(SIDE)
W = W_to_graph(W_dense)
n = SIDE * SIDE
print(f"n = {n}, W shape = {W_dense.shape}")

from neighbayes.dgp.cross_sectional import (
    simulate_sar,
    simulate_sdem,
    simulate_sdm,
    simulate_sem,
)

rng = np.random.default_rng(SEED)

# SAR data
sar_out = simulate_sar(W=W_dense, rho=0.6, sigma=1.0, rng=rng)
y_sar, X_sar = sar_out["y"], sar_out["X"]

# SEM data
sem_out = simulate_sem(W=W_dense, lam=0.6, sigma=1.0, rng=rng)
y_sem, X_sem = sem_out["y"], sem_out["X"]

# SDM data
sdm_out = simulate_sdm(W=W_dense, rho=0.4, sigma=1.0, rng=rng)
y_sdm, X_sdm = sdm_out["y"], sdm_out["X"]

# SDEM data
sdem_out = simulate_sdem(W=W_dense, lam=0.4, sigma=1.0, rng=rng)
y_sdem, X_sdem = sdem_out["y"], sdem_out["X"]

print(f"SAR:  y={y_sar.shape}, X={X_sar.shape}")
print(f"SEM:  y={y_sem.shape}, X={X_sem.shape}")
print(f"SDM:  y={y_sdm.shape}, X={X_sdm.shape}")
print(f"SDEM: y={y_sdem.shape}, X={X_sdem.shape}")

def profile_model(model, sampler_config: dict, label: str) -> dict:
    """Fit a model with a given sampler config and return timing + diagnostics."""
    t0 = time.perf_counter()
    idata = model.fit(**sampler_config)
    elapsed = time.perf_counter() - t0

    # Extract spatial parameter name (rho for SAR/SDM, lam for SEM/SDEM)
    spatial_vars = [p for p in idata.posterior.data_vars if p in ("rho", "lam")]
    spatial_param = spatial_vars[0] if spatial_vars else None

    result = {
        "label": label,
        "elapsed_s": round(elapsed, 2),
    }

    # Posterior means (scalar parameters only)
    # beta is (chain, draw, k) — report intercept (first coefficient)
    result["intercept_mean"] = float(
        idata.posterior["beta"].mean(dim=["chain", "draw"]).values[0]
    )
    result["sigma_mean"] = float(idata.posterior["sigma"].mean())
    if spatial_param:
        result[f"{spatial_param}_mean"] = float(idata.posterior[spatial_param].mean())

    # Diagnostics — include all parameters for thorough R-hat / ESS checks
    diag_vars = ["beta", "sigma"] + ([spatial_param] if spatial_param else [])
    summary = az.summary(idata, var_names=diag_vars)
    result["rhat_max"] = round(float(summary["r_hat"].max()), 3)
    result["ess_bulk_min"] = int(summary["ess_bulk"].min())

    return result

COMMON = dict(
    draws=DRAWS, tune=TUNE, chains=CHAINS, random_seed=SEED, progressbar=False
)

configs = {
    "NUTS": dict(sampler="nuts", target_accept=0.9, **COMMON),
    "Gibbs-NumPy": dict(sampler="gibbs", gibbs_backend="numpy", n_jobs=-1, **COMMON),
}

if HAS_JAX:
    configs["Gibbs-JAX"] = dict(sampler="gibbs", gibbs_backend="jax", **COMMON)

print(f"Configs: {list(configs.keys())}")

results = []

models_and_data = [
    ("SAR", SAR, y_sar, X_sar),
    ("SEM", SEM, y_sem, X_sem),
    ("SDM", SDM, y_sdm, X_sdm),
    ("SDEM", SDEM, y_sdem, X_sdem),
]

for model_name, ModelClass, y, X in models_and_data:
    print(f"\n{'=' * 60}")
    print(f"Model: {model_name}")
    print(f"{'=' * 60}")

    for config_name, config in configs.items():
        label = f"{model_name}/{config_name}"
        print(f"  Running {label}...", end=" ", flush=True)

        model = ModelClass(y=y, X=X, W=W)
        try:
            result = profile_model(model, config, label)
            print(
                f"{result['elapsed_s']:.1f}s  ρ̂={result.get('rho_mean', result.get('lam_mean', 'N/A')):.4f}  r̂={result['rhat_max']:.3f}  ESS={result['ess_bulk_min']}"
            )
        except Exception as e:
            result = {"label": label, "elapsed_s": None, "error": str(e)}
            print(f"FAILED: {e}")

        results.append(result)

print(f"\nDone! {len(results)} configurations profiled.")

# Build a clean results DataFrame
rows = []
for r in results:
    if "error" in r:
        rows.append(
            {"Model/Sampler": r["label"], "Time (s)": None, "Error": r["error"]}
        )
        continue
    model_name, config_name = r["label"].split("/")
    spatial_key = (
        "rho_mean" if "rho_mean" in r else ("lam_mean" if "lam_mean" in r else None)
    )
    rows.append(
        {
            "Model": model_name,
            "Sampler": config_name,
            "Time (s)": r["elapsed_s"],
            "Intercept": round(r["intercept_mean"], 4),
            "σ̂": round(r["sigma_mean"], 4),
            "ρ̂/λ̂": round(r.get(spatial_key, float("nan")), 4) if spatial_key else None,
            "max R̂": r["rhat_max"],
            "min ESS": r["ess_bulk_min"],
        }
    )

df = pd.DataFrame(rows)
# Drop duplicates in case cells were re-run
df = df.drop_duplicates(subset=["Model", "Sampler"], keep="last")
df.style.format(
    {"Time (s)": "{:.1f}", "Intercept": "{:.4f}", "σ̂": "{:.4f}", "ρ̂/λ̂": "{:.4f}"}
)

# Pivot to compare times (use drop_duplicates to guard against re-runs)
if len(df) > 0 and "Time (s)" in df.columns:
    pivot = df.drop_duplicates(subset=["Model", "Sampler"], keep="last").pivot(
        index="Model", columns="Sampler", values="Time (s)"
    )
    if "NUTS" in pivot.columns:
        for col in pivot.columns:
            if col != "NUTS":
                pivot[f"{col}/NUTS"] = pivot[col] / pivot["NUTS"]
        print("Time relative to NUTS (< 1 means faster):")
        speedup_cols = [c for c in pivot.columns if "/NUTS" in c]
        display(
            pivot[
                ["NUTS"]
                + [c for c in pivot.columns if c != "NUTS" and "/NUTS" not in c]
                + speedup_cols
            ].round(2)
        )
    else:
        display(pivot.round(2))
else:
    print("No results to display.")

import matplotlib.pyplot as plt

# Re-fit all models and store idatas for comparison plots
idatas = {}
for model_name, ModelClass, y, X in models_and_data:
    idatas[model_name] = {}
    for config_name, config in configs.items():
        label = f"{model_name}/{config_name}"
        model = ModelClass(y=y, X=X, W=W)
        try:
            idata = model.fit(**config)
            idatas[model_name][config_name] = idata
        except Exception:
            pass

# Plot comparison for each model
spatial_param_map = {"SAR": "rho", "SEM": "lam", "SDM": "rho", "SDEM": "lam"}

fig, axes = plt.subplots(2, 4, figsize=(16, 8))

for col, model_name in enumerate(["SAR", "SEM", "SDM", "SDEM"]):
    param = spatial_param_map[model_name]
    for row, param_name in enumerate([param, "sigma"]):
        ax = axes[row, col]
        for config_name in configs:
            if config_name in idatas.get(model_name, {}):
                idata = idatas[model_name][config_name]
                samples = idata.posterior[param_name].values.flatten()
                ax.hist(samples, bins=50, alpha=0.5, density=True, label=config_name)
        ax.set_title(f"{model_name}: {param_name}")
        if row == 0:
            ax.legend(fontsize=8)

plt.tight_layout()
plt.suptitle("Posterior Comparison: Gibbs vs NUTS", y=1.02, fontsize=14)
plt.show()

SCALE_DRAWS = 2000
SCALE_TUNE = 1000
SCALE_CHAINS = 2
SIDES = [10, 20, 50, 75]  # n = 16, 36, 64, 100

scale_results = []

scale_configs = {
    "NUTS": dict(
        sampler="nuts",
        target_accept=0.9,
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
    ),
    "Gibbs-NumPy": dict(
        sampler="gibbs",
        gibbs_backend="numpy",
        n_jobs=-1,
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
    ),
}
if HAS_JAX:
    scale_configs["Gibbs-JAX"] = dict(
        sampler="gibbs",
        gibbs_backend="jax",
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
    )
    scale_configs["NUTS-JAX"] = dict(
        sampler="nuts",
        gibbs_backend="jax",
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
        nuts_sampler="blackjax",
    )

for side in SIDES:
    n = side * side
    W_dense_s = make_rook_W(side)
    W_s = W_to_graph(W_dense_s)
    rng_s = np.random.default_rng(SEED)
    out_s = simulate_sar(W=W_dense_s, rho=0.5, sigma=1.0, rng=rng_s)
    y_s, X_s = out_s["y"], out_s["X"]

    print(f"\nn = {n} (side={side})")
    for config_name, config in scale_configs.items():
        label = f"n={n}/{config_name}"
        model = SAR(y=y_s, X=X_s, W=W_s)
        try:
            result = profile_model(model, config, label)
            print(f"  {config_name}: {result['elapsed_s']:.1f}s")
            result["n"] = n
            result["sampler"] = config_name
            scale_results.append(result)
        except Exception as e:
            print(f"  {config_name}: FAILED ({e})")

print("\nScaling benchmark complete.")

# Plot scaling results
if scale_results:
    scale_df = pd.DataFrame(scale_results)
    fig, ax = plt.subplots(figsize=(8, 5))
    for sampler in scale_df["sampler"].unique():
        sub = scale_df[scale_df["sampler"] == sampler].sort_values("n")
        ax.plot(sub["n"], sub["elapsed_s"], marker="o", label=sampler)
    ax.set_xlabel("Number of spatial units (n)")
    ax.set_ylabel("Wall-clock time (seconds)")
    ax.set_title(
        f"SAR Scaling: Gibbs vs NUTS (draws={SCALE_DRAWS}, tune={SCALE_TUNE}, chains={SCALE_CHAINS})"
    )
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()
else:
    print("No scaling results to plot.")

SCALE_DRAWS = 2000
SCALE_TUNE = 1000
SCALE_CHAINS = 2
SIDES = [10, 20, 50, 75]  # n = 16, 36, 64, 100

scale_results = []

scale_configs = {
    "NUTS": dict(
        sampler="nuts",
        target_accept=0.9,
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
    ),
    "Gibbs-NumPy": dict(
        sampler="gibbs",
        gibbs_backend="numpy",
        n_jobs=-1,
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
    ),
}
if HAS_JAX:
    scale_configs["Gibbs-JAX"] = dict(
        sampler="gibbs",
        gibbs_backend="jax",
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
    )
    scale_configs["NUTS-JAX"] = dict(
        sampler="nuts",
        gibbs_backend="jax",
        draws=SCALE_DRAWS,
        tune=SCALE_TUNE,
        chains=SCALE_CHAINS,
        random_seed=SEED,
        progressbar=False,
        nuts_sampler="blackjax",
    )

for side in SIDES:
    n = side * side
    W_dense_s = make_rook_W(side)
    W_s = W_to_graph(W_dense_s)
    rng_s = np.random.default_rng(SEED)
    out_s = simulate_sem(W=W_dense_s, lam=0.5, sigma=1.0, rng=rng_s)
    y_s, X_s = out_s["y"], out_s["X"]

    print(f"\nn = {n} (side={side})")
    for config_name, config in scale_configs.items():
        label = f"n={n}/{config_name}"
        model = SEM(y=y_s, X=X_s, W=W_s)
        try:
            result = profile_model(model, config, label)
            print(f"  {config_name}: {result['elapsed_s']:.1f}s")
            result["n"] = n
            result["sampler"] = config_name
            scale_results.append(result)
        except Exception as e:
            print(f"  {config_name}: FAILED ({e})")

print("\nScaling benchmark complete.")

# Plot scaling results
if scale_results:
    scale_df = pd.DataFrame(scale_results)
    fig, ax = plt.subplots(figsize=(8, 5))
    for sampler in scale_df["sampler"].unique():
        sub = scale_df[scale_df["sampler"] == sampler].sort_values("n")
        ax.plot(sub["n"], sub["elapsed_s"], marker="o", label=sampler)
    ax.set_xlabel("Number of spatial units (n)")
    ax.set_ylabel("Wall-clock time (seconds)")
    ax.set_title(
        f"SAR Scaling: Gibbs vs NUTS (draws={SCALE_DRAWS}, tune={SCALE_TUNE}, chains={SCALE_CHAINS})"
    )
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.show()
else:
    print("No scaling results to plot.")