Logdet Method Profiling Across Matrix Sizes¶

import time
from dataclasses import dataclass

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pytensor
import pytensor.tensor as pt
from libpysal import graph

from bayespecon import dgp
from bayespecon.logdet import make_logdet_fn

def make_grid_w(n_side: int) -> np.ndarray:
    """Create a row-standardized rook-contiguity matrix from an n_side x n_side polygon grid.

    Uses ``dgp.simulate_sar`` with ``create_gdf=True`` to generate the polygon
    geometry, then builds a contiguity graph from the returned GeoDataFrame.
    """
    gdf = dgp.simulate_sar(n=n_side, create_gdf=True)
    W = (
        graph.Graph.build_contiguity(gdf, rook=True)
        .transform("r")
        .sparse.toarray()
        .astype(np.float64)
    )
    return W


def compile_logdet_callable(W: np.ndarray, method: str, rho_min: float, rho_max: float):
    """Return a compiled callable f(rho) and its setup time in seconds."""
    t0 = time.perf_counter()
    rho = pt.scalar("rho")
    expr = make_logdet_fn(W, method=method, rho_min=rho_min, rho_max=rho_max)(rho)
    fn = pytensor.function([rho], expr)
    setup_s = time.perf_counter() - t0
    return fn, setup_s


def bench_eval_seconds(fn, rhos: np.ndarray, repeats: int = 5) -> float:
    """Median per-call evaluation latency in microseconds."""
    run_times = []
    for _ in range(repeats):
        t0 = time.perf_counter()
        for r in rhos:
            _ = fn(float(r))
        elapsed = time.perf_counter() - t0
        run_times.append(elapsed / len(rhos))
    return float(np.median(run_times))

@dataclass
class ProfileConfig:
    # Grid side lengths; obs count = n_side². e.g. 6→36, 9→81, 11→121, 13→169, 16→256.
    sizes: tuple[int, ...] = (10, 20, 25, 40, 50, 75)
    methods: tuple[str, ...] = (
        "exact",
        "eigenvalue",
        "grid_dense",
        "grid_sparse",
        "sparse_spline",
        "trace_mc",
        "grid_mc",
        "grid_ilu",
        "chebyshev",
    )
    method_max_n: dict = None
    eval_points: int = 80
    eval_repeats: int = 3
    seed: int = 2026

    def __post_init__(self):
        if self.method_max_n is None:
            self.method_max_n = {
                "exact": 20,
                "eigenvalue": 75,
                "grid_dense": 40,
                "grid_sparse": 75,
                "sparse_spline": 75,
                "trace_mc": 75,
                "grid_ilu": 75,
                "chebyshev": 75,
                "grid_mc": 75,
            }


cfg = ProfileConfig()

# spline/mc are defined on nonnegative rho ranges; others can use symmetric ranges.
method_rho_bounds = {
    "exact": (-0.95, 0.95),
    "eigenvalue": (-0.95, 0.95),
    "grid_dense": (-0.95, 0.95),
    "grid_sparse": (-0.95, 0.95),
    "grid_ilu": (-0.95, 0.95),
    "sparse_spline": (1e-5, 0.95),
    "trace_mc": (1e-5, 0.95),
    "chebyshev": (-0.95, 0.95),
    "grid_mc": (1e-5, 0.95),
}

results = []
skipped = []

for n in cfg.sizes:
    W = make_grid_w(n_side=n)
    print(f"Profiling n_side={n} ({n * n} obs)...")
    for method in cfg.methods:
        if n > cfg.method_max_n[method]:
            skipped.append(
                {
                    "n_side": n,
                    "n_obs": n * n,
                    "method": method,
                    "reason": "above method_max_n cap",
                }
            )
            continue

        rho_min, rho_max = method_rho_bounds[method]
        rho_grid = np.linspace(rho_min, rho_max, cfg.eval_points)

        try:
            fn, setup_s = compile_logdet_callable(
                W, method, rho_min=rho_min, rho_max=rho_max
            )
            eval_s = bench_eval_seconds(fn, rho_grid, repeats=cfg.eval_repeats)
            results.append(
                {
                    "n_side": n,
                    "n_obs": n * n,
                    "method": method,
                    "rho_min": rho_min,
                    "rho_max": rho_max,
                    "setup_ms": 1e3 * setup_s,
                    "eval_us": 1e6 * eval_s,
                }
            )
        except Exception as exc:
            skipped.append(
                {
                    "n_side": n,
                    "n_obs": n * n,
                    "method": method,
                    "reason": f"failed: {type(exc).__name__}: {exc}",
                }
            )

