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
import scipy.sparse as sp
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,
    trace_estimator: str | None = None,
    trace_k: int | None = None,
):
    """Return a compiled callable f(rho) and its setup time in seconds."""
    trace_kwargs = {}
    if trace_estimator is not None:
        trace_kwargs["trace_estimator"] = trace_estimator
    if trace_k is not None:
        trace_kwargs["trace_k"] = trace_k

    t0 = time.perf_counter()
    rho = pt.scalar("rho")
    expr = make_logdet_fn(
        W,
        method=method,
        rho_min=rho_min,
        rho_max=rho_max,
        **trace_kwargs,
    )(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. 10→100, 20→400, 75→5625.
    sizes: tuple[int, ...] = (10, 20, 25, 40, 50, 75)
    method_specs: tuple[dict, ...] = (
        {
            "label": "eigenvalue",
            "method": "eigenvalue",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "grid_sparse",
            "method": "grid_sparse",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "sparse_spline",
            "method": "sparse_spline",
            "rho_min": 1e-5,
            "rho_max": 0.95,
        },
        {"label": "grid_ilu", "method": "grid_ilu", "rho_min": -0.95, "rho_max": 0.95},
        {
            "label": "chebyshev",
            "method": "chebyshev",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "chebyshev+hutchinson",
            "method": "chebyshev",
            "trace_estimator": "hutchinson",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "chebyshev+hutchpp",
            "method": "chebyshev",
            "trace_estimator": "hutchpp",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "chebyshev+xtrace",
            "method": "chebyshev",
            "trace_estimator": "xtrace",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
    )
    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 = {spec["label"]: 75 for spec in self.method_specs}


cfg = ProfileConfig()

results = []
skipped = []

for n in cfg.sizes:
    W = make_grid_w(n_side=n)
    print(f"Profiling n_side={n} ({n * n} obs)...")
    for spec in cfg.method_specs:
        label = spec["label"]
        method = spec["method"]
        trace_estimator = spec.get("trace_estimator")
        trace_k = spec.get("trace_k")

        if n > cfg.method_max_n[label]:
            skipped.append(
                {
                    "n_side": n,
                    "n_obs": n * n,
                    "method": label,
                    "reason": "above method_max_n cap",
                }
            )
            continue

        rho_min = spec["rho_min"]
        rho_max = spec["rho_max"]
        rho_grid = np.linspace(rho_min, rho_max, cfg.eval_points)

        try:
            fn, setup_s = compile_logdet_callable(
                W,
                method=method,
                rho_min=rho_min,
                rho_max=rho_max,
                trace_estimator=trace_estimator,
                trace_k=trace_k,
            )
            eval_s = bench_eval_seconds(fn, rho_grid, repeats=cfg.eval_repeats)
            results.append(
                {
                    "n_side": n,
                    "n_obs": n * n,
                    "method": label,
                    "logdet_method": method,
                    "trace_estimator": trace_estimator or "",
                    "trace_k": trace_k,
                    "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": label,
                    "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.")

# Distinct color + marker per profiled configuration.
method_styles = {
    "eigenvalue": {"color": "#ff7f0e", "marker": "s", "linestyle": "-"},
    "grid_sparse": {"color": "#d62728", "marker": "v", "linestyle": "-"},
    "sparse_spline": {"color": "#9467bd", "marker": "D", "linestyle": "-"},
    "grid_ilu": {"color": "#8c564b", "marker": "P", "linestyle": "-"},
    "chebyshev": {"color": "#000000", "marker": "*", "linestyle": "-"},
    "chebyshev+hutchinson": {"color": "#17becf", "marker": "o", "linestyle": "--"},
    "chebyshev+hutchpp": {"color": "#bcbd22", "marker": "s", "linestyle": "--"},
    "chebyshev+xtrace": {"color": "#7f7f7f", "marker": "^", "linestyle": "--"},
}

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

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

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)

# Single shared legend to the right of the figure so it doesn't crowd the panels.
handles, labels = axes[0].get_legend_handles_labels()
fig.legend(
    handles,
    labels,
    loc="center left",
    bbox_to_anchor=(1.0, 0.5),
    frameon=False,
    title="configuration",
)

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._logdet import make_logdet_numpy_fn

# Use the same W from the last profiling run (or regenerate)
W_accuracy = make_grid_w(n_side=25)
W_sp = sp.csr_matrix(W_accuracy)

# Exact logdet via eigenvalues
W_dense = W_sp.toarray()
eigs = np.linalg.eigvals(W_dense)
rho_grid_acc = np.linspace(-0.95, 0.95, 200)
exact_logdet = np.array([np.sum(np.log(1 - r * eigs)) for r in rho_grid_acc])

trace_configs = {
    "chebyshev+hutchinson": "hutchinson",
    "chebyshev+hutchpp": "hutchpp",
    "chebyshev+xtrace": "xtrace",
}

acc_rows = []
for name, estimator in trace_configs.items():
    fn = make_logdet_numpy_fn(
        W_sp,
        None,
        method="chebyshev",
        rho_min=-0.95,
        rho_max=0.95,
        trace_estimator=estimator,
    )
    approx = np.array([fn(r) for r in rho_grid_acc])
    err = np.abs(approx - exact_logdet)
    acc_rows.append(
        {
            "method": name,
            "logdet_method": "chebyshev",
            "trace_estimator": estimator,
            "max_err": err.max(),
            "mean_err": err.mean(),
            "rmse": np.sqrt((err**2).mean()),
        }
    )

acc_df = pd.DataFrame(acc_rows).sort_values("rmse").reset_index(drop=True)
display(acc_df)

fig, ax = plt.subplots(figsize=(10, 5))
for name, estimator in trace_configs.items():
    fn = make_logdet_numpy_fn(
        W_sp,
        None,
        method="chebyshev",
        rho_min=-0.95,
        rho_max=0.95,
        trace_estimator=estimator,
    )
    approx = np.array([fn(r) for r in rho_grid_acc])
    ax.plot(rho_grid_acc, approx - exact_logdet, label=name, alpha=0.7)
ax.axhline(0, color="black", linestyle="--", linewidth=0.8)
ax.set_xlabel("rho")
ax.set_ylabel("approximation error")
ax.set_title("Logdet Approximation Error vs Exact (eigenvalue)")
ax.legend()
ax.grid(True, alpha=0.3)
plt.show()

from bayespecon import SAR


def simulate_sar_data(n_side: int = 25, 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,
    label: str | None = None,
    trace_estimator: str | None = None,
    trace_k: int | None = None,
    draws: int = 1000,
    tune: int = 1000,
    seed: int = 2026,
):
    """Fit SAR with a specific logdet configuration and return posterior means + runtime."""
    trace_kwargs = {}
    if trace_estimator is not None:
        trace_kwargs["trace_estimator"] = trace_estimator
    if trace_k is not None:
        trace_kwargs["trace_k"] = trace_k

    t0 = time.perf_counter()

    model = SAR(
        y=y,
        X=X,
        W=W,
        logdet_method=method,
        **trace_kwargs,
    )
    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": label or method,
        "logdet_method": method,
        "trace_estimator": trace_estimator or "",
        "trace_k": trace_k,
        "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 = [
    {"label": "eigenvalue", "method": "eigenvalue"},
    {"label": "chebyshev", "method": "chebyshev"},
    {
        "label": "chebyshev+hutchinson",
        "method": "chebyshev",
        "trace_estimator": "hutchinson",
    },
    {"label": "chebyshev+hutchpp", "method": "chebyshev", "trace_estimator": "hutchpp"},
    {"label": "chebyshev+xtrace", "method": "chebyshev", "trace_estimator": "xtrace"},
]
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 spec in methods_for_model:
    label = spec["label"]
    print(f"Estimating SAR with {label}...")
    try:
        model_rows.append(
            fit_sar_for_method(
                y_model,
                X_model,
                W_model,
                method=spec["method"],
                label=label,
                trace_estimator=spec.get("trace_estimator"),
                trace_k=spec.get("trace_k"),
            )
        )
    except Exception as exc:
        model_rows.append(
            {
                "method": label,
                "logdet_method": spec["method"],
                "trace_estimator": spec.get("trace_estimator", ""),
                "trace_k": spec.get("trace_k"),
                "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]

coef_compare

import arviz as az
import pandas as pd

# Refit each logdet configuration and compare effective sample sizes for rho.
ess_rows = []
for spec in methods_for_model:
    label = spec["label"]
    trace_kwargs = {}
    if spec.get("trace_estimator") is not None:
        trace_kwargs["trace_estimator"] = spec["trace_estimator"]
    if spec.get("trace_k") is not None:
        trace_kwargs["trace_k"] = spec["trace_k"]

    model = SAR(
        y=y_model,
        X=X_model,
        W=W_model,
        logdet_method=spec["method"],
        priors={"rho_lower": -1, "rho_upper": 1},
        **trace_kwargs,
    )
    idata = model.fit(
        draws=1000,
        tune=1000,
        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": label,
            "logdet_method": spec["method"],
            "trace_estimator": spec.get("trace_estimator", ""),
            "trace_k": spec.get("trace_k"),
            "ess_rho": ess,
        }
    )
    az.plot_trace(idata, var_names=["rho"], compact=True, legend=False)
    plt.suptitle(label)

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 Configuration"
)