Bayesian Spatial Flow (Origin–Destination) Models¶

import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from bayespecon import SARFlow, SARFlowSeparable, flow_design_matrix
from bayespecon.dgp.flows import generate_flow_data, generate_flow_data_separable

warnings.filterwarnings("ignore")
az.style.use("arviz-white")
rng = np.random.default_rng(42)

# True parameters (matching the table above)
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

# Generate synthetic flow data (auto-builds KNN graph + distance matrix)
data = generate_flow_data(
    n=12,
    rho_d=RHO_D,
    rho_o=RHO_O,
    rho_w=RHO_W,
    beta_d=BETA_D,
    beta_o=BETA_O,
    sigma=SIGMA,
    seed=42,
)

n = 12
N = n * n
G = data["G"]
gdf = data["gdf"]

print(f"Generated {N} O-D flows on {n} spatial units")
print(f"Design matrix: {data['X'].shape}, columns: {data['col_names']}")
print(f"Design: k_d={data['design'].k_d}, k_o={data['design'].k_o}")

y_obs = data["y_vec"]              # observed (positive) flows
y_vec = np.log(y_obs)              # latent SAR scale (DGP default is lognormal)
X = data["X"]              # (N, p) design matrix (incl. log_distance column)
col_names = data["col_names"]
G = data["G"]              # row-standardised KNN graph synthesised by the DGP
gdf = data["gdf"]          # synthetic point GeoDataFrame

print(f"Flow observations: N = {N}  ({n}×{n})")
print(f"Design matrix shape: {X.shape}")
print(f"Column names: {col_names}")
print(f"Last column ({col_names[-1]}) coefficient (gamma_dist) = {data['gamma_dist']}")
print(
    f"\ny (observed) summary:  min={y_obs.min():.2f}  mean={y_obs.mean():.2f}  max={y_obs.max():.2f}"
)
print(
    f"y (log scale) summary:  min={y_vec.min():.2f}  mean={y_vec.mean():.2f}  max={y_vec.max():.2f}"
)

# Visualise the flow matrix (observed scale) and the latent log-scale histogram
fig, axes = plt.subplots(1, 2, figsize=(10, 4))

im = axes[0].imshow(y_obs.reshape(n, n), cmap="YlOrRd")
axes[0].set_title("Simulated flow matrix $Y$ (origin × destination)")
axes[0].set_xlabel("Destination")
axes[0].set_ylabel("Origin")
plt.colorbar(im, ax=axes[0])

axes[1].hist(y_vec, bins=30, edgecolor="white")
axes[1].set_title(r"Distribution of $\log y_{ij}$ (latent SAR scale)")
axes[1].set_xlabel("log flow")
axes[1].set_ylabel("Count")

# plt.tight_layout()
plt.show()

sar_flow = SARFlow(
    y_vec,
    G,
    X,
    col_names=col_names,
    logdet_method="traces",  # Barry-Pace stochastic trace (default)
    restrict_positive=True,  # Dirichlet stability prior
    miter=20,  # trace polynomial order (increase for better accuracy)
    trace_seed=0,
)

idata_sar = sar_flow.fit(
    draws=1000,
    tune=1000,
    chains=4,
    target_accept=0.9,
    random_seed=42,
    progressbar=True,
)

# Posterior summary for the spatial autoregressive parameters
summary_rho = sar_flow.summary(var_names=["rho_d", "rho_o", "rho_w"])
print("=== SARFlow: spatial autoregressive parameters ===")
print(summary_rho.to_string())

# Compare to true values
print(f"\nTrue values:  rho_d={RHO_D}  rho_o={RHO_O}  rho_w={RHO_W}")

# Full coefficient summary
summary_beta = sar_flow.summary(var_names=["beta", "sigma"])
print("=== SARFlow: regression coefficients ===")
print(summary_beta.to_string())

# Trace plots for the three ρ parameters
az.plot_trace(
    idata_sar,
    var_names=["rho_d", "rho_o", "rho_w"],
    figsize=(10, 6),
)
plt.suptitle("SARFlow — posterior traces", y=1.01)
plt.tight_layout()
plt.show()

# Posterior densities with true-value markers
fig, axes = plt.subplots(1, 3, figsize=(12, 3))
params = [("rho_d", RHO_D), ("rho_o", RHO_O), ("rho_w", RHO_W)]

for ax, (param, true_val) in zip(axes, params):
    az.plot_posterior(
        idata_sar,
        var_names=[param],
        ax=ax,
        ref_val=true_val,
        hdi_prob=0.94,
    )
    # ax.set_title(f"${param.replace('_', r'\\_')}$ (true = {true_val})")

plt.suptitle("SARFlow — posterior distributions (red line = true value)", y=1.02)
plt.tight_layout()
plt.show()

# True parameters for the separable model (asymmetric for identifiability)
RHO_D_SEP = 0.40
RHO_O_SEP = 0.30
RHO_W_SEP = -RHO_D_SEP * RHO_O_SEP  # = -0.12 (constraint)

