nancyH/token_evaluation

1"""2Ablation analysis — mirrors all_chrom_evaluation.ipynb but for:3  - Ablation_no_partition, Ablation_no_priority, Ablation_no_length4  - Plus Baseline_bpe_5120 and Merged_uni_len2_5120 as references5Saves all plots to eval_outputs_ablation/plots/6"""7 8import os9import pandas as pd10import numpy as np11import matplotlib12matplotlib.use("Agg")13import matplotlib.pyplot as plt14import seaborn as sns15 16# ── Directories ──────────────────────────────────────────────────────────────17ABLATION_DIR = "/home/n5huang/dna_token/tokenizer_evaluation/eval_outputs_ablation"18MERGED_DIR   = "/home/n5huang/dna_token/tokenizer_evaluation/eval_outputs_all_chrom"19BASELINE_DIR = "/home/n5huang/dna_token/tokenizer_evaluation/eval_outputs_all_chrom_baseline_bpe"20PLOT_DIR     = os.path.join(ABLATION_DIR, "plots")21os.makedirs(PLOT_DIR, exist_ok=True)22 23pd.set_option("display.max_columns", 20)24pd.set_option("display.width", 200)25pd.set_option("display.float_format", "{:.6f}".format)26 27# ── Helper: compute summary row from an agg CSV ─────────────────────────────28def summarize_agg(agg: pd.DataFrame, name: str, scope: str, region: str = None):29    row = {30        "tokenizer": name,31        "scope": scope,32        "mean_of_mean_phyloP": agg["mean_mean"].mean(),33        "median_of_mean_phyloP": agg["mean_mean"].median(),34        "pct_mean_phyloP_above_0": (agg["mean_mean"] > 0).mean() * 100,35        "mean_of_variance": agg["mean_var"].mean(),36        "median_of_variance": agg["mean_var"].median(),37        "pct_variance_below_0.1": (agg["mean_var"] < 0.1).mean() * 100,38        "num_tokens": len(agg),39        "mean_token_count": agg["count"].mean(),40        "median_token_count": agg["count"].median(),41    }42    if region:43        row["region"] = region44    return row45 46 47# ── 1. Window-Based Summary ──────────────────────────────────────────────────48print("=" * 80)49print("1. WINDOW-BASED SUMMARY")50print("=" * 80)51 52ablation_names = ["Ablation_no_partition", "Ablation_no_priority", "Ablation_no_length"]53ref_names = {54    "Baseline_bpe_5120":    os.path.join(BASELINE_DIR, "agg_Baseline_bpe_5120.csv"),55    "Merged_uni_len2_5120": os.path.join(MERGED_DIR,   "agg_Merged_uni_len2_5120.csv"),56}57 58# Load all agg data (window)59agg_data = {}60window_summaries = []61 62for name in ablation_names:63    path = os.path.join(ABLATION_DIR, f"agg_{name}.csv")64    agg = pd.read_csv(path)65    agg_data[name] = agg66    window_summaries.append(summarize_agg(agg, name, "all_chrom_window"))67 68for name, path in ref_names.items():69    agg = pd.read_csv(path)70    agg_data[name] = agg71    window_summaries.append(summarize_agg(agg, name, "all_chrom_window"))72 73df_window = pd.DataFrame(window_summaries)74cols = ["tokenizer", "mean_of_mean_phyloP", "median_of_mean_phyloP",75        "pct_mean_phyloP_above_0", "mean_of_variance", "median_of_variance",76        "pct_variance_below_0.1", "num_tokens", "mean_token_count", "median_token_count"]77df_window = df_window[cols]78print(df_window.to_string(index=False))79df_window.to_csv(os.path.join(ABLATION_DIR, "summary_ablation_window.csv"), index=False)80 81 82# ── 2. Region-Based Summary ─────────────────────────────────────────────────83print("\n" + "=" * 80)84print("2. REGION-BASED SUMMARY")85print("=" * 80)86 87agg_region = {}88region_summaries = []89 90for name in ablation_names:91    for region in ["conserved", "neutral", "accelerated"]:92        path = os.path.join(ABLATION_DIR, f"agg_{name}_{region}.csv")93        if os.path.exists(path):94            agg = pd.read_csv(path)95            agg_region[(name, region)] = agg96            region_summaries.append(summarize_agg(agg, name, "all_chrom_region", region))97 98for name, base_path in ref_names.items():99    base_dir = os.path.dirname(base_path)100    for region in ["conserved", "neutral", "accelerated"]:101        path = os.path.join(base_dir, f"agg_{name}_{region}.csv")102        if os.path.exists(path):103            agg = pd.read_csv(path)104            agg_region[(name, region)] = agg105            region_summaries.append(summarize_agg(agg, name, "all_chrom_region", region))106 107df_region = pd.DataFrame(region_summaries)108cols_region = ["tokenizer", "region", "mean_of_mean_phyloP", "median_of_mean_phyloP",109               "pct_mean_phyloP_above_0", "mean_of_variance", "median_of_variance",110               "pct_variance_below_0.1", "num_tokens", "mean_token_count", "median_token_count"]111df_region = df_region[cols_region]112 113for region in ["conserved", "neutral", "accelerated"]:114    print(f"\n=== {region.upper()} ===")115    print(df_region[df_region["region"] == region].to_string(index=False))116 117df_region.to_csv(os.path.join(ABLATION_DIR, "summary_ablation_all_regions.csv"), index=False)118 119 120# ── Color scheme ─────────────────────────────────────────────────────────────121COLORS = {122    "Baseline_bpe_5120":      "steelblue",123    "Merged_uni_len2_5120":   "tomato",124    "Ablation_no_partition":  "#2ca02c",   # green125    "Ablation_no_priority":   "#9467bd",   # purple126    "Ablation_no_length":     "#ff7f0e",   # orange127}128 129 130# ── 3. Scatter: Mean phyloP vs Variance (Window) ────────────────────────────131print("\n" + "=" * 80)132print("3. SCATTER PLOTS: Mean phyloP vs Variance")133print("=" * 80)134 135fig, axes = plt.subplots(1, 2, figsize=(18, 8))136 137for ax_idx, (xcol, xlabel) in enumerate([("mean_mean", "Mean phyloP"), ("median_mean", "Median Mean phyloP")]):138    ax = axes[ax_idx]139    for name, color in COLORS.items():140        if name in agg_data:141            d = agg_data[name]142            ax.scatter(d[xcol], d["mean_var"], alpha=0.4, s=20, label=name, color=color)143    ax.set_xlabel(xlabel)144    ax.set_ylabel("Mean Variance")145    ax.set_title(f"Token Conservation: {xlabel} vs Variance")146    ax.legend(fontsize=8)147 148plt.tight_layout()149plt.savefig(os.path.join(PLOT_DIR, "scatter_mean_vs_var.png"), dpi=150, bbox_inches="tight")150plt.close()151print(f"  Saved scatter_mean_vs_var.png")152 153 154# ── 4. Histograms: Mean phyloP and Variance distributions ───────────────────155print("\n" + "=" * 80)156print("4. HISTOGRAMS")157print("=" * 80)158 159fig, axes = plt.subplots(1, 2, figsize=(18, 6))160 161ax = axes[0]162for name, color in COLORS.items():163    if name in agg_data:164        ax.hist(agg_data[name]["mean_mean"], bins=50, alpha=0.5, label=name, color=color)165ax.set_xlabel("Mean phyloP")166ax.set_ylabel("Token Count")167ax.set_title("Mean Conservation Distribution")168ax.legend(fontsize=8)169 170ax = axes[1]171for name, color in COLORS.items():172    if name in agg_data:173        ax.hist(agg_data[name]["mean_var"], bins=50, alpha=0.5, label=name, color=color)174ax.set_xlabel("Mean Variance")175ax.set_ylabel("Token Count")176ax.set_title("Variance Distribution")177ax.legend(fontsize=8)178 179plt.tight_layout()180plt.savefig(os.path.join(PLOT_DIR, "histograms_mean_var.png"), dpi=150, bbox_inches="tight")181plt.close()182print(f"  Saved histograms_mean_var.png")183 184 185# ── 5. Region scatter: by region ─────────────────────────────────────────────186print("\n" + "=" * 80)187print("5. REGION SCATTER PLOTS")188print("=" * 80)189 190fig, axes = plt.subplots(1, 3, figsize=(21, 7))191 192for i, region in enumerate(["conserved", "neutral", "accelerated"]):193    ax = axes[i]194    for name, color in COLORS.items():195        key = (name, region)196        if key in agg_region:197            d = agg_region[key]198            ax.scatter(d["mean_mean"], d["mean_var"], alpha=0.4, s=20, label=name, color=color)199    ax.set_xlabel("Mean phyloP")200    ax.set_ylabel("Mean Variance")201    ax.set_title(f"{region.capitalize()} Regions")202    ax.legend(fontsize=7, loc="upper left")203 204plt.suptitle("Region-Based: Ablation vs References", fontsize=14)205plt.tight_layout()206plt.savefig(os.path.join(PLOT_DIR, "scatter_by_region.png"), dpi=150, bbox_inches="tight")207plt.close()208print(f"  Saved scatter_by_region.png")209 210 211# ── 6. Bar charts: Summary metrics by tokenizer and region ──────────────────212print("\n" + "=" * 80)213print("6. BAR CHARTS: Summary metrics")214print("=" * 80)215 216metrics = ["mean_of_mean_phyloP", "pct_mean_phyloP_above_0", "mean_of_variance", "pct_variance_below_0.1"]217metric_labels = ["Mean of Mean phyloP", "% Mean phyloP > 0", "Mean of Variance", "% Variance < 0.1"]218 219fig, axes = plt.subplots(2, 2, figsize=(18, 12))220 221for idx, (metric, label) in enumerate(zip(metrics, metric_labels)):222    ax = axes[idx // 2, idx % 2]223    for region in ["conserved", "neutral", "accelerated"]:224        subset = df_region[df_region["region"] == region].copy()225        subset = subset.sort_values("tokenizer")226        x = range(len(subset))227        offset = {"conserved": -0.25, "neutral": 0, "accelerated": 0.25}[region]228        color = {"conserved": "forestgreen", "neutral": "gray", "accelerated": "firebrick"}[region]229        ax.bar([xi + offset for xi in x], subset[metric], width=0.25,230               label=region, color=color, alpha=0.8)231    ax.set_xticks(range(len(subset)))232    ax.set_xticklabels(subset["tokenizer"], rotation=45, ha="right", fontsize=8)233    ax.set_ylabel(label)234    ax.set_title(label)235    ax.legend(fontsize=8)236 237plt.suptitle("Ablation: Region-Based Metrics", fontsize=14)238plt.tight_layout()239plt.savefig(os.path.join(PLOT_DIR, "bar_region_metrics.png"), dpi=150, bbox_inches="tight")240plt.close()241print(f"  Saved bar_region_metrics.png")242 243 244# ── 7. Grouped bar: Ablation comparison (window metrics) ────────────────────245print("\n" + "=" * 80)246print("7. GROUPED BAR: Window metrics comparison")247print("=" * 80)248 249order = ["Baseline_bpe_5120", "Ablation_no_partition", "Ablation_no_priority",250         "Ablation_no_length", "Merged_uni_len2_5120"]251df_w_ordered = df_window.set_index("tokenizer").loc[order].reset_index()252 253fig, axes = plt.subplots(1, 3, figsize=(20, 6))254 255# Mean phyloP256ax = axes[0]257bars = ax.barh(df_w_ordered["tokenizer"], df_w_ordered["mean_of_mean_phyloP"],258               color=[COLORS[n] for n in df_w_ordered["tokenizer"]])259ax.set_xlabel("Mean of Mean phyloP")260ax.set_title("Conservation Signal (Higher is Better)")261ax.invert_yaxis()262for bar, val in zip(bars, df_w_ordered["mean_of_mean_phyloP"]):263    ax.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height()/2,264            f"{val:.4f}", va="center", fontsize=9)265 266# % above 0267ax = axes[1]268bars = ax.barh(df_w_ordered["tokenizer"], df_w_ordered["pct_mean_phyloP_above_0"],269               color=[COLORS[n] for n in df_w_ordered["tokenizer"]])270ax.set_xlabel("% Mean phyloP > 0")271ax.set_title("Positive Conservation Rate (Higher is Better)")272ax.invert_yaxis()273for bar, val in zip(bars, df_w_ordered["pct_mean_phyloP_above_0"]):274    ax.text(bar.get_width() + 0.1, bar.get_y() + bar.get_height()/2,275            f"{val:.1f}%", va="center", fontsize=9)276 277# Mean variance278ax = axes[2]279bars = ax.barh(df_w_ordered["tokenizer"], df_w_ordered["mean_of_variance"],280               color=[COLORS[n] for n in df_w_ordered["tokenizer"]])281ax.set_xlabel("Mean of Variance")282ax.set_title("Internal Variance (Lower is Better)")283ax.invert_yaxis()284for bar, val in zip(bars, df_w_ordered["mean_of_variance"]):285    ax.text(bar.get_width() + 0.001, bar.get_y() + bar.get_height()/2,286            f"{val:.4f}", va="center", fontsize=9)287 288plt.suptitle("Ablation: Window-Based Summary Metrics", fontsize=14)289plt.tight_layout()290plt.savefig(os.path.join(PLOT_DIR, "bar_window_comparison.png"), dpi=150, bbox_inches="tight")291plt.close()292print(f"  Saved bar_window_comparison.png")293 294 295# ── 8. Print final comparison table ─────────────────────────────────────────296print("\n" + "=" * 80)297print("8. FINAL COMPARISON TABLE")298print("=" * 80)299print("\nWindow-based:")300print(df_w_ordered[["tokenizer", "mean_of_mean_phyloP", "pct_mean_phyloP_above_0",301                     "mean_of_variance", "pct_variance_below_0.1", "num_tokens"]].to_string(index=False))302 303print("\nRegion-based (conserved only):")304cons = df_region[df_region["region"] == "conserved"].copy()305cons = cons[cons["tokenizer"].isin(order)]306print(cons[["tokenizer", "mean_of_mean_phyloP", "pct_mean_phyloP_above_0",307            "mean_of_variance", "num_tokens"]].to_string(index=False))308 309print(f"\nAll plots saved to: {PLOT_DIR}/")310print("All summary CSVs saved to:", ABLATION_DIR)311