res = pd.DataFrame(results).sort_values(["method", "n_obs"]).reset_index(drop=True)
if not res.empty:
    res["total_ms"] = res["setup_ms"] + (res["eval_us"] * cfg.eval_points / 1e3)
res

if res.empty:
    raise RuntimeError("No profiling results were generated.")

fig, axes = plt.subplots(1, 3, figsize=(16, 4.8), constrained_layout=True)

for method, grp in res.groupby("method"):
    grp = grp.sort_values("n_obs")
    axes[0].plot(grp["n_obs"], grp["setup_ms"], marker="o", label=method)
    axes[1].plot(grp["n_obs"], grp["eval_us"], marker="o", label=method)
    axes[2].plot(grp["n_obs"], grp["total_ms"], marker="o", label=method)

axes[0].set_title("Setup Time vs n")
axes[0].set_xlabel("observations")
axes[0].set_ylabel("setup time (ms)")
axes[0].set_yscale("log")
axes[0].grid(True, alpha=0.3)

axes[1].set_title("Evaluation Time vs n")
axes[1].set_xlabel("observations")
axes[1].set_ylabel("time per rho eval (us)")
axes[1].set_yscale("log")
axes[1].grid(True, alpha=0.3)

axes[2].set_title(f"Total Time vs n (setup + {cfg.eval_points} evals)")
axes[2].set_xlabel("observations")
axes[2].set_ylabel("total time (ms)")
axes[2].set_yscale("log")
axes[2].grid(True, alpha=0.3)

for ax in axes:
    ax.legend()

plt.show()

summary = res.pivot_table(
    index="n_obs", columns="method", values=["setup_ms", "eval_us", "total_ms"]
).sort_index()
display(summary)

if skipped:
    skipped_df = (
        pd.DataFrame(skipped).sort_values(["n_side", "method"]).reset_index(drop=True)
    )
    print("Skipped combinations (due to safety caps or failures):")
    display(skipped_df)

from bayespecon import SAR


def simulate_sar_data(n_side: int = 16, seed: int = 2026):
    """Simulate SAR data on an n_side x n_side polygon grid using the DGP module."""
    rng = np.random.default_rng(seed)
    beta_true = np.array([1.0, 0.8, -0.5], dtype=np.float64)
    rho_true = 0.35
    sigma_true = 0.7

    gdf = dgp.simulate_sar(
        n=n_side,
        rho=rho_true,
        beta=beta_true,
        sigma=sigma_true,
        rng=rng,
        create_gdf=True,
    )
    # Keep the SAR parameterization consistent with model assumptions.
    W_graph = graph.Graph.build_contiguity(gdf, rook=True).transform("r")
    y = gdf["y"].to_numpy()
    X_cols = [c for c in gdf.columns if c.startswith("X_")]
    X = gdf[X_cols].to_numpy()
    return gdf, y, X, W_graph, rho_true, beta_true


