How to estimate panel flow models

Origin–destination flows observed over several periods — trade by year, commuting by quarter. Four classes: SARFlowPanel and SARFlowSeparablePanel for continuous flows, SARNegBinFlowPanel and SARNegBinFlowSeparablePanel for counts.

Two things differ from the cross-sectional case and both will bite if missed: you must pass T, and the Gaussian classes expect the outcome on the latent SAR scale, so a lognormal outcome goes in as log(y). Equations and constructor arguments are in Supported Models.

import contextlib
import io
import time

import arviz as az
import numpy as np
import pandas as pd

from neighbayes.dgp.flows import (
    generate_panel_flow_data,
    generate_panel_flow_data_separable,
    generate_panel_negbin_flow_data,
    generate_panel_negbin_flow_data_separable,
)
from neighbayes.models import (
    SARFlowPanel,
    SARFlowSeparablePanel,
    SARNegBinFlowPanel,
    SARNegBinFlowSeparablePanel,
)

T = 3
RHO_D, RHO_O, RHO_W = 0.20, 0.15, 0.05
DGP = dict(n=12, T=T, seed=0, rho_d=RHO_D, rho_o=RHO_O, beta_d=[1.0], beta_o=[0.4])
FIT = dict(draws=800, tune=800, chains=4, random_seed=0, progressbar=False)


def fit_timed(model, **kw):
    start = time.perf_counter()
    with (
        contextlib.redirect_stdout(io.StringIO()),
        contextlib.redirect_stderr(io.StringIO()),
    ):
        idata = model.fit(**kw)
    return idata, time.perf_counter() - start


def report(label, idata, secs, params=("rho_d", "rho_o")):
    ess = az.ess(idata, var_names=list(params))
    rhat = az.rhat(idata, var_names=list(params))
    row = {"seconds": secs}
    for p in params:
        row[p] = float(idata.posterior[p].mean())
    row["min ess"] = float(min(float(ess[p]) for p in params))
    row["max rhat"] = float(max(float(rhat[p]) for p in params))
    return pd.Series(row, name=label)

/home/runner/micromamba/envs/test/lib/python3.14/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Fit continuous flows

The DGP returns lognormal flows, so the model gets log(y). effects=0 fits a pooled panel; the flow panel classes take the same effects argument as the other panel families.

data = generate_panel_flow_data(rho_w=RHO_W, sigma=0.5, **DGP)
G, X, names = data["G"], data["X"], data["col_names"]

print(f"periods      : {T}")
print(f"flows/period : {len(data['y']) // T}")
print(f"true rho     : rho_d={RHO_D}, rho_o={RHO_O}, rho_w={RHO_W}")

model_free = SARFlowPanel(
    y=np.log(data["y"]), W=G, X=X, T=T, col_names=names, effects=0
)
idata_free, secs_free = fit_timed(model_free, **FIT)

az.summary(idata_free, var_names=["rho_d", "rho_o", "rho_w", "sigma"]).round(3)

periods      : 3
flows/period : 144
true rho     : rho_d=0.2, rho_o=0.15, rho_w=0.05
mean sd hdi_3% hdi_97% mcse_mean mcse_sd ess_bulk ess_tail r_hat
rho_d 0.203 0.053 0.108 0.305 0.006 0.002 82.0 204.0 1.07
rho_o 0.239 0.060 0.136 0.356 0.009 0.004 42.0 150.0 1.09
rho_w 0.036 0.031 0.000 0.098 0.002 0.002 238.0 226.0 1.02
sigma 0.689 0.024 0.644 0.733 0.000 0.000 3145.0 2730.0 1.00

Look at ess_bulk and r_hat on the three ρ before anything else. They will be poor, and that is the model rather than the sampler: with \(\rho_d\), \(\rho_o\) and \(\rho_w\) all free they trade off along a nearly flat ridge that one-at-a-time updates cannot traverse.

Prefer the separable model

SARFlowSeparablePanel pins \(\rho_w = -\rho_d\rho_o\), which removes the ridge outright. Same data, same draws.

data_sep = generate_panel_flow_data_separable(sigma=0.5, **DGP)
model_sep = SARFlowSeparablePanel(
    y=np.log(data_sep["y"]),
    W=data_sep["G"],
    X=data_sep["X"],
    T=T,
    col_names=data_sep["col_names"],
    effects=0,
)
idata_sep, secs_sep = fit_timed(model_sep, **FIT)

