add tree-coloring appearance-tuning parameters

CHANGELOG.md

@@ -22,6 +22,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (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.
+- `tree_color_scale` (default `None`): override the default coloring
+  with an explicit `{category: color}` mapping. Keys must match the
+  tree's categories one-to-one and the legend order follows the
+  user's key order. CLI form: `"value1=#hex1,value2=#hex2,..."`.
+- `tree_color_legend_format` (default `None`): pass any subset of
+  Vega-Lite's
+  [Legend properties](https://vega.github.io/vega-lite/docs/legend.html#properties)
+  as a dict to style the tree's color legend (`orient`, `direction`,
+  `columns`, `padding`, `labelFontSize`, `titleFontSize`, …). When
+  `orient` is `"left"` or `"right"` and the user has not set
+  `columns` or `direction`, `columns=1` is forced so entries stack
+  vertically. CLI form: a JSON object string.
+- `tree_color_legend_show` (default `True`): set to `False` to hide
+  the tree's color legend entirely while still coloring the tree.
+- `scale_bar_font_size` (default `10`): font size for the tree's
+  scale bar label.
 
 ### Changed
 

docs/examples.md

@@ -170,6 +170,98 @@ so renders three states (N, K, D):
 CLI flag: `--color-tree-by genotype:HA1:158`. In Python:
 `color_tree_by="genotype:HA1:158"`.
 
+### Customizing the colors and legend
+
+Four optional knobs let you override the defaults:
+
+- `tree_color_scale` — supply your own `{category: color}` mapping
+  instead of the Auspice palette. Keys must match the tree's
+  categories one-to-one (extras or misses are an error). The legend
+  order follows the order of keys you pass. `"unknown"` is always
+  gray and must not be specified.
+- `tree_color_legend_format` — pass any subset of Vega-Lite's
+  [Legend properties](https://vega.github.io/vega-lite/docs/legend.html#properties)
+  as a dict to style the legend (`orient`, `direction`, `columns`,
+  `padding`, `offset`, `labelFontSize`, `titleFontSize`, …). Smart
+  default: when `orient` is `"left"` or `"right"` and you have not
+  set `columns` or `direction`, entries stack vertically
+  (`columns=1` is forced).
+- `tree_color_legend_show` — set to `False` to hide the legend
+  entirely while keeping the tree colored.
+- `scale_bar_font_size` — font size for the scale bar's label.
+
+The exact category strings depend on the form of `color_tree_by`:
+
+| `color_tree_by` | Example category strings |
+|---|---|
+| node attribute (e.g. `"subclade"`) | `"K"`, `"J.2"`, `"J.2.4"` |
+| `"genotype:HA1:158"` (single site) | `"K158"`, `"R158"`, `"E158"` |
+| `"genotype:HA1:158,189"` (both sites vary) | `"K158/E189"`, `"R158/E189"` |
+| `"genotype:HA1:158,189"` (189 invariant) | `"K158"`, `"R158"` (invariant site dropped) |
+
+If you pass keys that don't match, the error message lists the
+tree's actual categories so you can copy them.
+
+The H3N2 example below puts all four knobs to work: an explicit
+6-color palette (a colorblind-safe Okabe–Ito-inspired set) ordered
+K → J.2.4 → J.2.3 → J.2.2 → J.2 → G.1.3.1, the legend moved to the
+**left** of the combined plot at 14-pt, and a matching 14-pt
+scale-bar label.
+
+![H3N2 combined chart with custom colors and legend](images/h3n2_combined_custom_colors.svg)
+
+[Open the interactive chart in a new tab →](charts/h3n2_combined_custom_colors.html){target="_blank"}
+
+```python
+out = tree_annotated_plot.plot(
+    "examples/data/flu-seqneut-2025to2026_H3N2.json",
+    chart,
+    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,
+)
+```
+
+`--tree-color-scale` on the CLI is a comma-separated list of
+`key=color` pairs and `--tree-color-legend-format` is a JSON object
+string (quote the whole argument in both cases so the shell doesn't
+interpret `#`, braces, or quotes):
+
+```bash
+tree-annotated-plot \
+    --tree examples/data/flu-seqneut-2025to2026_H3N2.json \
+    --chart examples/data/flu-seqneut-2025to2026_H3N2_titers.json \
+    --chart-strain-field axis_label \
+    --tree-strain-field derived_haplotype \
+    --branch-length div \
+    --tree-size 140 \
+    --scale-bar \
+    --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 \
+    --output examples/data/h3n2_combined_custom_colors.html
+```
+
 ### Reproduce — command line
 
 ```bash

scripts/generate_docs_assets.py

@@ -183,6 +183,45 @@ def _render_kikawa() -> None:
     )
     _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/`."""

src/tree_annotated_plot/_color.py

@@ -119,6 +119,7 @@ def compute_node_color_values(
     root: TreeNode,
     color_spec: str,
     auspice_meta: dict | None = None,
+    tree_color_scale: dict[str, str] | None = None,
 ) -> ColorMapping:
     """Walk the tree and resolve per-node color categories + scale arrays.
 
@@ -135,6 +136,11 @@ def compute_node_color_values(
         is available (caller passed a pre-built `TreeNode`). Used only to
         consult ``meta.colorings[<key>].scale`` and ``.title`` for node-attr
         specs; ignored for genotype specs.
+    tree_color_scale
+        Optional user-supplied {category: color} mapping. When provided, its
+        keys must match the tree's real categories one-to-one (mismatch is a
+        ``ValueError``); the legend order is the dict's insertion order.
+        ``"unknown"`` is always gray and must not be specified.
 
     Returns
     -------
@@ -154,7 +160,10 @@ def compute_node_color_values(
         values_by_node = _color_by_genotype(root, gene, sites)
 
     categories = _ordered_categories(values_by_node.values())
-    domain, range_ = _resolve_scale(categories, parsed, auspice_meta)
+    if tree_color_scale is not None:
+        domain, range_ = _apply_user_scale(categories, tree_color_scale)
+    else:
+        domain, range_ = _resolve_scale(categories, parsed, auspice_meta)
     legend_title = _resolve_legend_title(color_spec, parsed, auspice_meta)
     legend_values = _resolve_legend_values(domain, values_by_node, root)
     return ColorMapping(
@@ -166,6 +175,54 @@ def compute_node_color_values(
     )
 
 
+def _apply_user_scale(
+    categories: list[str],
+    user_scale: dict[str, str],
+) -> tuple[list[str], list[str]]:
+    """Build (domain, range_) from a user-supplied {category: color} dict.
+
+    The user's key order is the legend order (Python dicts preserve insertion
+    order). Keys must match the tree's real categories one-to-one — extras or
+    misses raise ``ValueError`` listing the actual tree categories so the
+    user can copy-paste the correct keys (especially relevant for
+    genotype/haplotype categories like ``"K158"`` or ``"K158/E189"`` whose
+    exact form depends on the data). ``"unknown"`` is always gray and is
+    appended automatically when present in ``categories`` — the user must
+    not include it.
+    """
+    real_categories = [c for c in categories if c != _UNKNOWN]
+    user_keys = list(user_scale.keys())
+
+    if _UNKNOWN in user_keys:
+        raise ValueError(
+            f"tree_color_scale must not contain key {_UNKNOWN!r}; "
+            "missing values are always rendered gray."
+        )
+
+    tree_set = set(real_categories)
+    user_set = set(user_keys)
+    missing = sorted(tree_set - user_set)
+    extra = sorted(user_set - tree_set)
+    if missing or extra:
+        msg = (
+            "tree_color_scale keys don't match the tree's categories.\n"
+            f"  Tree categories: {real_categories!r}\n"
+            f"  Provided keys:   {user_keys!r}"
+        )
+        if missing:
+            msg += f"\n  Missing from your scale: {missing!r}"
+        if extra:
+            msg += f"\n  Unexpected in your scale: {extra!r}"
+        raise ValueError(msg)
+
+    domain = list(user_keys)
+    range_ = [user_scale[k] for k in user_keys]
+    if _UNKNOWN in categories:
+        domain.append(_UNKNOWN)
+        range_.append(_GRAY)
+    return domain, range_
+
+
 def _resolve_legend_values(
     domain: list[str],
     values_by_node: dict[str, str],

src/tree_annotated_plot/_config.py

@@ -18,7 +18,7 @@
 import dataclasses
 import textwrap
 import typing
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 TreeLocation = Literal["left", "right", "top", "bottom"]
 
@@ -165,6 +165,49 @@ class PlotConfig:
         "leaves the tree black.",
     ] = None
 
+    tree_color_scale: Annotated[
+        dict[str, str] | None,
+        "Hardcoded color scale that overrides the default coloring. Keys "
+        "are category labels, values are colors (any Vega-Lite-compatible "
+        "string — e.g. hex codes). The legend order follows the order keys "
+        "appear here. The keys must match the tree's categories one-to-one "
+        '(extra or missing keys are an error). "unknown" is always gray '
+        "and must not be specified. For genotype/haplotype colorings, the "
+        'category strings include the site number (e.g. "K158" or '
+        '"K158/E189"); a mismatch error lists the actual tree categories. '
+        'CLI form: "value1=#hex1,value2=#hex2,...".',
+    ] = None
+
+    tree_color_legend_format: Annotated[
+        dict[str, Any] | None,
+        "Vega-Lite Legend properties to apply to the tree-coloring legend. "
+        "Pass any subset of the keys at "
+        "https://vega.github.io/vega-lite/docs/legend.html#properties as a "
+        "Python dict (e.g. "
+        '`{"orient": "left", "labelFontSize": 13, "titleFontSize": 13}`). '
+        'Common keys: "orient" (default "bottom"), "direction", "columns", '
+        '"padding", "offset", "labelFontSize", "titleFontSize". Smart '
+        'default: when "orient" is "left" or "right" and you have not set '
+        '"columns" or "direction", "columns" is forced to 1 so entries '
+        "stack vertically. "
+        "None (default) leaves Vega-Lite's defaults. Has no effect when "
+        "`color_tree_by` is None. CLI form: a JSON object string (quote the "
+        'whole argument), e.g. \'{"orient":"left","labelFontSize":13}\'.',
+    ] = None
+
+    tree_color_legend_show: Annotated[
+        bool,
+        "Whether to render the tree-coloring legend. On (default) shows it. "
+        "Off hides the legend entirely while keeping the tree colored. Has "
+        "no effect when `color_tree_by` is None.",
+    ] = True
+
+    scale_bar_font_size: Annotated[
+        float,
+        "Font size (px) for the tree's scale bar label. Default 10. Has no "
+        "effect when `scale_bar` is off.",
+    ] = 10.0
+
 
 # Sidecar for Python-docstring-only prose, keyed by PlotConfig field name.
 # Empty by default — add an entry when a field's docstring entry needs more