supp_tokenizer_bench / code /plot_combined_panel_a_b.py

Add Supp. Fig. 2b JS-distance analysis: code, region/background BEDs, FASTAs, results, baseline tokenizers

89dec6b verified 17 days ago

5.66 kB

#!/usr/bin/env python3
"""Combined figure: (a) Motif Preservation (A1 + A2 stacked) and (b) JS heatmap, in one row."""

import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib as mpl
import seaborn as sns
from matplotlib.colors import LinearSegmentedColormap
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec

# ── Data ──
df_summary = pd.read_csv(
    "/data/n5huang/dna_token/tokenizer_evaluation/motif_eval_outputs_new/summary_by_tokenizer.csv"
)
df_js = pd.read_csv(
    "/data/n5huang/dna_token/enrichment_analysis/js_vs_random_matrix.csv",
    index_col=0,
)

# ── Method config (shared) ──
METHODS = ["DNAbert2", "Grover", "cCRE_region_BPE", "motif_region_BPE", "DNAMotifTokenizer"]
DISPLAY_NAMES = {
    "DNAbert2": "DNABERT-2",
    "Grover": "Grover",
    "cCRE_region_BPE": "cCRE-Region BPE",
    "motif_region_BPE": "Motif-Region BPE",
    "DNAMotifTokenizer": "DNAMotif",
}
METHOD_COLORS = {
    "DNAbert2": "#6C5CE7",          # purple (from C)
    "Grover": "#E74C3C",            # red (from T)
    "cCRE_region_BPE": "#52BE80",   # green (from A)
    "motif_region_BPE": "#F4D03F",  # yellow (Token Characterization)
    "DNAMotifTokenizer": "#E67E22", # orange (Fine-tuning Ability)
}
HEATMAP_REGIONS = ["promoter", "enhancer", "exon"]
JS_INDEX_DISPLAY = [DISPLAY_NAMES.get(n, n) for n in df_js.index]

# ── Plot style ──
mpl.rcParams.update({
    "font.family": "sans-serif",
    "font.size": 58,
    "axes.labelsize": 64,
    "axes.titlesize": 66,
    "axes.titleweight": "bold",
    "xtick.labelsize": 58,
    "ytick.labelsize": 58,
})


def get_value(method, source_df, col_name):
    row = source_df[source_df["Tokenizer"] == method]
    if len(row) == 0:
        return float("nan")
    return row[col_name].values[0]


def draw_bar_panel(ax, title, ylabel, col_name, higher_better, title_pad=30, ylabel_y=0.5):
    vals = [get_value(m, df_summary, col_name) for m in METHODS]
    colors = [METHOD_COLORS[m] for m in METHODS]
    x_labels = [DISPLAY_NAMES[m] for m in METHODS]
    x = np.arange(len(METHODS)) * 0.3
    bars = ax.bar(x, vals, color=colors, edgecolor="white", linewidth=0.5, alpha=0.85, width=0.2)

    for i, v in enumerate(vals):
        if np.isnan(v):
            continue
        y_off = max(vals) * 0.02
        fmt = f"{v:.2f}" if abs(v) >= 1 else f"{v:.4f}"
        ax.text(x[i], v + y_off, fmt, ha="center", va="bottom", fontsize=62, fontweight="bold")

    ax.set_xticks(x)
    ax.set_xticklabels(x_labels, rotation=15, ha="right")
    ax.set_xlim(x[0] - 0.2, x[-1] + 0.2)
    ax.set_ylabel(ylabel)
    # Pin ylabel position in axes coords so A1 and A2 line up regardless of tick-label widths.
    ax.yaxis.set_label_coords(-0.08, ylabel_y)
    ax.set_title(title, pad=title_pad)
    ax.tick_params(axis="x", pad=10)
    ax.tick_params(axis="y", pad=15)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(axis="y", alpha=0.3, linewidth=0.8)

    valid = [v for v in vals if not np.isnan(v)]
    if valid:
        best = (np.nanargmax if higher_better else np.nanargmin)(vals)
        bars[best].set_edgecolor("black")
        bars[best].set_linewidth(2)


# ── Figure layout: 1 row, 2 top-level columns ──
fig = plt.figure(figsize=(50, 26))
gs = GridSpec(1, 2, figure=fig, width_ratios=[1, 0.9], wspace=0.30)

# Left column (panel a): 2 stacked subplots
gs_a = GridSpecFromSubplotSpec(2, 1, subplot_spec=gs[0, 0], hspace=0.8)
ax_a1 = fig.add_subplot(gs_a[0, 0])
ax_a2 = fig.add_subplot(gs_a[1, 0])

# Right column (panel b): heatmap
ax_b = fig.add_subplot(gs[0, 1])

# ── Draw panel a (A1 + A2) ──
draw_bar_panel(ax_a1, "Perfect Match Rate", "Perfect match rate (%)",
               "Perfect_Match_Rate", higher_better=True, title_pad=80, ylabel_y=0.35)
