plot_ntroot_lai.py

#!/usr/bin/env python3
"""
plot_ntroot_lai.py
RLW/LC 06/2026
"""

import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
import os
import argparse
import numpy as np

# Define parameters / assumes GRCh38
# Can be overriden using --fai option to read chromosome lengths from a FASTA index file
CHR_LENGTHS = {
    "1":248956422,"2":242193529,"3":198295559,"4":190214555,
    "5":181538259,"6":170805979,"7":159345973,"8":145138636,
    "9":138394717,"10":133797422,"11":135086622,"12":133275309,
    "13":114364328,"14":107043718,"15":101991189,"16":90338345,
    "17":83257441,"18":80373285,"19":58617616,"20":64444167,
    "21":46709983,"22":50818468
}

# Define colours
COLOR_MAP = {
    "AFR": "orange",
    "EUR": "green",
    "EAS": "red",
    "AMR": "black",
    "SAS": "blue",
    "CAS": "violet",
    "OCE": "lightgreen",
    "NA": "#c000c0",
    "REF": "#d3d3d3"
}

# Define cluster colors
UNSUPERVISED_PALETTE = [
    "#4E79A7",  # blue
    "#F28E2B",  # orange
    "#E15759",  # red
    "#76B7B2",  # teal
    "#59A14F",  # green
    "#EDC948",  # yellow
    "#B07AA1",  # purple
    "#FF9DA7",  # pink
    "#9C755F",  # brown
    "#BAB0AC",  # gray
    "#8dd3c7",  # pastels V
    "#ffffb3",
    "#bebada",
    "#fb8072",
    "#80b1d3",
    "#fdb462",
    "#b3de69",
    "#fccde5",
    "#d9d9d9",
    "#bc80bd"
]

ANCESTRY_LABELS = {
    "AFR": "African",
    "AMR": "Admixed American",
    "EUR": "European",
    "EAS": "East Asian",
    "CAS": "Central Asian",
    "OCE": "Oceania",
    "NA": "Not available",
    "SAS": "South Asian"
}

#-----------------------------
def read_fai(fai_path):
    """
    Read chromosome lengths from a FASTA index (.fai) file.
    This function updates the global CHR_LENGTHS dictionary with values
    from the provided .fai file, overriding the default GRCh38 lengths.
    """
    global CHR_LENGTHS
    CHR_LENGTHS = {}
    
    with open(fai_path, "r") as f:
        for line in f:
            parts = line.strip().split("\t")
            if len(parts) < 2:
                continue
            chrom, length_str = parts[0], parts[1]
            CHR_LENGTHS[chrom] = int(length_str)

# ----------------------------
def get_cluster_color(cluster_id, unique_clusters):
    """
    Define unsupervised color
    """
    palette = UNSUPERVISED_PALETTE
    cluster_to_idx = {c: i for i, c in enumerate(unique_clusters)}
    idx = cluster_to_idx[cluster_id]
    return palette[idx % len(palette)]

# ----------------------------
def build_dynamic_ancestry_colors(ancestries, color_map, palette):
    """
    Define dynamic colour
    """
    dynamic = {}
    used = set(color_map.keys())

    # stable ordering prevents nondeterministic coloring
    unknowns = sorted([a for a in ancestries if a not in used])

    for i, anc in enumerate(unknowns):
        dynamic[anc] = palette[i % len(palette)]

    return dynamic

# -----------------------------
def parse_args():
    """
    Parse command-line arguments.
    """

    parser = argparse.ArgumentParser(
        description="Generate an interactive Local Ancestry (LAI) visualization."
    )

    parser.add_argument(
        "input_tsv",
        help="Tile-resolution ancestry prediction TSV generated by ntRoot/korus."
    )

    parser.add_argument(
        "output_html",
        help="Output interactive HTML visualization."
    )

    parser.add_argument(
        "predictor",
        nargs="?",
        default="ntRoot",
        help="Predictor name displayed in the plot title (default: ntRoot)."
    )
    
    parser.add_argument(
        "--fai", help="Optional FASTA index file to specify chromosome lengths (overrides built-in GRCh38).",
        required=False, type=str
    )

    return parser.parse_args()

