supp_tokenizer_bench / code /js_vs_random_heatmap.py

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

89dec6b verified 17 days ago

8.4 kB

"""
Simplified region-vs-random JS-distance heatmap.
For each tokenizer, compute token-frequency distributions for
promoter, enhancer, exon, and a random background (genomic shuffle
excluding those three region sets, total bp matched to exon).
Metric: Jensen-Shannon distance (sqrt of JS divergence) of each
region distribution vs the random-background distribution.
Output: 5 (tokenizers) x 3 (regions) heatmap.
"""

from __future__ import annotations

import argparse
import subprocess
from collections import Counter
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib.colors import LinearSegmentedColormap
from scipy.spatial.distance import jensenshannon
from tokenizers import Tokenizer

ROOT = Path("/data/n5huang/dna_token")
GENOME = ROOT / "hg38.fa"
FAI = ROOT / "hg38.fa.fai"
OUTDIR = ROOT / "enrichment_analysis"

REGION_BEDS = {
    "promoter": ROOT / "promoters_2kb.clean.merged.bed",
    "enhancer": ROOT / "enhancers_dels.clean.merged.bed",
    "exon":     ROOT / "exon.clean.merged.bed",
}

TOKENIZER_PATHS = {
    "DNAbert2":   ROOT / "pretrain/models/DNAbert2_Pretrained/tokenizer.json",
    "Grover":     ROOT / "pretrain/models/Grover_Pretrained/tokenizer.json",
    "cCRE_region_BPE":  ROOT / "tokenizer_files/cCRE_region_BPE_tokenizer.json",
    "motif_region_BPE": ROOT / "tokenizer_files/motif_region_BPE_tokenizer.json",
}

CUSTOM_TOKENIZER_PATHS = {
    "DNAMotifTokenizer": ROOT / "tokenizer_files/DNAMotifTokenizer_official",
}

REGIONS = ["promoter", "enhancer", "exon"]

TITLE_MAP = {
    "merge_len2": "EvoLen (5120)",
    "DNAbert2": "DNABERT-2",
    "Grover": "Grover",
    "cCRE_region_BPE": "cCRE-Region BPE",
    "motif_region_BPE": "Motif-Region BPE",
    "DNAMotifTokenizer": "DNAMotif",
}


def make_chrom_sizes(fai: Path, out: Path) -> None:
    if out.exists() and out.stat().st_size > 0:
        return
    with fai.open() as fh, out.open("w") as oh:
        for line in fh:
            name, size, *_ = line.rstrip("\n").split("\t")
            # keep only canonical autosomes + X/Y (skip alt/random/unlocalized)
            if "_" in name or name in ("chrM", "chrEBV"):
                continue
            oh.write(f"{name}\t{size}\n")


def union_bed(beds: list[Path], out: Path) -> None:
    if out.exists() and out.stat().st_size > 0:
        return
    cmd = "cat " + " ".join(str(b) for b in beds) + \
          f" | cut -f1-3 | bedtools sort -i - | bedtools merge -i - > {out}"
    subprocess.run(cmd, shell=True, check=True)


def shuffle_to_random(exon_bed: Path, excl_bed: Path, chrom_sizes: Path, out: Path) -> None:
    """Shuffle exon intervals into genome excluding region union -> same bp as exon."""
    if out.exists() and out.stat().st_size > 0:
        return
    subprocess.run(
        f"bedtools shuffle -i {exon_bed} -g {chrom_sizes} -excl {excl_bed} -seed 42 "
        f"| bedtools sort -i - | bedtools merge -i - > {out}",
        shell=True, check=True,
    )


def get_fasta(bed: Path, out_fa: Path) -> None:
    if out_fa.exists() and out_fa.stat().st_size > 0:
        return
    subprocess.run(
        f"bedtools getfasta -fi {GENOME} -bed {bed} -fo {out_fa}",
        shell=True, check=True,
    )