print(
    f"Separable true values:  rho_d={RHO_D_SEP}  rho_o={RHO_O_SEP}"
    f"  rho_w=-rho_d*rho_o={RHO_W_SEP:.4f}"
)

data_sep = generate_flow_data_separable(
    n,
    G,
    rho_d=RHO_D_SEP,
    rho_o=RHO_O_SEP,
    beta_d=BETA_D,
    beta_o=BETA_O,
    sigma=SIGMA,
    seed=7,
)

y_sep = np.log(data_sep["y_vec"])  # latent SAR scale (lognormal default)
X_sep = data_sep["X"]
cn_sep = data_sep["col_names"]

print(
    f"\ny_sep summary:  min={y_sep.min():.2f}  mean={y_sep.mean():.2f}  max={y_sep.max():.2f}"
)

sep_flow = SARFlowSeparable(
    y_sep,
    G,
    X_sep,
    col_names=cn_sep,
    # logdet_method="eigenvalue" is the default for SARFlowSeparable
    trace_seed=0,
)

idata_sep = sep_flow.fit(
    draws=1000,
    tune=1000,
    chains=4,
    target_accept=0.9,
    random_seed=42,
    progressbar=True,
)

summary_sep = sep_flow.summary(var_names=["rho_d", "rho_o"])
print("=== SARFlowSeparable: spatial autoregressive parameters ===")
print(summary_sep.to_string())
print(
    f"\nTrue values:  rho_d={RHO_D_SEP}  rho_o={RHO_O_SEP}"
    f"  (rho_w = -rho_d*rho_o = {RHO_W_SEP:.4f} is derived)"
)

az.plot_trace(
    idata_sep,
    var_names=["rho_d", "rho_o"],
    figsize=(10, 4),
)
plt.suptitle("SARFlowSeparable — posterior traces", y=1.01)
plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(8, 3))
for ax, (param, true_val) in zip(axes, [("rho_d", RHO_D_SEP), ("rho_o", RHO_O_SEP)]):
    az.plot_posterior(
        idata_sep, var_names=[param], ax=ax, ref_val=true_val, hdi_prob=0.94
    )
    ax.set_title(f"${param.replace('_', r'_')}$ (true = {true_val})")

plt.suptitle("SARFlowSeparable — posterior distributions", y=1.02)
plt.tight_layout()
plt.show()

# Build a design matrix from scratch using regional attributes
X_regional = rng.standard_normal((n, 2))  # n × k attribute matrix
dm = flow_design_matrix(X_regional, col_names=["income", "pop"])

print(f"Regional attribute matrix: {X_regional.shape}  (n × k)")
print(f"Full O-D design matrix:    {dm.combined.shape}  (N × p)")
print(f"\nColumn names:\n  {dm.feature_names}")

pd.DataFrame(dm.combined[:6], columns=dm.feature_names).round(3)

# Combined effects (default for symmetric Xo == Xd designs).
effects_df = sar_flow.spatial_effects(mode="combined")
effects_df.round(4)

# Small sampler settings used for the demo fits below.
SAMPLE_KWARGS_DEMO = dict(draws=200, tune=200, chains=2,
                          random_seed=0, progressbar=False)

from bayespecon.dgp.flows import generate_flow_data

# Asymmetric: 2 destination attributes, 1 origin attribute
data_asym = generate_flow_data(
    n=5,
    rho_d=0.2, rho_o=0.15, rho_w=0.05,
    beta_d=[1.0, -0.5],   # k_d = 2
    beta_o=[0.5],          # k_o = 1
    sigma=0.5, seed=11,
)
G_asym = data_asym["G"]
print(f"k_d={data_asym['design'].k_d}, k_o={data_asym['design'].k_o}")
print(f"Columns: {data_asym['col_names']}")

model_asym = SARFlow(
    np.log(data_asym["y_vec"]), G_asym, data_asym["X"],
    col_names=data_asym["col_names"],
)
print("symmetric_xo_xd =", model_asym._symmetric_xo_xd)
model_asym.fit(**SAMPLE_KWARGS_DEMO)
model_asym.spatial_effects().round(4)  # auto -> separate

# You can also build asymmetric designs directly with flow_design_matrix_asymmetric
from bayespecon.graph import flow_design_matrix_asymmetric

rng_asym = np.random.default_rng(7)
n_demo = 5
Xd = rng_asym.standard_normal((n_demo, 2))  # 2 destination variables
Xo = rng_asym.standard_normal((n_demo, 1))  # 1 origin variable

dm = flow_design_matrix_asymmetric(Xd, Xo)
print(f"Design shape: {dm.combined.shape}")
print(f"k_d={dm.k_d}, k_o={dm.k_o}")
print(f"Feature names: {dm.feature_names}")

from bayespecon import OLSFlow

