573 lines
20 KiB
Python
573 lines
20 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Benchmark visualization for per-feature KS / JSD / lag-1 diff + overall averages.
|
||
|
||
- Uses dummy data generators (AR(1) for continuous; Markov chain for discrete)
|
||
- Replace `real_df` and `gen_df` with your real/generated samples.
|
||
- Outputs publication-quality figures into ./benchmark_figs/
|
||
|
||
Dependencies:
|
||
pip install numpy pandas matplotlib scipy
|
||
(Optional for nicer label wrapping)
|
||
pip install textwrap3
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
from dataclasses import dataclass
|
||
from typing import Dict, List, Tuple
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
import matplotlib.pyplot as plt
|
||
from scipy.stats import ks_2samp
|
||
from scipy.spatial.distance import jensenshannon
|
||
|
||
|
||
# ----------------------------
|
||
# 1) Dummy data generation
|
||
# ----------------------------
|
||
|
||
def ar1_series(T: int, phi: float, sigma: float, mu: float = 0.0, seed: int | None = None) -> np.ndarray:
|
||
rng = np.random.default_rng(seed)
|
||
x = np.zeros(T, dtype=np.float64)
|
||
eps = rng.normal(0.0, sigma, size=T)
|
||
x[0] = mu + eps[0]
|
||
for t in range(1, T):
|
||
x[t] = mu + phi * (x[t - 1] - mu) + eps[t]
|
||
return x
|
||
|
||
def markov_chain(T: int, P: np.ndarray, pi0: np.ndarray, seed: int | None = None) -> np.ndarray:
|
||
"""Return integer states 0..K-1."""
|
||
rng = np.random.default_rng(seed)
|
||
K = P.shape[0]
|
||
states = np.zeros(T, dtype=np.int64)
|
||
states[0] = rng.choice(K, p=pi0)
|
||
for t in range(1, T):
|
||
states[t] = rng.choice(K, p=P[states[t - 1]])
|
||
return states
|
||
|
||
def make_dummy_real_gen(
|
||
T: int = 8000,
|
||
n_cont: int = 18,
|
||
n_disc: int = 6,
|
||
disc_vocab_sizes: List[int] | None = None,
|
||
seed: int = 7,
|
||
) -> Tuple[pd.DataFrame, pd.DataFrame, Dict[str, str], Dict[str, List[str]]]:
|
||
"""
|
||
Returns:
|
||
real_df, gen_df, feature_types, disc_vocab_map
|
||
where:
|
||
feature_types[col] in {"continuous","discrete"}
|
||
disc_vocab_map[col] = list of token strings (for discrete vars)
|
||
"""
|
||
rng = np.random.default_rng(seed)
|
||
|
||
cont_cols = [f"c{i:02d}" for i in range(n_cont)]
|
||
disc_cols = [f"d{i:02d}" for i in range(n_disc)]
|
||
|
||
if disc_vocab_sizes is None:
|
||
base = [3, 4, 5, 3, 6, 4]
|
||
disc_vocab_sizes = (base + [4] * max(0, n_disc - len(base)))[:n_disc]
|
||
elif len(disc_vocab_sizes) < n_disc:
|
||
# If you pass a shorter list, extend it safely
|
||
disc_vocab_sizes = list(disc_vocab_sizes) + [disc_vocab_sizes[-1]] * (n_disc - len(disc_vocab_sizes))
|
||
|
||
|
||
real = {}
|
||
gen = {}
|
||
|
||
# Continuous: AR(1) with slightly different parameters for generated
|
||
for i, col in enumerate(cont_cols):
|
||
phi_real = rng.uniform(0.75, 0.98)
|
||
sigma_real = rng.uniform(0.3, 1.2)
|
||
mu_real = rng.uniform(-1.0, 1.0)
|
||
|
||
# Generated is close but not identical
|
||
phi_gen = np.clip(phi_real + rng.normal(0, 0.02), 0.60, 0.995)
|
||
sigma_gen = max(0.05, sigma_real * rng.uniform(0.85, 1.15))
|
||
mu_gen = mu_real + rng.normal(0, 0.15)
|
||
|
||
real[col] = ar1_series(T, phi_real, sigma_real, mu_real, seed=seed + 100 + i)
|
||
gen[col] = ar1_series(T, phi_gen, sigma_gen, mu_gen, seed=seed + 200 + i)
|
||
|
||
# Discrete: Markov chains with slightly perturbed transition matrices
|
||
disc_vocab_map: Dict[str, List[str]] = {}
|
||
for j, col in enumerate(disc_cols):
|
||
K = disc_vocab_sizes[j]
|
||
disc_vocab_map[col] = [f"s{k}" for k in range(K)]
|
||
|
||
# Random but stable transition matrix for real
|
||
A = rng.uniform(0.1, 1.0, size=(K, K))
|
||
P_real = A / A.sum(axis=1, keepdims=True)
|
||
pi0 = rng.uniform(0.1, 1.0, size=K)
|
||
pi0 = pi0 / pi0.sum()
|
||
|
||
# Perturb for generated
|
||
noise = rng.normal(0, 0.06, size=(K, K))
|
||
P_gen = np.clip(P_real + noise, 1e-6, None)
|
||
P_gen = P_gen / P_gen.sum(axis=1, keepdims=True)
|
||
|
||
real_states = markov_chain(T, P_real, pi0, seed=seed + 300 + j)
|
||
gen_states = markov_chain(T, P_gen, pi0, seed=seed + 400 + j)
|
||
|
||
# Store as tokens (strings) to emphasize discrete nature
|
||
real[col] = np.array([disc_vocab_map[col][s] for s in real_states], dtype=object)
|
||
gen[col] = np.array([disc_vocab_map[col][s] for s in gen_states], dtype=object)
|
||
|
||
real_df = pd.DataFrame(real)
|
||
gen_df = pd.DataFrame(gen)
|
||
|
||
feature_types = {**{c: "continuous" for c in cont_cols}, **{d: "discrete" for d in disc_cols}}
|
||
return real_df, gen_df, feature_types, disc_vocab_map
|
||
|
||
|
||
# ----------------------------
|
||
# 2) Metric computation
|
||
# ----------------------------
|
||
|
||
def lag1_autocorr(x: np.ndarray) -> float:
|
||
"""Lag-1 Pearson autocorrelation (robust to constant arrays)."""
|
||
x = np.asarray(x, dtype=np.float64)
|
||
if x.size < 3:
|
||
return np.nan
|
||
x0 = x[:-1]
|
||
x1 = x[1:]
|
||
s0 = np.std(x0)
|
||
s1 = np.std(x1)
|
||
if s0 < 1e-12 or s1 < 1e-12:
|
||
return 0.0
|
||
return float(np.corrcoef(x0, x1)[0, 1])
|
||
|
||
def jsd_discrete(real_tokens: np.ndarray, gen_tokens: np.ndarray, vocab: List[str], base: float = 2.0) -> float:
|
||
"""Jensen–Shannon divergence for discrete distributions over vocab."""
|
||
# empirical pmf with Laplace smoothing to avoid zeros
|
||
eps = 1e-12
|
||
r = pd.Series(real_tokens).value_counts()
|
||
g = pd.Series(gen_tokens).value_counts()
|
||
p = np.array([r.get(v, 0) for v in vocab], dtype=np.float64)
|
||
q = np.array([g.get(v, 0) for v in vocab], dtype=np.float64)
|
||
p = p + eps
|
||
q = q + eps
|
||
p = p / p.sum()
|
||
q = q / q.sum()
|
||
# scipy returns sqrt(JS); square it to get JS divergence
|
||
return float(jensenshannon(p, q, base=base) ** 2)
|
||
|
||
@dataclass
|
||
class Metrics:
|
||
feature: str
|
||
ftype: str
|
||
ks: float | np.nan
|
||
jsd: float | np.nan
|
||
lag1_diff: float | np.nan
|
||
|
||
def compute_metrics(
|
||
real_df: pd.DataFrame,
|
||
gen_df: pd.DataFrame,
|
||
feature_types: Dict[str, str],
|
||
disc_vocab_map: Dict[str, List[str]] | None = None,
|
||
) -> pd.DataFrame:
|
||
rows: List[Metrics] = []
|
||
|
||
for col in real_df.columns:
|
||
ftype = feature_types[col]
|
||
r = real_df[col].to_numpy()
|
||
g = gen_df[col].to_numpy()
|
||
|
||
if ftype == "continuous":
|
||
# KS on raw values
|
||
ks_val = float(ks_2samp(r.astype(np.float64), g.astype(np.float64), alternative="two-sided").statistic)
|
||
# lag-1 diff
|
||
lag_r = lag1_autocorr(r)
|
||
lag_g = lag1_autocorr(g)
|
||
lag_diff = float(abs(lag_r - lag_g))
|
||
rows.append(Metrics(col, ftype, ks_val, np.nan, lag_diff))
|
||
|
||
elif ftype == "discrete":
|
||
# JSD on categorical pmf
|
||
if disc_vocab_map is None or col not in disc_vocab_map:
|
||
vocab = sorted(list(set(r).union(set(g))))
|
||
else:
|
||
vocab = disc_vocab_map[col]
|
||
jsd_val = jsd_discrete(r, g, vocab=vocab, base=2.0)
|
||
|
||
# optional: lag1 diff on integer encoding (captures “stickiness”; still interpretable)
|
||
# if you prefer lag1 only for continuous, set lag_diff=np.nan here.
|
||
mapping = {tok: i for i, tok in enumerate(vocab)}
|
||
r_int = np.array([mapping[t] for t in r], dtype=np.float64)
|
||
g_int = np.array([mapping[t] for t in g], dtype=np.float64)
|
||
lag_r = lag1_autocorr(r_int)
|
||
lag_g = lag1_autocorr(g_int)
|
||
lag_diff = float(abs(lag_r - lag_g))
|
||
|
||
rows.append(Metrics(col, ftype, np.nan, jsd_val, lag_diff))
|
||
else:
|
||
raise ValueError(f"Unknown feature type: {ftype}")
|
||
|
||
df = pd.DataFrame([r.__dict__ for r in rows])
|
||
|
||
# Overall (separate by metric meaning)
|
||
overall = {
|
||
"ks_avg_cont": float(df.loc[df.ftype == "continuous", "ks"].mean(skipna=True)),
|
||
"jsd_avg_disc": float(df.loc[df.ftype == "discrete", "jsd"].mean(skipna=True)),
|
||
"lag1_avg_all": float(df["lag1_diff"].mean(skipna=True)),
|
||
}
|
||
print("Overall averages:", overall)
|
||
return df
|
||
|
||
|
||
# ----------------------------
|
||
# 3) Fancy but honest visualizations
|
||
# ----------------------------
|
||
|
||
def _ensure_dir(p: str) -> None:
|
||
os.makedirs(p, exist_ok=True)
|
||
|
||
def plot_heatmap_metrics(metrics_df: pd.DataFrame, outdir: str, title: str = "Per-feature benchmark (normalized)") -> None:
|
||
"""
|
||
Heatmap of per-feature metrics.
|
||
We normalize each metric column to [0,1] so KS/JSD/lag1 are comparable visually.
|
||
"""
|
||
_ensure_dir(outdir)
|
||
|
||
df = metrics_df.copy()
|
||
|
||
# Create a unified table with columns present for all features
|
||
# For missing metric values (e.g., ks for discrete), keep NaN; normalize per metric ignoring NaNs.
|
||
show_cols = ["ks", "jsd", "lag1_diff"]
|
||
|
||
# sort features by "worst" normalized score (max across available metrics)
|
||
norm = df[show_cols].copy()
|
||
for c in show_cols:
|
||
col = norm[c]
|
||
if col.notna().sum() == 0:
|
||
continue
|
||
mn = float(col.min(skipna=True))
|
||
mx = float(col.max(skipna=True))
|
||
if abs(mx - mn) < 1e-12:
|
||
norm[c] = 0.0
|
||
else:
|
||
norm[c] = (col - mn) / (mx - mn)
|
||
df["_severity"] = norm.max(axis=1, skipna=True)
|
||
df = df.sort_values("_severity", ascending=False).drop(columns=["_severity"])
|
||
norm = norm.loc[df.index]
|
||
|
||
# Plot heatmap with matplotlib (no seaborn)
|
||
fig, ax = plt.subplots(figsize=(9.5, max(4.0, 0.28 * len(df))))
|
||
im = ax.imshow(norm.to_numpy(), aspect="auto", interpolation="nearest")
|
||
|
||
ax.set_yticks(np.arange(len(df)))
|
||
ax.set_yticklabels(df["feature"].tolist(), fontsize=8)
|
||
|
||
ax.set_xticks(np.arange(len(show_cols)))
|
||
ax.set_xticklabels(show_cols, fontsize=10)
|
||
|
||
ax.set_title(title, fontsize=13, pad=12)
|
||
cbar = fig.colorbar(im, ax=ax, fraction=0.035, pad=0.02)
|
||
cbar.set_label("min–max normalized (per metric)", rotation=90)
|
||
|
||
# Light gridlines for readability
|
||
ax.set_xticks(np.arange(-.5, len(show_cols), 1), minor=True)
|
||
ax.set_yticks(np.arange(-.5, len(df), 1), minor=True)
|
||
ax.grid(which="minor", color="white", linestyle="-", linewidth=0.6)
|
||
ax.tick_params(which="minor", bottom=False, left=False)
|
||
|
||
fig.tight_layout()
|
||
fig.savefig(os.path.join(outdir, "heatmap_metrics_normalized.png"), dpi=300)
|
||
plt.close(fig)
|
||
|
||
def plot_topk_lollipop(metrics_df: pd.DataFrame, metric: str, outdir: str, k: int = 15) -> None:
|
||
"""Fancy top-k lollipop chart for a metric (e.g., 'ks' or 'jsd' or 'lag1_diff')."""
|
||
_ensure_dir(outdir)
|
||
df = metrics_df.dropna(subset=[metric]).sort_values(metric, ascending=False).head(k)
|
||
|
||
fig, ax = plt.subplots(figsize=(9.5, 0.45 * len(df) + 1.5))
|
||
y = np.arange(len(df))[::-1]
|
||
x = df[metric].to_numpy()
|
||
|
||
# stems
|
||
ax.hlines(y=y, xmin=0, xmax=x, linewidth=2.0, alpha=0.85)
|
||
# heads
|
||
ax.plot(x, y, "o", markersize=7)
|
||
|
||
ax.set_yticks(y)
|
||
ax.set_yticklabels(df["feature"].tolist(), fontsize=9)
|
||
ax.set_xlabel(metric)
|
||
ax.set_title(f"Top-{k} features by {metric}", fontsize=13, pad=10)
|
||
|
||
# subtle x-grid
|
||
ax.grid(axis="x", linestyle="--", alpha=0.4)
|
||
fig.tight_layout()
|
||
fig.savefig(os.path.join(outdir, f"topk_{metric}_lollipop.png"), dpi=300)
|
||
plt.close(fig)
|
||
|
||
def plot_metric_distributions(metrics_df: pd.DataFrame, outdir: str) -> None:
|
||
"""Histogram distributions of each metric across features."""
|
||
_ensure_dir(outdir)
|
||
metrics = ["ks", "jsd", "lag1_diff"]
|
||
|
||
fig, ax = plt.subplots(figsize=(9.5, 5.2))
|
||
# overlay histograms (alpha) for compact comparison
|
||
for m in metrics:
|
||
vals = metrics_df[m].dropna().to_numpy()
|
||
if len(vals) == 0:
|
||
continue
|
||
ax.hist(vals, bins=20, alpha=0.55, density=True, label=m)
|
||
|
||
ax.set_title("Distribution of per-feature benchmark metrics", fontsize=13, pad=10)
|
||
ax.set_xlabel("metric value")
|
||
ax.set_ylabel("density")
|
||
ax.grid(axis="y", linestyle="--", alpha=0.35)
|
||
ax.legend()
|
||
fig.tight_layout()
|
||
fig.savefig(os.path.join(outdir, "metric_distributions.png"), dpi=300)
|
||
plt.close(fig)
|
||
|
||
def plot_overall_summary(metrics_df: pd.DataFrame, outdir: str) -> None:
|
||
"""Simple overall averages (separate meaning per metric)."""
|
||
_ensure_dir(outdir)
|
||
ks_avg = metrics_df.loc[metrics_df.ftype == "continuous", "ks"].mean(skipna=True)
|
||
jsd_avg = metrics_df.loc[metrics_df.ftype == "discrete", "jsd"].mean(skipna=True)
|
||
lag_avg = metrics_df["lag1_diff"].mean(skipna=True)
|
||
|
||
labels = ["KS(avg cont)", "JSD(avg disc)", "Lag1 diff(avg)"]
|
||
values = [ks_avg, jsd_avg, lag_avg]
|
||
|
||
fig, ax = plt.subplots(figsize=(7.8, 4.6))
|
||
ax.bar(labels, values)
|
||
ax.set_title("Overall benchmark summary", fontsize=13, pad=10)
|
||
ax.set_ylabel("value")
|
||
ax.grid(axis="y", linestyle="--", alpha=0.35)
|
||
fig.tight_layout()
|
||
fig.savefig(os.path.join(outdir, "overall_summary_bar.png"), dpi=300)
|
||
plt.close(fig)
|
||
|
||
|
||
# ----------------------------
|
||
# 4) continuous trends
|
||
# ----------------------------
|
||
|
||
|
||
def _smooth_ma(x: np.ndarray, w: int) -> np.ndarray:
|
||
"""Simple moving average smoothing; w>=1."""
|
||
if w <= 1:
|
||
return x.astype(np.float64, copy=True)
|
||
s = pd.Series(x.astype(np.float64))
|
||
return s.rolling(w, center=True, min_periods=max(2, w // 10)).mean().to_numpy()
|
||
|
||
def _rolling_quantile(x: np.ndarray, w: int, q: float) -> np.ndarray:
|
||
"""Rolling quantile for a local envelope."""
|
||
if w <= 1:
|
||
return x.astype(np.float64, copy=True)
|
||
s = pd.Series(x.astype(np.float64))
|
||
return s.rolling(w, center=True, min_periods=max(10, w // 5)).quantile(q).to_numpy()
|
||
|
||
def plot_trend_small_multiples(
|
||
real_df: pd.DataFrame,
|
||
gen_df: pd.DataFrame,
|
||
metrics_df: pd.DataFrame,
|
||
feature_types: Dict[str, str],
|
||
outdir: str,
|
||
k: int = 12,
|
||
rank_by: str = "lag1_diff", # or "ks"
|
||
smooth_w: int = 200,
|
||
env_w: int = 600,
|
||
downsample: int = 1,
|
||
) -> None:
|
||
"""
|
||
Small multiples of continuous-feature trends.
|
||
- Picks top-k continuous features by rank_by (ks or lag1_diff).
|
||
- Plots: faint raw + bold smoothed + rolling 10–90% envelope for real and gen.
|
||
"""
|
||
_ensure_dir(outdir)
|
||
|
||
# pick top-k continuous features
|
||
dfc = metrics_df[metrics_df["ftype"] == "continuous"].dropna(subset=[rank_by]).copy()
|
||
dfc = dfc.sort_values(rank_by, ascending=False).head(k)
|
||
feats = dfc["feature"].tolist()
|
||
if not feats:
|
||
print("[trend] No continuous features found to plot.")
|
||
return
|
||
|
||
n = len(feats)
|
||
ncols = 3
|
||
nrows = int(np.ceil(n / ncols))
|
||
|
||
fig = plt.subplots(figsize=(12, 3.2 * nrows))[0]
|
||
gs = fig.add_gridspec(nrows, ncols, hspace=0.35, wspace=0.22)
|
||
|
||
for idx, feat in enumerate(feats):
|
||
r = real_df[feat].to_numpy(dtype=np.float64)
|
||
g = gen_df[feat].to_numpy(dtype=np.float64)
|
||
|
||
if downsample > 1:
|
||
r = r[::downsample]
|
||
g = g[::downsample]
|
||
|
||
t = np.arange(len(r), dtype=np.int64)
|
||
|
||
# smoothed trend
|
||
r_tr = _smooth_ma(r, smooth_w)
|
||
g_tr = _smooth_ma(g, smooth_w)
|
||
|
||
# local envelope (10–90%)
|
||
r_lo = _rolling_quantile(r, env_w, 0.10)
|
||
r_hi = _rolling_quantile(r, env_w, 0.90)
|
||
g_lo = _rolling_quantile(g, env_w, 0.10)
|
||
g_hi = _rolling_quantile(g, env_w, 0.90)
|
||
|
||
ax = fig.add_subplot(gs[idx // ncols, idx % ncols])
|
||
|
||
# faint raw “watermark”
|
||
ax.plot(t, r, linewidth=0.6, alpha=0.18, label="real (raw)" if idx == 0 else None)
|
||
ax.plot(t, g, linewidth=0.6, alpha=0.18, label="gen (raw)" if idx == 0 else None)
|
||
|
||
# envelopes
|
||
ax.fill_between(t, r_lo, r_hi, alpha=0.18, label="real (10–90%)" if idx == 0 else None)
|
||
ax.fill_between(t, g_lo, g_hi, alpha=0.18, label="gen (10–90%)" if idx == 0 else None)
|
||
|
||
# bold trends
|
||
ax.plot(t, r_tr, linewidth=2.0, alpha=0.9, label="real (trend)" if idx == 0 else None)
|
||
ax.plot(t, g_tr, linewidth=2.0, alpha=0.9, label="gen (trend)" if idx == 0 else None)
|
||
|
||
# title with metric
|
||
mval = float(dfc.loc[dfc["feature"] == feat, rank_by].iloc[0])
|
||
ax.set_title(f"{feat} | {rank_by}={mval:.3f}", fontsize=10)
|
||
|
||
# cosmetics
|
||
ax.grid(axis="y", linestyle="--", alpha=0.25)
|
||
ax.tick_params(labelsize=8)
|
||
|
||
# shared legend (single)
|
||
handles, labels = fig.axes[0].get_legend_handles_labels()
|
||
fig.legend(handles, labels, loc="upper center", ncol=3, frameon=False, bbox_to_anchor=(0.5, 0.99))
|
||
fig.suptitle("Trend comparison (continuous features): raw watermark + smoothed trend + local envelope",
|
||
fontsize=14, y=1.02)
|
||
|
||
fig.tight_layout()
|
||
fig.savefig(os.path.join(outdir, f"trend_small_multiples_{rank_by}.png"), dpi=300, bbox_inches="tight")
|
||
plt.close(fig)
|
||
|
||
|
||
# ----------------------------
|
||
# 4) discrete trends
|
||
# ----------------------------
|
||
|
||
|
||
def plot_discrete_state_trends(
|
||
real_df: pd.DataFrame,
|
||
gen_df: pd.DataFrame,
|
||
metrics_df: pd.DataFrame,
|
||
feature_types: Dict[str, str],
|
||
disc_vocab_map: Dict[str, List[str]] | None,
|
||
outdir: str,
|
||
k: int = 6,
|
||
rank_by: str = "jsd",
|
||
win: int = 400,
|
||
step: int = 40,
|
||
) -> None:
|
||
"""
|
||
For top-k discrete features by JSD, plot rolling state occupancy over time:
|
||
p_t(v) = fraction of state v in window centered at t
|
||
"""
|
||
_ensure_dir(outdir)
|
||
|
||
dfd = metrics_df[metrics_df["ftype"] == "discrete"].dropna(subset=[rank_by]).copy()
|
||
dfd = dfd.sort_values(rank_by, ascending=False).head(k)
|
||
feats = dfd["feature"].tolist()
|
||
if not feats:
|
||
print("[trend] No discrete features found to plot.")
|
||
return
|
||
|
||
for feat in feats:
|
||
r = real_df[feat].to_numpy()
|
||
g = gen_df[feat].to_numpy()
|
||
vocab = disc_vocab_map.get(feat) if disc_vocab_map else sorted(list(set(r).union(set(g))))
|
||
|
||
# timeline centers
|
||
centers = np.arange(win // 2, len(r) - win // 2, step, dtype=np.int64)
|
||
t = centers
|
||
|
||
# occupancy matrices: shape [len(vocab), len(t)]
|
||
R = np.zeros((len(vocab), len(t)), dtype=np.float64)
|
||
G = np.zeros((len(vocab), len(t)), dtype=np.float64)
|
||
|
||
for ti, c in enumerate(centers):
|
||
lo = c - win // 2
|
||
hi = c + win // 2
|
||
rw = r[lo:hi]
|
||
gw = g[lo:hi]
|
||
|
||
# counts
|
||
rc = pd.Series(rw).value_counts()
|
||
gc = pd.Series(gw).value_counts()
|
||
for vi, v in enumerate(vocab):
|
||
R[vi, ti] = rc.get(v, 0) / max(1, len(rw))
|
||
G[vi, ti] = gc.get(v, 0) / max(1, len(gw))
|
||
|
||
# Plot as two stacked areas: Real vs Gen (same vocab order)
|
||
fig, ax = plt.subplots(figsize=(11.5, 4.8))
|
||
ax.stackplot(t, R, alpha=0.35, labels=[f"{v} (real)" for v in vocab])
|
||
ax.stackplot(t, G, alpha=0.35, labels=[f"{v} (gen)" for v in vocab])
|
||
|
||
mval = float(dfd.loc[dfd["feature"] == feat, rank_by].iloc[0])
|
||
ax.set_title(f"Discrete trend via rolling state occupancy — {feat} | {rank_by}={mval:.3f}", fontsize=13, pad=10)
|
||
ax.set_xlabel("time index (window centers)")
|
||
ax.set_ylabel("occupancy probability")
|
||
ax.grid(axis="y", linestyle="--", alpha=0.25)
|
||
|
||
# Put legend outside to keep plot clean
|
||
ax.legend(loc="center left", bbox_to_anchor=(1.01, 0.5), frameon=False, fontsize=8)
|
||
|
||
fig.tight_layout()
|
||
fig.savefig(os.path.join(outdir, f"disc_state_trend_{feat}.png"), dpi=300, bbox_inches="tight")
|
||
plt.close(fig)
|
||
|
||
|
||
|
||
# ----------------------------
|
||
# 4) Main entry
|
||
# ----------------------------
|
||
|
||
def main() -> None:
|
||
outdir = "benchmark_figs"
|
||
_ensure_dir(outdir)
|
||
real_df, gen_df, feature_types, disc_vocab_map = make_dummy_real_gen(
|
||
T=10000, n_cont=20, n_disc=8, seed=42
|
||
)
|
||
|
||
metrics_df = compute_metrics(real_df, gen_df, feature_types, disc_vocab_map)
|
||
metrics_df.to_csv(os.path.join(outdir, "metrics_per_feature.csv"), index=False)
|
||
|
||
plot_heatmap_metrics(metrics_df, outdir)
|
||
plot_metric_distributions(metrics_df, outdir)
|
||
plot_overall_summary(metrics_df, outdir)
|
||
|
||
# Top-k charts: choose what you emphasize
|
||
plot_topk_lollipop(metrics_df, metric="ks", outdir=outdir, k=15)
|
||
plot_topk_lollipop(metrics_df, metric="jsd", outdir=outdir, k=15)
|
||
plot_topk_lollipop(metrics_df, metric="lag1_diff", outdir=outdir, k=15)
|
||
|
||
# Fancy trend figures (new)
|
||
plot_trend_small_multiples(
|
||
real_df, gen_df, metrics_df, feature_types, outdir,
|
||
k=12, rank_by="lag1_diff", smooth_w=250, env_w=800, downsample=1
|
||
)
|
||
|
||
plot_trend_small_multiples(
|
||
real_df, gen_df, metrics_df, feature_types, outdir,
|
||
k=12, rank_by="ks", smooth_w=250, env_w=800, downsample=1
|
||
)
|
||
|
||
plot_discrete_state_trends(
|
||
real_df, gen_df, metrics_df, feature_types, disc_vocab_map, outdir,
|
||
k=6, rank_by="jsd", win=500, step=50
|
||
)
|
||
|
||
print(f"Saved figures + CSV under: ./{outdir}/")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|