tree-annotated-plot 0.2.1__tar.gz → 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (43) hide show
  1. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/CHANGELOG.md +38 -0
  2. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/PKG-INFO +5 -3
  3. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/README.md +1 -1
  4. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/docs/examples.md +13 -4
  5. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/docs/index.md +2 -2
  6. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/pyproject.toml +6 -4
  7. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/scripts/generate_docs_assets.py +6 -0
  8. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/src/tree_annotated_plot/_config.py +8 -3
  9. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/src/tree_annotated_plot/_plot.py +197 -11
  10. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/src/tree_annotated_plot/cli.py +1 -0
  11. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_cli.py +8 -0
  12. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_color_tree.py +34 -0
  13. tree_annotated_plot-0.3.0/tests/test_line_order.py +258 -0
  14. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_reconciliation.py +94 -3
  15. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/.github/workflows/docs.yml +0 -0
  16. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/.github/workflows/release.yml +0 -0
  17. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/.github/workflows/tests.yml +0 -0
  18. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/.gitignore +0 -0
  19. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/CLAUDE.md +0 -0
  20. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/LICENSE +0 -0
  21. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/docs/cli.md +0 -0
  22. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/docs/python-api.md +0 -0
  23. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/examples/fetch_auspice_data.py +0 -0
  24. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/examples/flu-seqneut-2025to2026_titer_charts.py +0 -0
  25. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/examples/synthetic_example.py +0 -0
  26. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/mkdocs.yml +0 -0
  27. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/scripts/build_docs.sh +0 -0
  28. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/scripts/check.sh +0 -0
  29. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/src/tree_annotated_plot/__init__.py +0 -0
  30. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/src/tree_annotated_plot/_color.py +0 -0
  31. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/src/tree_annotated_plot/_tree.py +0 -0
  32. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_appearance_tuning.py +0 -0
  33. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_branch_length.py +0 -0
  34. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_chart_loader.py +0 -0
  35. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_examples_run.py +0 -0
  36. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_introspection.py +0 -0
  37. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_label_connection.py +0 -0
  38. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_plot.py +0 -0
  39. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_plot_config.py +0 -0
  40. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_real_data.py +0 -0
  41. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_scale_bar.py +0 -0
  42. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_style.py +0 -0
  43. {tree_annotated_plot-0.2.1 → tree_annotated_plot-0.3.0}/tests/test_tree_location.py +0 -0
@@ -5,6 +5,44 @@ All notable changes to this project will be documented in this file.
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
6
6
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
7
 