def tokenize_fasta(fasta: Path, tok: Tokenizer, batch_size: int = 2000) -> Counter:
    counts: Counter = Counter()
    sequences: list[str] = []
    seq_parts: list[str] = []
    with fasta.open() as f:
        for line in f:
            line = line.strip()
            if line.startswith(">"):
                if seq_parts:
                    sequences.append("".join(seq_parts))
                    seq_parts = []
            else:
                seq_parts.append(line.upper().replace("N", ""))
    if seq_parts:
        sequences.append("".join(seq_parts))
    sequences = [s for s in sequences if s]

    for i in range(0, len(sequences), batch_size):
        batch = sequences[i:i + batch_size]
        for enc in tok.encode_batch(batch):
            counts.update(enc.tokens)
    return counts


def js_distance(p_counts: Counter, q_counts: Counter) -> float:
    vocab = sorted(set(p_counts) | set(q_counts))
    p = np.array([p_counts.get(t, 0) for t in vocab], dtype=float)
    q = np.array([q_counts.get(t, 0) for t in vocab], dtype=float)
    p_sum = p.sum() or 1.0
    q_sum = q.sum() or 1.0
    p /= p_sum
    q /= q_sum
    # scipy's jensenshannon already returns the distance (sqrt of divergence), base=2 below
    return float(jensenshannon(p, q, base=2))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--plot-only", action="store_true")
    args = parser.parse_args()

    OUTDIR.mkdir(exist_ok=True)
    beds_dir = OUTDIR / "beds_simple"
    fasta_dir = OUTDIR / "fastas_simple"
    beds_dir.mkdir(exist_ok=True)
    fasta_dir.mkdir(exist_ok=True)

    matrix_csv = OUTDIR / "js_vs_random_matrix.csv"

    if args.plot_only:
        df = pd.read_csv(matrix_csv, index_col=0)
    else:
        chrom_sizes = beds_dir / "chrom.sizes"
        make_chrom_sizes(FAI, chrom_sizes)

        # Build region union to use as exclusion
        excl = beds_dir / "region_union.bed"
        union_bed(list(REGION_BEDS.values()), excl)

        # Shuffle exon into non-region genome -> "random" background
        random_bed = beds_dir / "random.bed"
        shuffle_to_random(REGION_BEDS["exon"], excl, chrom_sizes, random_bed)

        # Extract FASTA for each set
        fastas: dict[str, Path] = {}
        for name, bed in REGION_BEDS.items():
            fa = fasta_dir / f"{name}.fa"
            get_fasta(bed, fa)
            fastas[name] = fa
        fa_rand = fasta_dir / "random.fa"
        get_fasta(random_bed, fa_rand)
        fastas["random"] = fa_rand

        rows = []
        tok_names_ordered = list(TOKENIZER_PATHS.keys()) + list(CUSTOM_TOKENIZER_PATHS.keys())
        # Load HF tokenizers
        loaded = {n: Tokenizer.from_file(str(p)) for n, p in TOKENIZER_PATHS.items()}
        # Load custom tokenizers
        import sys
        sys.path.insert(0, str(ROOT / "tokenizer_evaluation"))
        from dna_motif_tokenizer import DNAMotifTokenizer, Trie, TrieNode
        sys.modules["__main__"].Trie = Trie
        sys.modules["__main__"].TrieNode = TrieNode
        for n, p in CUSTOM_TOKENIZER_PATHS.items():
            loaded[n] = DNAMotifTokenizer(str(p), seed=42)

        for tok_name in tok_names_ordered:
            print(f"\nTokenizing with {tok_name}...")
            tok = loaded[tok_name]
            counts = {}
            for name, fa in fastas.items():
                counts[name] = tokenize_fasta(fa, tok)
                print(f"  {name}: {sum(counts[name].values()):,} tokens")
            row = {r: js_distance(counts[r], counts["random"]) for r in REGIONS}
            rows.append(row)
        df = pd.DataFrame(rows, index=tok_names_ordered, columns=REGIONS)
        df.to_csv(matrix_csv)
        print(f"\nSaved matrix: {matrix_csv}")
        print(df.round(4))

    # Plot
    display_index = [TITLE_MAP.get(n, n) for n in df.index]
    df_display = df.copy()
    df_display.index = display_index

    cmap = LinearSegmentedColormap.from_list("evo_heat", ["white", "#7EDCC5", "#636E72"])
    fig, ax = plt.subplots(figsize=(22, 32), constrained_layout=True)
    sns.heatmap(
        df_display,
        ax=ax,
        annot=True,
        fmt=".4f",
        annot_kws={"size": 56, "fontweight": "bold"},
        cmap=cmap,
        vmin=0.0,
        vmax=float(df_display.values.max()) or 0.001,
        linewidths=0.5,
        cbar_kws={"label": "Jensen-Shannon distance"},
    )
    ax.set_title("Region vs random-genome token-frequency separation",
                 fontsize=60, fontweight="bold", pad=30)
    ax.set_xlabel("Genomic region", fontsize=56)
    ax.set_ylabel("Tokenizer", fontsize=56)
    ax.tick_params(axis="x", labelsize=48)
    ax.tick_params(axis="y", labelsize=48, rotation=0)
    cbar = ax.collections[0].colorbar
    cbar.ax.tick_params(labelsize=48)
    cbar.set_label("Jensen-Shannon distance", size=56)

    out_pdf = OUTDIR / "js_vs_random_heatmap.pdf"
    out_png = OUTDIR / "js_vs_random_heatmap.png"
    fig.savefig(out_pdf, dpi=300, bbox_inches="tight")
    fig.savefig(out_png, dpi=300, bbox_inches="tight")
    print(f"Saved: {out_pdf}")
    print(f"Saved: {out_png}")


