Gibbs Sampler for Gaussian Spatial Models¶

%load_ext autoreload
%autoreload 2

import arviz as az
import geopandas as gpd
import libpysal
import numpy as np

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

# Load Columbus dataset
gdf = gpd.read_file(libpysal.examples.get_path("columbus.shp"))

# Response: crime rate
y = gdf["CRIME"].values.astype(float)

# Covariates: income + housing value
X = np.column_stack(
    [
        np.ones(len(y)),
        gdf["INC"].values.astype(float),
        gdf["HOVAL"].values.astype(float),
    ]
)

# Spatial weights: queen contiguity
W = libpysal.graph.Graph.build_contiguity(gdf)
W = W.transform("r")  # row-standardise

print(f"n = {len(y)}, k = {X.shape[1]}")

# --- NUTS  ---
model_nuts = SAR(y=y, X=X, W=W)
idata_nuts = model_nuts.fit(
    sampler="nuts",
    draws=2000,
    tune=1000,
    chains=4,
    target_accept=0.9,
    random_seed=42,
    idata_kwargs={"log_likelihood": True},
)

# --- Gibbs (numpy path) ---
model_gibbs = SAR(y=y, X=X, W=W)
idata_gibbs = model_gibbs.fit(
    sampler="gibbs",
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=42,
    n_jobs=-1,
    idata_kwargs={"log_likelihood": True},
)

print("NUTS posterior means:")
print(f"  rho   = {float(idata_nuts.posterior['rho'].mean()):.4f}")
print(f"  beta  = {idata_nuts.posterior['beta'].mean(dim=['chain', 'draw']).values}")
print(f"  sigma = {float(idata_nuts.posterior['sigma'].mean()):.4f}")
print()
print("Gibbs posterior means:")
print(f"  rho   = {float(idata_gibbs.posterior['rho'].mean()):.4f}")
print(f"  beta  = {idata_gibbs.posterior['beta'].mean(dim=['chain', 'draw']).values}")
print(f"  sigma = {float(idata_gibbs.posterior['sigma'].mean()):.4f}")

az.summary(idata_nuts)["ess_bulk"]

az.summary(idata_gibbs)["ess_bulk"]

model_gibbs.spatial_diagnostics_decision()

# JAX Gibbs with MALA (gradient-guided proposals)
model_jax_mala = SAR(y=y, X=X, W=W)
idata_jax_mala = model_jax_mala.fit(
    sampler="gibbs",
    gibbs_method="jax",
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=42,
    use_mala=True,
    use_slice=False,  # MALA instead of slice
)

# JAX Gibbs with slice sampling (default for JAX path)
model_jax_slice = SAR(y=y, X=X, W=W)
idata_jax_slice = model_jax_slice.fit(
    sampler="gibbs",
    gibbs_method="jax",  # slice sampling is the default
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=42,
)

print("JAX MALA posterior means:")
print(f"  rho   = {float(idata_jax_mala.posterior['rho'].mean()):.4f}")
print(f"  sigma = {float(idata_jax_mala.posterior['sigma'].mean()):.4f}")
print()
print("JAX Slice posterior means:")
print(f"  rho   = {float(idata_jax_slice.posterior['rho'].mean()):.4f}")
print(f"  sigma = {float(idata_jax_slice.posterior['sigma'].mean()):.4f}")

# Chebyshev logdet with JAX Gibbs
model_cheb = SAR(y=y, X=X, W=W, logdet_method="chebyshev")
idata_cheb = model_cheb.fit(
    sampler="gibbs",
    gibbs_method="jax",
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=42,
)

print(
    f"Chebyshev logdet + JAX Gibbs: rho = {float(idata_cheb.posterior['rho'].mean()):.4f}"
)

# SEM Gibbs
model_sem = SEM(y=y, X=X, W=W)
idata_sem = model_sem.fit(
    sampler="gibbs",
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=42,
    n_jobs=-1,
)
print(
    f"SEM Gibbs: lam = {float(idata_sem.posterior['lam'].mean()):.4f}, "
    f"sigma = {float(idata_sem.posterior['sigma'].mean()):.4f}"
)

# SDM Gibbs
model_sdm = SDM(y=y, X=X, W=W)
idata_sdm = model_sdm.fit(
    sampler="gibbs",
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=42,
    n_jobs=-1,
)
print(
    f"SDM Gibbs: rho = {float(idata_sdm.posterior['rho'].mean()):.4f}, "
    f"beta shape = {idata_sdm.posterior['beta'].shape}"
)

# SDEM Gibbs
model_sdem = SDEM(y=y, X=X, W=W)
idata_sdem = model_sdem.fit(
    sampler="gibbs",
    draws=2000,
    tune=1000,
    chains=4,
    random_seed=42,
    n_jobs=-1,
    chain_method="parallel",
)
print(
    f"SDEM Gibbs: lam = {float(idata_sdem.posterior['lam'].mean()):.4f}, "
    f"beta shape = {idata_sdem.posterior['beta'].shape}"
)

az.summary(idata_sdm)

# ArviZ diagnostics work with Gibbs-produced InferenceData
print("Groups:", idata_gibbs.groups())
print()

# LOO cross-validation
loo = az.loo(idata_gibbs)
print(f"LOO: elpd = {loo.elpd_loo:.2f}, SE = {loo.se:.2f}")
print()

# WAIC
waic = az.waic(idata_gibbs)
print(f"WAIC: elpd = {waic.elpd_waic:.2f}, SE = {waic.se:.2f}")
print()

# Summary
az.summary(idata_gibbs, var_names=["rho", "sigma"])

az.plot_trace(idata_gibbs)

# Run spatial diagnostics on the Gibbs-fitted SAR model
report = model_gibbs.spatial_diagnostics()
print(report)