Bayesian Spatial Diagnostics¶

import warnings

import libpysal
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from spreg import OLS as SpregOLS
from spreg import LMtests as SpregLMtests

from bayespecon import OLS, SAR, SDEM, SDM, SEM, SLX
from bayespecon.diagnostics.lmtests import (
    bayesian_lm_error_sdm_test,
    bayesian_lm_error_test,
    bayesian_lm_lag_sdem_test,
    bayesian_lm_lag_test,
    bayesian_lm_sdm_joint_test,
    bayesian_lm_slx_error_joint_test,
    bayesian_lm_wx_sem_test,
    bayesian_lm_wx_test,
    bayesian_robust_lm_error_sdem_test,
    bayesian_robust_lm_error_test,
    bayesian_robust_lm_lag_sdm_test,
    bayesian_robust_lm_lag_test,
    bayesian_robust_lm_wx_test,
)

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)

import geopandas as gpd
from libpysal.graph import Graph
from libpysal.weights import Rook

# Generate data under H0: no spatial effects
np.random.seed(42)

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

# Create a Graph (modern libpysal API) for bayespecon models
# Row-standardize the graph so spatial models work correctly
g = Graph.build_contiguity(gdf, rook=True).transform("r")
n = g.n

# Legacy W for spreg comparison
w_spreg = Rook.from_shapefile(columbus_path)
w_spreg.transform = "r"

# Get sparse and dense W matrices
W_sparse = g.sparse.tocsr().astype(np.float64)
W_dense = np.array(W_sparse.todense())

# Design matrix
k = 3
X = np.column_stack([np.ones(n), np.random.normal(size=(n, k - 1))])
beta_true = np.array([1.0, 2.0, -1.5])
y = X @ beta_true + np.random.normal(scale=1.0, size=n)

print(f"W shape: {W_sparse.shape}, nnz: {W_sparse.nnz}")

# Fit OLS model (null for joint tests)
ols_model = OLS(y=y, X=X, W=g)
ols_model.fit(draws=5000, tune=5000, chains=4, random_seed=42)

# Fit SAR model (null for LM-WX and robust LM-WX)
sar_model = SAR(y=y, X=X, W=g)
sar_model.fit(draws=5000, tune=5000, chains=4, random_seed=42)

# Fit SLX model (null for robust LM-Lag-SDM and robust LM-Error-SDEM)
slx_model = SLX(y=y, X=X, W=g)
slx_model.fit(draws=5000, tune=5000, chains=4, random_seed=42)

# Fit the SEM null required by LM-WX-SEM (Bayesian analogue of Koley–Bera 2024).
sem_model = SEM(y=y, X=X, W=g)
sem_model.fit(draws=2500, tune=2500, chains=2, random_seed=42)

# Non-robust Bayesian LM tests evaluated at the appropriate null model.
non_robust_results = pd.DataFrame(
    [
        bayesian_lm_lag_test(ols_model).to_series(),
        bayesian_lm_error_test(ols_model).to_series(),
        bayesian_lm_wx_test(sar_model).to_series(),
        bayesian_lm_wx_sem_test(sem_model).to_series(),
        bayesian_lm_sdm_joint_test(ols_model).to_series(),
        bayesian_lm_slx_error_joint_test(ols_model).to_series(),
    ],
    index=[
        "LM-Lag",
        "LM-Error",
        "LM-WX",
        "LM-WX-SEM",
        "LM-SDM Joint",
        "LM-SLX-Error Joint",
    ],
)

non_robust_results

# Robust Bayesian LM tests (Neyman orthogonal score)
robust_results = pd.DataFrame(
    [
        bayesian_robust_lm_lag_test(ols_model).to_series(),
        bayesian_robust_lm_error_test(ols_model).to_series(),
        bayesian_robust_lm_lag_sdm_test(slx_model).to_series(),
        bayesian_robust_lm_wx_test(sar_model).to_series(),
        bayesian_robust_lm_error_sdem_test(slx_model).to_series(),
    ],
    index=[
        "Robust LM-Lag",
        "Robust LM-Error",
        "Robust LM-Lag-SDM",
        "Robust LM-WX",
        "Robust LM-Error-SDEM",
    ],
)

robust_results

from scipy import stats as sp_stats

# Show six representative posterior LM distributions (a mix of non-robust and
# robust variants).  Each panel overlays the chi-squared 95% reference and the
# posterior mean.
panels = [
    ("LM-Lag", bayesian_lm_lag_test(ols_model)),
    ("LM-Error", bayesian_lm_error_test(ols_model)),
    ("LM-WX", bayesian_lm_wx_test(sar_model)),
    ("LM-SDM Joint", bayesian_lm_sdm_joint_test(ols_model)),
    ("Robust LM-WX", bayesian_robust_lm_wx_test(sar_model)),
    ("Robust LM-Lag-SDM", bayesian_robust_lm_lag_sdm_test(slx_model)),
]

