generate_docs_assets.py

"""Render the chart assets the docs site embeds.

For each example we ship two artifacts:

  - `docs/images/<name>.svg` — embedded inline on the docs page so the
    chart is visible without the user clicking anything.
  - `docs/charts/<name>.html` — full interactive Altair-rendered chart
    that the docs page links to ("Open the interactive chart →").

Both are gitignored. This script regenerates them from the example
modules. Run before `mkdocs build`; both `scripts/build_docs.sh` and
`.github/workflows/docs.yml` invoke it.

Idempotent and safe to re-run. If `examples/data/` is missing the two
upstream Auspice JSONs, the script invokes
`examples/fetch_auspice_data.py`'s `main()` to download them first.

Adding a new example: add a clause to `_render_examples()`.
"""

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path
from types import ModuleType

import altair as alt

import tree_annotated_plot

REPO = Path(__file__).resolve().parent.parent
EXAMPLES = REPO / "examples"
DATA_DIR = EXAMPLES / "data"
DOCS_IMAGES = REPO / "docs" / "images"
DOCS_CHARTS = REPO / "docs" / "charts"


def _import_path(name: str, path: Path) -> ModuleType:
    """Import a Python file by path. Used for the hyphenated-filename
    chart-builder module; same trick `tests/test_real_data.py` uses."""
    spec = importlib.util.spec_from_file_location(name, path)
    assert spec is not None and spec.loader is not None
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def _save_pair(chart, basename: str) -> None:
    """Save `chart` as both SVG (inline thumbnail) and HTML (interactive)."""
    DOCS_IMAGES.mkdir(parents=True, exist_ok=True)
    DOCS_CHARTS.mkdir(parents=True, exist_ok=True)
    svg_path = DOCS_IMAGES / f"{basename}.svg"
    html_path = DOCS_CHARTS / f"{basename}.html"
    chart.save(str(svg_path))
    chart.save(str(html_path))
    print(f"wrote {svg_path.relative_to(REPO)} ({svg_path.stat().st_size:,} B)")
    print(f"wrote {html_path.relative_to(REPO)} ({html_path.stat().st_size:,} B)")


def _render_synthetic() -> None:
    """The 8-tip synthetic example: minimum end-to-end."""
    syn = _import_path("syn_example", EXAMPLES / "synthetic_example.py")
    chart = syn.build_chart(syn.synthetic_titers())
    out = tree_annotated_plot.plot(
        syn.synthetic_auspice(),
        chart,
        chart_strain_field="strain",
        tree_strain_field="name",
        branch_length="div",
        tree_size=140,
    )
    _save_pair(out, "synthetic_example")


def _ensure_kikawa_auspice() -> None:
    """Download the H3N2 / H1N1 Auspice JSONs into examples/data/ if absent."""
    needed = [
        DATA_DIR / "flu-seqneut-2025to2026_H3N2.json",
        DATA_DIR / "flu-seqneut-2025to2026_H1N1.json",
    ]
    if all(p.exists() for p in needed):
        return
    fetcher = _import_path("fetcher", EXAMPLES / "fetch_auspice_data.py")
    fetcher.main()