# -----------------------------
def extract_sample_name(path):
    """
    Extract a sample identifier from an ntRoot/korus LAI TSV filename.

    The sample name is assumed to be the leading token before the first
    underscore ('_'). If no underscore exists, the leading token before
    the first period ('.') is used. This value is displayed beneath the
    plot title as a subtitle.

    Examples
    --------
    HG002_ntedit_k50_variants.vcf_ancestry.tsv -> HG002
    sample123.tsv                              -> sample123
    """

    name = os.path.basename(path)

    if "_" in name:
        return name.split("_", 1)[0]

    if "." in name:
        return name.split(".", 1)[0]

    return name

# -----------------------------
def load_data(path):
    """
    Load and validate an ntRoot/korus LAI tile-resolution TSV file.

    The input table is expected to contain:
      - chrom
      - start
      - end
      - ancestry_prediction

    Chromosome identifiers are converted to strings and restricted to
    autosomes present in CHR_LENGTHS. Coordinate columns are converted
    to numeric types to support downstream angular calculations.

    Returns
    -------
    pandas.DataFrame
        Cleaned dataframe ready for visualization.
    """

    df = pd.read_csv(path, sep="\t")
    df["chrom"] = df["chrom"].astype(str)
    df = df[df["chrom"].isin(CHR_LENGTHS.keys())].copy()
    df["start"] = pd.to_numeric(df["start"])
    df["end"] = pd.to_numeric(df["end"])
    return df

# -----------------------------
def build_angles():
    """
    Calculate chromosome angular coordinates for the circular genome plot.

    Chromosomes are allocated angular spans proportional to their GRCh38
    lengths. A fixed inter-chromosomal gap is inserted between adjacent
    chromosomes. Chromosome 1 begins at 0 degrees ("true north"), with
    the final rotation applied later during Plotly layout configuration.

    Returns
    -------
    dict
        Chromosome -> (start_angle, end_angle)
    """

    total = sum(CHR_LENGTHS.values())
    gap = 2.0
    usable = 360 - gap * len(CHR_LENGTHS)

    angles = {}
    current = 0.0   # chr1 starts at 0

    for c, L in CHR_LENGTHS.items():
        span = usable * L / total
        angles[c] = (current, current + span)
        current += span + gap

    return angles