fig, axes = plt.subplots(2, 3, figsize=(15, 8))
for ax, (name, res) in zip(axes.flat, panels):
    ax.hist(res.lm_samples, bins=50, density=True, alpha=0.7, color="steelblue")
    chi2_ref = sp_stats.chi2.ppf(0.95, res.df)
    ax.axvline(
        chi2_ref,
        color="red",
        linestyle="--",
        label=f"$\\chi^2_{{0.95,\\,df={res.df}}}$ = {chi2_ref:.2f}",
    )
    ax.axvline(res.mean, color="black", linestyle="-", label=f"Mean = {res.mean:.2f}")
    ax.set_title(name)
    ax.legend(fontsize=8)

plt.suptitle("Posterior Distributions of Bayesian LM Statistics", fontsize=14)
plt.tight_layout()
plt.show()

# Classical spreg LM tests
ols_spreg = SpregOLS(y, X)
lm_spreg = SpregLMtests(ols_spreg, w_spreg)

spreg_results = {
    "LM-Lag": lm_spreg.lml,
    "LM-Error": lm_spreg.lme,
    "LM-WX": lm_spreg.lmwx,
    "LM-SDM Joint": lm_spreg.lmspdurbin,
    "LM-SLX-Error Joint": lm_spreg.lmslxerr,
    "Robust LM-Lag": lm_spreg.rlml,
    "Robust LM-Error": lm_spreg.rlme,
    "Robust LM-Lag-SDM": lm_spreg.rlmdurlag,
    "Robust LM-WX": lm_spreg.rlmwx,
}

all_results = pd.concat([non_robust_results, robust_results])
comparison_rows = []
for bname in all_results.index:
    row = all_results.loc[bname]
    if bname in spreg_results:
        s_stat, s_pval = spreg_results[bname]
    else:
        s_stat, s_pval = np.nan, np.nan
    comparison_rows.append(
        {
            "bayes_mean": row["mean"],
            "spreg_stat": s_stat,
            "bayes_pvalue": row["bayes_pvalue"],
            "spreg_pvalue": s_pval,
        }
    )

comparison_df = pd.DataFrame(comparison_rows, index=all_results.index)
comparison_df

# Method-based API: one call returns all wired tests as a DataFrame
ols_model.spatial_diagnostics()

# The decision routine walks the appropriate tests and returns a recommendation
print("OLS recommends:")
ols_model.spatial_diagnostics_decision()

print("SAR recommends:")
sar_model.spatial_diagnostics_decision()

print("SLX recommends:")
slx_model.spatial_diagnostics_decision()

# SDM/SDEM-aware tests require fitted SDM / SDEM models so the residuals
# include the correct spatial filters.
sdm_model = SDM(y=y, X=X, W=g)
sdm_model.fit(draws=2000, tune=2000, chains=2, random_seed=42)

sdem_model = SDEM(y=y, X=X, W=g)
sdem_model.fit(draws=2000, tune=2000, chains=2, random_seed=42)

pd.DataFrame(
    [
        bayesian_lm_error_sdm_test(sdm_model).to_series(),
        bayesian_lm_lag_sdem_test(sdem_model).to_series(),
    ],
    index=["LM-Error-SDM", "LM-Lag-SDEM"],
)

from bayespecon.dgp.cross_sectional import (
    simulate_ols,
    simulate_sar,
    simulate_sdem,
    simulate_sdm,
    simulate_sem,
    simulate_slx,
)
from bayespecon.dgp.utils import rook_grid_weights

ALPHA = 0.05
SAMPLE_KW = dict(draws=600, tune=600, chains=2, random_seed=7, progressbar=False)

# 12x12 rook grid (n=144), large enough for stable LM tests but quick to fit.
N_SIDE = 12
W_dense_grid, W_graph = rook_grid_weights(N_SIDE)

beta = np.array([1.0, 2.0])
beta1 = np.array([1.0, 2.0])
beta2 = np.array([1.5])  # WX coefficient — large so LM-WX detects it

COMMON = dict(W=W_graph, seed=42, sigma=1.0)

scenarios = {
    "OLS": simulate_ols(beta=beta, **COMMON),
    "SAR": simulate_sar(rho=0.6, beta=beta, **COMMON),
    "SEM": simulate_sem(lam=0.6, beta=beta, **COMMON),
    "SLX": simulate_slx(beta1=beta1, beta2=beta2, **COMMON),
    "SDM": simulate_sdm(rho=0.5, beta1=beta1, beta2=beta2, **COMMON),
    "SDEM": simulate_sdem(lam=0.5, beta1=beta1, beta2=beta2, **COMMON),
}

