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 neighbayes import dgp
from neighbayes._logdet import make_logdet_fn, make_logdet_numpy_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. 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": "cheb_stochastic",
            "method": "cheb_stochastic",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "slq",
            "method": "slq",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "chebyshev",
            "method": "chebyshev",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "cheb_cholesky",
            "method": "cheb_cholesky",
            "rho_min": -0.95,
            "rho_max": 0.95,
        },
        {
            "label": "aaa",
            "method": "aaa",
            "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"]

        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,
            )
            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,
                    "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": "-"},
    "cheb_stochastic": {"color": "#1f77b4", "marker": "D", "linestyle": "-"},
    "slq": {"color": "#2ca02c", "marker": "o", "linestyle": "-"},
    "chebyshev": {"color": "#000000", "marker": "*", "linestyle": "-"},
    "cheb_cholesky": {"color": "#d62728", "marker": "^", "linestyle": "-"},
    "aaa": {"color": "#9467bd", "marker": "v", "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 neighbayes._logdet import (
    slq_logdet_eval,
    slq_logdet_precompute,
)
from neighbayes._logdet._aaa import (
    aaa_logdet_eval,
    aaa_logdet_precompute,
)
from neighbayes._logdet._chol_cheb import (
    chol_cheb_logdet_eval,
    chol_cheb_logdet_precompute,
)

# 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(np.abs(1 - r * eigs))) for r in rho_grid_acc])

# Chebyshev approximation (coefficients built from exact eigenvalues when available)
cheb_fn = make_logdet_numpy_fn(
    W_sp,
    eigs,
    method="chebyshev",
    rho_min=-0.95,
    rho_max=0.95,
)
cheb_approx = np.array([cheb_fn(r) for r in rho_grid_acc])
cheb_err = np.abs(cheb_approx - exact_logdet)

# SLQ approximation (Arnoldi-based, ρ-independent precompute)
slq_pre = slq_logdet_precompute(
    W_sp, n_probes=20, lanczos_deg=30, rng=np.random.default_rng(0)
)
slq_approx = np.array([slq_logdet_eval(slq_pre, r) for r in rho_grid_acc])
slq_err = np.abs(slq_approx - exact_logdet)

# Cholesky-Chebyshev (exact via sparse Cholesky at Chebyshev nodes, SPD symmetrized)
# Use the empirical range [0.1, 0.8] where the method is most accurate
rho_grid_emp = np.linspace(0.1, 0.8, 200)
exact_logdet_emp = np.array(
    [np.sum(np.log(np.abs(1 - r * eigs))) for r in rho_grid_emp]
)
chol_pre = chol_cheb_logdet_precompute(W_sp, order=None, rho_min=0.1, rho_max=0.8)
chol_approx = np.array([chol_cheb_logdet_eval(chol_pre, r) for r in rho_grid_emp])
chol_err = np.abs(chol_approx - exact_logdet_emp)

# AAA rational approximation (sparse LU at adaptively-selected support points)
aaa_pre = aaa_logdet_precompute(W_sp, rho_min=0.1, rho_max=0.8)
aaa_approx = np.array([aaa_logdet_eval(aaa_pre, r) for r in rho_grid_emp])
aaa_err = np.abs(aaa_approx - exact_logdet_emp)

acc_df = pd.DataFrame(
    [
        {
            "method": "chebyshev",
            "logdet_method": "chebyshev",
            "max_err": cheb_err.max(),
            "mean_err": cheb_err.mean(),
            "rmse": np.sqrt((cheb_err**2).mean()),
        },
        {
            "method": "slq",
            "logdet_method": "slq",
            "max_err": slq_err.max(),
            "mean_err": slq_err.mean(),
            "rmse": np.sqrt((slq_err**2).mean()),
        },
        {
            "method": "cheb_cholesky",
            "logdet_method": "cheb_cholesky",
            "max_err": chol_err.max(),
            "mean_err": chol_err.mean(),
            "rmse": np.sqrt((chol_err**2).mean()),
        },
        {
            "method": "aaa",
            "logdet_method": "aaa",
            "max_err": aaa_err.max(),
            "mean_err": aaa_err.mean(),
            "rmse": np.sqrt((aaa_err**2).mean()),
        },
    ]
)
display(acc_df)

fig, axes = plt.subplots(1, 2, figsize=(16, 5), constrained_layout=True)

# Full range [-0.95, 0.95]
axes[0].plot(rho_grid_acc, cheb_approx - exact_logdet, label="chebyshev", alpha=0.7)
axes[0].plot(rho_grid_acc, slq_approx - exact_logdet, label="slq", alpha=0.7)
axes[0].axhline(0, color="black", linestyle="--", linewidth=0.8)
axes[0].set_xlabel("rho")
axes[0].set_ylabel("approximation error")
axes[0].set_title("Logdet Error vs Exact (full range [-0.95, 0.95])")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# Empirical range [0.1, 0.8]
axes[1].plot(
    rho_grid_emp,
    chol_approx - exact_logdet_emp,
    label="cheb_cholesky",
    alpha=0.7,
    color="#d62728",
)
axes[1].plot(
    rho_grid_emp, aaa_approx - exact_logdet_emp, label="aaa", alpha=0.7, color="#9467bd"
)
axes[1].axhline(0, color="black", linestyle="--", linewidth=0.8)
axes[1].set_xlabel("rho")
axes[1].set_ylabel("approximation error")
axes[1].set_title("Logdet Error vs Exact (empirical range [0.1, 0.8])")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.show()

from neighbayes.models 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,
    draws: int = 1000,
    tune: int = 1000,
    seed: int = 2026,
):
    """Fit SAR with a specific logdet configuration and return posterior means + runtime."""
    t0 = time.perf_counter()

    model = SAR(
        y=y,
        X=X,
        W=W,
        logdet_method=method,
    )
    idata = model.fit(
        draws=draws,
        tune=tune,
        chains=2,
        random_seed=seed,
        progressbar=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,
        "beta_mean": beta_mean,
        "rho_mean": rho_mean,
        "fit_seconds": elapsed_s,
    }


# Build the dataset once and sweep the methods we actually compare.
gdf_model, y_model, X_model, W_model, rho_true, beta_true = simulate_sar_data(n_side=25)

methods_for_model = [
    {"label": "eigenvalue", "method": "eigenvalue"},
    {"label": "slq", "method": "slq"},
    {"label": "chebyshev", "method": "chebyshev"},
    {"label": "cheb_stochastic", "method": "cheb_stochastic"},
    {"label": "cheb_cholesky", "method": "cheb_cholesky"},
    {"label": "aaa", "method": "aaa"},
]

fit_rows = []
for spec in methods_for_model:
    print(f"Fitting SAR with logdet_method={spec['method']} ...")
    fit_rows.append(
        fit_sar_for_method(
            y_model,
            X_model,
            W_model,
            method=spec["method"],
            label=spec["label"],
        )
    )


coef_df = pd.DataFrame(fit_rows)

display(coef_df)

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"]
    model = SAR(
        y=y_model,
        X=X_model,
        W=W_model,
        logdet_method=spec["method"],
        priors={"rho_lower": -1, "rho_upper": 1},
    )
    idata = model.fit(
        draws=1000,
        tune=1000,
        chains=2,
        random_seed=2026,
        progressbar=False,
    )
    summary = az.summary(idata, var_names=["rho"])
    ess = summary.loc["rho", "ess_bulk"]
    ess_rows.append(
        {
            "method": label,
            "logdet_method": spec["method"],
            "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"
)