ols_flow = OLSFlow(y_vec, G, X, col_names=col_names)
ols_flow.fit(**SAMPLE_KWARGS_DEMO)
ols_flow.spatial_effects(mode="combined").round(4)

for label, idata in [("SARFlow", idata_sar), ("SARFlowSeparable", idata_sep)]:
    diag = az.summary(idata, var_names=["rho_d", "rho_o"], stat_focus="mean")
    rhat_ok = (diag["r_hat"] < 1.01).all()
    ess_ok = (diag["ess_bulk"] > 400).all()
    print(f"{label:25s}  r_hat<1.01: {rhat_ok}   ess_bulk>400: {ess_ok}")
    print(diag[["mean", "sd", "hdi_3%", "hdi_97%", "ess_bulk", "r_hat"]].to_string())
    print()

from bayespecon import spatial_mcmc_diagnostic

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

# Energy plot — checks that HMC is exploring the posterior efficiently
fig, axes = plt.subplots(1, 2, figsize=(10, 3))
az.plot_energy(idata_sar, ax=axes[0])
axes[0].set_title("SARFlow energy")
az.plot_energy(idata_sep, ax=axes[1])
axes[1].set_title("SARFlowSeparable energy")
plt.tight_layout()
plt.show()

# Fit unrestricted SARFlow on the same separable data for a fair comparison
sar_flow_on_sep = SARFlow(
    y_sep,
    G,
    X_sep,
    col_names=cn_sep,
    logdet_method="traces",
    miter=20,
    trace_seed=0,
)
idata_sar_on_sep = sar_flow_on_sep.fit(
    draws=1000,
    tune=1000,
    chains=4,
    target_accept=0.9,
    random_seed=42,
    progressbar=True,
)

# LOO-CV comparison (requires log-likelihood in idata; falls back to WAIC if unavailable)
try:
    loo_sar = az.loo(idata_sar_on_sep)
    loo_sep = az.loo(idata_sep)
    comparison = az.compare(
        {"SARFlow (3 ρ)": idata_sar_on_sep, "SARFlowSeparable": idata_sep}
    )
    print("LOO-CV model comparison:")
    print(
        comparison[
            ["elpd_loo", "p_loo", "elpd_diff", "weight", "se", "warning"]
        ].to_string()
    )
except Exception as e:
    print(f"LOO not available (log-likelihood not stored): {e}")
    print("Tip: pass compute_log_likelihood=True to fit() to enable LOO/WAIC.")

# Fit with restrict_positive=False to allow negative rho values
sar_flow_neg = SARFlow(
    y_vec,
    G,
    X,
    col_names=col_names,
    logdet_method="traces",
    restrict_positive=False,  # Uniform(-1,1) priors + stability potential
    miter=20,
    trace_seed=0,
)
idata_neg = sar_flow_neg.fit(
    draws=800,
    tune=1000,
    chains=4,
    target_accept=0.95,
    random_seed=42,
    progressbar=True,
)

summary_neg = sar_flow_neg.summary()
print("=== SARFlow (restrict_positive=False) ===")
print(summary_neg[["mean", "sd", "hdi_3%", "hdi_97%", "r_hat"]].to_string())

from bayespecon import (
    bayesian_lm_flow_dest_test,
    bayesian_lm_flow_orig_test,
    bayesian_lm_flow_network_test,
    bayesian_lm_flow_joint_test,
    bayesian_lm_flow_intra_test,
)

# All marginal + joint + intra tests, evaluated at the OLSFlow null draws.
flow_lm_tests = {
    "rho_d (dest)":     bayesian_lm_flow_dest_test(ols_flow),
    "rho_o (orig)":     bayesian_lm_flow_orig_test(ols_flow),
    "rho_w (network)":  bayesian_lm_flow_network_test(ols_flow),
    "joint (d,o,w)":    bayesian_lm_flow_joint_test(ols_flow),
    "intra block":      bayesian_lm_flow_intra_test(ols_flow),
}

import pandas as pd
pd.DataFrame(
    [(k, r.df, r.mean, r.bayes_pvalue) for k, r in flow_lm_tests.items()],
    columns=["test", "df", "LM mean", "Bayes p"],
).round(4)

from bayespecon import (
    bayesian_robust_lm_flow_dest_test,
    bayesian_robust_lm_flow_orig_test,
    bayesian_robust_lm_flow_network_test,
)

robust_tests = {
    "rho_d | (rho_o, rho_w)": bayesian_robust_lm_flow_dest_test(sar_flow),
    "rho_o | (rho_d, rho_w)": bayesian_robust_lm_flow_orig_test(sar_flow),
    "rho_w | (rho_d, rho_o)": bayesian_robust_lm_flow_network_test(sar_flow),
}

pd.DataFrame(
    [(k, r.mean, r.bayes_pvalue) for k, r in robust_tests.items()],
    columns=["test", "LM* mean", "Bayes p"],
).round(4)