draw_bar_panel(ax_a2, "Motif Fragmentation", "Avg tokens per motif",
               "Avg_Tokens_Per_Motif", higher_better=False, title_pad=80)

# ── Draw panel b (heatmap) ──
df_disp = df_js.copy()
df_disp.index = JS_INDEX_DISPLAY
cmap = LinearSegmentedColormap.from_list("evo_heat", ["white", "#AED6F1", "#5DADE2", "#2E86C1"])
sns.heatmap(
    df_disp,
    ax=ax_b,
    annot=True,
    fmt=".4f",
    annot_kws={"size": 62, "fontweight": "bold"},  # matches bar-value annotation size
    cmap=cmap,
    vmin=0.0,
    vmax=float(df_disp.values.max()) or 0.001,
    linewidths=0.5,
    cbar_kws={"label": "Jensen-Shannon distance"},
)
ax_b.set_title("Region vs random-genome separation",
               fontsize=66, fontweight="bold", pad=80)
ax_b.set_xlabel("Genomic region", fontsize=64, labelpad=15)
ax_b.set_ylabel("")
ax_b.tick_params(axis="x", labelsize=58)
ax_b.tick_params(axis="y", labelsize=58, rotation=0)
cbar = ax_b.collections[0].colorbar
cbar.ax.tick_params(labelsize=58)
cbar.set_label("Jensen-Shannon distance", size=64)

fig.subplots_adjust(left=0.06, right=0.98, top=0.88, bottom=0.10)

# ── Panel labels (a), (b) on same line as titles ──
# Use axes-relative y so they line up with titles (which sit just above each axes top).
ax_a1.text(-0.08, 1.234, "(a)", transform=ax_a1.transAxes,
           fontsize=66, fontweight="bold", va="center", ha="center")
ax_b.text(-0.30, 1.083, "(b)", transform=ax_b.transAxes,
          fontsize=66, fontweight="bold", va="center", ha="center")

