Spatial panel models: a walkthrough¶

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

from neighbayes import dgp
from neighbayes.models import (
    OLSPanelFE,
    OLSPanelRE,
    SARPanelFE,
    SARPanelRE,
    SDEMPanelFE,
    SDMPanelFE,
    SEMPanelFE,
    SEMPanelRE,
)

az.style.use("arviz-white")

# Generate base geometry via cross-sectional DGP, then simulate panel data via panel DGP.
rng = np.random.default_rng(2026)
xcols = ["poverty", "rev_rating", "num_spots", "crowded"]
ycol = "price_pp"

base_gdf = dgp.simulate_sar(n=20, seed=2026, create_gdf=True, geometry_type="polygon")
W = Graph.build_contiguity(base_gdf, rook=False).transform("r")
N = len(base_gdf)
T = 4
beta = np.array([1.5, -0.8, 0.6, 0.4, -0.5], dtype=float)

panel_sim = dgp.simulate_panel_sar_fe(
    N=N,
    T=T,
    rho=0.35,
    beta=beta,
    sigma=1.0,
    sigma_alpha=0.5,
    rng=rng,
    W=W,
)

panel = pd.DataFrame(
    {
        "unit": panel_sim["unit"],
        "time": panel_sim["time"],
        ycol: panel_sim["y"],
    }
)
for j, name in enumerate(xcols, start=1):
    panel[name] = panel_sim["X"][:, j]

def fit_panel_model(
    model_cls, formula, data, W, effects=3, draws=2000, tune=1000, chains=2, seed=42
):
    """Fit a panel model and return (model, summary, effects_df).

    Uses minimal MCMC settings for pedagogical demonstration.
    Increase draws/tune for real analyses.
    """
    m = model_cls(
        formula=formula,
        data=data,
        W=W,
        unit_col="unit",
        time_col="time",
        effects=effects,
    )
    m.fit(
        draws=draws,
        tune=tune,
        chains=chains,
        random_seed=seed,
        progressbar=True,
        # Stored only on request (as in PyMC); az.loo / az.waic need it.
        idata_kwargs={"log_likelihood": True},
    )
    summary = m.summary(round_to=3)
    effects_df = pd.DataFrame(m.spatial_effects())
    return m, summary, effects_df


def diagnostics_table(idata, var_names):
    """Show key MCMC diagnostics for the given parameters."""
    cols = ["mean", "sd", "ess_bulk", "ess_tail", "r_hat"]
    return az.summary(idata, var_names=var_names, round_to=3)[cols]


def show_trace(idata, var_names, title):
    """Plot trace plots for the given parameters."""
    az.plot_trace(idata, var_names=var_names)
    plt.suptitle(title, y=1.02)
    plt.tight_layout()
    plt.show()

formula = "price_pp ~ poverty + rev_rating + num_spots + crowded"
ols_panel, summary_ols, effects_ols = fit_panel_model(
    OLSPanelFE, formula, panel, W, effects=3
)
display(summary_ols)
display(effects_ols)
display(diagnostics_table(ols_panel.inference_data, ["beta", "sigma"]))
show_trace(ols_panel.inference_data, ["sigma"], "OLSPanelFE trace")

sar_panel, summary_sar, effects_sar = fit_panel_model(
    SARPanelFE, formula, panel, W, effects=3
)
display(summary_sar)
display(effects_sar)
display(diagnostics_table(sar_panel.inference_data, ["rho", "beta", "sigma"]))
show_trace(sar_panel.inference_data, ["rho", "sigma"], "SARPanelFE trace")

sem_panel, summary_sem, effects_sem = fit_panel_model(
    SEMPanelFE, formula, panel, W, effects=3
)
display(summary_sem)
display(effects_sem)
display(diagnostics_table(sem_panel.inference_data, ["lam", "beta", "sigma"]))
show_trace(sem_panel.inference_data, ["lam", "sigma"], "SEMPanelFE trace")

sdm_panel, summary_sdm, effects_sdm = fit_panel_model(
    SDMPanelFE, formula, panel, W, effects=3
)
display(summary_sdm)
display(effects_sdm)
display(diagnostics_table(sdm_panel.inference_data, ["rho", "beta", "sigma"]))
show_trace(sdm_panel.inference_data, ["rho", "sigma"], "SDMPanelFE trace")

sdem_panel, summary_sdem, effects_sdem = fit_panel_model(
    SDEMPanelFE, formula, panel, W, effects=3
)
display(summary_sdem)
display(effects_sdem)
display(diagnostics_table(sdem_panel.inference_data, ["lam", "beta", "sigma"]))
show_trace(sdem_panel.inference_data, ["lam", "sigma"], "SDEMPanelFE trace")

from neighbayes.diagnostics import spatial_mcmc_diagnostic

spatial_mcmc_diagnostic(sdm_panel, emit_warnings=False).to_frame()

ols_panel_re, summary_ols_re, effects_ols_re = fit_panel_model(
    OLSPanelRE, formula, panel, W, effects=0
)
display(summary_ols_re)
display(effects_ols_re)
display(
    diagnostics_table(ols_panel_re.inference_data, ["beta", "sigma", "sigma_alpha"])
)
show_trace(ols_panel_re.inference_data, ["sigma", "sigma_alpha"], "OLSPanelRE trace")

