JGSY2024 Replication with bayespecon¶

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

from bayespecon import SDMPanelFE

az.style.use("arviz-white")
pd.options.display.float_format = lambda x: f"{x:,.6f}"

# Paths and constants from JGSY2024
static_data_dir = "../_static/jgsy2024"

N = 266
T = 17
NT = N * T

# Coerce mixed-type spreadsheets to numeric and drop label rows/cols.
A_df = pd.read_excel(f"{static_data_dir}/Data.xlsx", header=None)
A_num = (
    A_df.apply(pd.to_numeric, errors="coerce")
    .dropna(axis=0, how="all")
    .dropna(axis=1, how="all")
)
A = A_num.to_numpy(dtype=float)

D_df = pd.read_excel(f"{static_data_dir}/distance.xlsx", header=None)
D_num = (
    D_df.apply(pd.to_numeric, errors="coerce")
    .dropna(axis=0, how="all")
    .dropna(axis=1, how="all")
)
D = D_num.to_numpy(dtype=float)

print(f"Loaded A shape: {A.shape}")
print(f"Loaded distance shape: {D.shape}")

# Reproduce data construction from JGSY2024_MainResults.m
# MATLAB indices are 1-based; Python is 0-based.
A_work = A.copy()

y = np.zeros(NT)
ylag = np.zeros(NT)
for t in range(T):
    t1 = t * N
    t2 = (t + 1) * N
    y[t1:t2] = A_work[t1 + 2 * N : t2 + 2 * N, 3] - A_work[t1 + N : t2 + N, 3]
    ylag[t1:t2] = A_work[t1 + N : t2 + N, 3] - A_work[t1:t2, 3]

lnpopg = A_work[:, 7].copy()
popg = np.exp(lnpopg) - 1.0
popg[popg <= 0] = 0.0001
A_work[:, 7] = np.log(popg)

x_core = np.column_stack(
    [
        A_work[2 * N : 2 * N + NT, 4],
        A_work[2 * N : 2 * N + NT, 7],
        A_work[N : N + NT, 3],
        ylag,
    ]
)

X_core = pd.DataFrame(x_core, columns=["inv", "pop", "initial_gdp", "lagged_growth"])

print("Correlation matrix of core regressors (matches script intent):")
display(X_core.corr())

def graph_from_dense(W_np: np.ndarray) -> Graph:
    rows, cols = np.where(W_np > 0)
    vals = W_np[rows, cols].astype(float)
    return Graph.from_arrays(rows, cols, vals)


def build_knn6_weights(distance_matrix: np.ndarray, m: int = 6) -> np.ndarray:
    Dm = distance_matrix.copy().astype(float)
    np.fill_diagonal(Dm, 0.0)
    order = np.argsort(Dm, axis=1)
    kth = np.take_along_axis(Dm, order[:, [m]], axis=1).ravel()

    W = np.zeros_like(Dm, dtype=float)
    for i in range(Dm.shape[0]):
        mask = (Dm[i] > 0) & (Dm[i] <= kth[i])
        W[i, mask] = 1.0

    row_sum = W.sum(axis=1, keepdims=True)
    row_sum[row_sum == 0] = 1.0
    return W / row_sum


def build_ed_weights(distance_matrix: np.ndarray) -> np.ndarray:
    DD0 = distance_matrix.astype(float) / 100.0
    W = np.exp(-DD0)
    np.fill_diagonal(W, 0.0)
    row_sum = W.sum(axis=1, keepdims=True)
    row_sum[row_sum == 0] = 1.0
    return W / row_sum


W_knn6 = build_knn6_weights(D, m=6)
W_ed = build_ed_weights(D)

print("W_knn6 row sums:", np.min(W_knn6.sum(axis=1)), np.max(W_knn6.sum(axis=1)))
print("W_ed row sums:", np.min(W_ed.sum(axis=1)), np.max(W_ed.sum(axis=1)))

def fit_sdm_fe(y_vec, X_df, W_np, draws=600, tune=600, chains=1, seed=42):
    model = SDMPanelFE(
        y=y_vec,
        X=X_df,
        W=graph_from_dense(W_np),
        N=N,
        T=T,
        model=3,
        logdet_method="eigenvalue",
    )

    idata = model.fit(
        draws=draws,
        tune=tune,
        chains=chains,
        cores=1,
        target_accept=0.9,
        random_seed=seed,
        progressbar=False,
        idata_kwargs={"log_likelihood": True},
    )
    return model, idata


model_knn6, idata_knn6 = fit_sdm_fe(y, X_core, W_knn6)
model_ed, idata_ed = fit_sdm_fe(y, X_core, W_ed)

print("Fit complete for both specifications.")

