Pólya-Gamma Model Backend Profiling: numpy vs JAX (sparse)¶

import time
import warnings
from dataclasses import dataclass

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

from neighbayes import dgp
from neighbayes.models import SARLogit, SARNegBin

warnings.filterwarnings("ignore")

def build_w(gdf, kind: str, k: int = 6):
    """Row-standardized W. ``kind="sym"`` = rook contiguity (D-symmetrizable →
    Cholesky logdet); ``kind="asym"`` = directed KNN on centroids
    (non-D-symmetrizable → LU/aaa logdet)."""
    if kind == "sym":
        return graph.Graph.build_contiguity(gdf, rook=True).transform("r")
    pts = gdf.copy()
    pts["geometry"] = gdf.geometry.centroid
    return graph.Graph.build_knn(pts, k=k).transform("r")


def make_dataset(model: str, n_side: int, kind: str, seed: int = 2026):
    """Simulate one PG dataset on an ``n_side × n_side`` grid with the chosen W."""
    gdf = dgp.simulate_sar(n=n_side, create_gdf=True, seed=seed)
    W = build_w(gdf, kind)
    sim = (dgp.simulate_sar_logit if model == "logit" else dgp.simulate_sar_negbin)(
        W=W, gdf=gdf, rho=0.4, seed=seed
    )
    y = sim["y"].to_numpy() if hasattr(sim["y"], "to_numpy") else np.asarray(sim["y"])
    X_cols = [c for c in sim.columns if c.startswith("X_")]
    X = sim[X_cols].to_numpy()
    return y, X, W


def fit_and_time(model: str, y, X, W, backend: str, draws, tune, chains, seed=2026):
    """Fit one PG model on one backend; return timing + ESS + resolved logdet method."""
    Model = SARLogit if model == "logit" else SARNegBin
    m = Model(y=y, X=X, W=W)
    t0 = time.perf_counter()
    idata = m.fit(
        draws=draws,
        tune=tune,
        chains=chains,
        random_seed=seed,
        progressbar=False,
        gibbs_backend=backend,
        sampler="gibbs",
    )
    elapsed = time.perf_counter() - t0
    ess = float(az.summary(idata, var_names=["rho"]).loc["rho", "ess_bulk"])
    rho_hat = float(idata.posterior["rho"].mean())
    return {
        "total_s": elapsed,
        "ess_rho": ess,
        "ess_per_s": ess / elapsed,
        "rho_hat": rho_hat,
        "logdet_method": m._resolved_logdet_method,
    }

@dataclass
class Config:
    # grid sides -> obs = side²: 20→400, 30→900, 45→2025, 71→5041 (≥5000)
    sides: tuple = (20, 30, 45, 71)
    models: tuple = ("logit", "nb")
    w_kinds: tuple = ("sym", "asym")
    backends: tuple = ("numpy", "jax")
    draws: int = 400
    tune: int = 400
    chains: int = 2
    seed: int = 2026


cfg = Config()
rows = []
for side in cfg.sides:
    n = side * side
    for model in cfg.models:
        for kind in cfg.w_kinds:
            y, X, W = make_dataset(model, side, kind, seed=cfg.seed)
            for backend in cfg.backends:
                print(
                    f"n={n:5d} | {model:5s} | W={kind:4s} | backend={backend} ...",
                    flush=True,
                )
                try:
                    r = fit_and_time(
                        model,
                        y,
                        X,
                        W,
                        backend,
                        cfg.draws,
                        cfg.tune,
                        cfg.chains,
                        cfg.seed,
                    )
                    r.update(n_obs=n, model=model, w_kind=kind, backend=backend)
                    rows.append(r)
                except Exception as exc:  # keep the sweep going
                    print(f"    FAILED: {type(exc).__name__}: {exc}")

res = pd.DataFrame(rows)
res

# Confirm the two W regimes land on the two exact logdet methods.
if not res.empty:
    method_map = (
        res.groupby(["w_kind"])["logdet_method"].unique().apply(sorted).to_dict()
    )
    print("Resolved logdet method by W regime:")
    for kind, methods in method_map.items():
        print(f"  W={kind:4s} -> {methods}")
    display(
        res.pivot_table(
            index=["model", "n_obs"], columns=["w_kind", "backend"], values="total_s"
        ).round(2)
    )

if not res.empty:
    fig, axes = plt.subplots(2, 2, figsize=(13, 9), constrained_layout=True)
    styles = {"numpy": dict(marker="o", ls="-"), "jax": dict(marker="s", ls="--")}
    colors = {"sym": "#1f77b4", "asym": "#d62728"}
    for j, model in enumerate(("logit", "nb")):
        for backend in cfg.backends:
            for kind in cfg.w_kinds:
                sub = res[
                    (res.model == model)
                    & (res.backend == backend)
                    & (res.w_kind == kind)
                ].sort_values("n_obs")
                if sub.empty:
                    continue
                lbl = f"{backend}·{kind}"
                axes[0, j].plot(
                    sub.n_obs,
                    sub.total_s,
                    color=colors[kind],
                    alpha=0.9 if backend == "jax" else 0.5,
                    label=lbl,
                    **styles[backend],
                )
                axes[1, j].plot(
                    sub.n_obs,
                    sub.ess_per_s,
                    color=colors[kind],
                    alpha=0.9 if backend == "jax" else 0.5,
                    label=lbl,
                    **styles[backend],
                )
        axes[0, j].set(
            title=f"{model}: total time vs n",
            xlabel="observations",
            ylabel="total time (s)",
            yscale="log",
        )
        axes[1, j].set(
            title=f"{model}: ESS(rho)/sec vs n",
            xlabel="observations",
            ylabel="ESS per second",
            yscale="log",
        )
        for ax in (axes[0, j], axes[1, j]):
            ax.grid(True, alpha=0.3)
            ax.legend(fontsize=8)
    plt.show()

# Correctness: numpy and jax must agree (within MC noise) on the posterior.
# Deterministic side-check of the JAX exact logdet is in logdet_profiling.ipynb;
# here we confirm the end-to-end samplers agree on rho.
if not res.empty:
    agree = res.pivot_table(
        index=["model", "w_kind", "n_obs"], columns="backend", values="rho_hat"
    ).dropna()
    agree["abs_diff"] = (agree["numpy"] - agree["jax"]).abs()
    print("Posterior mean rho: numpy vs jax (should agree within MC noise)")
    display(agree.round(3))