for name, d in scenarios.items():
    print(
        f"{name:5s}  n={len(d['y']):3d}  X.shape={d['X'].shape}  "
        f"true params: {list(d['params_true'].keys())}"
    )

def to_frame(X):
    """Wrap design matrix in a DataFrame with intercept + x1, x2, ... names."""
    cols = ["intercept"] + [f"x{i}" for i in range(1, X.shape[1])]
    return pd.DataFrame(X, columns=cols)


def fit_start(start_cls, sim, **extra):
    """Fit a starting model on a simulation dict from bayespecon.dgp."""
    Xf = to_frame(sim["X"])
    yf = sim["y"]
    model = start_cls(y=yf, X=Xf, W=W_graph, logdet_method="eigenvalue", **extra)
    model.fit(**SAMPLE_KW)
    return model


def diagnose(model, alpha=ALPHA):
    """Return (recommended_model, diagnostics_df)."""
    diag = model.spatial_diagnostics()
    rec = model.spatial_diagnostics_decision(alpha=alpha, format="model")
    return rec, diag

experiments = [
    ("OLS", OLS, "OLS"),
    ("SAR", OLS, "SAR"),
    ("SEM", OLS, "SEM"),
    ("SLX", SLX, "SLX"),
    ("SDM", SAR, "SDM"),
    ("SDM", SLX, "SDM"),
    ("SDEM", SEM, "SDEM"),
    ("SDEM", SLX, "SDEM"),
]

results = []
fitted = {}
for dgp_name, start_cls, expected in experiments:
    sim = scenarios[dgp_name]
    model = fit_start(start_cls, sim)
    rec, diag = diagnose(model)
    fitted[(dgp_name, start_cls.__name__)] = (model, diag)
    results.append(
        {
            "DGP": dgp_name,
            "Starting model": start_cls.__name__,
            "Recommended": rec,
            "Expected": expected,
            "Match": "yes" if rec == expected else "no",
        }
    )

summary = pd.DataFrame(results)
summary

for (dgp_name, start_name), (_, diag) in fitted.items():
    rec = summary.query("DGP == @dgp_name and `Starting model` == @start_name").iloc[0]
    print(
        f"\n=== DGP={dgp_name} | start={start_name} | "
        f"recommended={rec['Recommended']} ({rec['Match']} expected {rec['Expected']}) ==="
    )
    print(diag[["statistic", "df", "p_value"]].round(4).to_string())

# ASCII rendering of the full decision tree with the traversed path highlighted.
model_sdm_from_slx, _ = fitted[("SDM", "SLX")]
print(model_sdm_from_slx.spatial_diagnostics_decision(alpha=ALPHA, format="ascii"))

from bayespecon import OLSPanelFE
from bayespecon.dgp.panel_fe import simulate_panel_sar_fe
from bayespecon.dgp.utils import rook_grid_weights

# 9x9 rook grid (N=81), T=4 — small enough to fit quickly but large enough for
# the LM tests to behave well.
N_panel, T_panel = 81, 4
W_panel_dense, W_panel_graph = rook_grid_weights(int(np.sqrt(N_panel)))

# Simulate from a SAR-FE DGP with moderate spatial dependence so the LM-Lag
# direction is clearly significant.
panel_sim = simulate_panel_sar_fe(
    N=N_panel,
    T=T_panel,
    rho=0.4,
    beta=np.array([1.0, 2.0]),
    W=W_panel_graph,
    seed=11,
)

panel_model = OLSPanelFE(
    y=panel_sim["y"],
    X=panel_sim["X"],
    W=W_panel_graph,
    N=N_panel,
    T=T_panel,
    model=3,  # two-way fixed effects (unit + time)
)
panel_model.fit(draws=600, tune=600, chains=2, random_seed=11, progressbar=False)

panel_model.spatial_diagnostics()

print(
    f"Panel decision tree (DGP = SAR-FE) recommends: {panel_model.spatial_diagnostics_decision(alpha=0.05, format='model')}"
)

panel_model.spatial_diagnostics_decision(
    alpha=0.05,
)

from bayespecon import (
    OLSPanelFE,
    SARPanelFE,
    SEMPanelFE,
    SLXPanelFE,
)
from bayespecon.dgp.panel_fe import (
    simulate_panel_ols_fe,
    simulate_panel_sar_fe,
    simulate_panel_sdem_fe,
    simulate_panel_sdm_fe,
    simulate_panel_sem_fe,
    simulate_panel_slx_fe,
)