def _render_kikawa() -> None:
    """The Kikawa flu-seqneut H3N2 (vertical) and H1N1 (horizontal)
    end-to-ends — real Auspice JSON, real chart with VConcat(Facet(Layer))
    structure, scale bar enabled."""
    _ensure_kikawa_auspice()
    builder = _import_path(
        "kikawa_builder", EXAMPLES / "flu-seqneut-2025to2026_titer_charts.py"
    )
    titers, viruses, sera = builder.load_data()
    metadata = builder.build_metadata(sera)
    all_cohorts = ["All"] + sorted(sera["cohort"].unique())

    for subtype, chart_type, basename in [
        ("H3N2", "iqr", "h3n2"),
        ("H1N1", "lines", "h1n1"),
    ]:
        chart = builder.make_chart(
            subtype=subtype,
            chart_type=chart_type,
            titers=titers,
            viruses=viruses,
            metadata=metadata,
            all_cohorts=all_cohorts,
        )
        # Render the bare chart (no tree) so the docs page can show what
        # the chart looks like before tree-annotated-plot wraps it.
        _save_pair(chart, f"{basename}_chart_only")
        plot_kwargs = dict(
            chart_strain_field="axis_label",
            tree_strain_field="derived_haplotype",
            branch_length="div",
            tree_size=140,
            scale_bar=True,
            branch_length_units="substitutions",
        )
        if subtype == "H3N2":
            # Color H3N2 by subclade so the docs SVG matches what users see
            # on Nextstrain. The Auspice JSON's meta.colorings.subclade has
            # no `scale` defined, so colors come from the default palette.
            # Place the subclade legend on the left so it doesn't compete
            # with the cohort-selection legend that already sits at the
            # bottom of the chart.
            plot_kwargs["color_tree_by"] = "subclade"
            plot_kwargs["tree_color_legend_format"] = {"orient": "left"}
        out = tree_annotated_plot.plot(
            DATA_DIR / f"flu-seqneut-2025to2026_{subtype}.json",
            chart,
            **plot_kwargs,
        )
        _save_pair(out, f"{basename}_combined")

    # H3N2 again, with `connect_leader_to_label=True` and a 9-point label
    # font: opt-in label connection where leaders run flush into the
    # labels rendered alongside the tree (the default keeps the chart's
    # natural strain-axis labels untouched).
    h3n2_chart = builder.make_chart(
        subtype="H3N2",
        chart_type="iqr",
        titers=titers,
        viruses=viruses,
        metadata=metadata,
        all_cohorts=all_cohorts,
    )
    out = tree_annotated_plot.plot(
        DATA_DIR / "flu-seqneut-2025to2026_H3N2.json",
        h3n2_chart,
        chart_strain_field="axis_label",
        tree_strain_field="derived_haplotype",
        branch_length="div",
        tree_size=140,
        scale_bar=True,
        branch_length_units="substitutions",
        connect_leader_to_label=True,
        strain_label_font_size=9,
        shift_tree_loc=60,
        color_tree_by="subclade",
        tree_color_legend_format={"orient": "left"},
    )
    _save_pair(out, "h3n2_combined_label_connect")

    # The same label-connect H3N2 chart, with the recent vaccine strains'
    # labels drawn in bold red while a checkbox is ticked: `strain_label_styles`
    # with a `when` param. The checkbox starts ticked so the static SVG shows
    # the emphasis.
    h3n2_chart_styles = builder.make_chart(
        subtype="H3N2",
        chart_type="iqr",
        titers=titers,
        viruses=viruses,
        metadata=metadata,
        all_cohorts=all_cohorts,
    )
    vaccine_strains = alt.param(
        name="vaccine_strains",
        value=True,
        bind=alt.binding_checkbox(name="highlight vaccine strains "),
    )
    out = tree_annotated_plot.plot(
        DATA_DIR / "flu-seqneut-2025to2026_H3N2.json",
        h3n2_chart_styles,
        chart_strain_field="axis_label",
        tree_strain_field="derived_haplotype",
        branch_length="div",
        tree_size=140,
        scale_bar=True,
        branch_length_units="substitutions",
        connect_leader_to_label=True,
        strain_label_font_size=9,
        shift_tree_loc=60,
        color_tree_by="subclade",
        tree_color_legend_format={"orient": "left"},
        strain_label_styles=[
            {
                "strains": builder.vaccine_axis_labels(viruses, "H3N2"),
                "style": {"fill": "#d62728", "fontWeight": "bold"},
                "when": vaccine_strains,
            }
        ],
    )
    _save_pair(out, "h3n2_combined_label_styles")

    # H3N2 once more, colored by genotype at HA1 site 158: same chart and
    # default layout as `h3n2_combined`, with `color_tree_by` switched to
    # the genotype form. Site 158 has two mutations (N158K, N158D) in the
    # tree, so this renders three states (N, K, D).
    h3n2_chart_genotype = builder.make_chart(
        subtype="H3N2",
        chart_type="iqr",
        titers=titers,
        viruses=viruses,
        metadata=metadata,
        all_cohorts=all_cohorts,
    )
    out = tree_annotated_plot.plot(
        DATA_DIR / "flu-seqneut-2025to2026_H3N2.json",
        h3n2_chart_genotype,
        chart_strain_field="axis_label",
        tree_strain_field="derived_haplotype",
        branch_length="div",
        tree_size=140,
        scale_bar=True,
        branch_length_units="substitutions",
        color_tree_by="genotype:HA1:158",
        tree_color_legend_format={"orient": "left"},
    )
    _save_pair(out, "h3n2_combined_genotype_158")

    # H3N2 once more, demonstrating all four appearance-tuning knobs:
    # an explicit `tree_color_scale` (Okabe-Ito-inspired colorblind-safe
    # palette, ordered K, J.2.4, J.2.3, J.2.2, J.2, G.1.3.1), a 14-pt
    # legend on the left, and a 14-pt scale-bar label.
    h3n2_chart_custom = builder.make_chart(
        subtype="H3N2",
        chart_type="iqr",
        titers=titers,
        viruses=viruses,
        metadata=metadata,
        all_cohorts=all_cohorts,
    )
    out = tree_annotated_plot.plot(
        DATA_DIR / "flu-seqneut-2025to2026_H3N2.json",
        h3n2_chart_custom,
        chart_strain_field="axis_label",
        tree_strain_field="derived_haplotype",
        branch_length="div",
        tree_size=140,
        scale_bar=True,
        branch_length_units="substitutions",
        color_tree_by="subclade",
        tree_color_scale={
            "K": "#0072B2",
            "J.2.4": "#009E73",
            "J.2.3": "#D55E00",
            "J.2.2": "#CC79A7",
            "J.2": "#56B4E9",
            "G.1.3.1": "#E69F00",
        },
        tree_color_legend_format={
            "orient": "left",
            "labelFontSize": 14,
            "titleFontSize": 14,
        },
        scale_bar_font_size=14,
    )
    _save_pair(out, "h3n2_combined_custom_colors")


def main() -> None:
    """Render every example to SVG + interactive HTML under `docs/`."""
    _render_synthetic()
    _render_kikawa()


if __name__ == "__main__":
    try:
        main()
    except Exception as exc:
        print(f"asset generation failed: {exc}", file=sys.stderr)
        raise