# -----------------------------
def build(df, out_html, input_tsv, predictor):
    """
    Construct the interactive ntRoot/korus Local Ancestry (LAI) visualization.

    The figure contains four major components:

    1. Chromosome backbone
       Circular ideogram-like representation of chromosomes with lengths
       scaled according to GRCh38 chromosome sizes.

    2. Local ancestry tracks
       Tile-resolution ancestry assignments displayed as colored radial
       segments positioned according to genomic coordinates.

    3. Global ancestry summary
       Central donut chart summarizing ancestry proportions across all
       analyzed genomic tiles.

    4. Interactive chromosome labels
       Hoverable chromosome labels that dynamically highlight ancestry
       assignments from a selected chromosome.

    Additional JavaScript is injected into the exported HTML to support
    interactive highlighting between ancestry tracks, chromosome labels,
    and donut-chart ancestry categories.

    Parameters
    ----------
    df : pandas.DataFrame
        Parsed ntRoot/korus LAI tile-resolution table.

    out_html : str
        Output HTML file path.

    input_tsv : str
        Original TSV filename used to derive the displayed sample name.
    """

    sample_name = extract_sample_name(input_tsv)

    ancestries = sorted(df["ancestry_prediction"].dropna().unique())

    dynamic_color_map = build_dynamic_ancestry_colors(
        ancestries,
        COLOR_MAP,
        UNSUPERVISED_PALETTE
    )

    PLOT_Y0 = 0.09
    PLOT_Y1 = 0.97

    center_y = (PLOT_Y0 + PLOT_Y1) / 2

    pie_height = 0.3   # adjust size
    pie_width  = 0.3

    def anc_color(anc):
        if anc in COLOR_MAP:
            return COLOR_MAP[anc]
        return dynamic_color_map.get(anc, "#999999")

    angles = build_angles()

    fig = go.Figure()

    # =============================
    # BACKBONE CHROMOSOMES
    # =============================
    for c in CHR_LENGTHS:

        # -----------------------------
        # CHROMOSOME BOUNDARY SEPARATORS (dark radial lines)
        # -----------------------------
        a0, _ = angles[c]

        fig.add_trace(
                go.Scatterpolar(
                    r=[0.55, 1.05],
                    theta=[a0, a0],
                    mode="lines",
                    line=dict(color="rgba(80,80,80,0.35)", width=1),
                    hoverinfo="skip",
                    showlegend=False,
                    meta="BACKBONE_SEPARATOR"
                )
        )

        a0, a1 = angles[c]

        sub = df[df["chrom"] == c]
        summary_raw = sub.groupby("ancestry_prediction")["end"].count()
        summary_pct = (summary_raw / summary_raw.sum() * 100).sort_values(ascending=False)

        hover_txt = (
            f"chromosome{c}<br>"
            f"{CHR_LENGTHS[c]/1e6:.1f} Mb<br><br>"
            + "<br>".join([f"{k}: {v:.1f}%" for k, v in summary_pct.items()])
        )

        fig.add_trace(
            go.Barpolar(
                r=[1.0],
                #r=[0.2,1.0],
                theta=[(a0 + a1) / 2],
                width=[a1 - a0],
                base=0.0,
                marker=dict(color="#e0e0e0", line=dict(color="white", width=1)),
                hoverinfo="skip",
                showlegend=False,
                meta=f"BACKBONE|{c}"
            )
        )

    for anc in ancestries:

        if anc in ANCESTRY_LABELS:
            legend_name = f"{ANCESTRY_LABELS[anc]} ({anc})"
        else:
            legend_name = anc

        fig.add_trace(
            go.Scatterpolar(
                r=[None],
                theta=[None],
                mode="markers",
                marker=dict(color=anc_color(anc), size=12),
                name=legend_name,
                legendgroup=anc,
                showlegend=True
            )
        )

    # =============================
    # LAI TRACKS
    # =============================
    for anc in ancestries:
        for c in CHR_LENGTHS:

            sub = df[
                (df["ancestry_prediction"] == anc) &
                (df["chrom"] == c)
            ]

            t, w, custom = [], [], []

            for _, r in sub.iterrows():
                c = str(r["chrom"])
                a0, a1 = angles[c]
                span = a1 - a0
                clen = CHR_LENGTHS[c]

                t0 = a0 + span * (r["start"] / clen)
                t1 = a0 + span * (r["end"] / clen)

                t.append((t0 + t1) / 2)
                w.append(t1 - t0) 
                custom.append([c, int(r["start"]), int(r["end"]), ANCESTRY_LABELS.get(anc, anc)])

            fig.add_trace(
                go.Barpolar(
                    r=[0.45]*len(t),
                    theta=t,
                    width=w,
                    base=0.55,
                    marker=dict(
                        color=anc_color(anc),
                        line=dict(width=0)
                    ),
                    legendgroup=anc,
                    showlegend=False,
                    customdata=custom,
                    hovertemplate=
                        "Chromosome %{customdata[0]}<br>"
                        "%{customdata[1]:,} - %{customdata[2]:,}<br>"
                        "%{customdata[3]}<extra></extra>",
                    meta=f"{anc}|{c}"
                )
            )

    # =============================
    # PIE
    # =============================
    df["tile_len"] = df["end"] - df["start"] + 1
    frac_bp = df.groupby("ancestry_prediction")["tile_len"].sum()
    frac_mb = (frac_bp / 1e6).round(1)
    frac_pct = frac_bp / frac_bp.sum() * 100

    pie_text = []

    for anc in frac_mb.index:
        pct = frac_pct[anc]

        if pct >= 10:
            pie_text.append(f"{anc}<br>{pct:.1f}%")
        else:
            pie_text.append("")

    fig.add_trace(
        go.Pie(
            labels=list(frac_mb.index),
            values=list(frac_mb.values),
            hole=0.55,
            sort=False,
            marker=dict(colors=[anc_color(a) for a in frac_mb.index]),
            text=pie_text,
            textinfo="text",
            textposition="inside",
            insidetextorientation="horizontal", 
            textfont=dict(size=11, color="white"),
            showlegend=False,
            domain=dict(
                x=[0.5 - pie_width/2, 0.5 + pie_width/2],
                y=[center_y - pie_height/2, center_y + pie_height/2]
            ),
            customdata=[ANCESTRY_LABELS.get(a, a) for a in frac_mb.index],
            hovertemplate="%{customdata}<br>%{percent:.1%}<br>%{value} Mb<extra></extra>",
        )
    )
    # =============================
    # CHR LABELS
    # =============================
    label_theta = []
    label_text = []
    hover_texts = []

    for c in CHR_LENGTHS:
        a0, a1 = angles[c]

        sub = df[df["chrom"] == c]
        summary_raw = sub.groupby("ancestry_prediction")["end"].count()
        summary_pct = (summary_raw / summary_raw.sum() * 100).sort_values(ascending=False)

        hover_txt = (
            f"{CHR_LENGTHS[c]/1e6:.1f} Mb<br><br>"
            + "<br>".join([f"{k}: {v:.1f}%" for k, v in summary_pct.items()])
        )

        label_theta.append((a0 + a1) / 2)
        label_text.append(f"{c}")
        hover_texts.append(hover_txt)

    fig.add_trace(
        go.Scatterpolar(
            r=[1.12]*len(label_theta),
            theta=label_theta,
            mode="text",
            text=label_text,
            customdata=[
                [c, hover_texts[i]] for i, c in enumerate(CHR_LENGTHS)
            ],
            hovertemplate=
                "Chromosome %{customdata[0]}<br>"
                "%{customdata[1]}<extra></extra>",
            showlegend=False,
            textfont=dict(size=14),
            meta="CHR_LABELS"
        )
    )

    # =============================
    # LAYOUT
    # =============================
    fig.update_layout(

        title=dict(
            text=f"{predictor} — Chromosomal Local Ancestry (LAI)<br><sup>{sample_name}</sup>",
            x=0.38,
            y=0.95,
            yanchor="top",
            xanchor="center",
            pad=dict(t=0, b=0),
            font=dict(size=18)
        ),      
        legend=dict(
            x=1.4,
            y=1.0,
            xanchor="right",
            yanchor="top",
            font=dict(size=12)
        ),
        polar=dict(
            bgcolor="white",
            radialaxis=dict(visible=False, range=[0, 1.2]),
            angularaxis=dict(
                rotation=90,        # chr1 at top
                direction="clockwise",
                showgrid=False,
                showticklabels=False  
            ),
            domain=dict(y=[PLOT_Y0, PLOT_Y1])
        ),
        margin=dict(l=40, r=140, t=40, b=20),
        autosize=True,
        height=750
    )

    # =============================
    # LINKING (LAI ↔ PIE ↔ BACKBONE)
    # =============================
    js = """
<script>
document.addEventListener("DOMContentLoaded", function() {

const plot = document.getElementsByClassName("plotly-graph-div")[0];

let selectedAnc = null;

plot.on("plotly_click", function(e){
    const pt = e.points[0];

    // PIE click
    if (pt.data.type === "pie") {
        selectedAnc = pt.label;
        highlight(null, selectedAnc);
        return;
    }

    // LAI segment click
    if (pt.data.meta && pt.data.meta.includes("|")) {
        const anc = pt.data.meta.split("|")[0];
        selectedAnc = anc;
        highlight(null, anc);
        return;
    }
});

function highlight(activeChrom, activeAnc){

    const op = [];

    const ancToUse = activeAnc || selectedAnc;

    const modeChrom = activeChrom !== null && ancToUse === null;
    const modeAnc = ancToUse !== null;

    for (let i = 0; i < plot.data.length; i++) {
        const tr = plot.data[i];

        if (tr.meta === "CHR_LABELS" || !tr.meta) {
            op.push(1.0);
            continue;
        }

        // -------------------------
        // BACKBONE CHROMOSOMES
        // -------------------------
        if (tr.meta.startsWith("BACKBONE")) {

           const backboneChr = tr.meta.split("|")[1];

           if (modeChrom) {
               op.push(backboneChr === activeChrom ? 1.0 : 0.15);
           } else {
               op.push(1.0);
           }

           continue;
        }


        const [anc, chr] = tr.meta.split("|");

        // -------------------------
        // CHR MODE
        // -------------------------
        if (modeChrom) {
            const keep = (chr === activeChrom);
            op.push(keep ? 1.0 : 0.05);
            continue;
        }

        // -------------------------
        // ANC MODE
        // -------------------------
        if (modeAnc) {
            const keep = (anc === ancToUse);
            op.push(keep ? 1.0 : 0.08);
            continue;
        }

        // -------------------------
        // DEFAULT
        // -------------------------
        op.push(1.0);
    }

    Plotly.restyle(plot, {
        opacity: op
    });
}

plot.on("plotly_doubleclick", function(){
    selectedAnc = null;
    Plotly.restyle(plot, {"opacity": plot.data.map(_ => 1.0)});
});

plot.on("plotly_legendclick", function(){
    selectedAnc = null;

    // allow Plotly's normal legend toggle to proceed
    setTimeout(function(){
        Plotly.restyle(
            plot,
            {"opacity": plot.data.map(_ => 1.0)}
        );
    }, 0);

    return true;
});

plot.on("plotly_legenddoubleclick", function(){
    selectedAnc = null;

    setTimeout(function(){
        Plotly.restyle(
            plot,
            {"opacity": plot.data.map(_ => 1.0)}
        );
    }, 0);

    return true;
});

plot.on("plotly_hover", function(e){
    const pt = e.points[0];

    // -------------------------
    // chromosome label hover
    // -------------------------
    if (pt.data.meta === "CHR_LABELS") {
        const chrom = pt.customdata[0];
        highlight(chrom, null);
        return;
    }

    // -------------------------
    // LAI hover → chromosome focus
    // -------------------------
    if (pt.data.meta && pt.data.meta.includes("|")) {
        const parts = pt.data.meta.split("|");
        const anc = parts[0];
        const chr = parts[1];

        highlight(chr, null);
        return;
    }

});

plot.on("plotly_unhover", function(){

    // keep persistent ancestry selection if user clicked
    if (selectedAnc !== null) {
        highlight(null, selectedAnc);
        return;
    }

    highlight(null, null);
});

});
</script>
"""

    html = fig.to_html(full_html=True, include_plotlyjs=True, config={"responsive": True}) 
    html = html.replace("</body>", js + "\n</body>")

    with open(out_html, "w") as f:
        f.write(html)

# -----------------------------
def main():
    """
    Command-line entry point.

    Expected arguments
    ------------------
    input.tsv
        Tile-resolution local ancestry prediction file.

    output.html
        Interactive Plotly visualization.

    [predictor]
        Optional predictor name displayed in the title
        (default: ntRoot).
        
    --fai
        Optional FASTA index file to specify chromosome lengths
        (overrides built-in GRCh38).

    The TSV is loaded, converted into an interactive circular LAI plot,
    and written as a standalone HTML document.
    """

    args = parse_args()
    
    if args.fai:
        # Override chromosome lengths if FASTA index provided
        read_fai(args.fai)

    df = load_data(args.input_tsv)

    build(
        df,
        args.output_html,
        args.input_tsv,
        args.predictor
    )

if __name__ == "__main__":
    main()