8
+ ## [Unreleased]
9
+
10
+ ## [0.3.0] - 2026-05-19
11
+
12
+ ### Added
13
+
14
+ - `prune_chart_to_tree` (default `False`): when on, chart rows whose
15
+ `chart_strain_field` value isn't a tree tip are filtered out before
16
+ drawing. CLI form:
17
+ `--prune-chart-to-tree / --no-prune-chart-to-tree`.
18
+ ([#6](https://github.com/jbloomlab/tree-annotated-plot/issues/6))
19
+ - `-h` is accepted as a short alias for `--help` on the
20
+ `tree-annotated-plot` CLI.
21
+
22
+ ### Changed
23
+
24
+ - Minimum supported Python lowered from 3.13 to 3.11. CI still
25
+ runs on 3.13.
26
+
27
+ ## [0.2.2] - 2026-05-09
28
+
29
+ ### Fixed
30
+
31
+ - `mark_line`, `mark_trail`, and `mark_area` now connect points in tree
32
+ tip order regardless of other encodings on the chart. Previously, a
33
+ user chart with an explicit categorical color-scale `domain` rendered
34
+ with crisscrossing line segments because Vega-Lite's default
35
+ connection-order heuristic ignored the strain-axis sort. The package
36
+ now attaches a `calculate` transform that derives a per-row tip rank
37
+ and points the `order` channel at that field on these marks. Any
38
+ user-supplied `order` is left in place.
39
+ - When `color_tree_by` is used together with a user chart that has its
40
+ own `color` encoding (e.g. a titer plot colored by `cell_line`), the
41
+ user chart's marks no longer disappear. The concat container now
42
+ resolves the `color` scale as `independent` so the tree's
43
+ `color_value:N` scale (with its tree-specific domain) is not merged
44
+ with the user chart's color scale.
45
+
8
46
  ## [0.2.1] - 2026-05-06
9
47
 
10
48
  ### Fixed
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: tree-annotated-plot
3
- Version: 0.2.1
3
+ Version: 0.3.0
4
4
  Summary: Annotate the axis of an Altair / Vega-Lite plot with a phylogenetic tree.
5
5
  Project-URL: Documentation, https://jbloomlab.github.io/tree-annotated-plot/
6
6
  Project-URL: Homepage, https://github.com/jbloomlab/tree-annotated-plot
@@ -30,10 +30,12 @@ License-File: LICENSE
30
30
  Classifier: License :: OSI Approved :: MIT License
31
31
  Classifier: Operating System :: OS Independent
32
32
  Classifier: Programming Language :: Python :: 3
33
+ Classifier: Programming Language :: Python :: 3.11
34
+ Classifier: Programming Language :: Python :: 3.12
33
35
  Classifier: Programming Language :: Python :: 3.13
34
36
  Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
35
37
  Classifier: Topic :: Scientific/Engineering :: Visualization
36
- Requires-Python: >=3.13
38
+ Requires-Python: >=3.11
37
39
  Requires-Dist: altair<7,>=6.0
38
40
  Requires-Dist: click>=8.0
39
41
  Requires-Dist: pandas>=2.0
@@ -66,7 +68,7 @@ See [https://jbloomlab.github.io/tree-annotated-plot/](https://jbloomlab.github.
66
68
 
67
69
  ## Installation
68
70
 
69
- Released on [PyPI](https://pypi.org/project/tree-annotated-plot/). Requires Python 3.13+.
71
+ Released on [PyPI](https://pypi.org/project/tree-annotated-plot/). Requires Python 3.11+.
70
72
 
71
73
  ```bash
72
74
  pip install tree-annotated-plot
@@ -15,7 +15,7 @@ See [https://jbloomlab.github.io/tree-annotated-plot/](https://jbloomlab.github.
15
15
 
16
16
  ## Installation
17
17
 
18
- Released on [PyPI](https://pypi.org/project/tree-annotated-plot/). Requires Python 3.13+.
18
+ Released on [PyPI](https://pypi.org/project/tree-annotated-plot/). Requires Python 3.11+.
19
19
 
20
20
  ```bash
21
21
  pip install tree-annotated-plot
@@ -105,8 +105,12 @@ edge.
105
105
 
106
106
  The H3N2 example above is rendered with `color_tree_by="subclade"`,
107
107
  which colors the tree's branches and tip circles by the
108
- `node_attrs.subclade` value at each node and adds a categorical legend
109
- below the plot. See "Color the tree" below for the full set of options.
108
+ `node_attrs.subclade` value at each node. The package's default places
109
+ the categorical legend at the bottom of the combined plot; here we pass
110
+ `tree_color_legend_format={"orient": "left"}` to push it to the left
111
+ side instead, so it doesn't compete with the cohort-selection legend
112
+ already sitting below the chart. See "Color the tree" below for the
113
+ full set of options.
110
114
 
111
115
  ### Optional: connect leaders all the way to the labels
112
116
 
@@ -156,8 +160,11 @@ Colors match what you'd see on the Nextstrain view of the same tree —
156
160
  either from the JSON's palette information when the build provides it,
157
161
  or from the same default palette Auspice uses when it doesn't.
158
162
  Categories are ordered by descending frequency in both cases. Missing
159
- values render in gray, and the legend is drawn at the bottom of the
160
- combined plot.
163
+ values render in gray. By default the legend is drawn at the bottom of
164
+ the combined plot; the H3N2 examples in this section pass
165
+ `tree_color_legend_format={"orient": "left"}` to move it to the left
166
+ instead, since the bottom edge already carries the cohort-selection
167
+ legend.
161
168
 
162
169
  The example below colors the same H3N2 chart by genotype at HA1
163
170
  site 158, which has two mutations in the tree (`N158K`, `N158D`) and
@@ -277,6 +284,7 @@ tree-annotated-plot \
277
284
  --scale-bar \
278
285
  --branch-length-units substitutions \
279
286
  --color-tree-by subclade \
287
+ --tree-color-legend-format '{"orient":"left"}' \
280
288
  --output examples/data/h3n2_combined.json
281
289
  ```
282
290
 
@@ -295,6 +303,7 @@ out = tree_annotated_plot.plot(
295
303
  scale_bar=True,
296
304
  branch_length_units="substitutions",
297
305
  color_tree_by="subclade",
306
+ tree_color_legend_format={"orient": "left"},
298
307
  )
299
308
  ```
300
309
 
@@ -27,7 +27,7 @@ sit on by default (left for y-axis, bottom for x-axis).
27
27
  ## Installation
28
28
 
29
29
  Released on [PyPI](https://pypi.org/project/tree-annotated-plot/).
30
- Requires Python 3.13+.
30
+ Requires Python 3.11+.
31
31
 
32
32
  ```bash
33
33
  pip install tree-annotated-plot
@@ -36,7 +36,7 @@ pip install tree-annotated-plot
36
36
  To pin a specific version:
37
37
 
38
38
  ```bash
39
- pip install tree-annotated-plot==0.0.1
39
+ pip install tree-annotated-plot==0.3.0
40
40
  ```
41
41
 
42
42
  To install the bleeding edge directly from the
@@ -4,15 +4,17 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "tree-annotated-plot"
7
- version = "0.2.1"
7
+ version = "0.3.0"
8
8
  description = "Annotate the axis of an Altair / Vega-Lite plot with a phylogenetic tree."
9
9
  readme = "README.md"
10
- requires-python = ">=3.13"
10
+ requires-python = ">=3.11"
11
11
  license = { file = "LICENSE" }
12
12
  authors = [{ name = "Jesse Bloom", email = "jbloom@fredhutch.org" }]
13
13
  classifiers = [
14
14
  "License :: OSI Approved :: MIT License",
15
15
  "Programming Language :: Python :: 3",
16
+ "Programming Language :: Python :: 3.11",
17
+ "Programming Language :: Python :: 3.12",
16
18
  "Programming Language :: Python :: 3.13",
17
19
  "Operating System :: OS Independent",
18
20
  "Topic :: Scientific/Engineering :: Bio-Informatics",
@@ -52,7 +54,7 @@ packages = ["src/tree_annotated_plot"]
52
54
  dev-mode-dirs = ["src"]
53
55
 
54
56
  [tool.black]
55
- target-version = ["py313"]
57
+ target-version = ["py311"]
56
58
 
57
59
  [tool.ruff]
58
- target-version = "py313"
60
+ target-version = "py311"
@@ -122,7 +122,11 @@ def _render_kikawa() -> None:
122
122
  # Color H3N2 by subclade so the docs SVG matches what users see
123
123
  # on Nextstrain. The Auspice JSON's meta.colorings.subclade has
124
124
  # no `scale` defined, so colors come from the default palette.
125
+ # Place the subclade legend on the left so it doesn't compete
126
+ # with the cohort-selection legend that already sits at the
127
+ # bottom of the chart.
125
128
  plot_kwargs["color_tree_by"] = "subclade"
129
+ plot_kwargs["tree_color_legend_format"] = {"orient": "left"}
126
130
  out = tree_annotated_plot.plot(
127
131
  DATA_DIR / f"flu-seqneut-2025to2026_{subtype}.json",
128
132
  chart,
@@ -155,6 +159,7 @@ def _render_kikawa() -> None:
155
159
  strain_label_font_size=9,
156
160
  shift_tree_loc=60,
157
161
  color_tree_by="subclade",
162
+ tree_color_legend_format={"orient": "left"},
158
163
  )
159
164
  _save_pair(out, "h3n2_combined_label_connect")
160
165
 
@@ -180,6 +185,7 @@ def _render_kikawa() -> None:
180
185
  scale_bar=True,
181
186
  branch_length_units="substitutions",
182
187
  color_tree_by="genotype:HA1:158",
188
+ tree_color_legend_format={"orient": "left"},
183
189
  )
184
190
  _save_pair(out, "h3n2_combined_genotype_158")
185
191
 
@@ -109,9 +109,14 @@ class PlotConfig:
109
109
  "When off (default), tree tips not present in the chart's strain "
110
110
  "set are a fatal error. When on, those tips (and any internal "
111
111
  "nodes whose subtrees become empty) are dropped before drawing, "
112
- "with single-child internals collapsed into their kept child. "
113
- "Chart strains not present in the tree are *always* fatal "
114
- "regardless of this flag — pruning would silently lose plot data.",
112
+ "with single-child internals collapsed into their kept child.",
113
+ ] = False
114
+
115
+ prune_chart_to_tree: Annotated[
116
+ bool,
117
+ "When off (default), chart strains not present in the tree are a "
118
+ "fatal error. When on, chart rows whose `chart_strain_field` "
119
+ "value isn't a tree tip are filtered out before drawing.",
115
120
  ] = False
116
121
 
117
122
  strict_version: Annotated[
@@ -55,6 +55,7 @@ def plot(
55
55
  scale_bar: bool = False,
56
56
  branch_length_units: str | None = None,
57
57
  prune_tree_to_chart: bool = False,
58
+ prune_chart_to_tree: bool = False,
58
59
  strict_version: bool = True,
59
60
  connect_leader_to_label: bool = False,
60
61
  strain_label_font_size: float = 10.0,
@@ -82,6 +83,7 @@ def plot(
82
83
  scale_bar=scale_bar,
83
84
  branch_length_units=branch_length_units,
84
85
  prune_tree_to_chart=prune_tree_to_chart,
86
+ prune_chart_to_tree=prune_chart_to_tree,
85
87
  strict_version=strict_version,
86
88
  connect_leader_to_label=connect_leader_to_label,
87
89
  strain_label_font_size=strain_label_font_size,
@@ -168,12 +170,33 @@ def _build(
168
170
 
169
171
  chart_strains = _extract_chart_strains(spec, axis_hits, config.chart_strain_field)
170
172
 
173
+ if config.prune_chart_to_tree and (set(chart_strains) - set(tip_names)):
174
+ _prune_chart_spec_to_strains(
175
+ spec,
176
+ chart_strain_field=config.chart_strain_field,
177
+ keep_strains=set(tip_names),
178
+ )
179
+ chart = alt.Chart.from_dict(spec)
180
+ chart_strains = _extract_chart_strains(
181
+ spec, axis_hits, config.chart_strain_field
182
+ )
183
+ if not chart_strains:
184
+ raise ValueError(
185
+ "prune_chart_to_tree=True dropped every chart row: no "
186
+ "chart strain matched any tree tip under "
187
+ f"chart_strain_field={config.chart_strain_field!r} / "
188
+ f"tree_strain_field={config.tree_strain_field!r}. Pruning "
189
+ "is meant for charts that bundle a superset of strains; "
190
+ "an empty intersection suggests a wrong field choice."
191
+ )
192
+
171
193
  _reconcile_tips_and_strains(
172
194
  tree_strains=tip_names,
173
195
  chart_strains=chart_strains,
174
196
  chart_strain_field=config.chart_strain_field,
175
197
  tree_strain_field=config.tree_strain_field,
176
198
  prune_tree_to_chart=config.prune_tree_to_chart,
199
+ prune_chart_to_tree=config.prune_chart_to_tree,
177
200
  chart_spec=spec,
178
201
  tree_source=tree,
179
202
  )
@@ -233,6 +256,28 @@ def _build(
233
256
  ch.axis = alt.Axis(labels=False, ticks=False, domain=False, title=None)
234
257
  n_hits += 1
235
258
  _check_walker_hits("strain-axis update", n_hits, len(axis_hits), axis)
259
+
260
+ # Pin line/trail/area connection order to tip order via Vega-Lite's
261
+ # `order` channel. Without this, an explicit categorical `color` scale
262
+ # `domain` (or other encoding choices) can shift Vega-Lite's
263
+ # connection-order heuristic away from the strain-axis sort, causing
264
+ # lines to crisscross. The order channel's `sort` only accepts
265
+ # ascending/descending, so we attach a `calculate` transform that
266
+ # computes a per-row tip rank and point `order` at that derived
267
+ # quantitative field. User-supplied `order` always wins.
268
+ rank_expr = (
269
+ f"indexof({json.dumps(list(tip_names))}, "
270
+ f"datum[{json.dumps(config.chart_strain_field)}])"
271
+ )
272
+ for node in _iter_connection_order_nodes(new_chart, config.chart_strain_field):
273
+ enc = _live_attr(node, "encoding")
274
+ if _live_attr(enc, "order") is not None:
275
+ continue
276
+ existing = _live_attr(node, "transform") or []
277
+ node.transform = list(existing) + [
278
+ {"calculate": rank_expr, "as": _TIP_ORDER_RANK_FIELD}
279
+ ]
280
+ enc.order = alt.Order(f"{_TIP_ORDER_RANK_FIELD}:Q")
236
281
  hoisted_config, hoisted_other = _pop_toplevel_only_attrs(new_chart)
237
282
 
238
283
  combined = _concat_for_location(
@@ -276,22 +321,32 @@ def _concat_for_location(
276
321
  user_chart: alt.TopLevelMixin,
277
322
  location: TreeLocation,
278
323
  ) -> alt.HConcatChart | alt.VConcatChart:
279
- """Concat tree and chart in the order implied by the tree's location."""
324
+ """Concat tree and chart in the order implied by the tree's location.
325
+
326
+ The strain axis is resolved independent so the tree and chart can use
327
+ different scales on that axis (the tree's branch length vs. the chart's
328
+ measurement value), while still sharing the orthogonal strain axis. The
329
+ `color` scale is also resolved independent: when ``color_tree_by`` is set
330
+ the tree panel emits a `color_value:N` color scale with a tree-specific
331
+ domain, and Vega-Lite's default of sharing color across concat views
332
+ would merge it with any color encoding on the user's chart, hiding
333
+ user-chart marks whose color values aren't in the tree's domain.
334
+ """
280
335
  if location == "left":
281
336
  return alt.hconcat(tree_chart, user_chart, spacing=0).resolve_scale(
282
- y="independent"
337
+ y="independent", color="independent"
283
338
  )
284
339
  if location == "right":
285
340
  return alt.hconcat(user_chart, tree_chart, spacing=0).resolve_scale(
286
- y="independent"
341
+ y="independent", color="independent"
287
342
  )
288
343
  if location == "top":
289
344
  return alt.vconcat(tree_chart, user_chart, spacing=0).resolve_scale(
290
- x="independent"
345
+ x="independent", color="independent"
291
346
  )
292
347
  if location == "bottom":
293
348
  return alt.vconcat(user_chart, tree_chart, spacing=0).resolve_scale(
294
- x="independent"
349
+ x="independent", color="independent"
295
350
  )
296
351
  raise ValueError(f"unreachable: tree_location={location!r}")
297
352
 
@@ -609,6 +664,70 @@ def _extract_chart_strains(
609
664
  return _extract_field_values_from_spec_data(spec, chart_strain_field)
610
665
 
611
666
 
667
+ def _prune_chart_spec_to_strains(
668
+ spec: dict, *, chart_strain_field: str, keep_strains: set[str]
669
+ ) -> None:
670
+ """Filter a Vega-Lite spec in place to drop rows outside `keep_strains`.
671
+
672
+ Mutates three kinds of structure:
673
+ - top-level `datasets` entries (each a list of row-dicts).
674
+ - inline `data.values` lists anywhere in the spec tree.
675
+ - explicit `sort` lists on encoding channels bound to
676
+ `chart_strain_field` (any other `sort` is left alone).
677
+
678
+ Rows that don't carry `chart_strain_field` at all are preserved
679
+ (we have no signal to drop them). URL-backed data raises — we
680
+ can't fetch + filter at plot time, mirroring `_extract_chart_strains`.
681
+ """
682
+
683
+ def row_kept(row: Any) -> bool:
684
+ if not isinstance(row, dict):
685
+ return True
686
+ if chart_strain_field not in row:
687
+ return True
688
+ return row[chart_strain_field] in keep_strains
689
+
690
+ datasets = spec.get("datasets") if isinstance(spec, dict) else None
691
+ if isinstance(datasets, dict):
692
+ for name, rows in list(datasets.items()):
693
+ if isinstance(rows, list):
694
+ datasets[name] = [row for row in rows if row_kept(row)]
695
+
696
+ def walk(node: Any) -> None:
697
+ if isinstance(node, dict):
698
+ data = node.get("data")
699
+ if isinstance(data, dict):
700
+ if "url" in data:
701
+ raise ValueError(
702
+ f"chart references data via URL ({data['url']!r}); "
703
+ "URL data is not supported, so prune_chart_to_tree "
704
+ "cannot filter it. Materialize the data inline "
705
+ "(via alt.Chart(df) with a pandas DataFrame) before "
706
+ "saving the chart."
707
+ )
708
+ if "values" in data and isinstance(data["values"], list):
709
+ data["values"] = [row for row in data["values"] if row_kept(row)]
710
+ encoding = node.get("encoding")
711
+ if isinstance(encoding, dict):
712
+ for channel in encoding.values():
713
+ if not isinstance(channel, dict):
714
+ continue
715
+ if channel.get("field") != chart_strain_field:
716
+ continue
717
+ sort = channel.get("sort")
718
+ if isinstance(sort, list):
719
+ channel["sort"] = [s for s in sort if s in keep_strains]
720
+ for k, v in node.items():
721
+ if k in ("data", "datasets"):
722
+ continue
723
+ walk(v)
724
+ elif isinstance(node, list):
725
+ for item in node:
726
+ walk(item)
727
+
728
+ walk(spec)
729
+
730
+
612
731
  def _extract_field_values_from_spec_data(spec: dict, field: str) -> list[str]:
613
732
  """Walk spec for inline / named data and return distinct values of `field`.
614
733
 
@@ -676,13 +795,16 @@ def _reconcile_tips_and_strains(
676
795
  chart_strain_field: str,
677
796
  tree_strain_field: str,
678
797
  prune_tree_to_chart: bool,
798
+ prune_chart_to_tree: bool,
679
799
  chart_spec: dict,
680
800
  tree_source: Any,
681
801
  ) -> None:
682
802
  """Verify tree strains and chart strains are reconcilable.
683
803
 
684
804
  Three asymmetries:
685
- - chart strains not in tree → always fatal.
805
+ - chart strains not in tree → fatal unless `prune_chart_to_tree=True`
806
+ (in which case the chart spec has already been pre-filtered upstream
807
+ and this set is expected to be empty by the time we get here).
686
808
  - tree tips not in chart → fatal unless `prune_tree_to_chart=True`.
687
809
  - (duplicate tree_strain_field values across tips → handled by the
688
810
  separate `_check_no_duplicate_tip_strains`.)
@@ -696,7 +818,9 @@ def _reconcile_tips_and_strains(
696
818
  chart_minus_tree = chart_set - tree_set
697
819
  tree_minus_chart = tree_set - chart_set
698
820
 
699
- if not chart_minus_tree and (not tree_minus_chart or prune_tree_to_chart):
821
+ chart_ok = not chart_minus_tree
822
+ tree_ok = not tree_minus_chart or prune_tree_to_chart
823
+ if chart_ok and tree_ok:
700
824
  return
701
825
 
702
826
  hints = _candidate_field_hints(
@@ -716,6 +840,7 @@ def _reconcile_tips_and_strains(
716
840
  chart_minus_tree=chart_minus_tree,
717
841
  tree_minus_chart=tree_minus_chart,
718
842
  prune_tree_to_chart=prune_tree_to_chart,
843
+ prune_chart_to_tree=prune_chart_to_tree,
719
844
  hints=hints,
720
845
  )
721
846
  )
@@ -730,14 +855,16 @@ def _format_strain_mismatch(
730
855
  chart_minus_tree: set[str],
731
856
  tree_minus_chart: set[str],
732
857
  prune_tree_to_chart: bool,
858
+ prune_chart_to_tree: bool,
733
859
  hints: list[str],
734
860
  ) -> str:
735
861
  parts: list[str] = []
736
- if chart_minus_tree:
862
+ if chart_minus_tree and not prune_chart_to_tree:
737
863
  parts.append(
738
864
  f"{len(chart_minus_tree)} chart strain(s) are not present in the "
739
- "tree (these would be silently dropped if we pruned, so this is "
740
- "always fatal)."
865
+ "tree. Pass `prune_chart_to_tree=True` to drop the offending "
866
+ "chart rows automatically (use with care — this discards plot "
867
+ "data)."
741
868
  )
742
869
  if tree_minus_chart and not prune_tree_to_chart:
743
870
  parts.append(
@@ -751,7 +878,7 @@ def _format_strain_mismatch(
751
878
  )
752
879
  parts.append("Sample chart_strain_field values: " f"{sorted(chart_strains)[:5]}")
753
880
  parts.append("Sample tree_strain_field values: " f"{sorted(tree_strains)[:5]}")
754
- if chart_minus_tree:
881
+ if chart_minus_tree and not prune_chart_to_tree:
755
882
  parts.append(f"Sample chart-only values: {sorted(chart_minus_tree)[:5]}")
756
883
  if tree_minus_chart and not prune_tree_to_chart:
757
884
  parts.append(f"Sample tree-only values: {sorted(tree_minus_chart)[:5]}")
@@ -999,6 +1126,65 @@ def _live_attr(obj: Any, name: str) -> Any:
999
1126
  return v
1000
1127
 
1001
1128
 
1129
+ # Vega-Lite marks whose drawing order along the discrete axis is determined
1130
+ # by the order channel (rule 1 of the connection-order fallback chain). For
1131
+ # these marks pinning the order channel to a per-row tip-order rank ensures
1132
+ # the line / trail / area connects strains in tip order regardless of how
1133
+ # Vega-Lite would otherwise resolve the fallback (e.g. when an explicit
1134
+ # color-scale `domain` shifts it away from the y-axis sort).
1135
+ _CONNECTION_ORDER_MARKS = frozenset({"line", "trail", "area"})
1136
+
1137
+ # Derived field name appended to each connection-order mark's transform
1138
+ # pipeline. Chosen to be unlikely to collide with user data; private prefix
1139
+ # makes the intent clear if anyone inspects the rendered spec.
1140
+ _TIP_ORDER_RANK_FIELD = "_tap_strain_order_idx"
1141
+
1142
+
1143
+ def _mark_type(node: Any) -> str | None:
1144
+ """Return the mark type string for a chart node, or None.
1145
+
1146
+ Handles altair's two mark forms: the plain string (e.g. ``"line"`` from
1147
+ ``chart.mark_line()``) and the ``MarkDef`` object (e.g. from
1148
+ ``chart.mark_line(point=True)``). Container nodes (LayerChart,
1149
+ HConcatChart, etc.) have no mark and return None.
1150
+ """
1151
+ mark = _live_attr(node, "mark")
1152
+ if mark is None:
1153
+ return None
1154
+ if isinstance(mark, str):
1155
+ return mark
1156
+ type_ = _live_attr(mark, "type")
1157
+ return type_ if isinstance(type_, str) else None
1158
+
1159
+
1160
+ def _iter_connection_order_nodes(node: Any, chart_strain_field: str) -> Iterator[Any]:
1161
+ """Yield the live chart node for every node whose mark is in
1162
+ ``_CONNECTION_ORDER_MARKS`` and whose ``x``/``y`` strain encoding
1163
+ matches ``chart_strain_field``.
1164
+
1165
+ Mirrors `_iter_strain_axis_channels`'s traversal (hconcat / vconcat /
1166
+ concat / layer / spec descent). Yields the node itself (not just the
1167
+ encoding) so the caller can attach both a ``calculate`` transform and
1168
+ the ``order`` channel.
1169
+ """
1170
+ if _mark_type(node) in _CONNECTION_ORDER_MARKS:
1171
+ enc = _live_attr(node, "encoding")
1172
+ if enc is not None:
1173
+ for channel in ("x", "y"):
1174
+ ch = _live_attr(enc, channel)
1175
+ if ch is not None and _channel_field(ch) == chart_strain_field:
1176
+ yield node
1177
+ break
1178
+ for attr in ("hconcat", "vconcat", "concat", "layer"):
1179
+ sub = _live_attr(node, attr)
1180
+ if isinstance(sub, list):
1181
+ for s in sub:
1182
+ yield from _iter_connection_order_nodes(s, chart_strain_field)
1183
+ spec = _live_attr(node, "spec")
1184
+ if spec is not None:
1185
+ yield from _iter_connection_order_nodes(spec, chart_strain_field)
1186
+
1187
+
1002
1188
  def _pop_toplevel_only_attrs(
1003
1189
  chart: alt.TopLevelMixin,
1004
1190
  ) -> tuple[Any, dict]:
@@ -263,6 +263,7 @@ def _stack_config_options(command: Any) -> Any:
263
263
  "Save the result as HTML / JSON / PNG / SVG / PDF "
264
264
  "(format is dispatched on --output's extension)."
265
265
  ),
266
+ context_settings={"help_option_names": ["-h", "--help"]},
266
267
  )
267
268
  @click.option(
268
269
  "--tree",
@@ -27,6 +27,14 @@ def test_help_lists_data_options() -> None:
27
27
  assert "--output" in result.output
28
28
 
29
29
 
30
+ def test_short_help_flag_is_alias_for_long() -> None:
31
+ """`-h` should produce identical output to `--help`."""
32
+ long_form = _runner().invoke(main, ["--help"])
33
+ short_form = _runner().invoke(main, ["-h"])
34
+ assert short_form.exit_code == 0
35
+ assert short_form.output == long_form.output
36
+
37
+
30
38
  def test_help_lists_auto_generated_options() -> None:
31
39
  """A handful of PlotConfig-derived options should appear in --help."""
32
40
  result = _runner().invoke(main, ["--help"])
@@ -663,6 +663,40 @@ def test_color_tree_by_legend_hides_internal_only_unknown():
663
663
  assert "unknown" in enc["scale"]["domain"]
664
664
 
665
665
 
666
+ def test_color_tree_by_color_scale_independent_from_user_chart():
667
+ """When the tree is colored, the concat container must resolve the color
668
+ scale as `independent` so the tree's `color_value:N` scale doesn't merge
669
+ with a user-chart color encoding (which would hide user marks whose
670
+ color values aren't in the tree's domain)."""
671
+ df = pd.DataFrame(
672
+ {
673
+ "strain": ["A", "B", "C", "D"] * 2,
674
+ "titer": [1.0, 2.0, 4.0, 8.0, 1.5, 2.5, 4.5, 8.5],
675
+ "cell_line": ["X", "X", "X", "X", "Y", "Y", "Y", "Y"],
676
+ }
677
+ )
678
+ user_chart = (
679
+ alt.Chart(df)
680
+ .mark_line(point=True)
681
+ .encode(
682
+ x="titer:Q",
683
+ y=alt.Y("strain:N"),
684
+ color=alt.Color("cell_line:N"),
685
+ )
686
+ .properties(width=200, height=200)
687
+ )
688
+ out = tree_annotated_plot.plot(
689
+ _attr_auspice(),
690
+ user_chart,
691
+ **_kw(),
692
+ color_tree_by="subclade",
693
+ )
694
+ spec = out.to_dict()
695
+ resolve = spec.get("resolve") or {}
696
+ scale = resolve.get("scale") or {}
697
+ assert scale.get("color") == "independent"
698
+
699
+
666
700
  # -----------------------------------------------------------------------------
667
701
  # CLI
668
702
  # -----------------------------------------------------------------------------
@@ -0,0 +1,258 @@
1
+ """Connection order on `mark_line` / `mark_trail` / `mark_area` follows tip order.
2
+
3
+ Regression: when the user chart had `mark_line` with strain on a discrete
4
+ axis, lines connected points in tip order most of the time — but adding an
5
+ explicit categorical color-scale `domain` shifted Vega-Lite's default
6
+ connection-order heuristic away from the strain-axis `sort`, causing lines
7
+ to crisscross. The package now sets the `order` channel to a per-row
8
+ tip-rank field on these marks (unless the user has set their own `order`),
9
+ pinning rule 1 of the Vega-Lite fallback chain.
10
+
11
+ The tip-rank is computed via a `calculate` transform that runs `indexof`
12
+ against the tip-names list — Vega-Lite's `order` channel `sort` only
13
+ accepts ascending/descending, so a derived quantitative field is the
14
+ direct way to express custom ordering.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ from typing import Any
20
+
21
+ import altair as alt
22
+ import pandas as pd
23
+ import pytest
24
+
25
+ import tree_annotated_plot
26
+ from tree_annotated_plot._plot import _TIP_ORDER_RANK_FIELD
27
+
28
+
29
+ def _synthetic_auspice() -> dict:
30
+ """4-tip Auspice tree; tip layout order is [A, B, C, D]."""
31
+ return {
32
+ "version": "v2",
33
+ "meta": {},
34
+ "tree": {
35
+ "name": "ROOT",
36
+ "node_attrs": {"div": 0.0},
37
+ "children": [
38
+ {
39
+ "name": "int1",
40
+ "node_attrs": {"div": 0.01},
41
+ "children": [
42
+ {"name": "A", "node_attrs": {"div": 0.04}},
43
+ {"name": "B", "node_attrs": {"div": 0.05}},
44
+ ],
45
+ },
46
+ {
47
+ "name": "int2",
48
+ "node_attrs": {"div": 0.02},
49
+ "children": [
50
+ {"name": "C", "node_attrs": {"div": 0.03}},
51
+ {"name": "D", "node_attrs": {"div": 0.06}},
52
+ ],
53
+ },
54
+ ],
55
+ },
56
+ }
57
+
58
+
59
+ TIP_ORDER = ["A", "B", "C", "D"]
60
+
61
+
62
+ def _df() -> pd.DataFrame:
63
+ return pd.DataFrame(
64
+ [
65
+ {"strain": s, "serum": serum, "titer": v}
66
+ for s, serum, v in [
67
+ ("A", "s1", 100),
68
+ ("B", "s1", 200),
69
+ ("C", "s1", 400),
70
+ ("D", "s1", 800),
71
+ ("A", "s2", 150),
72
+ ("B", "s2", 250),
73
+ ("C", "s2", 350),
74
+ ("D", "s2", 700),
75
+ ]
76
+ ]
77
+ )
78
+
79
+
80
+ def _plot_user_panel(chart: alt.TopLevelMixin) -> dict:
81
+ """Run plot() with default settings and return the user-panel spec dict."""
82
+ out = tree_annotated_plot.plot(
83
+ _synthetic_auspice(),
84
+ chart,
85
+ chart_strain_field="strain",
86
+ tree_strain_field="name",
87
+ branch_length="div",
88
+ )
89
+ # Default tree_location for y-axis strain is "left" — user panel is hconcat[1].
90
+ return out.to_dict()["hconcat"][1]
91
+
92
+
93
+ def _find_order_encodings(spec: Any) -> list[dict]:
94
+ """Walk a spec dict and return every `encoding.order` dict found."""
95
+ found: list[dict] = []
96
+
97
+ def walk(node: Any) -> None:
98
+ if isinstance(node, dict):
99
+ enc = node.get("encoding")
100
+ if isinstance(enc, dict) and "order" in enc:
101
+ found.append(enc["order"])
102
+ for k, v in node.items():
103
+ if k != "encoding":
104
+ walk(v)
105
+ elif isinstance(node, list):
106
+ for item in node:
107
+ walk(item)
108
+
109
+ walk(spec)
110
+ return found
111
+
112
+
113
+ def _find_tip_rank_transforms(spec: Any) -> list[dict]:
114
+ """Walk a spec dict and return every calculate transform that computes
115
+ the package's tip-rank derived field."""
116
+ found: list[dict] = []
117
+
118
+ def walk(node: Any) -> None:
119
+ if isinstance(node, dict):
120
+ transforms = node.get("transform")
121
+ if isinstance(transforms, list):
122
+ for t in transforms:
123
+ if isinstance(t, dict) and t.get("as") == _TIP_ORDER_RANK_FIELD:
124
+ found.append(t)
125
+ for v in node.values():
126
+ walk(v)
127
+ elif isinstance(node, list):
128
+ for item in node:
129
+ walk(item)
130
+
131
+ walk(spec)
132
+ return found
133
+
134
+
135
+ def _assert_tip_rank_wired(spec: dict) -> None:
136
+ """Spec contains both an order-channel pointing at the tip-rank field
137
+ and a calculate transform that derives that field from the tip names."""
138
+ orders = _find_order_encodings(spec)
139
+ assert len(orders) == 1, f"expected one order encoding, got {orders}"
140
+ assert orders[0]["field"] == _TIP_ORDER_RANK_FIELD
141
+ assert orders[0]["type"] == "quantitative"
142
+ transforms = _find_tip_rank_transforms(spec)
143
+ assert len(transforms) == 1, f"expected one tip-rank transform, got {transforms}"
144
+ expr = transforms[0]["calculate"]
145
+ assert "indexof" in expr
146
+ # The expression embeds the tip names list verbatim.
147
+ for tip in TIP_ORDER:
148
+ assert f'"{tip}"' in expr, f"tip {tip!r} not in calculate expression {expr!r}"
149
+
150
+
151
+ def test_mark_line_with_color_domain_gets_order_in_tip_order():
152
+ """The original bug: explicit color-scale `domain` no longer crisscrosses lines."""
153
+ chart = (
154
+ alt.Chart(_df())
155
+ .mark_line(point=True)
156
+ .encode(
157
+ x=alt.X("titer:Q", scale=alt.Scale(type="log")),
158
+ y=alt.Y("strain:N"),
159
+ color=alt.Color(
160
+ "serum:N",
161
+ scale=alt.Scale(domain=["s1", "s2"], range=["red", "blue"]),
162
+ ),
163
+ )
164
+ .properties(width=300, height=200)
165
+ )
166
+ _assert_tip_rank_wired(_plot_user_panel(chart))
167
+
168
+
169
+ def test_user_supplied_order_is_preserved():
170
+ """User-supplied `order` always wins; the package never overwrites it."""
171
+ chart = (
172
+ alt.Chart(_df())
173
+ .mark_line()
174
+ .encode(
175
+ x=alt.X("titer:Q"),
176
+ y=alt.Y("strain:N"),
177
+ color="serum:N",
178
+ order=alt.Order("titer:Q", sort="ascending"),
179
+ )
180
+ .properties(width=300, height=200)
181
+ )
182
+ spec = _plot_user_panel(chart)
183
+ orders = _find_order_encodings(spec)
184
+ assert len(orders) == 1
185
+ assert orders[0]["field"] == "titer"
186
+ # The package did not attach its tip-rank transform either.
187
+ assert _find_tip_rank_transforms(spec) == []
188
+
189
+
190
+ @pytest.mark.parametrize("mark_method", ["mark_line", "mark_area", "mark_trail"])
191
+ def test_connection_order_marks_get_order_in_tip_order(mark_method: str):
192
+ """All three connection-order marks (line / area / trail) get `order` injected."""
193
+ base = alt.Chart(_df()).encode(
194
+ x=alt.X("titer:Q"),
195
+ y=alt.Y("strain:N"),
196
+ color="serum:N",
197
+ )
198
+ chart = getattr(base, mark_method)().properties(width=300, height=200)
199
+ _assert_tip_rank_wired(_plot_user_panel(chart))
200
+
201
+
202
+ def test_layered_chart_only_line_layer_gets_order():
203
+ """Mixed-mark layers: line layer gets `order`, circle layer is untouched."""
204
+ base = alt.Chart(_df()).encode(
205
+ x=alt.X("titer:Q"),
206
+ y=alt.Y("strain:N"),
207
+ color="serum:N",
208
+ )
209
+ layered = (base.mark_line() + base.mark_circle()).properties(width=300, height=200)
210
+ _assert_tip_rank_wired(_plot_user_panel(layered))
211
+
212
+
213
+ def test_no_connection_order_marks_means_no_order_injected():
214
+ """A chart with only point/circle/etc. marks is a no-op for the walker."""
215
+ chart = (
216
+ alt.Chart(_df())
217
+ .mark_circle()
218
+ .encode(
219
+ x=alt.X("titer:Q"),
220
+ y=alt.Y("strain:N"),
221
+ color="serum:N",
222
+ )
223
+ .properties(width=300, height=200)
224
+ )
225
+ spec = _plot_user_panel(chart)
226
+ assert _find_order_encodings(spec) == []
227
+ assert _find_tip_rank_transforms(spec) == []
228
+
229
+
230
+ def test_faceted_chart_gets_order_injection():
231
+ """The walker descends into the inner spec of a faceted chart."""
232
+ df_facet = pd.DataFrame(
233
+ [
234
+ {"strain": s, "serum": serum, "panel": p, "titer": v}
235
+ for s, serum, p, v in [
236
+ ("A", "s1", "P1", 100),
237
+ ("B", "s1", "P1", 200),
238
+ ("C", "s1", "P1", 400),
239
+ ("D", "s1", "P1", 800),
240
+ ("A", "s1", "P2", 150),
241
+ ("B", "s1", "P2", 250),
242
+ ("C", "s1", "P2", 350),
243
+ ("D", "s1", "P2", 700),
244
+ ]
245
+ ]
246
+ )
247
+ base = (
248
+ alt.Chart(df_facet)
249
+ .mark_line()
250
+ .encode(
251
+ x=alt.X("titer:Q"),
252
+ y=alt.Y("strain:N"),
253
+ color="serum:N",
254
+ )
255
+ .properties(width=300, height=200)
256
+ )
257
+ faceted = base.facet(facet="panel:N", columns=2)
258
+ _assert_tip_rank_wired(_plot_user_panel(faceted))
@@ -73,7 +73,7 @@ def _chart_for_strains(strains: list[str], *, height: int = 200) -> alt.Chart:
73
73
  )
74
74
 
75
75
 
76
- # ---------- chart-not-in-tree (always fatal) ----------
76
+ # ---------- chart-not-in-tree (fatal unless prune_chart_to_tree) ----------
77
77
 
78
78
 
79
79
  def test_chart_strain_not_in_tree_is_fatal_default() -> None:
@@ -90,9 +90,9 @@ def test_chart_strain_not_in_tree_is_fatal_default() -> None:
90
90
  )
91
91
 
92
92
 
93
- def test_chart_strain_not_in_tree_is_fatal_even_with_prune() -> None:
93
+ def test_chart_strain_not_in_tree_is_fatal_with_only_tree_prune() -> None:
94
94
  """`prune_tree_to_chart=True` only drops *tree* tips. A chart strain
95
- not in the tree still raises — pruning would silently lose plot data."""
95
+ not in the tree still raises — the two flags are orthogonal."""
96
96
  chart = _chart_for_strains(["A1", "A2", "X"])
97
97
  with pytest.raises(ValueError, match="not present in the tree"):
98
98
  tree_annotated_plot.plot(
@@ -105,6 +105,97 @@ def test_chart_strain_not_in_tree_is_fatal_even_with_prune() -> None:
105
105
  )
106
106
 
107
107
 
108
+ def test_chart_strain_not_in_tree_succeeds_with_prune_chart_to_tree() -> None:
109
+ """`prune_chart_to_tree=True` filters chart rows whose strain isn't a
110
+ tree tip. The resulting chart's strain-axis sort matches the kept tree
111
+ tip order, and the dropped strain is gone from the chart's data."""
112
+ chart = _chart_for_strains(["A1", "A2", "A3", "B1", "B2", "X"])
113
+ out = tree_annotated_plot.plot(
114
+ _auspice_two_clades(),
115
+ chart,
116
+ chart_strain_field="strain",
117
+ tree_strain_field="name",
118
+ branch_length="div",
119
+ prune_chart_to_tree=True,
120
+ )
121
+ assert isinstance(out, alt.HConcatChart)
122
+ spec = out.to_dict()
123
+ # The strain-axis sort on the user-chart panel should be exactly the
124
+ # tree's tip order — no `X`.
125
+ sorts = []
126
+ for ch in out.hconcat:
127
+ ch_spec = ch.to_dict()
128
+ enc = ch_spec.get("encoding", {})
129
+ for channel in enc.values():
130
+ if isinstance(channel, dict) and channel.get("field") == "strain":
131
+ if isinstance(channel.get("sort"), list):
132
+ sorts.append(channel["sort"])
133
+ assert sorts, "expected at least one strain-axis encoding with a sort"
134
+ for s in sorts:
135
+ assert "X" not in s
136
+ assert set(s) <= {"A1", "A2", "A3", "B1", "B2"}
137
+ # And the user-chart data (now in `datasets`) should not contain X rows.
138
+ for rows in (spec.get("datasets") or {}).values():
139
+ if isinstance(rows, list) and rows and isinstance(rows[0], dict):
140
+ if "strain" in rows[0]:
141
+ assert all(r["strain"] != "X" for r in rows)
142
+
143
+
144
+ def test_prune_chart_to_tree_zero_overlap_raises() -> None:
145
+ """If pruning would drop every chart row (no overlap with tree tips),
146
+ raise a clear error rather than producing an empty plot."""
147
+ chart = _chart_for_strains(["X", "Y", "Z"])
148
+ with pytest.raises(ValueError, match="dropped every chart row"):
149
+ tree_annotated_plot.plot(
150
+ _auspice_two_clades(),
151
+ chart,
152
+ chart_strain_field="strain",
153
+ tree_strain_field="name",
154
+ branch_length="div",
155
+ prune_chart_to_tree=True,
156
+ )
157
+
158
+
159
+ def test_prune_chart_and_tree_combined() -> None:
160
+ """Both flags may be on at once: the chart has extras (X) and the tree
161
+ has tips the chart lacks (B1, B2). Pruning is bidirectional."""
162
+ chart = _chart_for_strains(["A1", "A2", "A3", "X"])
163
+ out = tree_annotated_plot.plot(
164
+ _auspice_two_clades(),
165
+ chart,
166
+ chart_strain_field="strain",
167
+ tree_strain_field="name",
168
+ branch_length="div",
169
+ prune_tree_to_chart=True,
170
+ prune_chart_to_tree=True,
171
+ )
172
+ assert isinstance(out, alt.HConcatChart)
173
+ # Final intersection is {A1, A2, A3}: chart loses X, tree loses B1/B2.
174
+ found_sort = None
175
+ for ch in out.hconcat:
176
+ enc = ch.to_dict().get("encoding", {})
177
+ for channel in enc.values():
178
+ if isinstance(channel, dict) and channel.get("field") == "strain":
179
+ if isinstance(channel.get("sort"), list):
180
+ found_sort = channel["sort"]
181
+ assert found_sort is not None
182
+ assert set(found_sort) == {"A1", "A2", "A3"}
183
+
184
+
185
+ def test_prune_chart_to_tree_default_error_mentions_flag() -> None:
186
+ """When chart has strains not in tree and the user hasn't opted in,
187
+ the error message should suggest `prune_chart_to_tree=True`."""
188
+ chart = _chart_for_strains(["A1", "A2", "A3", "B1", "B2", "X"])
189
+ with pytest.raises(ValueError, match="prune_chart_to_tree=True"):
190
+ tree_annotated_plot.plot(
191
+ _auspice_two_clades(),
192
+ chart,
193
+ chart_strain_field="strain",
194
+ tree_strain_field="name",
195
+ branch_length="div",
196
+ )
197
+
198
+
108
199
  # ---------- tree-not-in-chart (fatal unless prune) ----------
109
200
 
110
201