unitheory.py

import numpy as np
import random
import matplotlib.pyplot as plt
import scipy.sparse as sp
from pygsp import graphs

from gvg import GVG
from theory import GVG_Theory

np.random.seed(0)
random.seed(0)

# ============================================================
# Config
# ============================================================
R = 5
T = 1           # (optional) empirical MC draws; currently unused
N_SAMP = 2048   # endpoint samples for theory p_st estimate

dists = ["gaussian", "laplace", "uniform", "exponential"]

# Paper-ish style (keep it simple, deterministic)
plt.rcParams.update({
    "font.family": "serif",
    "font.serif": ["Times New Roman", "Times", "DejaVu Serif"],
    "mathtext.fontset": "cm",
    "font.size": 14,
    "axes.labelsize": 17,
    "axes.titlesize": 14,
    "legend.fontsize": 12,
    "xtick.labelsize": 12,
    "ytick.labelsize": 12,
    "lines.linewidth": 2.2,
    "lines.markersize": 6,
    "axes.grid": True,
    "savefig.dpi": 300,
})

color_map = {
    "gaussian": "#1f77b4",
    "laplace": "#ff7f0e",
    "uniform": "#2ca02c",
    "exponential": "#d62728",
}
marker_map = {
    "gaussian": "o",
    "laplace": "s",
    "uniform": "^",
    "exponential": "D",
}
label_map = {
    "gaussian": "Gaussian",
    "laplace": "Laplace",
    "uniform": "Uniform",
    "exponential": "Exponential",
}

# ============================================================
# Helpers
# ============================================================
def sample_full_signal(dist_name, N, seed):
    rng = np.random.default_rng(seed)

    if dist_name == "gaussian":
        return rng.standard_normal(N)

    elif dist_name == "laplace":
        b = 1.0 / np.sqrt(2.0)   # unit-variance Laplace
        return rng.laplace(0.0, b, size=N)

    elif dist_name == "uniform":
        c = np.sqrt(3.0)         # unit-variance Uniform[-c,c]
        return rng.uniform(-c, c, size=N)

    elif dist_name == "exponential":
        # Exp(1): mean=1, var=1 (keep as-is; affine invariance anyway)
        return rng.exponential(scale=1.0, size=N)

    else:
        raise ValueError(dist_name)


def theory_Edeg_and_hop_probs(gvg_th):
    """
    Compute:
      - Edeg[s] = sum_{v: d(s,v)<=R} p_{sv}
      - hop_mean[h] = average p_{sv} over all ordered pairs (s,v) with d(s,v)=h
    """
    dist = gvg_th.comp_dijkstra()
    N = dist.shape[0]

    xs, xv = gvg_th.sample_endp()

    Edeg = np.zeros(N, dtype=np.float64)
    hop_sums = np.zeros(gvg_th.R + 1, dtype=np.float64)
    hop_counts = np.zeros(gvg_th.R + 1, dtype=np.int64)

    for s in range(N):
        sv = np.where(np.isfinite(dist[s, :]) & (dist[s, :] > 0) & (dist[s, :] <= gvg_th.R))[0]
        if sv.size == 0:
            continue

        degs = 0.0
        for v in sv:
            h = int(dist[s, v])
            p = gvg_th.P_sv(dist, s, v, xs, xv)

            degs += p
            hop_sums[h] += p
            hop_counts[h] += 1

        Edeg[s] = degs

    hop_mean = np.zeros_like(hop_sums)
    for h in range(1, gvg_th.R + 1):
        if hop_counts[h] > 0:
            hop_mean[h] = hop_sums[h] / float(hop_counts[h])
        else:
            hop_mean[h] = np.nan

    return Edeg, float(np.mean(Edeg)), hop_mean, hop_counts, dist


def make_graph(graph_name, seed=0):
    """
    Returns: (G, W_bin_csr, title_str)
    Keep graphs unweighted & comparable (binary adjacency).
    """
    rng = np.random.default_rng(seed)

    if graph_name == "sensor_knn":
        # kNN Sensor network: N=1000, k=10
        N = 1000
        G = graphs.Sensor(N, k=5, seed=seed)
        title = r"Sensor $k$-NN Graph ($n=1,000, k=5$)"

    elif graph_name == "grid2d":
        # Use a moderate grid so shortest-path computations don't explode
        # N = m*m
        m = 50
        G = graphs.Grid2d(m, m)
        title = rf"2-D Grid Graph ($n_1={m}, n_2={m}$)"

    elif graph_name == "erdos_renyi":
        # ER: N=1000, p=0.01
        N = 1000
        p = 0.01
        G = graphs.ErdosRenyi(N, p, seed=seed)
        title = rf"Erdős-Rényi Graph ($n=1,000, p={p}$)"

    else:
        raise ValueError(graph_name)

    G.compute_laplacian()

    W = G.W.tocsr()
    W = (W > 0).astype(np.int8).tocsr()
    return G, W, title