if __name__ == "__main__":
    main()

	"""
	Simplified region-vs-random JS-distance heatmap.

	For each tokenizer, compute token-frequency distributions for
	promoter, enhancer, exon, and a random background (genomic shuffle
	excluding those three region sets, total bp matched to exon).

	Metric: Jensen-Shannon distance (sqrt of JS divergence) of each
	region distribution vs the random-background distribution.

	Output: 5 (tokenizers) x 3 (regions) heatmap.
	"""

	from __future__ import annotations

	import argparse
	import subprocess
	from collections import Counter
	from pathlib import Path

	import matplotlib.pyplot as plt
	import numpy as np
	import pandas as pd
	import seaborn as sns
	from matplotlib.colors import LinearSegmentedColormap
	from scipy.spatial.distance import jensenshannon
	from tokenizers import Tokenizer

	ROOT = Path("/data/n5huang/dna_token")
	GENOME = ROOT / "hg38.fa"
	FAI = ROOT / "hg38.fa.fai"
	OUTDIR = ROOT / "enrichment_analysis"

	REGION_BEDS = {
	"promoter": ROOT / "promoters_2kb.clean.merged.bed",
	"enhancer": ROOT / "enhancers_dels.clean.merged.bed",
	"exon": ROOT / "exon.clean.merged.bed",
	}

	TOKENIZER_PATHS = {
	"DNAbert2": ROOT / "pretrain/models/DNAbert2_Pretrained/tokenizer.json",
	"Grover": ROOT / "pretrain/models/Grover_Pretrained/tokenizer.json",
	"cCRE_region_BPE": ROOT / "tokenizer_files/cCRE_region_BPE_tokenizer.json",
	"motif_region_BPE": ROOT / "tokenizer_files/motif_region_BPE_tokenizer.json",
	}

	CUSTOM_TOKENIZER_PATHS = {
	"DNAMotifTokenizer": ROOT / "tokenizer_files/DNAMotifTokenizer_official",
	}

	REGIONS = ["promoter", "enhancer", "exon"]

	TITLE_MAP = {
	"merge_len2": "EvoLen (5120)",
	"DNAbert2": "DNABERT-2",
	"Grover": "Grover",
	"cCRE_region_BPE": "cCRE-Region BPE",
	"motif_region_BPE": "Motif-Region BPE",
	"DNAMotifTokenizer": "DNAMotif",
	}


	def make_chrom_sizes(fai: Path, out: Path) -> None:
	if out.exists() and out.stat().st_size > 0:
	return
	with fai.open() as fh, out.open("w") as oh:
	for line in fh:
	name, size, *_ = line.rstrip("\n").split("\t")
	# keep only canonical autosomes + X/Y (skip alt/random/unlocalized)
	if "_" in name or name in ("chrM", "chrEBV"):
	continue
	oh.write(f"{name}\t{size}\n")


	def union_bed(beds: list[Path], out: Path) -> None:
	if out.exists() and out.stat().st_size > 0:
	return
	cmd = "cat " + " ".join(str(b) for b in beds) + \
	f" \| cut -f1-3 \| bedtools sort -i - \| bedtools merge -i - > {out}"
	subprocess.run(cmd, shell=True, check=True)


	def shuffle_to_random(exon_bed: Path, excl_bed: Path, chrom_sizes: Path, out: Path) -> None:
	"""Shuffle exon intervals into genome excluding region union -> same bp as exon."""
	if out.exists() and out.stat().st_size > 0:
	return
	subprocess.run(
	f"bedtools shuffle -i {exon_bed} -g {chrom_sizes} -excl {excl_bed} -seed 42 "
	f"\| bedtools sort -i - \| bedtools merge -i - > {out}",
	shell=True, check=True,
	)


	def get_fasta(bed: Path, out_fa: Path) -> None:
	if out_fa.exists() and out_fa.stat().st_size > 0:
	return
	subprocess.run(
	f"bedtools getfasta -fi {GENOME} -bed {bed} -fo {out_fa}",
	shell=True, check=True,
	)


	def tokenize_fasta(fasta: Path, tok: Tokenizer, batch_size: int = 2000) -> Counter:
	counts: Counter = Counter()
	sequences: list[str] = []
	seq_parts: list[str] = []
	with fasta.open() as f:
	for line in f:
	line = line.strip()
	if line.startswith(">"):
	if seq_parts:
	sequences.append("".join(seq_parts))
	seq_parts = []
	else:
	seq_parts.append(line.upper().replace("N", ""))
	if seq_parts:
	sequences.append("".join(seq_parts))
	sequences = [s for s in sequences if s]

	for i in range(0, len(sequences), batch_size):
	batch = sequences[i:i + batch_size]
	for enc in tok.encode_batch(batch):
	counts.update(enc.tokens)
	return counts


	def js_distance(p_counts: Counter, q_counts: Counter) -> float:
	vocab = sorted(set(p_counts) \| set(q_counts))
	p = np.array([p_counts.get(t, 0) for t in vocab], dtype=float)
	q = np.array([q_counts.get(t, 0) for t in vocab], dtype=float)
	p_sum = p.sum() or 1.0
	q_sum = q.sum() or 1.0
	p /= p_sum
	q /= q_sum
	# scipy's jensenshannon already returns the distance (sqrt of divergence), base=2 below
	return float(jensenshannon(p, q, base=2))


	def main() -> None:
	parser = argparse.ArgumentParser()
	parser.add_argument("--plot-only", action="store_true")
	args = parser.parse_args()

	OUTDIR.mkdir(exist_ok=True)
	beds_dir = OUTDIR / "beds_simple"
	fasta_dir = OUTDIR / "fastas_simple"
	beds_dir.mkdir(exist_ok=True)
	fasta_dir.mkdir(exist_ok=True)

	matrix_csv = OUTDIR / "js_vs_random_matrix.csv"

	if args.plot_only:
	df = pd.read_csv(matrix_csv, index_col=0)
	else:
	chrom_sizes = beds_dir / "chrom.sizes"
	make_chrom_sizes(FAI, chrom_sizes)

	# Build region union to use as exclusion
	excl = beds_dir / "region_union.bed"
	union_bed(list(REGION_BEDS.values()), excl)

	# Shuffle exon into non-region genome -> "random" background
	random_bed = beds_dir / "random.bed"
	shuffle_to_random(REGION_BEDS["exon"], excl, chrom_sizes, random_bed)

	# Extract FASTA for each set
	fastas: dict[str, Path] = {}
	for name, bed in REGION_BEDS.items():
	fa = fasta_dir / f"{name}.fa"
	get_fasta(bed, fa)
	fastas[name] = fa
	fa_rand = fasta_dir / "random.fa"
	get_fasta(random_bed, fa_rand)
	fastas["random"] = fa_rand

	rows = []
	tok_names_ordered = list(TOKENIZER_PATHS.keys()) + list(CUSTOM_TOKENIZER_PATHS.keys())
	# Load HF tokenizers
	loaded = {n: Tokenizer.from_file(str(p)) for n, p in TOKENIZER_PATHS.items()}
	# Load custom tokenizers
	import sys
	sys.path.insert(0, str(ROOT / "tokenizer_evaluation"))
	from dna_motif_tokenizer import DNAMotifTokenizer, Trie, TrieNode
	sys.modules["__main__"].Trie = Trie
	sys.modules["__main__"].TrieNode = TrieNode
	for n, p in CUSTOM_TOKENIZER_PATHS.items():
	loaded[n] = DNAMotifTokenizer(str(p), seed=42)

	for tok_name in tok_names_ordered:
	print(f"\nTokenizing with {tok_name}...")
	tok = loaded[tok_name]
	counts = {}
	for name, fa in fastas.items():
	counts[name] = tokenize_fasta(fa, tok)
	print(f" {name}: {sum(counts[name].values()):,} tokens")
	row = {r: js_distance(counts[r], counts["random"]) for r in REGIONS}
	rows.append(row)
	df = pd.DataFrame(rows, index=tok_names_ordered, columns=REGIONS)
	df.to_csv(matrix_csv)
	print(f"\nSaved matrix: {matrix_csv}")
	print(df.round(4))

	# Plot
	display_index = [TITLE_MAP.get(n, n) for n in df.index]
	df_display = df.copy()
	df_display.index = display_index

	cmap = LinearSegmentedColormap.from_list("evo_heat", ["white", "#7EDCC5", "#636E72"])
	fig, ax = plt.subplots(figsize=(22, 32), constrained_layout=True)
	sns.heatmap(
	df_display,
	ax=ax,
	annot=True,
	fmt=".4f",
	annot_kws={"size": 56, "fontweight": "bold"},
	cmap=cmap,
	vmin=0.0,
	vmax=float(df_display.values.max()) or 0.001,
	linewidths=0.5,
	cbar_kws={"label": "Jensen-Shannon distance"},
	)
	ax.set_title("Region vs random-genome token-frequency separation",
	fontsize=60, fontweight="bold", pad=30)
	ax.set_xlabel("Genomic region", fontsize=56)
	ax.set_ylabel("Tokenizer", fontsize=56)
	ax.tick_params(axis="x", labelsize=48)
	ax.tick_params(axis="y", labelsize=48, rotation=0)
	cbar = ax.collections[0].colorbar
	cbar.ax.tick_params(labelsize=48)
	cbar.set_label("Jensen-Shannon distance", size=56)

	out_pdf = OUTDIR / "js_vs_random_heatmap.pdf"
	out_png = OUTDIR / "js_vs_random_heatmap.png"
	fig.savefig(out_pdf, dpi=300, bbox_inches="tight")
	fig.savefig(out_png, dpi=300, bbox_inches="tight")
	print(f"Saved: {out_pdf}")
	print(f"Saved: {out_png}")


	if __name__ == "__main__":
	main()