""" 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()