Blame for code/js_vs_random_heatmap.py

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