pd.DataFrame(
    [
        report("SARFlowPanel (3 free rho)", idata_free, secs_free),
        report("SARFlowSeparablePanel", idata_sep, secs_sep),
    ]
).round(3)
seconds rho_d rho_o min ess max rhat
SARFlowPanel (3 free rho) 15.595 0.203 0.239 41.869 1.095
SARFlowSeparablePanel 16.988 0.230 0.270 1921.145 1.001

The separable model buys well over an order of magnitude in effective samples for less time. Reach for it by default, and use the unrestricted model only when \(\rho_w\) is itself the object of interest — in which case expect to run far longer chains and to check convergence carefully.

You can confirm the restriction is exactly enforced rather than merely encouraged:

identity_error = float(
    np.max(
        np.abs(
            idata_sep.posterior["rho_w"].values
            + idata_sep.posterior["rho_d"].values * idata_sep.posterior["rho_o"].values
        )
    )
)
print(
    f"max |rho_w + rho_d*rho_o| = {identity_error:.2e}   (deterministic, not sampled)"
)

max |rho_w + rho_d*rho_o| = 0.00e+00   (deterministic, not sampled)

Fit flow counts

The negative binomial panel classes take integer counts directly — no log transform — and add a dispersion parameter alpha. They default to the reduced-form Pólya–Gamma Gibbs sampler, and are currently pooled-only (effects=0).

nb = generate_panel_negbin_flow_data(rho_w=RHO_W, alpha=2.0, **DGP)
nb_sep = generate_panel_negbin_flow_data_separable(alpha=2.0, **DGP)

model_nb = SARNegBinFlowPanel(
    y=nb["y"], W=nb["G"], X=nb["X"], T=T, col_names=nb["col_names"], effects=0
)
idata_nb, secs_nb = fit_timed(model_nb, **FIT)

model_nb_sep = SARNegBinFlowSeparablePanel(
    y=nb_sep["y"],
    W=nb_sep["G"],
    X=nb_sep["X"],
    T=T,
    col_names=nb_sep["col_names"],
    effects=0,
)
idata_nb_sep, secs_nb_sep = fit_timed(model_nb_sep, **FIT)

pd.DataFrame(
    [
        report("SARNegBinFlowPanel", idata_nb, secs_nb),
        report("SARNegBinFlowSeparablePanel", idata_nb_sep, secs_nb_sep),
    ]
).round(3)
seconds rho_d rho_o min ess max rhat
SARNegBinFlowPanel 48.217 0.127 0.122 1224.495 1.007
SARNegBinFlowSeparablePanel 25.377 0.180 -0.294 1509.175 1.002

What to check before trusting the output

r_hat and ess_bulk on every ρ, and — for the count models — on alpha too.

Warning

These panels are far too small to identify ρ The examples above use a 12-unit grid over 3 periods so the docs build in under a minute. At that size the spatial parameters of a flow model are not identified, and the point estimates below will not match the values the data were generated from — some come back with the wrong sign, and they move substantially from one random seed to the next while r_hat stays at 1.00 and ess_bulk runs into the thousands.

Good convergence diagnostics say the sampler explored the posterior faithfully. They say nothing about whether that posterior is concentrated near the truth. Read the credible intervals, notice how wide they are, and size your data accordingly — the count flow models want thousands of flows per period before ρ_d and ρ_o separate.

pd.DataFrame(
    [
        report("Gaussian, 3 free rho", idata_free, secs_free),
        report("Gaussian, separable", idata_sep, secs_sep),
        report("NegBin, 3 free rho", idata_nb, secs_nb),
        report("NegBin, separable", idata_nb_sep, secs_nb_sep),
    ]
).round(3)
seconds rho_d rho_o min ess max rhat
Gaussian, 3 free rho 15.595 0.203 0.239 41.869 1.095
Gaussian, separable 16.988 0.230 0.270 1921.145 1.001
NegBin, 3 free rho 48.217 0.127 0.122 1224.495 1.007
NegBin, separable 25.377 0.180 -0.294 1509.175 1.002

See also