def fit_sar_for_method(
    y,
    X,
    W,
    method: str,
    draws: int = 400,
    tune: int = 400,
    seed: int = 2026,
    rho_lower: float = 0.000001,
    rho_upper: float = 0.95,
):
    """Fit SAR with a specific logdet method and return posterior means + runtime."""
    t0 = time.perf_counter()

    model = SAR(
        y=y,
        X=X,
        W=W,
        logdet_method=method,
        priors={"rho_lower": rho_lower, "rho_upper": rho_upper},
    )
    idata = model.fit(
        draws=draws,
        tune=tune,
        chains=2,
        cores=1,
        random_seed=seed,
        target_accept=0.95,
        progressbar=False,
        compute_convergence_checks=False,
    )

    elapsed_s = time.perf_counter() - t0
    beta_mean = idata.posterior["beta"].mean(("chain", "draw")).to_numpy()
    rho_mean = float(idata.posterior["rho"].mean(("chain", "draw")).to_numpy())

    return {
        "method": method,
        "total_time_s": elapsed_s,
        "rho": rho_mean,
        "beta_0": float(beta_mean[0]),
        "beta_1": float(beta_mean[1]),
        "beta_2": float(beta_mean[2]),
    }


methods_for_model = [
    "eigenvalue",
    "sparse_spline",
    "trace_mc",
    "grid_mc",
    "grid_ilu",
    "chebyshev",
]
gdf_model, y_model, X_model, W_model, rho_true, beta_true = simulate_sar_data(n_side=25)

fig, ax = plt.subplots(1, 1, figsize=(7, 7))
gdf_model.plot(
    column="y", cmap="viridis", legend=True, linewidth=0.15, edgecolor="white", ax=ax
)
ax.set_title("Simulated y on 25x25 polygon grid")
ax.set_axis_off()
plt.show()

model_rows = []
for m in methods_for_model:
    print(f"Estimating SAR with method={m}...")
    try:
        model_rows.append(fit_sar_for_method(y_model, X_model, W_model, method=m))
    except Exception as exc:
        model_rows.append(
            {
                "method": m,
                "total_time_s": np.nan,
                "rho": np.nan,
                "beta_0": np.nan,
                "beta_1": np.nan,
                "beta_2": np.nan,
                "error": f"{type(exc).__name__}: {exc}",
            }
        )

coef_compare = pd.DataFrame(model_rows)
coef_compare = coef_compare.sort_values("total_time_s", na_position="last").reset_index(
    drop=True
)

coef_compare["rho_true"] = rho_true
coef_compare["abs_err_rho"] = (coef_compare["rho"] - rho_true).abs()
for j, btrue in enumerate(beta_true):
    coef_compare[f"beta_{j}_true"] = btrue
    coef_compare[f"abs_err_beta_{j}"] = (coef_compare[f"beta_{j}"] - btrue).abs()

if "eigenvalue" in coef_compare["method"].values:
    base = coef_compare.loc[
        coef_compare["method"] == "eigenvalue", ["rho", "beta_0", "beta_1", "beta_2"]
    ].iloc[0]
    for col in ["rho", "beta_0", "beta_1", "beta_2"]:
        coef_compare[f"delta_vs_eigen_{col}"] = coef_compare[col] - base[col]

display(coef_compare)

import arviz as az
import pandas as pd

# Assuming you have a list of inference data objects (idata) for each method,
# or you can refit/collect them as needed. If you only have the summary DataFrames,
# adapt the code to use those directly.

ess_rows = []
for i, m in enumerate(methods_for_model):
    model = SAR(
        y=y_model,
        X=X_model,
        W=W_model,
        logdet_method=m,
        priors={"rho_lower": 1e-5, "rho_upper": 0.95},
    )
    idata = model.fit(
        draws=400,
        tune=400,
        chains=2,
        cores=1,
        random_seed=2026,
        target_accept=0.95,
        progressbar=False,
        compute_convergence_checks=False,
    )
    summary = az.summary(idata, var_names=["rho"])
    ess = summary.loc["rho", "ess_bulk"]
    ess_rows.append({"method": m, "ess_rho": ess})
    az.plot_trace(idata, var_names=["rho"], compact=True, legend=False)
    plt.suptitle(m)

ess_df = pd.DataFrame(ess_rows)
display(ess_df)

ess_df.set_index("method")["ess_rho"].plot.bar(
    ylabel="ESS (rho)", title="Effective Sample Size for rho by Logdet Method"
)