def run_theory_for_graph(W):
    results = {}

    for dist_name in dists:
        print("\n" + "=" * 70)
        print("Distribution:", dist_name)
        print("=" * 70)

        gvg_th = GVG_Theory(W, R=R, x_dist=dist_name, n_samp=N_SAMP, seed=0)
        Edeg, avg_Edeg, hop_mean, hop_counts, dist = theory_Edeg_and_hop_probs(gvg_th)

        print("Theory avg E[deg] (mean over nodes):", avg_Edeg)
        print("Hop-wise mean visibility probabilities (ordered pairs s->t):")
        for h in range(1, R + 1):
            if hop_counts[h] > 0:
                print(f"  h={h}: mean p_h = {hop_mean[h]:.6f}   (pairs={hop_counts[h]})")
            else:
                print(f"  h={h}: mean p_h = NaN        (pairs=0)")

        results[dist_name] = {
            "Edeg": Edeg,
            "avg_Edeg": avg_Edeg,
            "hop_mean": hop_mean,
            "hop_counts": hop_counts,
        }

    return results


# ============================================================
# Main: run 3 graphs, make a 3-panel figure (for IEEE figure*)
# ============================================================
graph_specs = [
    ("sensor_knn",  "sensor"),
    ("grid2d",      "grid"),
    ("erdos_renyi", "er"),
]

all_results = {}
all_titles = {}

for graph_key, short in graph_specs:
    print("\n" + "#" * 80)
    print("Graph:", graph_key)
    print("#" * 80)

    G, W, title = make_graph(graph_key, seed=0)

    print("Num nodes:", G.N)
    print("Num edges:", W.nnz // 2)

    all_titles[graph_key] = title
    all_results[graph_key] = run_theory_for_graph(W)

# ============================================================
# Plot (3 horizontal panels, shared y, integer hops only)
# ============================================================
hs = np.arange(1, R + 1)

fig, axes = plt.subplots(
    1, 3,
    figsize=(12, 4),
    sharey=True
)

for ax, (graph_key, _) in zip(axes, graph_specs):
    results = all_results[graph_key]

    for dist_name in dists:
        hop_mean = results[dist_name]["hop_mean"]
        ax.plot(
            hs,
            hop_mean[1:R+1],
            marker=marker_map[dist_name],
            color=color_map[dist_name],
            label=label_map[dist_name],
        )

    ax.set_title(all_titles[graph_key])
    ax.set_xlim(1, R)
    ax.set_xticks(np.arange(1, R + 1))
    ax.set_ylim(0, 1.0)
    ax.set_yticks(np.arange(0.0, 1.01, 0.1))
    ax.tick_params(direction="in", top=True, right=True)

axes[0].set_ylabel(r"$\bar p(h)$")
fig.supxlabel(r"$h$")

# One legend for the whole figure
handles, labels = axes[0].get_legend_handles_labels()
fig.legend(
    handles, labels,
    loc="upper center",
    ncol=len(dists),
    frameon=True,
    fancybox=False,
    edgecolor="black",
    bbox_to_anchor=(0.5, 1.05)
)

# -----------------------------
# Refinements for paper layout
# -----------------------------
# 1) prevent endpoint markers from being clipped
for ax in axes.ravel():  # if you have axs array from subplots
    ax.set_xlim(0.85, R + 0.15)   # small padding left/right
    ax.set_ylim(-0.02, 1.02)      # tiny padding bottom/top
    ax.margins(x=0.02, y=0.02)    # extra safety (doesn't hurt)

    # 2) reduce tick label clash near bottom-left
    ax.tick_params(axis="x", pad=6)
    ax.tick_params(axis="y", pad=6)

# (if you want the range too)
# fig.supxlabel(r"Hop Distance ($1 \le h \le R$, $R=5$)", y=0.04)

# tighten layout and explicitly reserve a small bottom margin for supxlabel
fig.tight_layout()
fig.subplots_adjust(bottom=0.14, top=0.84)  # increase if still clipped; decrease if too much space

# save
fig.savefig("distr.pdf", bbox_inches="tight")
plt.show()