How to estimate origin–destination flow models¶

import warnings

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

from neighbayes.dgp.flows import generate_flow_data, generate_flow_data_separable
from neighbayes.graph import flow_design_matrix
from neighbayes.models import SARFlow, SARFlowSeparable

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=20,
    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 = 20
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"].transform("r")  # 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,
    X,
    G,
    col_names=col_names,
    restrict_positive=True,  # Dirichlet stability prior,
)

idata_sar = sar_flow.fit(
    sampler="nuts",
    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,
    X_sep,
    G,
    col_names=cn_sep,
    # logdet_method=None (auto) selects the separable Kronecker method
)

idata_sep = sep_flow.fit(
    draws=1000,
    tune=1000,
    chains=4,
    target_accept=0.9,
    random_seed=42,
    progressbar=True,
    # Stored only on request (as in PyMC); the LOO comparison below needs it.
    idata_kwargs={"log_likelihood": 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 neighbayes.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"]),
    data_asym["X"],
    G_asym,
    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 neighbayes.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 neighbayes.models import OLSFlow

ols_flow = OLSFlow(y_vec, X, G, 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 neighbayes.diagnostics import spatial_mcmc_diagnostic

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

# Fit unrestricted SARFlow on the same separable data for a fair comparison
sar_flow_on_sep = SARFlow(
    y_sep,
    X_sep,
    G,
    col_names=cn_sep,
)
idata_sar_on_sep = sar_flow_on_sep.fit(
    draws=1000,
    tune=1000,
    chains=4,
    target_accept=0.9,
    random_seed=42,
    progressbar=True,
    idata_kwargs={"log_likelihood": 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 idata_kwargs={"log_likelihood": True} to fit() for LOO/WAIC.')

# Fit with restrict_positive=False to allow negative rho values
sar_flow_neg = SARFlow(
    y_vec,
    X,
    G,
    col_names=col_names,
    restrict_positive=False,  # Uniform(-1,1) priors + stability potential
)
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 neighbayes.diagnostics.lmtests import (
    bayesian_robust_lm_flow_dest_test,
    bayesian_robust_lm_flow_network_test,
    bayesian_robust_lm_flow_orig_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)