How to choose and configure the Gibbs sampler¶

import time

import arviz as az
import geopandas as gpd
import libpysal
import numpy as np
import pandas as pd

from neighbayes.models import SAR, SDEM, SDM, SEM

gdf = gpd.read_file(libpysal.examples.get_path("columbus.shp"))
y = gdf["CRIME"].values.astype(float)
X = np.column_stack(
    [
        np.ones(len(y)),
        gdf["INC"].values.astype(float),
        gdf["HOVAL"].values.astype(float),
    ]
)
W = libpysal.graph.Graph.build_contiguity(gdf).transform("r")

COMMON = dict(draws=2000, tune=1000, chains=4, random_seed=42)


def fit_timed(model, **kw):
    start = time.perf_counter()
    idata = model.fit(progressbar=False, **kw)
    return idata, time.perf_counter() - start


def efficiency(label, idata, secs, params=("rho", "sigma")):
    ess = az.ess(idata, var_names=list(params))
    rhat = az.rhat(idata, var_names=list(params))
    row = {"seconds": secs}
    for p in params:
        row[f"ess {p}"] = float(ess[p])
        row[f"ess/sec {p}"] = float(ess[p]) / secs
    row["max rhat"] = float(max(float(rhat[p]) for p in params))
    return pd.Series(row, name=label)


print(f"n = {len(y)}, k = {X.shape[1]}")

idata_default, secs_default = fit_timed(SAR(y=y, X=X, W=W), **COMMON)

nuts_only = {"tree_depth", "diverging", "step_size"}
present = nuts_only & set(idata_default.sample_stats.data_vars)
print("sample_stats:", sorted(idata_default.sample_stats.data_vars))
print("ran NUTS" if present else "ran Gibbs")

idata_nuts, secs_nuts = fit_timed(
    SAR(y=y, X=X, W=W), sampler="nuts", target_accept=0.9, **COMMON
)

pd.DataFrame(
    {
        "Gibbs": [
            float(idata_default.posterior["rho"].mean()),
            float(idata_default.posterior["sigma"].mean()),
        ],
        "NUTS": [
            float(idata_nuts.posterior["rho"].mean()),
            float(idata_nuts.posterior["sigma"].mean()),
        ],
    },
    index=["rho", "sigma"],
).round(4)

pd.DataFrame(
    [
        efficiency("Gibbs", idata_default, secs_default),
        efficiency("NUTS", idata_nuts, secs_nuts),
    ]
).round(3)

try:
    SAR(y=y, X=X, W=W).fit(sampler="gibbs", target_accept=0.9, draws=10, tune=10)
except TypeError as err:
    print(f"TypeError: {err}")

pd.DataFrame(
    [
        efficiency(
            backend,
            *fit_timed(
                SAR(y=y, X=X, W=W), sampler="gibbs", gibbs_backend=backend, **COMMON
            ),
        )
        for backend in ("numpy", "jax")
    ]
).round(3)

idata_tuned, secs_tuned = fit_timed(
    SAR(y=y, X=X, W=W),
    sampler="gibbs",
    gibbs_backend="numpy",
    slice_width=0.25,
    n_jobs=-1,
    thin=2,
    **COMMON,
)
print(f"draws kept per chain with thin=2: {idata_tuned.posterior.sizes['draw']}")
efficiency("slice_width=0.25, thin=2", idata_tuned, secs_tuned).round(3)

# The constructor reports the valid names, so this cannot drift out of date.
try:
    SAR(y=y, X=X, W=W, logdet_method="list-them-please")
except ValueError as err:
    print(err)

pd.DataFrame(
    [
        pd.Series(
            {
                "rho mean": float(idata.posterior["rho"].mean()),
                "rho sd": float(idata.posterior["rho"].std()),
                "seconds": secs,
                "ess rho": float(az.ess(idata, var_names=["rho"])["rho"]),
            },
            name=method,
        )
        for method, (idata, secs) in (
            (
                m,
                fit_timed(
                    SAR(y=y, X=X, W=W, logdet_method=m), sampler="gibbs", **COMMON
                ),
            )
            for m in ("eigenvalue", "chebyshev", "cheb_stochastic")
        )
    ]
).round(4)

pd.DataFrame(
    [
        pd.Series(
            {
                "spatial parameter": param,
                "posterior mean": float(idata.posterior[param].mean()),
                "n coefficients": idata.posterior["beta"].sizes["coefficient"],
                "ess": float(az.ess(idata, var_names=[param])[param]),
                "rhat": float(az.rhat(idata, var_names=[param])[param]),
            },
            name=name,
        )
        for name, param, idata in (
            (n, p, fit_timed(cls(y=y, X=X, W=W), sampler="gibbs", **COMMON)[0])
            for n, cls, p in (
                ("SAR", SAR, "rho"),
                ("SDM", SDM, "rho"),
                ("SEM", SEM, "lam"),
                ("SDEM", SDEM, "lam"),
            )
        )
    ]
).round(4)

az.summary(idata_default, var_names=["rho", "sigma"]).round(4)

model_diag = SAR(y=y, X=X, W=W)
model_diag.fit(sampler="gibbs", progressbar=False, **COMMON)
model_diag.spatial_diagnostics_decision()