add color_tree_by parameter for Auspice-style tree coloring

CHANGELOG.md

@@ -16,6 +16,20 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `strain_label_font_size` (default `10`), `strain_label_font_weight`
   (default `"normal"`), and `shift_tree_loc` (default `0`) for tuning
   the size, weight, and placement of the connected labels.
+- `color_tree_by` (default `None`): color the tree's branches and tip
+  circles by an Auspice node attribute (e.g. `"subclade"`) or by the
+  inferred genotype state at one or more sites
+  (e.g. `"genotype:HA1:158"` or `"genotype:HA1:158,189"`). Colors,
+  category ordering, and the bottom-of-plot legend match the
+  Nextstrain view of the same tree.
+
+### Changed
+
+- Default `tree_line_width` bumped from `1.5` to `2`, default
+  `tree_node_size` from `28` to `45`. Tree branch lines and tip
+  circles are now drawn at full opacity. The thicker / fuller
+  defaults read better when the tree is colored (the prior values
+  were tuned for unicolor black trees).
 
 ## [0.1.0] - 2026-05-04
 

docs/examples.md

@@ -103,6 +103,11 @@ chart-builder wrote them (fonts, ticks, axis title, and all), and the
 tree's dashed leader lines stop at the tree panel's chart-facing
 edge.
 
+The H3N2 example above is rendered with `color_tree_by="subclade"`,
+which colors the tree's branches and tip circles by the
+`node_attrs.subclade` value at each node and adds a categorical legend
+below the plot. See "Color the tree" below for the full set of options.
+
 ### Optional: connect leaders all the way to the labels
 
 If you'd prefer the dashed leaders to run flush into the strain
@@ -131,6 +136,40 @@ CLI flags: `--connect-leader-to-label --strain-label-font-size 9
 --shift-tree-loc 60`. In Python:
 `connect_leader_to_label=True, strain_label_font_size=9, shift_tree_loc=60`.
 
+### Color the tree
+
+Pass `color_tree_by` to color the tree's branches and tip circles by
+any property the Auspice JSON exposes — broadly, anything that
+appears in the "Color By" dropdown on the Nextstrain view of the same
+tree. Two forms are supported:
+
+- A named attribute, e.g. `color_tree_by="subclade"` (used in the
+  example above). Common alternatives include `"clade_membership"`,
+  `"region"`, `"country"` — whichever the tree provides.
+- A genotype at one or more sites in a gene. For a single site,
+  `color_tree_by="genotype:HA1:158"` colors each tip by the amino
+  acid at HA1 site 158. A comma-separated list gives a haplotype:
+  `"genotype:HA1:158,189"`. Sites that don't vary in the tree are
+  dropped from the haplotype label.
+
+Colors match what you'd see on the Nextstrain view of the same tree —
+either from the JSON's palette information when the build provides it,
+or from the same default palette Auspice uses when it doesn't.
+Categories are ordered by descending frequency in both cases. Missing
+values render in gray, and the legend is drawn at the bottom of the
+combined plot.
+
+The example below colors the same H3N2 chart by genotype at HA1
+site 158, which has two mutations in the tree (`N158K`, `N158D`) and
+so renders three states (N, K, D):
+
+![H3N2 combined chart, colored by genotype HA1:158](images/h3n2_combined_genotype_158.svg)
+
+[Open the interactive chart in a new tab →](charts/h3n2_combined_genotype_158.html){target="_blank"}
+
+CLI flag: `--color-tree-by genotype:HA1:158`. In Python:
+`color_tree_by="genotype:HA1:158"`.
+
 ### Reproduce — command line
 
 ```bash
@@ -145,6 +184,7 @@ tree-annotated-plot \
     --tree-size 140 \
     --scale-bar \
     --branch-length-units substitutions \
+    --color-tree-by subclade \
     --output examples/data/h3n2_combined.json
 ```
 
@@ -162,6 +202,7 @@ out = tree_annotated_plot.plot(
     tree_size=140,
     scale_bar=True,
     branch_length_units="substitutions",
+    color_tree_by="subclade",
 )
 ```
 

scripts/generate_docs_assets.py

@@ -110,16 +110,24 @@ def _render_kikawa() -> None:
         # 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")
-        out = tree_annotated_plot.plot(
-            DATA_DIR / f"flu-seqneut-2025to2026_{subtype}.json",
-            chart,
+        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.
+            plot_kwargs["color_tree_by"] = "subclade"
+        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
@@ -146,9 +154,35 @@ def _render_kikawa() -> None:
         connect_leader_to_label=True,
         strain_label_font_size=9,
         shift_tree_loc=60,
+        color_tree_by="subclade",
     )
     _save_pair(out, "h3n2_combined_label_connect")
 
+    # 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",
+    )
+    _save_pair(out, "h3n2_combined_genotype_158")
+
 
 def main() -> None:
     """Render every example to SVG + interactive HTML under `docs/`."""