add color_tree_by parameter for Auspice-style tree coloring (#2)

Introduces a color_tree_by parameter that colors the tree's branches
and tip circles by:

- A node_attrs key like "subclade" — colored by node_attrs[<key>].value
  on each node.
- "genotype:<GENE>:<SITE>" — colored by the inferred amino-acid (or
  nucleotide) state at the site, computed from branch_attrs.mutations
  walked from the root.
- "genotype:<GENE>:<SITE1>,<SITE2>,..." — same but for a haplotype
  across sites; sites that don't vary in the tree drop out of the
  label; if every requested site is invariant, every node gets a single
  "<no variation>" category.

Color and ordering are chosen to match Nextstrain views for the same
tree: when the Auspice JSON defines meta.colorings[<key>].scale that
palette is used; otherwise the same per-N palette Auspice's frontend
uses (reproduced in _color.py with attribution to AGPL-licensed Auspice)
fills in. Categories are sorted by descending frequency (ties broken
alphabetically), matching Auspice's sortedDomain. The "unknown" category
renders in gray and is hidden from the legend when only internal nodes
lack the attribute.

Other changes bundled in:

- TreeNode now carries node_attrs and branch_attrs for downstream
  consumers; load_auspice_with_meta sibling helper exposes the JSON's
  top-level meta dict alongside the parsed root.
- Default tree_line_width bumped from 1.5 to 2 and tree_node_size from
  28 to 45 with explicit opacity=1, since the prior defaults were tuned
  for unicolor black trees and read poorly when colored.
- New "Color the tree" subsection in docs/examples.md with a second
  H3N2 example colored by genotype HA1:158.

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

Author: jbloom

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/`."""