out_dir = "/data/n5huang/dna_token/paper_figures"
os.makedirs(out_dir, exist_ok=True)
out = os.path.join(out_dir, "combined_panel_a_b.pdf")
fig.savefig(out, dpi=300)
fig.savefig(out.replace(".pdf", ".png"), dpi=80)
print(f"Saved: {out}")
print(f"Saved: {out.replace('.pdf', '.png')}")
plt.close(fig)

	#!/usr/bin/env python3
	"""Combined figure: (a) Motif Preservation (A1 + A2 stacked) and (b) JS heatmap, in one row."""

	import os
	import numpy as np
	import pandas as pd
	import matplotlib.pyplot as plt
	import matplotlib as mpl
	import seaborn as sns
	from matplotlib.colors import LinearSegmentedColormap
	from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec

	# ── Data ──
	df_summary = pd.read_csv(
	"/data/n5huang/dna_token/tokenizer_evaluation/motif_eval_outputs_new/summary_by_tokenizer.csv"
	)
	df_js = pd.read_csv(
	"/data/n5huang/dna_token/enrichment_analysis/js_vs_random_matrix.csv",
	index_col=0,
	)

	# ── Method config (shared) ──
	METHODS = ["DNAbert2", "Grover", "cCRE_region_BPE", "motif_region_BPE", "DNAMotifTokenizer"]
	DISPLAY_NAMES = {
	"DNAbert2": "DNABERT-2",
	"Grover": "Grover",
	"cCRE_region_BPE": "cCRE-Region BPE",
	"motif_region_BPE": "Motif-Region BPE",
	"DNAMotifTokenizer": "DNAMotif",
	}
	METHOD_COLORS = {
	"DNAbert2": "#6C5CE7", # purple (from C)
	"Grover": "#E74C3C", # red (from T)
	"cCRE_region_BPE": "#52BE80", # green (from A)
	"motif_region_BPE": "#F4D03F", # yellow (Token Characterization)
	"DNAMotifTokenizer": "#E67E22", # orange (Fine-tuning Ability)
	}
	HEATMAP_REGIONS = ["promoter", "enhancer", "exon"]
	JS_INDEX_DISPLAY = [DISPLAY_NAMES.get(n, n) for n in df_js.index]

	# ── Plot style ──
	mpl.rcParams.update({
	"font.family": "sans-serif",
	"font.size": 58,
	"axes.labelsize": 64,
	"axes.titlesize": 66,
	"axes.titleweight": "bold",
	"xtick.labelsize": 58,
	"ytick.labelsize": 58,
	})


	def get_value(method, source_df, col_name):
	row = source_df[source_df["Tokenizer"] == method]
	if len(row) == 0:
	return float("nan")
	return row[col_name].values[0]


	def draw_bar_panel(ax, title, ylabel, col_name, higher_better, title_pad=30, ylabel_y=0.5):
	vals = [get_value(m, df_summary, col_name) for m in METHODS]
	colors = [METHOD_COLORS[m] for m in METHODS]
	x_labels = [DISPLAY_NAMES[m] for m in METHODS]
	x = np.arange(len(METHODS)) * 0.3
	bars = ax.bar(x, vals, color=colors, edgecolor="white", linewidth=0.5, alpha=0.85, width=0.2)

	for i, v in enumerate(vals):
	if np.isnan(v):
	continue
	y_off = max(vals) * 0.02
	fmt = f"{v:.2f}" if abs(v) >= 1 else f"{v:.4f}"
	ax.text(x[i], v + y_off, fmt, ha="center", va="bottom", fontsize=62, fontweight="bold")

	ax.set_xticks(x)
	ax.set_xticklabels(x_labels, rotation=15, ha="right")
	ax.set_xlim(x[0] - 0.2, x[-1] + 0.2)
	ax.set_ylabel(ylabel)
	# Pin ylabel position in axes coords so A1 and A2 line up regardless of tick-label widths.
	ax.yaxis.set_label_coords(-0.08, ylabel_y)
	ax.set_title(title, pad=title_pad)
	ax.tick_params(axis="x", pad=10)
	ax.tick_params(axis="y", pad=15)
	ax.spines["top"].set_visible(False)
	ax.spines["right"].set_visible(False)
	ax.grid(axis="y", alpha=0.3, linewidth=0.8)

	valid = [v for v in vals if not np.isnan(v)]
	if valid:
	best = (np.nanargmax if higher_better else np.nanargmin)(vals)
	bars[best].set_edgecolor("black")
	bars[best].set_linewidth(2)


	# ── Figure layout: 1 row, 2 top-level columns ──
	fig = plt.figure(figsize=(50, 26))
	gs = GridSpec(1, 2, figure=fig, width_ratios=[1, 0.9], wspace=0.30)

	# Left column (panel a): 2 stacked subplots
	gs_a = GridSpecFromSubplotSpec(2, 1, subplot_spec=gs[0, 0], hspace=0.8)
	ax_a1 = fig.add_subplot(gs_a[0, 0])
	ax_a2 = fig.add_subplot(gs_a[1, 0])

	# Right column (panel b): heatmap
	ax_b = fig.add_subplot(gs[0, 1])

	# ── Draw panel a (A1 + A2) ──
	draw_bar_panel(ax_a1, "Perfect Match Rate", "Perfect match rate (%)",
	"Perfect_Match_Rate", higher_better=True, title_pad=80, ylabel_y=0.35)
	draw_bar_panel(ax_a2, "Motif Fragmentation", "Avg tokens per motif",
	"Avg_Tokens_Per_Motif", higher_better=False, title_pad=80)

	# ── Draw panel b (heatmap) ──
	df_disp = df_js.copy()
	df_disp.index = JS_INDEX_DISPLAY
	cmap = LinearSegmentedColormap.from_list("evo_heat", ["white", "#AED6F1", "#5DADE2", "#2E86C1"])
	sns.heatmap(
	df_disp,
	ax=ax_b,
	annot=True,
	fmt=".4f",
	annot_kws={"size": 62, "fontweight": "bold"}, # matches bar-value annotation size
	cmap=cmap,
	vmin=0.0,
	vmax=float(df_disp.values.max()) or 0.001,
	linewidths=0.5,
	cbar_kws={"label": "Jensen-Shannon distance"},
	)
	ax_b.set_title("Region vs random-genome separation",
	fontsize=66, fontweight="bold", pad=80)
	ax_b.set_xlabel("Genomic region", fontsize=64, labelpad=15)
	ax_b.set_ylabel("")
	ax_b.tick_params(axis="x", labelsize=58)
	ax_b.tick_params(axis="y", labelsize=58, rotation=0)
	cbar = ax_b.collections[0].colorbar
	cbar.ax.tick_params(labelsize=58)
	cbar.set_label("Jensen-Shannon distance", size=64)

	fig.subplots_adjust(left=0.06, right=0.98, top=0.88, bottom=0.10)

	# ── Panel labels (a), (b) on same line as titles ──
	# Use axes-relative y so they line up with titles (which sit just above each axes top).
	ax_a1.text(-0.08, 1.234, "(a)", transform=ax_a1.transAxes,
	fontsize=66, fontweight="bold", va="center", ha="center")
	ax_b.text(-0.30, 1.083, "(b)", transform=ax_b.transAxes,
	fontsize=66, fontweight="bold", va="center", ha="center")

	out_dir = "/data/n5huang/dna_token/paper_figures"
	os.makedirs(out_dir, exist_ok=True)
	out = os.path.join(out_dir, "combined_panel_a_b.pdf")
	fig.savefig(out, dpi=300)
	fig.savefig(out.replace(".pdf", ".png"), dpi=80)
	print(f"Saved: {out}")
	print(f"Saved: {out.replace('.pdf', '.png')}")
	plt.close(fig)