# Original JGSY2024 Table 1 values (from readme / script output)
orig_coef = {
    "knn6": {
        "inv": 0.002422,
        "pop": -0.008074,
        "initial_gdp": -0.093832,
        "lagged_growth": 0.013879,
        "W*inv": 0.008933,
        "W*pop": -0.002317,
        "W*initial_gdp": 0.065751,
        "W*lagged_growth": 0.191021,
        "rho": 0.534381,
    },
    "ed": {
        "inv": 0.001616,
        "pop": -0.007904,
        "initial_gdp": -0.088731,
        "lagged_growth": 0.011519,
        "W*inv": 0.009081,
        "W*pop": -0.004958,
        "W*initial_gdp": 0.070415,
        "W*lagged_growth": 0.194065,
        "rho": 0.650187,
    },
}

orig_effects = {
    "knn6": pd.DataFrame(
        {
            "feature": ["inv", "pop", "initial_gdp", "lagged_growth"],
            "direct": [0.0035, -0.0088, -0.0921, 0.0351],
            "indirect": [0.0212, -0.0135, 0.0322, 0.4029],
            "total": [0.0247, -0.0223, -0.0599, 0.4380],
        }
    ),
    "ed": pd.DataFrame(
        {
            "feature": ["inv", "pop", "initial_gdp", "lagged_growth"],
            "direct": [0.0026, -0.0089, -0.0876, 0.0291],
            "indirect": [0.0281, -0.0282, 0.0352, 0.5566],
            "total": [0.0307, -0.0371, -0.0523, 0.5857],
        }
    ),
}


def extract_bayes_coef(model, idata):
    beta_names = model._beta_names()
    beta_mean = idata.posterior["beta"].mean(("chain", "draw")).values
    out = dict(zip(beta_names, beta_mean))
    out["rho"] = float(idata.posterior["rho"].mean())
    keep = [
        "inv",
        "pop",
        "initial_gdp",
        "lagged_growth",
        "W*inv",
        "W*pop",
        "W*initial_gdp",
        "W*lagged_growth",
        "rho",
    ]
    return pd.Series({k: out[k] for k in keep}, name="bayespecon")


def compare_coef(spec_name, model, idata):
    b = extract_bayes_coef(model, idata).rename("bayespecon")
    o = pd.Series(orig_coef[spec_name], name="original")
    cmp = pd.concat([o, b], axis=1)
    cmp["diff"] = cmp["bayespecon"] - cmp["original"]
    cmp["abs_diff"] = cmp["diff"].abs()
    cmp["pct_diff_abs"] = (
        100 * cmp["abs_diff"] / cmp["original"].abs().replace(0, np.nan)
    )
    return cmp


def compare_effects(spec_name, model):
    eff = model.spatial_effects()
    be = pd.DataFrame(
        {
            "feature": eff.index.tolist(),
            "direct": eff["direct"].values,
            "indirect": eff["indirect"].values,
            "total": eff["total"].values,
        }
    )
    merged = orig_effects[spec_name].merge(
        be, on="feature", suffixes=("_orig", "_bayes")
    )
    for col in ["direct", "indirect", "total"]:
        merged[f"{col}_diff"] = merged[f"{col}_bayes"] - merged[f"{col}_orig"]
        merged[f"{col}_abs_diff"] = merged[f"{col}_diff"].abs()
    return merged


coef_cmp_knn6 = compare_coef("knn6", model_knn6, idata_knn6)
coef_cmp_ed = compare_coef("ed", model_ed, idata_ed)

eff_cmp_knn6 = compare_effects("knn6", model_knn6)
eff_cmp_ed = compare_effects("ed", model_ed)

print("Coefficient comparison: 6NN")
display(coef_cmp_knn6)
print("Coefficient comparison: ED")
display(coef_cmp_ed)

print("Effects comparison: 6NN")
display(eff_cmp_knn6)
print("Effects comparison: ED")
display(eff_cmp_ed)

def summarize_comparison(coef_cmp, eff_cmp, label):
    return pd.Series(
        {
            "spec": label,
            "coef_MAE": coef_cmp["abs_diff"].mean(),
            "coef_median_abs_pct_diff": np.nanmedian(
                coef_cmp["pct_diff_abs"].to_numpy()
            ),
            "direct_effect_MAE": eff_cmp["direct_abs_diff"].mean(),
            "indirect_effect_MAE": eff_cmp["indirect_abs_diff"].mean(),
            "total_effect_MAE": eff_cmp["total_abs_diff"].mean(),
        }
    )


summary_table = pd.DataFrame(
    [
        summarize_comparison(coef_cmp_knn6, eff_cmp_knn6, "Table 1 col[1] 6NN"),
        summarize_comparison(coef_cmp_ed, eff_cmp_ed, "Table 1 col[2] ED"),
    ]
)

display(summary_table)

print("Top absolute coefficient gaps (6NN):")
display(coef_cmp_knn6.sort_values("abs_diff", ascending=False).head(5))

print("Top absolute coefficient gaps (ED):")
display(coef_cmp_ed.sort_values("abs_diff", ascending=False).head(5))