sar_panel_re, summary_sar_re, effects_sar_re = fit_panel_model(
    SARPanelRE, formula, panel, W, effects=0
)
display(summary_sar_re)
display(effects_sar_re)
display(
    diagnostics_table(
        sar_panel_re.inference_data, ["rho", "beta", "sigma", "sigma_alpha"]
    )
)
show_trace(
    sar_panel_re.inference_data, ["rho", "sigma", "sigma_alpha"], "SARPanelRE trace"
)

sem_panel_re, summary_sem_re, effects_sem_re = fit_panel_model(
    SEMPanelRE, formula, panel, W, effects=0
)
display(summary_sem_re)
display(effects_sem_re)
display(
    diagnostics_table(
        sem_panel_re.inference_data, ["lam", "beta", "sigma", "sigma_alpha"]
    )
)
show_trace(
    sem_panel_re.inference_data, ["lam", "sigma", "sigma_alpha"], "SEMPanelRE trace"
)

# Collect all fitted panel models for comparison
panel_models = {
    "OLSPanelFE": ols_panel,
    "SARPanelFE": sar_panel,
    "SEMPanelFE": sem_panel,
    "SDMPanelFE": sdm_panel,
    "SDEMPanelFE": sdem_panel,
    "OLSPanelRE": ols_panel_re,
    "SARPanelRE": sar_panel_re,
    "SEMPanelRE": sem_panel_re,
}
idata_dict = {name: m.inference_data for name, m in panel_models.items()}

# WAIC and LOO comparison
for ic in ("waic", "loo"):
    try:
        cmp = az.compare(idata_dict, ic=ic, method="BB-pseudo-BMA")
        print(f"\n{ic.upper()} comparison")
        display(cmp)
    except Exception as e:
        print(f"{ic.upper()} comparison not available: {type(e).__name__}: {e}")

# Per-model IC table
rows = []
for name, model in panel_models.items():
    idata = model.inference_data
    row = {"model": name}
    try:
        waic_res = az.waic(idata)
        row["elpd_waic"] = float(waic_res.elpd_waic)
        row["p_waic"] = float(waic_res.p_waic)
    except Exception:
        row["elpd_waic"] = np.nan
        row["p_waic"] = np.nan
    try:
        loo_res = az.loo(idata)
        row["elpd_loo"] = float(loo_res.elpd_loo)
        row["p_loo"] = float(loo_res.p_loo)
    except Exception:
        row["elpd_loo"] = np.nan
        row["p_loo"] = np.nan
    rows.append(row)

pd.DataFrame(rows).sort_values("model").reset_index(drop=True)

from neighbayes.diagnostics.lmtests import (
    bayesian_panel_lm_error_test,
    bayesian_panel_lm_lag_test,
    bayesian_panel_lm_sdm_joint_test,
    bayesian_panel_lm_slx_error_joint_test,
    bayesian_panel_lm_wx_test,
    bayesian_panel_robust_lm_error_sdem_test,
    bayesian_panel_robust_lm_error_test,
    bayesian_panel_robust_lm_lag_sdm_test,
    bayesian_panel_robust_lm_lag_test,
    bayesian_panel_robust_lm_wx_test,
)

# Standard panel LM tests (from OLS null)
panel_standard = pd.DataFrame(
    [
        bayesian_panel_lm_lag_test(ols_panel).to_series(),
        bayesian_panel_lm_error_test(ols_panel).to_series(),
        bayesian_panel_lm_sdm_joint_test(ols_panel).to_series(),
        bayesian_panel_lm_slx_error_joint_test(ols_panel).to_series(),
    ],
    index=["LM-Lag", "LM-Error", "LM-SDM Joint", "LM-SLX-Error Joint"],
)

panel_standard

# Robust panel LM tests (from OLS null)
panel_robust = pd.DataFrame(
    [
        bayesian_panel_robust_lm_lag_test(ols_panel).to_series(),
        bayesian_panel_robust_lm_error_test(ols_panel).to_series(),
    ],
    index=["Robust LM-Lag", "Robust LM-Error"],
)

panel_robust

from neighbayes.models import SLXPanelFE

slx_panel, summary_slx, effects_slx = fit_panel_model(
    SLXPanelFE, formula, panel, W, effects=3
)
display(summary_slx)

# SDM/SDEM variant panel LM tests
panel_sdem = pd.DataFrame(
    [
        bayesian_panel_lm_wx_test(sar_panel).to_series(),
        bayesian_panel_robust_lm_lag_sdm_test(slx_panel).to_series(),
        bayesian_panel_robust_lm_wx_test(sar_panel).to_series(),
        bayesian_panel_robust_lm_error_sdem_test(slx_panel).to_series(),
    ],
    index=["LM-WX", "Robust LM-Lag-SDM", "Robust LM-WX", "Robust LM-Error-SDEM"],
)

panel_sdem

from scipy import stats as sp_stats

# All panel LM tests in a single table
all_panel = pd.concat([panel_standard, panel_robust, panel_sdem])

# Add chi-squared reference values
all_panel["chi2_ref_95"] = all_panel["df"].apply(lambda df: sp_stats.chi2.ppf(0.95, df))

all_panel

ols_panel.spatial_diagnostics()

print(
    "OLS panel recommends:",
)
ols_panel.spatial_diagnostics_decision()

print("SAR panel recommends:")
sar_panel.spatial_diagnostics_decision()