PANEL_N_SIDE = 10
_, W_panel_graph = rook_grid_weights(PANEL_N_SIDE)
PANEL_N = PANEL_N_SIDE * PANEL_N_SIDE
PANEL_T = 5
PANEL_COMMON = dict(N=PANEL_N, T=PANEL_T, W=W_panel_graph, seed=42, sigma=1.0)
PANEL_SAMPLE_KW = dict(draws=400, tune=400, chains=2, random_seed=7, progressbar=False)

panel_scenarios = {
    "OLS": simulate_panel_ols_fe(beta=beta, **PANEL_COMMON),
    "SAR": simulate_panel_sar_fe(rho=0.5, beta=beta, **PANEL_COMMON),
    "SEM": simulate_panel_sem_fe(lam=0.5, beta=beta, **PANEL_COMMON),
    "SLX": simulate_panel_slx_fe(beta1=beta1, beta2=beta2, **PANEL_COMMON),
    "SDM": simulate_panel_sdm_fe(rho=0.4, beta1=beta1, beta2=beta2, **PANEL_COMMON),
    "SDEM": simulate_panel_sdem_fe(lam=0.4, beta1=beta1, beta2=beta2, **PANEL_COMMON),
}

def fit_panel_start(start_cls, sim):
    Xf = to_frame(sim["X"])
    m = start_cls(
        y=sim["y"],
        X=Xf,
        W=W_panel_graph,
        N=PANEL_N,
        T=PANEL_T,
        model=3,  # two-way fixed effects
        logdet_method="eigenvalue",
    )
    m.fit(**PANEL_SAMPLE_KW)
    return m


panel_experiments = [
    ("OLS", OLSPanelFE, "OLSPanelFE"),
    ("SAR", OLSPanelFE, "SARPanelFE"),
    ("SEM", OLSPanelFE, "SEMPanelFE"),
    ("SLX", SLXPanelFE, "SLXPanelFE"),
    ("SDM", SARPanelFE, "SDMPanelFE"),
    ("SDM", SLXPanelFE, "SDMPanelFE"),
    ("SDEM", SEMPanelFE, "SDEMPanelFE"),
    ("SDEM", SLXPanelFE, "SDEMPanelFE"),
]

panel_results = []
panel_fitted = {}
for dgp_name, start_cls, expected in panel_experiments:
    m = fit_panel_start(start_cls, panel_scenarios[dgp_name])
    diag = m.spatial_diagnostics()
    rec = m.spatial_diagnostics_decision(alpha=ALPHA, format="model")
    panel_fitted[(dgp_name, start_cls.__name__)] = (m, diag)
    panel_results.append(
        {
            "DGP": dgp_name,
            "Starting model": start_cls.__name__,
            "Recommended": rec,
            "Expected": expected,
            "Match": "yes" if rec == expected else "no",
        }
    )

panel_summary = pd.DataFrame(panel_results)
panel_summary

for (dgp_name, start_name), (_, diag) in panel_fitted.items():
    rec = panel_summary.query(
        "DGP == @dgp_name and `Starting model` == @start_name"
    ).iloc[0]
    print(
        f"\n=== DGP={dgp_name} | start={start_name} | "
        f"recommended={rec['Recommended']} ({rec['Match']} expected {rec['Expected']}) ==="
    )
    print(diag[["statistic", "df", "p_value"]].round(4).to_string())

from bayespecon import OLSFlow, SARFlow
from bayespecon.dgp.flows import generate_flow_data

# Simulate a small SAR flow DGP on n=8 spatial units (N = 64 OD cells).
flow_data = generate_flow_data(
    n=8,
    rho_d=0.35,
    rho_o=0.25,
    rho_w=0.10,
    beta_d=[1.0, -0.5],
    beta_o=[0.5, 0.3],
    sigma=1.0,
    seed=42,
)

y_flow = np.log(flow_data["y_vec"])  # latent SAR scale (DGP default is lognormal)
X_flow = flow_data["X"]
G_flow = flow_data["G"]

ols_flow = OLSFlow(y_flow, G_flow, X_flow, col_names=flow_data["col_names"])
ols_flow.fit(draws=600, tune=600, chains=2, random_seed=42, progressbar=False)

ols_flow.spatial_diagnostics()

# Fit the SAR flow alternative; its diagnostic registry consists of the robust
# (Neyman-orthogonal) marginal-extension tests.
sar_flow = SARFlow(y_flow, G_flow, X_flow, col_names=flow_data["col_names"])
sar_flow.fit(draws=600, tune=600, chains=2, random_seed=42, progressbar=False)

sar_flow.spatial_diagnostics()