From 4061f38b88ff85f4047be3f7eb63e9e0d3d01ab4 Mon Sep 17 00:00:00 2001 From: ivan-aksamentov Date: Wed, 22 Jul 2026 12:45:20 +0200 Subject: [PATCH 1/4] refactor(io): make graph authoritative for tree output - Apply output topology preferences in place before every tree writer reads the final graph.\n- Remove TreeIR projections and command-specific format gates so adapters read command graph data directly. --- packages/treetime-graph/src/edge.rs | 2 +- packages/treetime-graph/src/graph.rs | 35 +- packages/treetime-graph/src/graph_traverse.rs | 56 +- packages/treetime-graph/src/node.rs | 2 +- packages/treetime-graph/src/topology_order.rs | 394 +++++---- packages/treetime-io/src/auspice.rs | 7 +- packages/treetime-io/src/graph.rs | 47 +- packages/treetime-io/src/lib.rs | 1 - packages/treetime-io/src/nex.rs | 12 +- packages/treetime-io/src/nwk.rs | 12 +- packages/treetime-io/src/phyloxml.rs | 133 +-- .../treetime-io/src/tree_ir/__tests__/mod.rs | 4 - .../src/tree_ir/__tests__/test_auspice.rs | 182 ---- .../src/tree_ir/__tests__/test_mutation.rs | 44 - .../src/tree_ir/__tests__/test_phyloxml.rs | 100 --- .../src/tree_ir/__tests__/test_usher.rs | 136 --- packages/treetime-io/src/tree_ir/auspice.rs | 342 -------- packages/treetime-io/src/tree_ir/mod.rs | 7 - packages/treetime-io/src/tree_ir/mutation.rs | 82 -- packages/treetime-io/src/tree_ir/phyloxml.rs | 268 ------ packages/treetime-io/src/tree_ir/types.rs | 160 ---- packages/treetime-io/src/tree_ir/usher.rs | 169 ---- packages/treetime-io/src/usher_mat.rs | 35 +- packages/treetime/src/ancestral/fitch.rs | 6 +- packages/treetime/src/clock/clock_graph.rs | 2 +- .../src/commands/ancestral/augur_node_data.rs | 12 +- .../treetime/src/commands/ancestral/result.rs | 44 +- .../treetime/src/commands/ancestral/run.rs | 82 +- packages/treetime/src/commands/clock/run.rs | 79 +- .../src/commands/mugration/augur_node_data.rs | 20 +- .../treetime/src/commands/mugration/run.rs | 30 +- .../src/commands/optimize/augur_node_data.rs | 8 +- .../treetime/src/commands/optimize/result.rs | 42 +- .../treetime/src/commands/optimize/run.rs | 70 +- .../treetime/src/commands/prune/result.rs | 29 +- packages/treetime/src/commands/prune/run.rs | 25 +- .../src/commands/shared/__tests__/mod.rs | 1 + .../__tests__/test_output_resolution.rs | 96 +- .../shared/__tests__/test_tree_output.rs | 298 +++++++ .../src/commands/shared/ir_projection.rs | 210 ----- packages/treetime/src/commands/shared/mod.rs | 2 +- .../treetime/src/commands/shared/output.rs | 131 ++- .../src/commands/shared/tree_output.rs | 822 ++++++++++++++++++ .../commands/timetree/output/__tests__/mod.rs | 1 - .../timetree/output/__tests__/test_ir.rs | 108 --- .../timetree/output/augur_node_data.rs | 10 +- .../src/commands/timetree/output/ir.rs | 84 -- .../src/commands/timetree/output/mod.rs | 1 - .../treetime/src/commands/timetree/result.rs | 61 +- .../treetime/src/commands/timetree/run.rs | 104 +-- packages/treetime/src/mugration/result.rs | 31 +- .../treetime/src/partition/marginal_core.rs | 8 +- .../treetime/src/partition/marginal_dense.rs | 6 +- .../treetime/src/partition/marginal_sparse.rs | 2 +- packages/treetime/src/partition/timetree.rs | 2 +- packages/treetime/src/partition/traits.rs | 6 +- packages/treetime/src/payload/ancestral.rs | 2 +- packages/treetime/src/prune/pipeline.rs | 3 + 58 files changed, 2069 insertions(+), 2599 deletions(-) delete mode 100644 packages/treetime-io/src/tree_ir/__tests__/mod.rs delete mode 100644 packages/treetime-io/src/tree_ir/__tests__/test_auspice.rs delete mode 100644 packages/treetime-io/src/tree_ir/__tests__/test_mutation.rs delete mode 100644 packages/treetime-io/src/tree_ir/__tests__/test_phyloxml.rs delete mode 100644 packages/treetime-io/src/tree_ir/__tests__/test_usher.rs delete mode 100644 packages/treetime-io/src/tree_ir/auspice.rs delete mode 100644 packages/treetime-io/src/tree_ir/mod.rs delete mode 100644 packages/treetime-io/src/tree_ir/mutation.rs delete mode 100644 packages/treetime-io/src/tree_ir/phyloxml.rs delete mode 100644 packages/treetime-io/src/tree_ir/types.rs delete mode 100644 packages/treetime-io/src/tree_ir/usher.rs create mode 100644 packages/treetime/src/commands/shared/__tests__/test_tree_output.rs delete mode 100644 packages/treetime/src/commands/shared/ir_projection.rs create mode 100644 packages/treetime/src/commands/shared/tree_output.rs delete mode 100644 packages/treetime/src/commands/timetree/output/__tests__/test_ir.rs delete mode 100644 packages/treetime/src/commands/timetree/output/ir.rs diff --git a/packages/treetime-graph/src/edge.rs b/packages/treetime-graph/src/edge.rs index 481a8265d..131a0cf9e 100644 --- a/packages/treetime-graph/src/edge.rs +++ b/packages/treetime-graph/src/edge.rs @@ -40,7 +40,7 @@ pub trait TimeLength { fn set_time_length(&mut self, length: Option); } -pub trait GraphEdge: Clone + Debug + Sync + Send {} +pub trait GraphEdge: Debug + Sync + Send {} /// Composite trait for edges that support ancestral reconstruction. /// Currently equivalent to `GraphEdge` for consistency with `NodeAncestralOps`. diff --git a/packages/treetime-graph/src/graph.rs b/packages/treetime-graph/src/graph.rs index bac937150..470628564 100644 --- a/packages/treetime-graph/src/graph.rs +++ b/packages/treetime-graph/src/graph.rs @@ -39,7 +39,7 @@ where pub(crate) edges: Vec>>>>, pub(crate) roots: Vec, pub(crate) leaves: Vec, - pub(crate) data: Arc>, + pub(crate) data: D, } impl Graph @@ -57,7 +57,7 @@ where edges: Vec::new(), roots: vec![], leaves: vec![], - data: Arc::new(RwLock::new(D::default())), + data: D::default(), } } @@ -67,16 +67,33 @@ where edges: Vec::new(), roots: vec![], leaves: vec![], - data: Arc::new(RwLock::new(data)), + data, } } - pub fn data(&self) -> Arc> { - Arc::clone(&self.data) + pub const fn data(&self) -> &D { + &self.data + } + + pub fn data_mut(&mut self) -> &mut D { + &mut self.data } pub fn set_data(&mut self, data: D) { - self.data = Arc::new(RwLock::new(data)); + self.data = data; + } + + pub fn map_data(self, data: T) -> Graph + where + T: Sync + Send, + { + Graph { + nodes: self.nodes, + edges: self.edges, + roots: self.roots, + leaves: self.leaves, + data, + } } /// Retrieve parent nodes of a given node and the corresponding edges. @@ -200,7 +217,7 @@ where } /// Iterates nodes synchronously and in unspecified order - pub fn for_each(&self, f: &mut dyn FnMut(GraphNodeSafe)) { + pub fn for_each(&self, f: &mut dyn FnMut(GraphNodeSafe)) { self .nodes .iter() @@ -211,7 +228,7 @@ where /// Iterates nodes synchronously and in unspecified order pub fn map(&self, mut f: F) -> Vec where - F: FnMut(GraphNodeSafe) -> T, + F: FnMut(GraphNodeSafe) -> T, { self .nodes @@ -224,7 +241,7 @@ where /// Iterates nodes synchronously and in unspecified order pub fn filter_map(&self, mut f: F) -> Vec where - F: FnMut(GraphNodeSafe) -> Option, + F: FnMut(GraphNodeSafe) -> Option, { self .nodes diff --git a/packages/treetime-graph/src/graph_traverse.rs b/packages/treetime-graph/src/graph_traverse.rs index 66d499cf4..264c0516a 100644 --- a/packages/treetime-graph/src/graph_traverse.rs +++ b/packages/treetime-graph/src/graph_traverse.rs @@ -15,7 +15,7 @@ use treetime_utils::collections::container::{get_exactly_one, get_exactly_one_mu /// Represents graph node during forward traversal #[must_use] #[derive(Debug)] -pub struct GraphNodeForward +pub struct GraphNodeForward where N: GraphNode, E: GraphEdge, @@ -28,16 +28,17 @@ where pub parent_keys: Vec<(GraphNodeKey, GraphEdgeKey)>, pub child_edges: Vec>, pub child_edge_keys: Vec, - pub data: Arc>, } -impl GraphNodeForward +impl GraphNodeForward where N: GraphNode, E: GraphEdge, - D: Sync + Send, { - pub fn new(graph: &Graph, node: &Node) -> Self { + pub fn new(graph: &Graph, node: &Node) -> Self + where + D: Sync + Send, + { let is_leaf = node.is_leaf(); let is_root = node.is_root(); let key = node.key(); @@ -68,8 +69,6 @@ where .map(|(_, edge)| edge.write_arc().payload().write_arc()) .collect_vec(); - let data = Arc::clone(&graph.data); - Self { is_root, is_leaf, @@ -79,7 +78,6 @@ where parent_keys, child_edges, child_edge_keys, - data, } } @@ -93,11 +91,10 @@ where /// Represents graph node during backwards traversal #[must_use] #[derive(Debug)] -pub struct GraphNodeBackward +pub struct GraphNodeBackward where N: GraphNode, E: GraphEdge, - D: Sync + Send, { pub is_root: bool, pub is_leaf: bool, @@ -107,16 +104,17 @@ where pub child_keys: Vec<(GraphNodeKey, GraphEdgeKey)>, pub parent_edges: Vec>, pub parent_edge_keys: Vec, - pub data: Arc>, } -impl GraphNodeBackward +impl GraphNodeBackward where N: GraphNode, E: GraphEdge, - D: Sync + Send, { - pub fn new(graph: &Graph, node: &Node) -> Self { + pub fn new(graph: &Graph, node: &Node) -> Self + where + D: Sync + Send, + { let is_leaf = node.is_leaf(); let is_root = node.is_root(); let key = node.key(); @@ -147,8 +145,6 @@ where .map(|(_, edge)| edge.write_arc().payload().write_arc()) .collect_vec(); - let data = Arc::clone(&graph.data); - Self { is_root, is_leaf, @@ -158,7 +154,6 @@ where child_keys, parent_edges, parent_edge_keys, - data, } } @@ -169,11 +164,10 @@ where /// Represents graph node during safe traversal #[derive(Debug)] -pub struct GraphNodeSafe +pub struct GraphNodeSafe where N: GraphNode, E: GraphEdge, - D: Sync + Send, { pub is_root: bool, pub is_leaf: bool, @@ -181,16 +175,17 @@ where pub payload: Arc>, pub children: Vec>, pub parents: Vec>, - pub data: Arc>, } -impl GraphNodeSafe +impl GraphNodeSafe where N: GraphNode, E: GraphEdge, - D: Sync + Send, { - pub fn from_node(graph: &Graph, node: &Arc>>) -> Self { + pub fn from_node(graph: &Graph, node: &Arc>>) -> Self + where + D: Sync + Send, + { let node = node.read(); let is_leaf = node.is_leaf(); let is_root = node.is_root(); @@ -198,8 +193,6 @@ where let payload = node.payload(); let parents = graph.parents_of(&node); let children = graph.children_of(&node); - let data = Arc::clone(&graph.data); - Self { is_root, is_leaf, @@ -207,7 +200,6 @@ where payload, children, parents, - data, } } } @@ -243,7 +235,7 @@ where /// before the traversal returns. Subsequent errors are discarded (first error wins). pub fn par_iter_breadth_first_forward(&self, explorer: F) -> Result<(), Report> where - F: Fn(GraphNodeForward) -> Result + Sync + Send, + F: Fn(GraphNodeForward) -> Result + Sync + Send, { let roots = self.roots.iter().filter_map(|idx| self.get_node(*idx)).collect_vec(); @@ -260,7 +252,7 @@ where /// See [`Self::par_iter_breadth_first_forward`] for callback semantics. pub fn par_iter_breadth_first_backward(&self, explorer: F) -> Result<(), Report> where - F: Fn(GraphNodeBackward) -> Result + Sync + Send, + F: Fn(GraphNodeBackward) -> Result + Sync + Send, { let leaves = self .leaves @@ -280,7 +272,7 @@ where /// Serial depth-first preorder forward traversal (roots to leaves, parents before children). pub fn iter_depth_first_preorder_forward(&self, mut explorer: F) -> Result<(), Report> where - F: FnMut(GraphNodeForward) -> Result<(), Report>, + F: FnMut(GraphNodeForward) -> Result<(), Report>, { let root = self .get_exactly_one_root() @@ -299,7 +291,7 @@ where /// Serial depth-first postorder forward traversal (children before parents). pub fn iter_depth_first_postorder_forward(&self, mut explorer: F) -> Result<(), Report> where - F: FnMut(GraphNodeBackward) -> Result<(), Report>, + F: FnMut(GraphNodeBackward) -> Result<(), Report>, { let root = self .get_exactly_one_root() @@ -331,7 +323,7 @@ where /// per-node work must capture mutable outer state, which a parallel callback cannot. pub fn iter_breadth_first_forward(&self, mut explorer: F) -> Result<(), Report> where - F: FnMut(GraphNodeForward) -> Result<(), Report>, + F: FnMut(GraphNodeForward) -> Result<(), Report>, { let root = self .get_exactly_one_root() @@ -352,7 +344,7 @@ where /// Serial breadth-first backward traversal (leaves to roots, against edge directions). pub fn iter_breadth_first_backward(&self, mut explorer: F) -> Result<(), Report> where - F: FnMut(GraphNodeBackward) -> Result<(), Report>, + F: FnMut(GraphNodeBackward) -> Result<(), Report>, { let root = self .get_exactly_one_root() diff --git a/packages/treetime-graph/src/node.rs b/packages/treetime-graph/src/node.rs index dfb511daa..efde590ae 100644 --- a/packages/treetime-graph/src/node.rs +++ b/packages/treetime-graph/src/node.rs @@ -48,7 +48,7 @@ pub trait TimeConstraint { fn set_bad_branch(&mut self, bad: bool); } -pub trait GraphNode: Clone + Debug + Sync + Send {} +pub trait GraphNode: Debug + Sync + Send {} /// Composite trait for nodes that support ancestral reconstruction pub trait NodeAncestralOps: GraphNode + Named + Described {} diff --git a/packages/treetime-graph/src/topology_order.rs b/packages/treetime-graph/src/topology_order.rs index baa5a63c3..8fc1eac58 100644 --- a/packages/treetime-graph/src/topology_order.rs +++ b/packages/treetime-graph/src/topology_order.rs @@ -1,14 +1,12 @@ -use crate::edge::{Edge, GraphEdge, GraphEdgeKey, HasBranchLength}; -use crate::graph::Graph; -use crate::node::{GraphNode, GraphNodeKey, Named, Node}; +use crate::edge::{GraphEdge, GraphEdgeKey, HasBranchLength}; +use crate::graph::{Graph, SafeNode}; +use crate::node::{GraphNode, GraphNodeKey, Named}; use eyre::Report; use itertools::Itertools; use ordered_float::OrderedFloat; -use parking_lot::RwLock; use serde::{Deserialize, Serialize}; use std::cmp::Ordering; use std::collections::{BTreeMap, VecDeque}; -use std::sync::Arc; use treetime_utils::{make_error, make_report}; #[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] @@ -47,42 +45,63 @@ impl TopologyOrderSpec { } } - pub fn plan(&self, graph: &Graph) -> Result + /// Apply the requested logical topology order to a graph. + /// + /// All fallible computation and validation completes before the graph is + /// mutated. Node and edge keys and their slot storage remain unchanged. + pub fn apply(&self, graph: &mut Graph) -> Result<(), Report> where N: GraphNode + Named, E: GraphEdge + HasBranchLength, D: Sync + Send, { - if self.preset == TopologyOrderPreset::Keep { - return build_plan_unordered(graph); - } - - let postorder = postorder_keys(graph)?; - let reverse = self.preset.is_reverse(); + let order = if self.preset == TopologyOrderPreset::Keep { + build_order_unmodified(graph) + } else { + let postorder = postorder_keys(graph)?; + let reverse = self.preset.is_reverse(); + match self.preset { + TopologyOrderPreset::Keep => unreachable!(), + TopologyOrderPreset::DescendantCount | TopologyOrderPreset::DescendantCountReverse => { + let keys = compute_descendant_counts(graph, &postorder); + build_order(graph, &keys, reverse) + }, + TopologyOrderPreset::Height | TopologyOrderPreset::HeightReverse => { + let keys = compute_heights(graph, &postorder); + build_order(graph, &keys, reverse) + }, + TopologyOrderPreset::Divergence | TopologyOrderPreset::DivergenceReverse => { + let keys = compute_divergences(graph, &postorder); + build_order(graph, &keys, reverse) + }, + TopologyOrderPreset::Label | TopologyOrderPreset::LabelReverse => { + let keys = compute_labels(graph, &postorder)?; + build_order(graph, &keys, reverse) + }, + TopologyOrderPreset::TargetOrder | TopologyOrderPreset::TargetOrderReverse => { + let keys = compute_target_scores(graph, &postorder, &self.target_order, self.target_aggregate)?; + build_order(graph, &keys, reverse) + }, + } + }?; + + let ordered_nodes: Vec<(SafeNode, Vec)> = order + .outbound_edges + .into_iter() + .map(|(node_key, outbound_edges)| { + graph + .get_node(node_key) + .map(|node| (node, outbound_edges)) + .ok_or_else(|| make_report!("Node {node_key} disappeared while applying topology order")) + }) + .try_collect()?; - match self.preset { - TopologyOrderPreset::Keep => unreachable!(), - TopologyOrderPreset::DescendantCount | TopologyOrderPreset::DescendantCountReverse => { - let keys = compute_descendant_counts(graph, &postorder); - build_plan(graph, &keys, reverse) - }, - TopologyOrderPreset::Height | TopologyOrderPreset::HeightReverse => { - let keys = compute_heights(graph, &postorder); - build_plan(graph, &keys, reverse) - }, - TopologyOrderPreset::Divergence | TopologyOrderPreset::DivergenceReverse => { - let keys = compute_divergences(graph, &postorder); - build_plan(graph, &keys, reverse) - }, - TopologyOrderPreset::Label | TopologyOrderPreset::LabelReverse => { - let keys = compute_labels(graph, &postorder)?; - build_plan(graph, &keys, reverse) - }, - TopologyOrderPreset::TargetOrder | TopologyOrderPreset::TargetOrderReverse => { - let keys = compute_target_scores(graph, &postorder, &self.target_order, self.target_aggregate)?; - build_plan(graph, &keys, reverse) - }, + for (node, outbound_edges) in ordered_nodes { + *node.write_arc().outbound_mut() = outbound_edges; } + graph.roots = order.roots; + graph.leaves = order.leaves; + Ok(()) } } @@ -126,71 +145,20 @@ pub enum TopologyOrderTargetAggregate { Median, } -#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)] -pub struct TopologyOrderPlan { - pub roots: Vec, - pub leaves: Vec, - pub outbound_edges: BTreeMap>, -} - -impl TopologyOrderPlan { - pub fn ordered_graph(&self, graph: &Graph) -> Result, Report> - where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - { - let mut ordered = Graph { - nodes: graph - .nodes - .iter() - .map(|node| { - node.as_ref().map(|node| { - let node = node.read_arc(); - let mut cloned = Node::new(node.key(), node.payload().read_arc().clone()); - *cloned.outbound_mut() = self - .outbound_edges - .get(&node.key()) - .cloned() - .unwrap_or_else(|| node.outbound().to_vec()); - *cloned.inbound_mut() = node.inbound().to_vec(); - Arc::new(RwLock::new(cloned)) - }) - }) - .collect(), - edges: graph - .edges - .iter() - .map(|edge| { - edge.as_ref().map(|edge| { - let edge = edge.read_arc(); - Arc::new(RwLock::new(Edge::new( - edge.key(), - edge.source(), - edge.target(), - edge.payload().read_arc().clone(), - ))) - }) - }) - .collect(), - roots: self.roots.clone(), - leaves: self.leaves.clone(), - data: graph.data(), - }; - ordered.build()?; - ordered.roots = self.roots.clone(); - ordered.leaves = self.leaves.clone(); - Ok(ordered) - } +#[derive(Debug, Eq, PartialEq)] +struct TopologyOrder { + roots: Vec, + leaves: Vec, + outbound_edges: BTreeMap>, } -fn build_plan_unordered(graph: &Graph) -> Result +fn build_order_unmodified(graph: &Graph) -> Result where N: GraphNode, E: GraphEdge, D: Sync + Send, { - Ok(TopologyOrderPlan { + Ok(TopologyOrder { roots: graph.roots.clone(), leaves: graph.leaves.clone(), outbound_edges: graph @@ -205,11 +173,11 @@ where }) } -fn build_plan( +fn build_order( graph: &Graph, keys: &BTreeMap, reverse: bool, -) -> Result +) -> Result where N: GraphNode, E: GraphEdge, @@ -247,7 +215,7 @@ where .collect_vec() }; - Ok(TopologyOrderPlan { + Ok(TopologyOrder { roots: sort_node_keys(&graph.roots), leaves: sort_node_keys(&graph.leaves), outbound_edges: graph @@ -398,6 +366,15 @@ where .map(|(i, label)| (label.as_str(), i)) .collect(); + if position_of.len() != target_order.len() { + let duplicate = target_order + .iter() + .duplicates() + .next() + .expect("target-order length mismatch guarantees a duplicate"); + return make_error!("When ordering topology: target order contains duplicate leaf label '{duplicate}'"); + } + validate_target_order(graph, &position_of)?; match aggregate { @@ -509,6 +486,7 @@ where E: GraphEdge, D: Sync + Send, { + let mut final_labels = BTreeMap::new(); for leaf in graph.get_leaves() { let leaf = leaf.read_arc(); let label = leaf @@ -520,6 +498,12 @@ where if !position_of.contains_key(label.as_str()) { return make_error!("When validating target order: leaf '{label}' is absent from target order"); } + if let Some(previous_key) = final_labels.insert(label.clone(), leaf.key()) { + return make_error!( + "When validating target order: final leaf label '{label}' is duplicated by nodes {previous_key} and {}", + leaf.key() + ); + } } Ok(()) } @@ -578,25 +562,26 @@ mod tests { #[test] fn topology_order_descendant_count_sorts_children_ascending() -> Result<(), Report> { - let graph = fixture_tree()?; - let plan = TopologyOrderSpec::default().plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + let mut graph = fixture_tree()?; + let original = child_names(&graph, "root")?; + TopologyOrderSpec::default().apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["A", "BC", "DEF"], actual); - assert_eq!(vec!["DEF", "A", "BC"], child_names(&graph, "root")?); + assert_eq!(vec!["DEF", "A", "BC"], original); Ok(()) } #[test] fn topology_order_descendant_count_reverse_sorts_children_descending() -> Result<(), Report> { - let graph = fixture_tree()?; - let plan = TopologyOrderSpec::descendant_count(true).plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + let mut graph = fixture_tree()?; + TopologyOrderSpec::descendant_count(true).apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["DEF", "BC", "A"], actual); @@ -605,11 +590,11 @@ mod tests { #[test] fn topology_order_keep_preserves_outbound_order() -> Result<(), Report> { - let graph = fixture_tree()?; - let plan = TopologyOrderSpec::keep().plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + let mut graph = fixture_tree()?; + TopologyOrderSpec::keep().apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["DEF", "A", "BC"], actual); @@ -632,10 +617,10 @@ mod tests { graph.add_edge(right, right_only, TestEdge::new())?; graph.build()?; - let plan = TopologyOrderSpec::default().plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + TopologyOrderSpec::default().apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["left", "right"], actual); @@ -644,7 +629,7 @@ mod tests { #[test] fn topology_order_target_order_uses_requested_tip_order() -> Result<(), Report> { - let graph = fixture_tree()?; + let mut graph = fixture_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::TargetOrder, target_order: vec!["D", "E", "F", "B", "C", "A"] @@ -653,10 +638,10 @@ mod tests { .collect(), target_aggregate: TopologyOrderTargetAggregate::Mean, }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["DEF", "BC", "A"], actual); @@ -675,7 +660,7 @@ mod tests { graph.add_edge(c, a, TestEdge::new())?; graph.build()?; - let err = TopologyOrderSpec::default().plan(&graph).unwrap_err(); + let err = TopologyOrderSpec::default().apply(&mut graph).unwrap_err(); assert!(err.to_string().contains("directed cycle")); @@ -684,15 +669,15 @@ mod tests { #[test] fn topology_order_height_sorts_by_subtree_depth() -> Result<(), Report> { - let graph = fixture_deep_tree()?; + let mut graph = fixture_deep_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::Height, ..TopologyOrderSpec::default() }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; // A is leaf (height 0), shallow has height 1, deep has height 2 assert_eq!(vec!["A", "shallow", "deep"], actual); @@ -702,15 +687,15 @@ mod tests { #[test] fn topology_order_height_reverse_sorts_deepest_first() -> Result<(), Report> { - let graph = fixture_deep_tree()?; + let mut graph = fixture_deep_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::HeightReverse, ..TopologyOrderSpec::default() }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["deep", "shallow", "A"], actual); @@ -723,15 +708,15 @@ mod tests { // short: max divergence = 0.1 + 0.1 = 0.2 // long: max divergence = 0.5 + 0.2 = 0.7 // D: leaf, divergence = 0.0 - let graph = fixture_branch_length_tree()?; + let mut graph = fixture_branch_length_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::Divergence, ..TopologyOrderSpec::default() }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["D", "short", "long"], actual); @@ -740,15 +725,15 @@ mod tests { #[test] fn topology_order_divergence_reverse_sorts_longest_first() -> Result<(), Report> { - let graph = fixture_branch_length_tree()?; + let mut graph = fixture_branch_length_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::DivergenceReverse, ..TopologyOrderSpec::default() }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["long", "short", "D"], actual); @@ -757,15 +742,15 @@ mod tests { #[test] fn topology_order_label_sorts_alphabetically() -> Result<(), Report> { - let graph = fixture_tree()?; + let mut graph = fixture_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::Label, ..TopologyOrderSpec::default() }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; // A has label "A", BC has min label "B", DEF has min label "D" assert_eq!(vec!["A", "BC", "DEF"], actual); @@ -775,15 +760,15 @@ mod tests { #[test] fn topology_order_label_reverse_sorts_descending() -> Result<(), Report> { - let graph = fixture_tree()?; + let mut graph = fixture_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::LabelReverse, ..TopologyOrderSpec::default() }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["DEF", "BC", "A"], actual); @@ -792,7 +777,7 @@ mod tests { #[test] fn topology_order_target_order_median_uses_median_position() -> Result<(), Report> { - let graph = fixture_tree()?; + let mut graph = fixture_tree()?; // Target order: D=0, E=1, F=2, B=3, C=4, A=5 // DEF median of [0,1,2] = 1, BC median of [3,4] = 3.5, A = 5 let spec = TopologyOrderSpec { @@ -803,10 +788,10 @@ mod tests { .collect(), target_aggregate: TopologyOrderTargetAggregate::Median, }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["DEF", "BC", "A"], actual); @@ -815,15 +800,15 @@ mod tests { #[test] fn topology_order_propagates_through_nested_levels() -> Result<(), Report> { - let graph = fixture_deep_tree()?; - let plan = TopologyOrderSpec::default().plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + let mut graph = fixture_deep_tree()?; + TopologyOrderSpec::default().apply(&mut graph)?; + let ordered = &graph; - let deep_children = child_names(&ordered, "deep")?; + let deep_children = child_names(ordered, "deep")?; // mid (2 leaves) vs D (1 leaf): D first assert_eq!(vec!["D", "mid"], deep_children); - let mid_children = child_names(&ordered, "mid")?; + let mid_children = child_names(ordered, "mid")?; assert_eq!(vec!["E", "F"], mid_children); Ok(()) @@ -831,7 +816,7 @@ mod tests { #[test] fn topology_order_target_order_reverse_inverts_order() -> Result<(), Report> { - let graph = fixture_tree()?; + let mut graph = fixture_tree()?; let spec = TopologyOrderSpec { preset: TopologyOrderPreset::TargetOrderReverse, target_order: vec!["D", "E", "F", "B", "C", "A"] @@ -840,10 +825,10 @@ mod tests { .collect(), target_aggregate: TopologyOrderTargetAggregate::Mean, }; - let plan = spec.plan(&graph)?; - let ordered = plan.ordered_graph(&graph)?; + spec.apply(&mut graph)?; + let ordered = &graph; - let actual = child_names(&ordered, "root")?; + let actual = child_names(ordered, "root")?; assert_eq!(vec!["A", "BC", "DEF"], actual); @@ -852,16 +837,89 @@ mod tests { #[test] fn topology_order_target_order_rejects_empty() { - let graph = fixture_tree().unwrap(); + let mut graph = fixture_tree().unwrap(); let spec = TopologyOrderSpec { preset: TopologyOrderPreset::TargetOrder, target_order: vec![], target_aggregate: TopologyOrderTargetAggregate::Mean, }; - let err = spec.plan(&graph).unwrap_err(); + let err = spec.apply(&mut graph).unwrap_err(); assert!(err.to_string().contains("non-empty target order")); } + #[test] + fn topology_order_target_order_rejects_duplicate_ranking_labels() { + let mut graph = fixture_tree().unwrap(); + let spec = TopologyOrderSpec { + preset: TopologyOrderPreset::TargetOrder, + target_order: vec!["A", "B", "B", "C", "D", "E", "F"] + .into_iter() + .map(str::to_owned) + .collect(), + target_aggregate: TopologyOrderTargetAggregate::Mean, + }; + + let error = spec.apply(&mut graph).unwrap_err(); + + assert!(error.to_string().contains("duplicate leaf label 'B'")); + } + + #[test] + fn topology_order_target_order_rejects_duplicate_final_leaf_labels() -> Result<(), Report> { + let mut graph = fixture_tree()?; + graph + .get_node(find_node(&graph, "C")?) + .expect("fixture node C must exist") + .write_arc() + .payload() + .write_arc() + .0 = "B".to_owned(); + let spec = TopologyOrderSpec { + preset: TopologyOrderPreset::TargetOrder, + target_order: vec!["A", "B", "D", "E", "F"].into_iter().map(str::to_owned).collect(), + target_aggregate: TopologyOrderTargetAggregate::Mean, + }; + + let error = spec.apply(&mut graph).unwrap_err(); + + assert!(error.to_string().contains("final leaf label 'B' is duplicated")); + Ok(()) + } + + #[test] + fn topology_order_target_order_ignores_absent_ranking_labels() -> Result<(), Report> { + let mut graph = fixture_tree()?; + let spec = TopologyOrderSpec { + preset: TopologyOrderPreset::TargetOrder, + target_order: vec!["removed", "D", "E", "F", "B", "C", "A"] + .into_iter() + .map(str::to_owned) + .collect(), + target_aggregate: TopologyOrderTargetAggregate::Mean, + }; + + spec.apply(&mut graph)?; + + assert_eq!(vec!["DEF", "BC", "A"], child_names(&graph, "root")?); + Ok(()) + } + + #[test] + fn topology_order_is_idempotent_and_preserves_graph_data_identity() -> Result<(), Report> { + let graph = fixture_tree()?; + let mut graph = graph.map_data(NonCloneData); + let data = std::ptr::from_ref(graph.data()); + + TopologyOrderSpec::default().apply(&mut graph)?; + let first = child_names_with_data(&graph, "root")?; + TopologyOrderSpec::default().apply(&mut graph)?; + let second = child_names_with_data(&graph, "root")?; + + assert_eq!(first, second); + assert!(std::ptr::eq(data, graph.data())); + Ok(()) + } + /// root -> [deep -> [D, mid -> [E, F]], shallow -> [B, C], A] fn fixture_deep_tree() -> Result, Report> { let mut graph = Graph::::new(); @@ -957,7 +1015,7 @@ mod tests { .ok_or_else(|| make_report!("Node '{name}' not found")) } - #[derive(Clone, Debug, Eq, PartialEq)] + #[derive(Debug, Eq, PartialEq)] struct TestNode(String); impl TestNode { @@ -978,7 +1036,7 @@ mod tests { } } - #[derive(Clone, Debug, PartialEq)] + #[derive(Debug, PartialEq)] struct TestEdge { branch_length: Option, } @@ -1006,4 +1064,26 @@ mod tests { self.branch_length = branch_length; } } + + #[derive(Debug)] + struct NonCloneData; + + fn child_names_with_data( + graph: &Graph, + parent_name: &str, + ) -> Result, Report> { + let parent_key = graph + .find_node(|node| node.0 == parent_name) + .ok_or_else(|| make_report!("Node '{parent_name}' not found"))?; + let parent = graph + .get_node(parent_key) + .ok_or_else(|| make_report!("Node {parent_key} not found"))?; + Ok( + graph + .children_of(&parent.read_arc()) + .into_iter() + .map(|(node, _)| node.read_arc().payload().read_arc().0.clone()) + .collect_vec(), + ) + } } diff --git a/packages/treetime-io/src/auspice.rs b/packages/treetime-io/src/auspice.rs index 64ba13bf2..54ef4dec2 100644 --- a/packages/treetime-io/src/auspice.rs +++ b/packages/treetime-io/src/auspice.rs @@ -7,7 +7,7 @@ use std::io::Cursor; use std::io::{Read, Write}; use std::path::Path; use std::sync::Arc; -use treetime_graph::edge::{Edge, GraphEdge}; +use treetime_graph::edge::{Edge, GraphEdge, GraphEdgeKey}; use treetime_graph::graph::Graph; use treetime_graph::node::{GraphNode, GraphNodeKey, Node}; use treetime_utils::io::file::create_file_or_stdout; @@ -155,6 +155,7 @@ where pub node: &'a N, pub parent_key: Option, pub parent: Option<&'a N>, + pub edge_key: Option, pub edge: Option<&'a E>, pub graph: &'a Graph, } @@ -214,6 +215,9 @@ where .as_ref() .map(|parent: &Arc>>| parent.read_arc().payload().read_arc()); let parent = parent.as_deref(); + let edge_key = current_edge + .as_ref() + .map(|edge: &Arc>>| edge.read_arc().key()); let edge = current_edge .as_ref() .map(|edge: &Arc>>| edge.read_arc().payload().read_arc()); @@ -223,6 +227,7 @@ where node, parent_key, parent, + edge_key, edge, graph, }; diff --git a/packages/treetime-io/src/graph.rs b/packages/treetime-io/src/graph.rs index 9ad7f85bc..6922c0f93 100644 --- a/packages/treetime-io/src/graph.rs +++ b/packages/treetime-io/src/graph.rs @@ -1,12 +1,9 @@ -use crate::auspice::auspice_write_file; +use crate::auspice::{AuspiceWrite, auspice_write_file}; use crate::graphviz::{EdgeToGraphviz, NodeToGraphviz, graphviz_write_file}; use crate::nex::{NexWriteOptions, nex_write_file_with}; use crate::nwk::{CommentProviders, EdgeToNwk, NodeToNwk, NwkWriteOptions, nwk_write_file_with}; -use crate::phyloxml::{PhyloxmlJsonOptions, phyloxml_json_write_file, phyloxml_write_file}; -use crate::tree_ir::auspice::TreeIrAuspiceWriter; -use crate::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode}; -use crate::tree_ir::usher::TreeIrUsherWriter; -use crate::usher_mat::{UsherMatJsonOptions, usher_mat_json_write_file, usher_mat_pb_write_file}; +use crate::phyloxml::{PhyloxmlFromGraph, PhyloxmlJsonOptions, phyloxml_json_write_file, phyloxml_write_file}; +use crate::usher_mat::{UsherMatJsonOptions, UsherWrite, usher_mat_json_write_file, usher_mat_pb_write_file}; use eyre::Report; use serde::Serialize; use std::collections::BTreeMap; @@ -15,13 +12,8 @@ use treetime_graph::edge::{GraphEdge, HasBranchLength}; use treetime_graph::graph::Graph; use treetime_graph::node::{GraphNode, Named}; use treetime_utils::io::json::{JsonPretty, json_write_file}; -use treetime_utils::make_report; use util_newick::NwkStyle; -/// Format-neutral intermediate representation graph for the formats that cannot be -/// derived from the domain graph alone (PhyloXML, Auspice v2, UShER MAT). -pub type TreeIrGraph = Graph; - /// Dispatch tag for tree output formats. Used as keys in the resolved output map. #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)] pub enum TreeWriteKind { @@ -51,23 +43,17 @@ pub struct NwkWriteSpec { pub style: NwkStyle, } -/// Write the requested tree outputs. -/// -/// Newick, Nexus, Graphviz, and graph-JSON are serialized directly from the domain -/// `graph`. PhyloXML, Auspice v2, and UShER MAT are serialized from the -/// format-neutral `ir` graph, which the command builds by projecting its analysis -/// result and partition data. When an IR format is requested but no `ir` graph is -/// provided, the command does not support that format and a clear error is returned. -pub fn write_tree_outputs( +/// Write every requested tree format from one authoritative graph. +pub fn write_tree_outputs( graph: &Graph, outputs: &BTreeMap, providers: &CommentProviders, - ir: Option<&TreeIrGraph>, ) -> Result<(), Report> where N: GraphNode + Named + NodeToNwk + NodeToGraphviz + Serialize, E: GraphEdge + HasBranchLength + EdgeToNwk + EdgeToGraphviz + Serialize, - D: Send + Sync + Default + Serialize, + D: Send + Sync + Serialize, + C: AuspiceWrite + PhyloxmlFromGraph + UsherWrite, { for (kind, path) in outputs { match kind { @@ -92,30 +78,21 @@ where graphviz_write_file(path, graph)?; }, TreeWriteKind::Auspice => { - let ir = require_ir(ir, "Auspice v2 JSON")?; - auspice_write_file::(path, ir)?; + auspice_write_file::(path, graph)?; }, TreeWriteKind::Phyloxml => { - let ir = require_ir(ir, "PhyloXML")?; - phyloxml_write_file(path, ir)?; + phyloxml_write_file::(path, graph)?; }, TreeWriteKind::PhyloxmlJson => { - let ir = require_ir(ir, "PhyloXML JSON")?; - phyloxml_json_write_file(path, ir, &PhyloxmlJsonOptions::default())?; + phyloxml_json_write_file::(path, graph, &PhyloxmlJsonOptions::default())?; }, TreeWriteKind::MatPb => { - let ir = require_ir(ir, "UShER MAT protobuf")?; - usher_mat_pb_write_file::(path, ir)?; + usher_mat_pb_write_file::(path, graph)?; }, TreeWriteKind::MatJson => { - let ir = require_ir(ir, "UShER MAT JSON")?; - usher_mat_json_write_file::(path, ir, &UsherMatJsonOptions::default())?; + usher_mat_json_write_file::(path, graph, &UsherMatJsonOptions::default())?; }, } } Ok(()) } - -fn require_ir<'a>(ir: Option<&'a TreeIrGraph>, format: &str) -> Result<&'a TreeIrGraph, Report> { - ir.ok_or_else(|| make_report!("{format} output was requested but is not available for this command")) -} diff --git a/packages/treetime-io/src/lib.rs b/packages/treetime-io/src/lib.rs index d5c150c9c..53c0919e9 100644 --- a/packages/treetime-io/src/lib.rs +++ b/packages/treetime-io/src/lib.rs @@ -14,7 +14,6 @@ pub mod nex; pub mod nwk; pub mod parse_delimited; pub mod phyloxml; -pub mod tree_ir; pub mod usher_mat; #[cfg(test)] diff --git a/packages/treetime-io/src/nex.rs b/packages/treetime-io/src/nex.rs index 1c7b8f585..47b8cad9c 100644 --- a/packages/treetime-io/src/nex.rs +++ b/packages/treetime-io/src/nex.rs @@ -31,7 +31,7 @@ pub fn nex_write_file( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let mut f = create_file_or_stdout(filepath)?; nex_write(&mut f, graph, options)?; @@ -48,7 +48,7 @@ pub fn nex_write_file_with( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let mut f = create_file_or_stdout(filepath)?; nex_write_with(&mut f, graph, options, providers)?; @@ -60,7 +60,7 @@ pub fn nex_write_str(graph: &Graph, options: &NexWriteOptions) where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let providers = CommentProviders::new(); nex_write_str_with(graph, options, &providers) @@ -75,7 +75,7 @@ pub fn nex_write_str_with( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let mut buf = Vec::new(); nex_write_with(&mut buf, graph, options, providers)?; @@ -86,7 +86,7 @@ pub fn nex_write(w: &mut impl Write, graph: &Graph, options: & where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let providers = CommentProviders::new(); nex_write_with(w, graph, options, &providers) @@ -102,7 +102,7 @@ pub fn nex_write_with( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let n_leaves = graph.num_leaves(); let leaf_names = graph diff --git a/packages/treetime-io/src/nwk.rs b/packages/treetime-io/src/nwk.rs index 674a4963b..3298fa0b2 100644 --- a/packages/treetime-io/src/nwk.rs +++ b/packages/treetime-io/src/nwk.rs @@ -128,7 +128,7 @@ pub fn nwk_write_file( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let mut f = create_file_or_stdout(filepath)?; nwk_write(&mut f, graph, options)?; @@ -145,7 +145,7 @@ pub fn nwk_write_file_with( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let mut f = create_file_or_stdout(filepath)?; nwk_write_with(&mut f, graph, options, providers)?; @@ -157,7 +157,7 @@ pub fn nwk_write_str(graph: &Graph, options: &NwkWriteOptions) where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let providers = CommentProviders::new(); nwk_write_str_with(graph, options, &providers) @@ -172,7 +172,7 @@ pub fn nwk_write_str_with( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let mut buf = Vec::new(); nwk_write_with(&mut buf, graph, options, providers)?; @@ -187,7 +187,7 @@ pub fn nwk_write( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let providers = CommentProviders::new(); nwk_write_with(writer, graph, options, &providers) @@ -205,7 +205,7 @@ pub fn nwk_write_with( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { let roots = graph.get_roots(); if roots.is_empty() { diff --git a/packages/treetime-io/src/phyloxml.rs b/packages/treetime-io/src/phyloxml.rs index 09130b863..21e5e8362 100644 --- a/packages/treetime-io/src/phyloxml.rs +++ b/packages/treetime-io/src/phyloxml.rs @@ -23,112 +23,113 @@ pub use util_phyloxml::{ PhyloxmlTaxonomy, PhyloxmlUri, Polygon, }; -pub fn phyloxml_read_file(filepath: impl AsRef) -> Result, Report> +pub fn phyloxml_read_file(filepath: impl AsRef) -> Result, Report> where N: GraphNode, E: GraphEdge, D: PhyloxmlDataToGraphData + Sync + Send, - (): PhyloxmlToGraph, + C: PhyloxmlToGraph, { let filepath = filepath.as_ref(); - phyloxml_read(open_file_or_stdin(&Some(filepath))?) + phyloxml_read::(open_file_or_stdin(&Some(filepath))?) .wrap_err_with(|| format!("When reading PhyloXML file '{}'", filepath.display())) } -pub fn phyloxml_read_str(s: impl AsRef) -> Result, Report> +pub fn phyloxml_read_str(s: impl AsRef) -> Result, Report> where N: GraphNode, E: GraphEdge, D: PhyloxmlDataToGraphData + Sync + Send, - (): PhyloxmlToGraph, + C: PhyloxmlToGraph, { - phyloxml_read(Cursor::new(s.as_ref())).wrap_err("When reading PhyloXML string") + phyloxml_read::(Cursor::new(s.as_ref())).wrap_err("When reading PhyloXML string") } -pub fn phyloxml_read(reader: impl Read) -> Result, Report> +pub fn phyloxml_read(reader: impl Read) -> Result, Report> where N: GraphNode, E: GraphEdge, D: PhyloxmlDataToGraphData + Sync + Send, - (): PhyloxmlToGraph, + C: PhyloxmlToGraph, { let pxml = util_phyloxml::phyloxml_read(reader).wrap_err("When reading PhyloXML")?; - phyloxml_to_graph(&pxml) + phyloxml_to_graph::(&pxml) } -pub fn phyloxml_json_read_file(filepath: impl AsRef) -> Result, Report> +pub fn phyloxml_json_read_file(filepath: impl AsRef) -> Result, Report> where N: GraphNode, E: GraphEdge, D: PhyloxmlDataToGraphData + Sync + Send, - (): PhyloxmlToGraph, + C: PhyloxmlToGraph, { let filepath = filepath.as_ref(); let tree = json_read_file(filepath).wrap_err_with(|| format!("When reading PhyloXML JSON file '{}'", filepath.display()))?; - phyloxml_to_graph(&tree) + phyloxml_to_graph::(&tree) } -pub fn phyloxml_json_read_str(s: impl AsRef) -> Result, Report> +pub fn phyloxml_json_read_str(s: impl AsRef) -> Result, Report> where N: GraphNode, E: GraphEdge, D: PhyloxmlDataToGraphData + Sync + Send, - (): PhyloxmlToGraph, + C: PhyloxmlToGraph, { let pxml = json_read_str(s).wrap_err("When reading PhyloXML JSON string")?; - phyloxml_to_graph(&pxml) + phyloxml_to_graph::(&pxml) } -pub fn phyloxml_json_read(reader: impl Read) -> Result, Report> +pub fn phyloxml_json_read(reader: impl Read) -> Result, Report> where N: GraphNode, E: GraphEdge, D: PhyloxmlDataToGraphData + Sync + Send, - (): PhyloxmlToGraph, + C: PhyloxmlToGraph, { let pxml = json_read(reader).wrap_err("When reading PhyloXML JSON")?; - phyloxml_to_graph(&pxml) + phyloxml_to_graph::(&pxml) } -pub fn phyloxml_write_file(filepath: impl AsRef, graph: &Graph) -> Result<(), Report> +pub fn phyloxml_write_file(filepath: impl AsRef, graph: &Graph) -> Result<(), Report> where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, + C: PhyloxmlFromGraph, { let filepath = filepath.as_ref(); let mut f = create_file_or_stdout(filepath)?; - phyloxml_write(&mut f, graph).wrap_err_with(|| format!("When writing PhyloXML file '{}'", filepath.display()))?; + phyloxml_write::(&mut f, graph) + .wrap_err_with(|| format!("When writing PhyloXML file '{}'", filepath.display()))?; writeln!(f)?; Ok(()) } -pub fn phyloxml_write_str(graph: &Graph) -> Result +pub fn phyloxml_write_str(graph: &Graph) -> Result where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, + C: PhyloxmlFromGraph, { let mut buf = Vec::new(); - phyloxml_write(&mut buf, graph).wrap_err("When writing PhyloXML string")?; + phyloxml_write::(&mut buf, graph).wrap_err("When writing PhyloXML string")?; String::from_utf8(buf).wrap_err("PhyloXML output is not valid UTF-8") } -pub fn phyloxml_write(writer: &mut impl Write, graph: &Graph) -> Result<(), Report> +pub fn phyloxml_write(writer: &mut impl Write, graph: &Graph) -> Result<(), Report> where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, + C: PhyloxmlFromGraph, { - let pxml = phyloxml_from_graph(graph)?; + let pxml = phyloxml_from_graph::(graph)?; util_phyloxml::phyloxml_write(writer, &pxml).wrap_err("When writing PhyloXML") } -pub fn phyloxml_json_write_file( +pub fn phyloxml_json_write_file( filepath: impl AsRef, graph: &Graph, options: &PhyloxmlJsonOptions, @@ -136,28 +137,31 @@ pub fn phyloxml_json_write_file( where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, + C: PhyloxmlFromGraph, { let filepath = filepath.as_ref(); - let pxml = phyloxml_from_graph(graph)?; + let pxml = phyloxml_from_graph::(graph)?; json_write_file(filepath, &pxml, JsonPretty(options.pretty)) .wrap_err_with(|| format!("When writing PhyloXML JSON file: '{}'", filepath.display()))?; Ok(()) } -pub fn phyloxml_json_write_str(graph: &Graph, options: &PhyloxmlJsonOptions) -> Result +pub fn phyloxml_json_write_str( + graph: &Graph, + options: &PhyloxmlJsonOptions, +) -> Result where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, + C: PhyloxmlFromGraph, { - let pxml = phyloxml_from_graph(graph)?; + let pxml = phyloxml_from_graph::(graph)?; json_write_str(&pxml, JsonPretty(options.pretty)).wrap_err("When writing PhyloXML JSON string") } -pub fn phyloxml_json_write( +pub fn phyloxml_json_write( writer: &mut impl Write, graph: &Graph, options: &PhyloxmlJsonOptions, @@ -165,10 +169,10 @@ pub fn phyloxml_json_write( where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, + C: PhyloxmlFromGraph, { - let pxml = phyloxml_from_graph(graph)?; + let pxml = phyloxml_from_graph::(graph)?; json_write(writer, &pxml, JsonPretty(options.pretty)).wrap_err("When writing PhyloXML JSON") } @@ -176,10 +180,6 @@ pub trait PhyloxmlDataToGraphData: Sized { fn phyloxml_data_to_graph_data(pxml: &Phyloxml) -> Result; } -pub trait PhyloxmlDataFromGraphData { - fn phyloxml_data_from_graph_data(&self) -> Result; -} - #[derive(SmartDefault)] pub struct PhyloxmlJsonOptions { #[default = true] @@ -207,12 +207,12 @@ pub struct PhyloxmlNodeImpl { } /// Convert PhyloXML to graph -pub fn phyloxml_to_graph(pxml: &Phyloxml) -> Result, Report> +pub fn phyloxml_to_graph(pxml: &Phyloxml) -> Result, Report> where N: GraphNode, E: GraphEdge, D: PhyloxmlDataToGraphData + Sync + Send, - (): PhyloxmlToGraph, + C: PhyloxmlToGraph, { let phylogeny = { let n_phylogeny = pxml.phylogeny.len(); @@ -226,8 +226,7 @@ where let mut graph = Graph::::with_data(D::phyloxml_data_to_graph_data(pxml)?); let mut queue: VecDeque<_> = phylogeny.clade.iter().map(|clade| (None, clade)).collect(); while let Some((parent_key, clade)) = queue.pop_front() { - let (graph_node, graph_edge) = - <() as PhyloxmlToGraph>::phyloxml_node_to_graph_components(&PhyloxmlContext { clade, pxml })?; + let (graph_node, graph_edge) = C::phyloxml_node_to_graph_components(&PhyloxmlContext { clade, pxml })?; let node_key = graph.add_node(graph_node); if let Some(parent_key) = parent_key { graph.add_edge(parent_key, node_key, graph_edge)?; @@ -245,10 +244,11 @@ pub struct PhyloxmlGraphContext<'a, N, E, D> where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, { + pub node_key: treetime_graph::node::GraphNodeKey, pub node: &'a N, + pub edge_key: Option, pub edge: Option<&'a E>, pub graph: &'a Graph, } @@ -258,19 +258,20 @@ pub trait PhyloxmlFromGraph where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, { + fn phyloxml_data_from_graph_data(graph: &Graph) -> Result; + fn phyloxml_node_from_graph_components(context: &PhyloxmlGraphContext) -> Result; } /// Convert graph to PhyloXML -pub fn phyloxml_from_graph(graph: &Graph) -> Result +pub fn phyloxml_from_graph(graph: &Graph) -> Result where N: GraphNode, E: GraphEdge, - D: PhyloxmlDataFromGraphData + Sync + Send, - (): PhyloxmlFromGraph, + D: Sync + Send, + C: PhyloxmlFromGraph, { let root = graph.get_exactly_one_root()?; @@ -282,16 +283,20 @@ where let current_node = current_node.read_arc(); let node = &*current_node.payload().read_arc(); + let edge_key = current_edge + .as_ref() + .map(|edge: &Arc>>| edge.read_arc().key()); let edge = current_edge .as_ref() .map(|edge: &Arc>>| edge.read_arc().payload().read_arc()); let edge = edge.as_deref(); - let current_tree_node = - <() as PhyloxmlFromGraph>::phyloxml_node_from_graph_components(&PhyloxmlGraphContext { - node, - edge, - graph, - })?; + let current_tree_node = C::phyloxml_node_from_graph_components(&PhyloxmlGraphContext { + node_key: current_node.key(), + node, + edge_key, + edge, + graph, + })?; for (child, edge) in graph.children_of(¤t_node) { queue.push_back((child, Some(edge))); } @@ -324,7 +329,7 @@ where } } - let mut data = graph.data().read_arc().phyloxml_data_from_graph_data()?; + let mut data = C::phyloxml_data_from_graph_data(graph)?; let root_key = root.read_arc().key(); let clade = node_map .remove(&root_key) diff --git a/packages/treetime-io/src/tree_ir/__tests__/mod.rs b/packages/treetime-io/src/tree_ir/__tests__/mod.rs deleted file mode 100644 index 5c98fb3b3..000000000 --- a/packages/treetime-io/src/tree_ir/__tests__/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -mod test_auspice; -mod test_mutation; -mod test_phyloxml; -mod test_usher; diff --git a/packages/treetime-io/src/tree_ir/__tests__/test_auspice.rs b/packages/treetime-io/src/tree_ir/__tests__/test_auspice.rs deleted file mode 100644 index 99debe6ea..000000000 --- a/packages/treetime-io/src/tree_ir/__tests__/test_auspice.rs +++ /dev/null @@ -1,182 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::auspice::{auspice_read_str, auspice_write_str}; - use crate::auspice_types::AuspiceTree; - use crate::tree_ir::auspice::{TreeIrAuspiceReader, TreeIrAuspiceWriter, format_number}; - use approx::assert_ulps_eq; - use eyre::Report; - use pretty_assertions::assert_eq; - - // Oracle: augur export_v2.format_number keeps `precision` significant figures in the - // fractional part while preserving integer digits. - #[test] - fn test_tree_ir_auspice_format_number_fractional_precision() { - assert_ulps_eq!(0.123457, format_number(0.12345678, 6), max_ulps = 0); - assert_ulps_eq!(123.456789, format_number(123.456789, 6), max_ulps = 0); - assert_ulps_eq!(0.0, format_number(0.0, 6), max_ulps = 0); - assert_ulps_eq!(2020.123, format_number(2020.1234567, 3), max_ulps = 0); - } - - #[test] - fn test_tree_ir_auspice_div_is_cumulative_and_rounded() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let json = auspice_write_str::(&graph)?; - let tree: AuspiceTree = serde_json::from_str(&json)?; - - assert_eq!("root", tree.tree.name); - assert_eq!(Some(0.0), tree.tree.node_attrs.div); - let child_a = tree.tree.children.iter().find(|c| c.name == "A").unwrap(); - assert_eq!(Some(0.5), child_a.node_attrs.div); - Ok(()) - } - - #[test] - fn test_tree_ir_auspice_mutations_grouped_by_gene() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let json = auspice_write_str::(&graph)?; - let tree: AuspiceTree = serde_json::from_str(&json)?; - - let child_a = tree.tree.children.iter().find(|c| c.name == "A").unwrap(); - assert_eq!(vec!["A123T".to_owned()], child_a.branch_attrs.mutations["nuc"]); - Ok(()) - } - - #[test] - fn test_tree_ir_auspice_num_date_with_confidence() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let json = auspice_write_str::(&graph)?; - let tree: AuspiceTree = serde_json::from_str(&json)?; - - let child_a = tree.tree.children.iter().find(|c| c.name == "A").unwrap(); - let num_date = child_a.node_attrs.num_date.as_ref().unwrap(); - assert_ulps_eq!(2020.5, num_date.value, max_ulps = 0); - assert_eq!(Some([2020.2, 2020.8]), num_date.confidence); - Ok(()) - } - - #[test] - fn test_tree_ir_auspice_bad_branch_categorical() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let json = auspice_write_str::(&graph)?; - let tree: AuspiceTree = serde_json::from_str(&json)?; - - let child_b = tree.tree.children.iter().find(|c| c.name == "B").unwrap(); - assert_eq!("Yes", child_b.node_attrs.bad_branch.as_ref().unwrap().value); - assert_eq!("No", tree.tree.node_attrs.bad_branch.as_ref().unwrap().value); - Ok(()) - } - - #[test] - fn test_tree_ir_auspice_colorings_present() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let json = auspice_write_str::(&graph)?; - let tree: AuspiceTree = serde_json::from_str(&json)?; - - let keys: Vec<&str> = tree.data.meta.colorings.iter().map(|c| c.key.as_str()).collect(); - assert!(keys.contains(&"num_date")); - assert!(keys.contains(&"bad_branch")); - assert!(keys.contains(&"gt")); - Ok(()) - } - - #[test] - fn test_tree_ir_auspice_roundtrip_is_idempotent() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let json1 = auspice_write_str::(&graph)?; - let graph2 = auspice_read_str::(&json1)?; - let json2 = auspice_write_str::(&graph2)?; - assert_eq!(json1, json2); - Ok(()) - } - - #[test] - fn test_tree_ir_auspice_non_finite_div_errors() -> Result<(), Report> { - let graph = helpers::graph_with_root_div(f64::NAN)?; - let err = auspice_write_str::(&graph).unwrap_err(); - assert!(format!("{err:?}").contains("non-finite div")); - Ok(()) - } - - mod helpers { - use crate::graph::TreeIrGraph; - use crate::tree_ir::mutation::{NUC_GENE, TreeIrSub}; - use crate::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode}; - use eyre::Report; - use treetime_primitives::AsciiChar; - - pub fn nuc(c: char) -> AsciiChar { - AsciiChar::try_from_char(c).unwrap() - } - - pub fn sub(position: usize, parent: char, child: char) -> TreeIrSub { - TreeIrSub { - gene: NUC_GENE.to_owned(), - position, - parent: nuc(parent), - child: nuc(child), - } - } - - pub fn data() -> TreeIrData { - TreeIrData { - has_dates: true, - has_bad_branch: true, - has_mutations: true, - ..TreeIrData::default() - } - } - - pub fn sample_graph() -> Result { - let mut graph = TreeIrGraph::with_data(data()); - let root = graph.add_node(TreeIrNode { - name: Some("root".to_owned()), - div: Some(0.0), - date: Some(2020.0), - ..TreeIrNode::default() - }); - let a = graph.add_node(TreeIrNode { - name: Some("A".to_owned()), - div: Some(0.5), - date: Some(2020.5), - date_confidence: Some([2020.2, 2020.8]), - ..TreeIrNode::default() - }); - let b = graph.add_node(TreeIrNode { - name: Some("B".to_owned()), - div: Some(1.0), - bad_branch: true, - ..TreeIrNode::default() - }); - graph.add_edge( - root, - a, - TreeIrEdge { - branch_length: Some(0.5), - mutations: vec![sub(123, 'A', 'T')], - ..TreeIrEdge::default() - }, - )?; - graph.add_edge( - root, - b, - TreeIrEdge { - branch_length: Some(1.0), - ..TreeIrEdge::default() - }, - )?; - graph.build()?; - Ok(graph) - } - - pub fn graph_with_root_div(div: f64) -> Result { - let mut graph = TreeIrGraph::with_data(data()); - graph.add_node(TreeIrNode { - name: Some("root".to_owned()), - div: Some(div), - ..TreeIrNode::default() - }); - graph.build()?; - Ok(graph) - } - } -} diff --git a/packages/treetime-io/src/tree_ir/__tests__/test_mutation.rs b/packages/treetime-io/src/tree_ir/__tests__/test_mutation.rs deleted file mode 100644 index d341a7448..000000000 --- a/packages/treetime-io/src/tree_ir/__tests__/test_mutation.rs +++ /dev/null @@ -1,44 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::tree_ir::mutation::{NUC_GENE, TreeIrSub}; - use pretty_assertions::assert_eq; - use treetime_primitives::AsciiChar; - - fn nuc(c: char) -> AsciiChar { - AsciiChar::try_from_char(c).unwrap() - } - - #[test] - fn test_tree_ir_mutation_nuc_string_roundtrip() { - let sub = TreeIrSub { - gene: NUC_GENE.to_owned(), - position: 123, - parent: nuc('A'), - child: nuc('T'), - }; - assert_eq!("A123T", sub.to_auspice_string()); - assert_eq!(sub, TreeIrSub::from_auspice_string(NUC_GENE, "A123T").unwrap()); - } - - #[test] - fn test_tree_ir_mutation_aa_string_roundtrip() { - let sub = TreeIrSub { - gene: "M".to_owned(), - position: 1, - parent: nuc('M'), - child: nuc('I'), - }; - assert_eq!("M1I", sub.to_auspice_string()); - assert_eq!(sub, TreeIrSub::from_auspice_string("M", "M1I").unwrap()); - } - - #[test] - fn test_tree_ir_mutation_too_short_is_error() { - let _error = TreeIrSub::from_auspice_string(NUC_GENE, "A").unwrap_err(); - } - - #[test] - fn test_tree_ir_mutation_non_numeric_position_is_error() { - let _error = TreeIrSub::from_auspice_string(NUC_GENE, "AxT").unwrap_err(); - } -} diff --git a/packages/treetime-io/src/tree_ir/__tests__/test_phyloxml.rs b/packages/treetime-io/src/tree_ir/__tests__/test_phyloxml.rs deleted file mode 100644 index 6c58ad589..000000000 --- a/packages/treetime-io/src/tree_ir/__tests__/test_phyloxml.rs +++ /dev/null @@ -1,100 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::graph::TreeIrGraph; - use crate::phyloxml::{phyloxml_read_str, phyloxml_write_str}; - use eyre::Report; - use pretty_assertions::assert_eq; - use treetime_graph::node::Named; - - #[test] - fn test_tree_ir_phyloxml_emits_name_branch_length_and_div_property() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let xml = phyloxml_write_str(&graph)?; - assert!(xml.contains("A")); - assert!(xml.contains("0.5")); - assert!(xml.contains("treetime:divergence")); - assert!(xml.contains("treetime:mutation")); - Ok(()) - } - - #[test] - fn test_tree_ir_phyloxml_roundtrip_is_idempotent() -> Result<(), Report> { - let graph = helpers::sample_graph()?; - let xml1 = phyloxml_write_str(&graph)?; - let graph2: TreeIrGraph = phyloxml_read_str(&xml1)?; - let xml2 = phyloxml_write_str(&graph2)?; - assert_eq!(xml1, xml2); - Ok(()) - } - - #[test] - fn test_tree_ir_phyloxml_escapes_xml_special_chars_in_name() -> Result<(), Report> { - let graph = helpers::single_node_named("a AsciiChar { - AsciiChar::try_from_char(c).unwrap() - } - - pub fn sample_graph() -> Result { - let mut graph = TreeIrGraph::with_data(TreeIrData::default()); - let root = graph.add_node(TreeIrNode { - name: Some("root".to_owned()), - div: Some(0.0), - ..TreeIrNode::default() - }); - let a = graph.add_node(TreeIrNode { - name: Some("A".to_owned()), - div: Some(0.5), - date: Some(2020.5), - ..TreeIrNode::default() - }); - graph.add_edge( - root, - a, - TreeIrEdge { - branch_length: Some(0.5), - mutations: vec![TreeIrSub { - gene: NUC_GENE.to_owned(), - position: 7, - parent: nuc('A'), - child: nuc('G'), - }], - ..TreeIrEdge::default() - }, - )?; - graph.build()?; - Ok(graph) - } - - pub fn single_node_named(name: &str) -> Result { - let mut graph = TreeIrGraph::with_data(TreeIrData::default()); - graph.add_node(TreeIrNode { - name: Some(name.to_owned()), - ..TreeIrNode::default() - }); - graph.build()?; - Ok(graph) - } - } -} diff --git a/packages/treetime-io/src/tree_ir/__tests__/test_usher.rs b/packages/treetime-io/src/tree_ir/__tests__/test_usher.rs deleted file mode 100644 index 507825648..000000000 --- a/packages/treetime-io/src/tree_ir/__tests__/test_usher.rs +++ /dev/null @@ -1,136 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::tree_ir::usher::{TreeIrUsherReader, TreeIrUsherWriter}; - use crate::usher_mat::{usher_from_graph, usher_to_graph}; - use eyre::Report; - use pretty_assertions::assert_eq; - - #[test] - fn test_tree_ir_usher_encodes_nuc_substitution() -> Result<(), Report> { - let graph = helpers::sample_graph(helpers::nuc_sub())?; - let tree = usher_from_graph::(&graph)?; - - let muts: Vec<_> = tree.node_mutations.iter().flat_map(|ml| &ml.mutation).collect(); - assert_eq!(1, muts.len()); - assert_eq!(123, muts[0].position); - assert_eq!(0, muts[0].par_nuc); // A - assert_eq!(vec![3], muts[0].mut_nuc); // T - assert_eq!(0, muts[0].ref_nuc); // no reference -> falls back to parent (A) - Ok(()) - } - - #[test] - fn test_tree_ir_usher_roundtrip_preserves_mutations_and_topology() -> Result<(), Report> { - let graph = helpers::sample_graph(helpers::nuc_sub())?; - let tree = usher_from_graph::(&graph)?; - let graph2 = usher_to_graph::(&tree)?; - let tree2 = usher_from_graph::(&graph2)?; - - assert_eq!(tree.newick, tree2.newick); - assert_eq!(tree.node_mutations, tree2.node_mutations); - Ok(()) - } - - #[test] - fn test_tree_ir_usher_drops_amino_acid_mutations() -> Result<(), Report> { - let graph = helpers::sample_graph(helpers::aa_sub())?; - let tree = usher_from_graph::(&graph)?; - let total: usize = tree.node_mutations.iter().map(|ml| ml.mutation.len()).sum(); - assert_eq!(0, total); - Ok(()) - } - - #[test] - fn test_tree_ir_usher_drops_indels() -> Result<(), Report> { - let graph = helpers::sample_graph_with_indel()?; - let tree = usher_from_graph::(&graph)?; - let total: usize = tree.node_mutations.iter().map(|ml| ml.mutation.len()).sum(); - assert_eq!(0, total); - Ok(()) - } - - mod helpers { - use crate::graph::TreeIrGraph; - use crate::tree_ir::mutation::{IndelKind, NUC_GENE, TreeIrIndel, TreeIrSub}; - use crate::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode}; - use eyre::Report; - use treetime_primitives::AsciiChar; - - fn nuc(c: char) -> AsciiChar { - AsciiChar::try_from_char(c).unwrap() - } - - pub fn nuc_sub() -> TreeIrSub { - TreeIrSub { - gene: NUC_GENE.to_owned(), - position: 123, - parent: nuc('A'), - child: nuc('T'), - } - } - - pub fn aa_sub() -> TreeIrSub { - TreeIrSub { - gene: "M".to_owned(), - position: 1, - parent: nuc('M'), - child: nuc('I'), - } - } - - fn named_node(name: &str) -> TreeIrNode { - TreeIrNode { - name: Some(name.to_owned()), - ..TreeIrNode::default() - } - } - - pub fn sample_graph(sub: TreeIrSub) -> Result { - let mut graph = TreeIrGraph::with_data(TreeIrData::default()); - let root = graph.add_node(named_node("root")); - let a = graph.add_node(named_node("A")); - let b = graph.add_node(named_node("B")); - graph.add_edge( - root, - a, - TreeIrEdge { - branch_length: Some(0.5), - mutations: vec![sub], - ..TreeIrEdge::default() - }, - )?; - graph.add_edge( - root, - b, - TreeIrEdge { - branch_length: Some(1.0), - ..TreeIrEdge::default() - }, - )?; - graph.build()?; - Ok(graph) - } - - pub fn sample_graph_with_indel() -> Result { - let mut graph = TreeIrGraph::with_data(TreeIrData::default()); - let root = graph.add_node(named_node("root")); - let a = graph.add_node(named_node("A")); - graph.add_edge( - root, - a, - TreeIrEdge { - branch_length: Some(0.5), - indels: vec![TreeIrIndel { - gene: NUC_GENE.to_owned(), - kind: IndelKind::Deletion, - start: 10, - seq: vec![nuc('A'), nuc('C')], - }], - ..TreeIrEdge::default() - }, - )?; - graph.build()?; - Ok(graph) - } - } -} diff --git a/packages/treetime-io/src/tree_ir/auspice.rs b/packages/treetime-io/src/tree_ir/auspice.rs deleted file mode 100644 index 06f88f31f..000000000 --- a/packages/treetime-io/src/tree_ir/auspice.rs +++ /dev/null @@ -1,342 +0,0 @@ -//! Auspice v2 JSON adapter for the format-neutral TreeIR graph. -//! -//! Output semantics follow augur's `export_v2.py`: -//! - `node_attrs.div` is cumulative divergence from the root (root = 0), precomputed -//! on each IR node, rounded with [`format_number`] at precision 6. -//! - `node_attrs.num_date` carries the numeric date with optional `[lower, upper]` -//! confidence, rounded at precision 3. -//! - `meta.colorings` are derived from which data classes the dataset populates: -//! `num_date` (continuous) when dates are present, `bad_branch` (categorical) when -//! the outlier flag is present, one categorical coloring per discrete trait, and -//! the genotype coloring `gt` when branch mutations are present. -//! - `branch_attrs.mutations` groups per-branch substitutions by gene track. - -use crate::auspice::{AuspiceGraphContext, AuspiceRead, AuspiceTreeContext, AuspiceWrite}; -use crate::auspice_types::{ - AuspiceColoring, AuspiceDisplayDefaults, AuspiceNumDate, AuspiceTree, AuspiceTreeBranchAttrs, AuspiceTreeData, - AuspiceTreeMeta, AuspiceTreeNode, AuspiceTreeNodeAttr, AuspiceTreeNodeAttrs, -}; -use crate::tree_ir::mutation::TreeIrSub; -use crate::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode, TreeIrTrait}; -use eyre::Report; -use serde_json::{Value, json}; -use std::collections::BTreeMap; -use treetime_graph::graph::Graph; -use treetime_graph::node::Named; -use treetime_utils::make_error; - -const COLORING_GENOTYPE: &str = "gt"; -const COLORING_NUM_DATE: &str = "num_date"; -const COLORING_BAD_BRANCH: &str = "bad_branch"; - -/// Round a float the way augur's `format_number` does: keep `precision` significant -/// digits in the fractional part while preserving all integer digits. -/// -/// Reimplements `augur.export_v2.format_number`: total significant figures equal the -/// number of integer digits plus `precision`. Formatting through scientific notation -/// and parsing back yields the same `f64` augur stores after `float(f"{n:.{k}g}")`. -pub fn format_number(n: f64, precision: i32) -> f64 { - if n == 0.0 || !n.is_finite() { - return n; - } - let integral = n.abs().trunc(); - let significand = if integral >= 1.0 { - integral.log10().floor() as i32 + 1 - } else { - 0 - }; - let sig_figs = (significand + precision).max(1) as usize; - format!("{n:.*e}", sig_figs - 1) - .parse() - .expect("a float formatted in scientific notation must parse back") -} - -/// Converter writing a TreeIR graph to Auspice v2 JSON. -pub struct TreeIrAuspiceWriter { - has_bad_branch: bool, -} - -impl AuspiceWrite for TreeIrAuspiceWriter { - fn new(graph: &Graph) -> Result { - let has_bad_branch = graph.data().read_arc().has_bad_branch; - Ok(Self { has_bad_branch }) - } - - fn auspice_data_from_graph_data( - &self, - graph: &Graph, - ) -> Result { - let data = graph.data().read_arc(); - - let mut colorings = vec![]; - if data.has_dates { - colorings.push(coloring(COLORING_NUM_DATE, "Date", "continuous")); - } - if data.has_bad_branch { - colorings.push(coloring(COLORING_BAD_BRANCH, "Excluded", "categorical")); - } - for attr in &data.trait_attrs { - colorings.push(coloring(attr, attr, "categorical")); - } - if data.has_mutations { - colorings.push(coloring(COLORING_GENOTYPE, "Genotype", "categorical")); - } - - // Default coloring: a discrete trait if present, else the outlier flag, else date. - let color_by = data - .trait_attrs - .first() - .cloned() - .or_else(|| data.has_bad_branch.then(|| COLORING_BAD_BRANCH.to_owned())) - .or_else(|| data.has_dates.then(|| COLORING_NUM_DATE.to_owned())); - - let mut filters = data.trait_attrs.clone(); - if data.has_bad_branch { - filters.push(COLORING_BAD_BRANCH.to_owned()); - } - - let root_sequence = (!data.root_sequence.is_empty()).then(|| data.root_sequence.clone()); - - Ok(AuspiceTreeData { - version: Some("v2".to_owned()), - meta: AuspiceTreeMeta { - title: data.title.clone().or_else(|| Some("TreeTime analysis".to_owned())), - description: data.description.clone(), - panels: vec!["tree".to_owned()], - colorings, - filters, - display_defaults: AuspiceDisplayDefaults { - color_by, - ..AuspiceDisplayDefaults::default() - }, - ..AuspiceTreeMeta::default() - }, - root_sequence, - other: Value::default(), - }) - } - - fn auspice_node_from_graph_components( - &mut self, - context: &AuspiceGraphContext, - ) -> Result { - let &AuspiceGraphContext { - node_key, node, edge, .. - } = context; - - let name = node - .name() - .map_or_else(|| format!("node_{}", node_key.as_usize()), |n| n.as_ref().to_owned()); - - let div = match node.div { - Some(div) => { - if !div.is_finite() { - return make_error!("Node '{name}' has non-finite div={div}"); - } - Some(format_number(div, 6)) - }, - None => None, - }; - - let num_date = match node.date { - Some(date) => { - if !date.is_finite() { - return make_error!("Node '{name}' has non-finite date={date}"); - } - let confidence = match node.date_confidence { - Some([lo, hi]) => { - if !lo.is_finite() || !hi.is_finite() { - return make_error!("Node '{name}' has non-finite date confidence [{lo}, {hi}]"); - } - Some([format_number(lo, 3), format_number(hi, 3)]) - }, - None => None, - }; - Some(AuspiceNumDate { - value: format_number(date, 3), - confidence, - }) - }, - None => None, - }; - - let bad_branch = self - .has_bad_branch - .then(|| AuspiceTreeNodeAttr::new(if node.bad_branch { "Yes" } else { "No" })); - - let other = build_trait_attrs(&node.traits); - - let mutations = edge.map_or_else(BTreeMap::new, |edge| group_mutations(&edge.mutations)); - - Ok(AuspiceTreeNode { - name, - branch_attrs: AuspiceTreeBranchAttrs { - mutations, - labels: None, - other: Value::default(), - }, - node_attrs: AuspiceTreeNodeAttrs { - div, - num_date, - bad_branch, - clade_membership: None, - region: None, - country: None, - division: None, - other, - }, - children: vec![], - other: Value::default(), - }) - } -} - -/// Converter reading a TreeIR graph from Auspice v2 JSON. -pub struct TreeIrAuspiceReader { - trait_attrs: Vec, -} - -impl AuspiceRead for TreeIrAuspiceReader { - fn new(tree: &AuspiceTree) -> Result { - Ok(Self { - trait_attrs: trait_attrs_from_meta(&tree.data.meta), - }) - } - - fn auspice_data_to_graph_data(&mut self, tree: &AuspiceTree) -> Result { - let meta = &tree.data.meta; - let keys: Vec<&str> = meta.colorings.iter().map(|c| c.key.as_str()).collect(); - Ok(TreeIrData { - title: meta.title.clone(), - description: meta.description.clone(), - root_sequence: tree.data.root_sequence.clone().unwrap_or_default(), - trait_attrs: self.trait_attrs.clone(), - has_dates: keys.contains(&COLORING_NUM_DATE), - has_bad_branch: keys.contains(&COLORING_BAD_BRANCH), - has_mutations: keys.contains(&COLORING_GENOTYPE), - }) - } - - fn auspice_node_to_graph_components( - &mut self, - context: &AuspiceTreeContext, - ) -> Result<(TreeIrNode, TreeIrEdge), Report> { - let node = context.node; - - let date = node.node_attrs.num_date.as_ref().map(|nd| nd.value); - let date_confidence = node.node_attrs.num_date.as_ref().and_then(|nd| nd.confidence); - let bad_branch = node - .node_attrs - .bad_branch - .as_ref() - .is_some_and(|attr| attr.value == "Yes"); - - let mut traits = BTreeMap::new(); - for attr in &self.trait_attrs { - if let Some(value) = node.node_attrs.other.get(attr) { - if let Some(trait_) = parse_trait(value) { - traits.insert(attr.clone(), trait_); - } - } - } - - let ir_node = TreeIrNode { - name: Some(node.name.clone()), - div: node.node_attrs.div, - date, - date_confidence, - bad_branch, - traits, - ..TreeIrNode::default() - }; - - let mut mutations = vec![]; - for (gene, muts) in &node.branch_attrs.mutations { - for m in muts { - mutations.push(TreeIrSub::from_auspice_string(gene, m)?); - } - } - - let ir_edge = TreeIrEdge { - branch_length: context.branch_length(), - mutations, - ..TreeIrEdge::default() - }; - - Ok((ir_node, ir_edge)) - } -} - -fn coloring(key: &str, title: &str, type_: &str) -> AuspiceColoring { - AuspiceColoring { - key: key.to_owned(), - title: title.to_owned(), - type_: type_.to_owned(), - ..AuspiceColoring::default() - } -} - -fn build_trait_attrs(traits: &BTreeMap) -> Value { - if traits.is_empty() { - return Value::default(); - } - let mut map = serde_json::Map::new(); - for (attr, trait_) in traits { - let mut obj = serde_json::Map::new(); - obj.insert("value".to_owned(), json!(trait_.value)); - if !trait_.confidence.is_empty() { - let conf: BTreeMap = trait_ - .confidence - .iter() - .map(|(k, v)| (k.clone(), format_number(*v, 3))) - .collect(); - obj.insert("confidence".to_owned(), json!(conf)); - } - if let Some(entropy) = trait_.entropy { - obj.insert("entropy".to_owned(), json!(format_number(entropy, 3))); - } - map.insert(attr.clone(), Value::Object(obj)); - } - Value::Object(map) -} - -fn group_mutations(subs: &[TreeIrSub]) -> BTreeMap> { - let mut grouped: BTreeMap> = BTreeMap::new(); - for sub in subs { - grouped - .entry(sub.gene.clone()) - .or_default() - .push(sub.to_auspice_string()); - } - grouped -} - -fn trait_attrs_from_meta(meta: &AuspiceTreeMeta) -> Vec { - let excluded = [COLORING_GENOTYPE, COLORING_NUM_DATE, COLORING_BAD_BRANCH, "author"]; - meta - .colorings - .iter() - .map(|c| c.key.clone()) - .filter(|key| !excluded.contains(&key.as_str())) - .collect() -} - -fn parse_trait(value: &Value) -> Option { - let value_str = value.get("value")?.as_str()?.to_owned(); - let confidence = value - .get("confidence") - .and_then(Value::as_object) - .map(|obj| { - obj - .iter() - .filter_map(|(k, v)| v.as_f64().map(|f| (k.clone(), f))) - .collect() - }) - .unwrap_or_default(); - let entropy = value.get("entropy").and_then(Value::as_f64); - Some(TreeIrTrait { - value: value_str, - confidence, - entropy, - }) -} diff --git a/packages/treetime-io/src/tree_ir/mod.rs b/packages/treetime-io/src/tree_ir/mod.rs deleted file mode 100644 index 93107f822..000000000 --- a/packages/treetime-io/src/tree_ir/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -#[cfg(test)] -mod __tests__; -pub mod auspice; -pub mod mutation; -pub mod phyloxml; -pub mod types; -pub mod usher; diff --git a/packages/treetime-io/src/tree_ir/mutation.rs b/packages/treetime-io/src/tree_ir/mutation.rs deleted file mode 100644 index f93c1ab38..000000000 --- a/packages/treetime-io/src/tree_ir/mutation.rs +++ /dev/null @@ -1,82 +0,0 @@ -use eyre::{Report, WrapErr}; -use serde::{Deserialize, Serialize}; -use treetime_primitives::AsciiChar; -use treetime_utils::make_error; - -/// Gene key for nucleotide substitutions. -/// -/// Matches the key augur uses in `branch_attrs.mutations` and in the -/// `root_sequence` map for the whole-genome nucleotide track. -pub const NUC_GENE: &str = "nuc"; - -/// A single substitution on a branch: ancestral state -> derived state at one position. -/// -/// `gene` is [`NUC_GENE`] for nucleotide substitutions or a CDS/gene name for -/// amino-acid substitutions. `position` is 1-based, matching both the Auspice -/// mutation-string convention (`A123T`) and the UShER protobuf position field. -/// `parent` and `child` are the ancestral and derived characters (`A`/`C`/`G`/`T` -/// for nucleotides, single-letter amino acids otherwise). -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct TreeIrSub { - pub gene: String, - pub position: usize, - pub parent: AsciiChar, - pub child: AsciiChar, -} - -impl TreeIrSub { - pub fn is_nuc(&self) -> bool { - self.gene == NUC_GENE - } - - /// Render in Auspice mutation-string form, e.g. `A123T`. - pub fn to_auspice_string(&self) -> String { - format!("{}{}{}", self.parent, self.position, self.child) - } - - /// Parse an Auspice mutation string (`A123T`) for a given gene track. - /// - /// The leading character is the ancestral state, the trailing character is the - /// derived state, and the digits in between are the 1-based position. - pub fn from_auspice_string(gene: &str, s: &str) -> Result { - let bytes = s.as_bytes(); - if bytes.len() < 3 { - return make_error!("Invalid mutation string '{s}': expected form like 'A123T'"); - } - let parent = AsciiChar::try_new(bytes[0])?; - let child = AsciiChar::try_new(bytes[bytes.len() - 1])?; - let position_bytes = &bytes[1..bytes.len() - 1]; - let position: usize = std::str::from_utf8(position_bytes) - .wrap_err_with(|| format!("Invalid mutation string '{s}': position is not valid UTF-8"))? - .parse() - .wrap_err_with(|| format!("Invalid mutation string '{s}': position is not a number"))?; - Ok(Self { - gene: gene.to_owned(), - position, - parent, - child, - }) - } -} - -/// Whether an indel is an insertion or a deletion relative to the parent. -#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub enum IndelKind { - Insertion, - Deletion, -} - -/// An insertion or deletion event on a branch. -/// -/// `start` is the 1-based start position; `seq` is the inserted (for -/// [`IndelKind::Insertion`]) or deleted (for [`IndelKind::Deletion`]) sequence. -/// Indels are not representable in every output format (UShER MAT carries only -/// nucleotide substitutions); adapters that cannot represent them drop indels -/// with a warning. -#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -pub struct TreeIrIndel { - pub gene: String, - pub kind: IndelKind, - pub start: usize, - pub seq: Vec, -} diff --git a/packages/treetime-io/src/tree_ir/phyloxml.rs b/packages/treetime-io/src/tree_ir/phyloxml.rs deleted file mode 100644 index e309527d6..000000000 --- a/packages/treetime-io/src/tree_ir/phyloxml.rs +++ /dev/null @@ -1,268 +0,0 @@ -//! PhyloXML adapter for the format-neutral TreeIR graph. -//! -//! PhyloXML carries name, branch length, and date natively. TreeTime-specific -//! data (divergence, discrete traits, mutations, indels, branch-rate multiplier, -//! bad-branch flag) has no native PhyloXML element, so it is projected into -//! generic `` elements with descriptive `ref` attributes and the XSD -//! `datatype`/`applies_to` qualifiers. The `treetime:` ref prefix is a placeholder -//! pending a dedicated property namespace (see -//! `kb/proposals/phyloxml-treetime-property-namespace.md`). - -use crate::phyloxml::{ - Phyloxml, PhyloxmlClade, PhyloxmlContext, PhyloxmlDataFromGraphData, PhyloxmlDataToGraphData, PhyloxmlDate, - PhyloxmlFromGraph, PhyloxmlGraphContext, PhyloxmlPhylogeny, PhyloxmlProperty, PhyloxmlToGraph, -}; -use crate::tree_ir::mutation::{IndelKind, TreeIrIndel, TreeIrSub}; -use crate::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode, TreeIrTrait}; -use eyre::{Report, WrapErr}; -use std::collections::BTreeMap; -use treetime_primitives::AsciiChar; -use treetime_utils::{make_error, make_report}; - -const BRANCH_LENGTH_UNIT: &str = "subs/site"; - -const REF_DIV: &str = "treetime:divergence"; -const REF_BAD_BRANCH: &str = "treetime:bad_branch"; -const REF_DATE_INFERRED: &str = "treetime:date_inferred"; -const REF_GAMMA: &str = "treetime:gamma"; -const REF_MUTATION: &str = "treetime:mutation"; -const REF_INDEL: &str = "treetime:indel"; -const REF_TRAIT_PREFIX: &str = "treetime:trait:"; - -const APPLIES_NODE: &str = "node"; -const APPLIES_BRANCH: &str = "parent_branch"; - -const DT_DOUBLE: &str = "xsd:double"; -const DT_STRING: &str = "xsd:string"; -const DT_BOOLEAN: &str = "xsd:boolean"; - -impl PhyloxmlDataFromGraphData for TreeIrData { - fn phyloxml_data_from_graph_data(&self) -> Result { - let phylogeny = PhyloxmlPhylogeny { - rooted: true, - rerootable: None, - branch_length_unit: Some(BRANCH_LENGTH_UNIT.to_owned()), - phylogeny_type: None, - name: self.title.clone(), - id: None, - description: self.description.clone(), - date: None, - confidence: vec![], - clade: None, - clade_relation: vec![], - sequence_relation: vec![], - property: vec![], - other: BTreeMap::new(), - }; - Ok(Phyloxml { - phylogeny: vec![phylogeny], - other: BTreeMap::new(), - }) - } -} - -impl PhyloxmlDataToGraphData for TreeIrData { - fn phyloxml_data_to_graph_data(pxml: &Phyloxml) -> Result { - let phylogeny = pxml - .phylogeny - .first() - .ok_or_else(|| make_report!("PhyloXML has no phylogeny"))?; - Ok(Self { - title: phylogeny.name.clone(), - description: phylogeny.description.clone(), - ..Self::default() - }) - } -} - -impl PhyloxmlFromGraph for () { - fn phyloxml_node_from_graph_components( - context: &PhyloxmlGraphContext, - ) -> Result { - let node = context.node; - let edge = context.edge; - - let mut property = vec![]; - - if let Some(div) = node.div { - property.push(make_property(REF_DIV, DT_DOUBLE, APPLIES_NODE, &div.to_string())); - } - if node.bad_branch { - property.push(make_property(REF_BAD_BRANCH, DT_BOOLEAN, APPLIES_NODE, "true")); - } - if node.date_is_inferred { - property.push(make_property(REF_DATE_INFERRED, DT_BOOLEAN, APPLIES_NODE, "true")); - } - for (attr, value) in &node.traits { - let ref_ = format!("{REF_TRAIT_PREFIX}{attr}"); - property.push(make_property(&ref_, DT_STRING, APPLIES_NODE, &value.value)); - } - - if let Some(edge) = edge { - if let Some(gamma) = edge.gamma { - property.push(make_property(REF_GAMMA, DT_DOUBLE, APPLIES_BRANCH, &gamma.to_string())); - } - for sub in &edge.mutations { - property.push(make_property( - REF_MUTATION, - DT_STRING, - APPLIES_BRANCH, - &encode_mutation(sub), - )); - } - for indel in &edge.indels { - property.push(make_property( - REF_INDEL, - DT_STRING, - APPLIES_BRANCH, - &encode_indel(indel), - )); - } - } - - let date = node.date.map(|value| { - let (minimum, maximum) = match node.date_confidence { - Some([lo, hi]) => (Some(lo), Some(hi)), - None => (None, None), - }; - PhyloxmlDate { - desc: None, - value: Some(value), - minimum, - maximum, - unit: Some("year".to_owned()), - } - }); - - Ok(PhyloxmlClade { - name: node.name.clone(), - branch_length_elem: edge.and_then(|edge| edge.branch_length), - branch_length_attr: None, - confidence: vec![], - width: None, - color: None, - node_id: None, - taxonomy: vec![], - sequence: vec![], - events: None, - binary_characters: None, - distribution: vec![], - date, - reference: vec![], - property, - clade: vec![], - other: BTreeMap::new(), - }) - } -} - -impl PhyloxmlToGraph for () { - fn phyloxml_node_to_graph_components(context: &PhyloxmlContext) -> Result<(TreeIrNode, TreeIrEdge), Report> { - let clade = context.clade; - - let mut node = TreeIrNode { - name: clade.name.clone(), - ..TreeIrNode::default() - }; - let mut edge = TreeIrEdge { - branch_length: clade.branch_length_elem.or(clade.branch_length_attr), - ..TreeIrEdge::default() - }; - - if let Some(date) = &clade.date { - node.date = date.value; - node.date_confidence = match (date.minimum, date.maximum) { - (Some(lo), Some(hi)) => Some([lo, hi]), - _ => None, - }; - } - - for prop in &clade.property { - match prop.ref_.as_str() { - REF_DIV => node.div = Some(parse_f64(&prop.value, REF_DIV)?), - REF_BAD_BRANCH => node.bad_branch = prop.value == "true", - REF_DATE_INFERRED => node.date_is_inferred = prop.value == "true", - REF_GAMMA => edge.gamma = Some(parse_f64(&prop.value, REF_GAMMA)?), - REF_MUTATION => edge.mutations.push(decode_mutation(&prop.value)?), - REF_INDEL => edge.indels.push(decode_indel(&prop.value)?), - other if other.starts_with(REF_TRAIT_PREFIX) => { - let attr = other.trim_start_matches(REF_TRAIT_PREFIX).to_owned(); - node.traits.insert( - attr, - TreeIrTrait { - value: prop.value.clone(), - ..TreeIrTrait::default() - }, - ); - }, - _ => {}, - } - } - - Ok((node, edge)) - } -} - -fn make_property(ref_: &str, datatype: &str, applies_to: &str, value: &str) -> PhyloxmlProperty { - PhyloxmlProperty { - value: value.to_owned(), - ref_: ref_.to_owned(), - unit: None, - datatype: datatype.to_owned(), - applies_to: applies_to.to_owned(), - id_ref: None, - } -} - -/// Encode a substitution as `gene:A123T` so the gene track survives the round-trip. -fn encode_mutation(sub: &TreeIrSub) -> String { - format!("{}:{}", sub.gene, sub.to_auspice_string()) -} - -fn decode_mutation(value: &str) -> Result { - let (gene, mut_str) = value - .split_once(':') - .ok_or_else(|| make_report!("Invalid mutation property '{value}': expected 'gene:A123T'"))?; - TreeIrSub::from_auspice_string(gene, mut_str) -} - -/// Encode an indel as `gene:ins|del:start:SEQ`. -fn encode_indel(indel: &TreeIrIndel) -> String { - let kind = match indel.kind { - IndelKind::Insertion => "ins", - IndelKind::Deletion => "del", - }; - let seq: String = indel.seq.iter().map(|c| char::from(*c)).collect(); - format!("{}:{}:{}:{}", indel.gene, kind, indel.start, seq) -} - -fn decode_indel(value: &str) -> Result { - let parts: Vec<&str> = value.splitn(4, ':').collect(); - if parts.len() != 4 { - return make_error!("Invalid indel property '{value}': expected 'gene:ins|del:start:SEQ'"); - } - let kind = match parts[1] { - "ins" => IndelKind::Insertion, - "del" => IndelKind::Deletion, - other => return make_error!("Invalid indel kind '{other}' in '{value}'"), - }; - let start: usize = parts[2] - .parse() - .wrap_err_with(|| format!("Invalid indel start in '{value}'"))?; - let seq = parts[3] - .bytes() - .map(AsciiChar::try_new) - .collect::, _>>()?; - Ok(TreeIrIndel { - gene: parts[0].to_owned(), - kind, - start, - seq, - }) -} - -fn parse_f64(value: &str, ref_: &str) -> Result { - value - .parse() - .wrap_err_with(|| format!("Invalid numeric value '{value}' for property '{ref_}'")) -} diff --git a/packages/treetime-io/src/tree_ir/types.rs b/packages/treetime-io/src/tree_ir/types.rs deleted file mode 100644 index 6873972da..000000000 --- a/packages/treetime-io/src/tree_ir/types.rs +++ /dev/null @@ -1,160 +0,0 @@ -use crate::graphviz::{EdgeToGraphviz, NodeToGraphviz}; -use crate::nwk::{EdgeFromNwk, EdgeToNwk, NodeFromNwk, NodeToNwk}; -use crate::tree_ir::mutation::{TreeIrIndel, TreeIrSub}; -use eyre::Report; -use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; -use treetime_graph::edge::{GraphEdge, HasBranchLength}; -use treetime_graph::node::{GraphNode, Named}; - -/// Format-neutral tree node. -/// -/// Carries the union of node-level data the supported tree formats can represent. -/// Each adapter projects the subset its format supports; data a format cannot carry -/// is dropped (with a warning where it is lossy and silent omission would surprise). -/// -/// `div` is cumulative divergence from the root (root = 0), precomputed at projection -/// time so every adapter emits a consistent value without re-deriving it. The -/// projection chooses the per-branch contribution source (mutation counts when -/// divergence is measured in mutations, branch length otherwise), matching augur's -/// `node_div` preference order. -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct TreeIrNode { - pub name: Option, - pub desc: Option, - pub div: Option, - /// Numeric (decimal-year) date or inferred time. - pub date: Option, - /// Date confidence interval `[lower, upper]`. - pub date_confidence: Option<[f64; 2]>, - /// Whether the date was inferred (internal nodes) rather than observed (tips). - pub date_is_inferred: bool, - /// Original date string for tips (augur `num_date.raw_value`). - pub date_raw_value: Option, - /// Whether the node is flagged as a bad branch / temporal outlier. - pub bad_branch: bool, - /// Discrete trait assignments keyed by attribute name (e.g. mugration `country`). - pub traits: BTreeMap, -} - -impl GraphNode for TreeIrNode {} - -impl Named for TreeIrNode { - fn name(&self) -> Option> { - self.name.as_deref() - } - - fn set_name(&mut self, name: Option>) { - self.name = name.map(|name| name.as_ref().to_owned()); - } -} - -impl NodeToNwk for TreeIrNode { - fn nwk_name(&self) -> Option> { - self.name.as_deref() - } -} - -impl NodeFromNwk for TreeIrNode { - fn from_nwk( - name: Option>, - _confidence: Option, - _comments: &BTreeMap, - ) -> Result { - Ok(Self { - name: name.map(|name| name.as_ref().to_owned()), - ..Self::default() - }) - } -} - -impl NodeToGraphviz for TreeIrNode { - fn to_graphviz_label(&self) -> Option> { - self.name.as_deref() - } -} - -/// Discrete trait assignment for a node. -/// -/// `value` is the assigned state; `confidence` maps candidate states to their -/// marginal probabilities; `entropy` is the Shannon entropy of that distribution. -/// Mirrors augur's `node_attrs.` object (`value`, `confidence`, `entropy`). -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct TreeIrTrait { - pub value: String, - #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] - pub confidence: BTreeMap, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub entropy: Option, -} - -/// Format-neutral tree edge (branch leading to the child node). -/// -/// `mutations` and `indels` describe the events on the branch above the child. -/// `branch_length` is in substitutions per site. `gamma` is the branch-rate -/// multiplier from timetree relaxed-clock inference. -#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)] -pub struct TreeIrEdge { - pub branch_length: Option, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub mutations: Vec, - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub indels: Vec, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub gamma: Option, -} - -impl GraphEdge for TreeIrEdge {} - -impl HasBranchLength for TreeIrEdge { - fn branch_length(&self) -> Option { - self.branch_length - } - - fn set_branch_length(&mut self, branch_length: Option) { - self.branch_length = branch_length; - } -} - -impl EdgeToNwk for TreeIrEdge { - fn nwk_weight(&self) -> Option { - self.branch_length - } -} - -impl EdgeFromNwk for TreeIrEdge { - fn from_nwk(weight: Option) -> Result { - Ok(Self { - branch_length: weight, - ..Self::default() - }) - } -} - -impl EdgeToGraphviz for TreeIrEdge {} - -/// Format-neutral graph-level data. -/// -/// Carries metadata that does not belong to any single node: dataset title and -/// description, per-gene root reference sequences, the set of discrete trait -/// attributes present (to build colorings), and flags recording which node/branch -/// data classes are populated (to decide which colorings and properties to emit). -#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] -pub struct TreeIrData { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub title: Option, - #[serde(default, skip_serializing_if = "Option::is_none")] - pub description: Option, - /// Root reference sequence per gene track (e.g. `"nuc" -> "ACGT..."`). - #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] - pub root_sequence: BTreeMap, - /// Names of discrete trait attributes present on nodes. - #[serde(default, skip_serializing_if = "Vec::is_empty")] - pub trait_attrs: Vec, - /// Whether nodes carry dates (drives the `num_date` coloring and node attribute). - pub has_dates: bool, - /// Whether nodes carry the bad-branch flag (drives the `bad_branch` coloring). - pub has_bad_branch: bool, - /// Whether branches carry mutations (drives the genotype `gt` coloring). - pub has_mutations: bool, -} diff --git a/packages/treetime-io/src/tree_ir/usher.rs b/packages/treetime-io/src/tree_ir/usher.rs deleted file mode 100644 index 46ab7b9d7..000000000 --- a/packages/treetime-io/src/tree_ir/usher.rs +++ /dev/null @@ -1,169 +0,0 @@ -//! UShER MAT (mutation-annotated tree) adapter for the format-neutral TreeIR graph. -//! -//! UShER MAT carries only nucleotide substitutions (encoded as integers -//! `0:A, 1:C, 2:G, 3:T`) plus an embedded Newick string for topology, names, and -//! branch lengths. Amino-acid substitutions, indels, dates, and discrete traits -//! have no representation in the format and are dropped with a warning. Ambiguous -//! or non-ACGT nucleotides cannot be encoded and are skipped with a warning. - -use crate::tree_ir::mutation::{NUC_GENE, TreeIrSub}; -use crate::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode}; -use crate::usher_mat::{ - UsherGraphContext, UsherMetadata, UsherMutation, UsherMutationList, UsherRead, UsherTree, UsherTreeContext, - UsherTreeNode, UsherWrite, -}; -use eyre::Report; -use log::warn; -use treetime_graph::graph::Graph; -use treetime_primitives::AsciiChar; -use treetime_utils::make_error; - -/// Converter writing a TreeIR graph to a UShER MAT. -pub struct TreeIrUsherWriter { - /// Reference (root) nucleotide sequence, used to fill `ref_nuc`. Falls back to the - /// parent nucleotide when the reference is unavailable at a position. - reference: Vec, - warned_non_nuc: bool, - warned_indel: bool, - warned_ambiguous: bool, -} - -impl UsherWrite for TreeIrUsherWriter { - fn new(graph: &Graph) -> Result { - let reference = graph - .data() - .read_arc() - .root_sequence - .get(NUC_GENE) - .map(|seq| seq.bytes().filter_map(|b| AsciiChar::try_new(b).ok()).collect()) - .unwrap_or_default(); - Ok(Self { - reference, - warned_non_nuc: false, - warned_indel: false, - warned_ambiguous: false, - }) - } - - fn usher_node_from_graph_components( - &mut self, - context: &UsherGraphContext, - ) -> Result<(UsherTreeNode, UsherMutationList, UsherMetadata), Report> { - let node_name = context.node.name.clone().unwrap_or_default(); - - let mut mutation = vec![]; - if let Some(edge) = context.edge { - if !edge.indels.is_empty() && !self.warned_indel { - warn!("UShER MAT format cannot represent indels; dropping them from output"); - self.warned_indel = true; - } - for sub in &edge.mutations { - if !sub.is_nuc() { - if !self.warned_non_nuc { - warn!("UShER MAT format carries only nucleotide substitutions; dropping amino-acid mutations"); - self.warned_non_nuc = true; - } - continue; - } - let (Some(par_nuc), Some(mut_nuc)) = (nuc_to_int(sub.parent), nuc_to_int(sub.child)) else { - if !self.warned_ambiguous { - warn!("UShER MAT format cannot encode non-ACGT nucleotides; skipping ambiguous substitutions"); - self.warned_ambiguous = true; - } - continue; - }; - let ref_nuc = self - .reference - .get(sub.position.saturating_sub(1)) - .and_then(|c| nuc_to_int(*c)) - .unwrap_or(par_nuc); - mutation.push(UsherMutation { - position: sub.position as i32, - ref_nuc, - par_nuc, - mut_nuc: vec![mut_nuc], - chromosome: String::new(), - }); - } - } - - let node = UsherTreeNode { - node_name, - condensed_leaves: vec![], - }; - let metadata = UsherMetadata { - clade_annotations: vec![], - }; - Ok((node, UsherMutationList { mutation }, metadata)) - } -} - -/// Converter reading a TreeIR graph from a UShER MAT. -pub struct TreeIrUsherReader; - -impl UsherRead for TreeIrUsherReader { - fn new(_tree: &UsherTree) -> Result { - Ok(Self) - } - - fn usher_data_to_graph_data(&mut self, tree: &UsherTree) -> Result { - let has_mutations = tree.node_mutations.iter().any(|ml| !ml.mutation.is_empty()); - Ok(TreeIrData { - has_mutations, - ..TreeIrData::default() - }) - } - - fn usher_node_to_graph_components(&mut self, context: &UsherTreeContext) -> Result<(TreeIrNode, TreeIrEdge), Report> { - let node = TreeIrNode { - name: context.node.name.clone(), - ..TreeIrNode::default() - }; - - let mut mutations = vec![]; - for m in &context.node.mutations { - let Some(&mut_nuc) = m.mut_nuc.first() else { - continue; - }; - mutations.push(TreeIrSub { - gene: NUC_GENE.to_owned(), - position: m.position as usize, - parent: int_to_nuc(m.par_nuc)?, - child: int_to_nuc(mut_nuc)?, - }); - } - - let edge = TreeIrEdge { - branch_length: (context.node.branch_length != 0.0).then_some(context.node.branch_length), - mutations, - ..TreeIrEdge::default() - }; - - Ok((node, edge)) - } -} - -/// Encode a nucleotide as the UShER integer code (`0:A, 1:C, 2:G, 3:T`). -/// -/// Returns `None` for ambiguous or non-ACGT characters, which UShER cannot represent. -fn nuc_to_int(c: AsciiChar) -> Option { - match char::from(c).to_ascii_uppercase() { - 'A' => Some(0), - 'C' => Some(1), - 'G' => Some(2), - 'T' => Some(3), - _ => None, - } -} - -/// Decode a UShER integer nucleotide code (`0:A, 1:C, 2:G, 3:T`) to a character. -fn int_to_nuc(code: i32) -> Result { - let ch = match code { - 0 => 'A', - 1 => 'C', - 2 => 'G', - 3 => 'T', - other => return make_error!("Invalid UShER nucleotide code {other}: expected 0..=3"), - }; - AsciiChar::try_from_char(ch) -} diff --git a/packages/treetime-io/src/usher_mat.rs b/packages/treetime-io/src/usher_mat.rs index af7c8f774..63754d1e0 100644 --- a/packages/treetime-io/src/usher_mat.rs +++ b/packages/treetime-io/src/usher_mat.rs @@ -4,9 +4,9 @@ use eyre::{Report, WrapErr}; use smart_default::SmartDefault; use std::io::{Read, Write}; use std::path::Path; -use treetime_graph::edge::{GraphEdge, HasBranchLength}; +use treetime_graph::edge::{GraphEdge, GraphEdgeKey, HasBranchLength}; use treetime_graph::graph::Graph; -use treetime_graph::node::{GraphNode, Named}; +use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; use treetime_utils::io::file::create_file_or_stdout; use treetime_utils::io::file::open_file_or_stdin; use treetime_utils::io::json::{ @@ -89,7 +89,7 @@ pub fn usher_mat_pb_write_file(filepath: impl AsRef, graph: &G where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, C: UsherWrite, { let filepath = filepath.as_ref(); @@ -104,7 +104,7 @@ pub fn usher_mat_pb_write_bytes(graph: &Graph) -> Result<() where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, C: UsherWrite, { let tree = usher_from_graph::(graph)?; @@ -116,7 +116,7 @@ pub fn usher_mat_pb_write(writer: &mut impl Write, graph: &Graph, { let tree = usher_from_graph::(graph)?; @@ -132,7 +132,7 @@ pub fn usher_mat_json_write_file( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, C: UsherWrite, { let filepath = filepath.as_ref(); @@ -150,7 +150,7 @@ pub fn usher_mat_json_write_str( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, C: UsherWrite, { let tree = usher_from_graph::(graph)?; @@ -166,7 +166,7 @@ pub fn usher_mat_json_write( where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, C: UsherWrite, { let tree = usher_from_graph::(graph)?; @@ -265,9 +265,11 @@ pub struct UsherGraphContext<'a, N, E, D> where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { + pub node_key: GraphNodeKey, pub node: &'a N, + pub edge_key: Option, pub edge: Option<&'a E>, pub graph: &'a Graph, } @@ -276,7 +278,7 @@ pub trait UsherWrite: Sized where N: GraphNode + NodeToNwk, E: GraphEdge + EdgeToNwk, - D: Sync + Send + Default, + D: Sync + Send, { fn new(graph: &Graph) -> Result; @@ -291,7 +293,7 @@ pub fn usher_from_graph(graph: &Graph) -> Result, { let mut node_mutations = vec![]; @@ -300,12 +302,19 @@ where let mut converter = C::new(graph)?; graph.iter_depth_first_preorder_forward(|node| { + let node_key = node.key; let edge = node.parents.first().map(|(_, edge)| edge.read_arc()); + let edge_key = node.parent_keys.first().map(|(_, edge_key)| *edge_key); let edge = edge.as_deref(); let node = &node.payload; - let (node, mutations, meta) = - converter.usher_node_from_graph_components(&UsherGraphContext { node, edge, graph })?; + let (node, mutations, meta) = converter.usher_node_from_graph_components(&UsherGraphContext { + node_key, + node, + edge_key, + edge, + graph, + })?; node_mutations.push(mutations); condensed_nodes.push(node); diff --git a/packages/treetime/src/ancestral/fitch.rs b/packages/treetime/src/ancestral/fitch.rs index e72b47144..11e862597 100644 --- a/packages/treetime/src/ancestral/fitch.rs +++ b/packages/treetime/src/ancestral/fitch.rs @@ -354,7 +354,7 @@ pub fn ancestral_reconstruction_fitch( graph: &GraphAncestral, include_leaves: bool, partitions: &[Arc>], - mut visitor: impl FnMut(&GraphNodeForward, &Seq) -> Result<(), Report>, + mut visitor: impl FnMut(&GraphNodeForward, &Seq) -> Result<(), Report>, ) -> Result<(), Report> { graph .iter_depth_first_preorder_forward(|node| run_fitch_reconstruction(include_leaves, partitions, &mut visitor, &node)) @@ -363,8 +363,8 @@ pub fn ancestral_reconstruction_fitch( fn run_fitch_reconstruction( include_leaves: bool, partitions: &[Arc>], - mut visitor: impl FnMut(&GraphNodeForward, &Seq) -> Result<(), Report>, - node: &GraphNodeForward, + mut visitor: impl FnMut(&GraphNodeForward, &Seq) -> Result<(), Report>, + node: &GraphNodeForward, ) -> Result<(), Report> { if !include_leaves && node.is_leaf { return Ok(()); diff --git a/packages/treetime/src/clock/clock_graph.rs b/packages/treetime/src/clock/clock_graph.rs index 3382077d0..1d9e3b9a6 100644 --- a/packages/treetime/src/clock/clock_graph.rs +++ b/packages/treetime/src/clock/clock_graph.rs @@ -9,7 +9,7 @@ use treetime_graph::node::{GraphNode, Named, Outlier}; use treetime_io::graphviz::{EdgeToGraphviz, NodeToGraphviz}; use treetime_io::nwk::{EdgeFromNwk, EdgeToNwk, NodeFromNwk, NodeToNwk}; -pub type GraphClock = Graph; +pub type GraphClock = Graph; #[derive(Debug, Default, Clone, Serialize, Deserialize)] pub struct NodeClock { diff --git a/packages/treetime/src/commands/ancestral/augur_node_data.rs b/packages/treetime/src/commands/ancestral/augur_node_data.rs index 62d87db21..3bde862e8 100644 --- a/packages/treetime/src/commands/ancestral/augur_node_data.rs +++ b/packages/treetime/src/commands/ancestral/augur_node_data.rs @@ -23,8 +23,8 @@ use util_augur_node_data_json::{ /// the per-position `mask` string, and per-node `muts` (mask-filtered) and /// `sequence` (masked positions set to the ambiguous character, matching augur's /// `collect_sequences`). -pub fn write_augur_node_data_json( - graph: &GraphAncestral, +pub fn write_augur_node_data_json( + graph: &GraphAncestral, partition: &dyn AugurNodeDataJsonAncestralPartition, mask: &[bool], path: &Path, @@ -32,8 +32,8 @@ pub fn write_augur_node_data_json( write_augur_node_data_json_with_aa(graph, partition, mask, None, path) } -pub fn build_augur_node_data_json( - graph: &GraphAncestral, +pub fn build_augur_node_data_json( + graph: &GraphAncestral, partition: &dyn AugurNodeDataJsonAncestralPartition, mask: &[bool], aa_node_data: Option<&AaNodeData>, @@ -127,8 +127,8 @@ pub fn build_augur_node_data_json( }) } -pub fn write_augur_node_data_json_with_aa( - graph: &GraphAncestral, +pub fn write_augur_node_data_json_with_aa( + graph: &GraphAncestral, partition: &dyn AugurNodeDataJsonAncestralPartition, mask: &[bool], aa_node_data: Option<&AaNodeData>, diff --git a/packages/treetime/src/commands/ancestral/result.rs b/packages/treetime/src/commands/ancestral/result.rs index 11f765e23..b0989751c 100644 --- a/packages/treetime/src/commands/ancestral/result.rs +++ b/packages/treetime/src/commands/ancestral/result.rs @@ -1,13 +1,51 @@ +use crate::ancestral::pipeline::AncestralPartition; use crate::gtr::get_gtr::GtrModelName; use crate::gtr::gtr::GTR; use crate::payload::ancestral::GraphAncestral; use serde::Serialize; -#[derive(Debug, Serialize)] -pub struct AncestralResult { +#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. +#[derive(Serialize)] +#[serde(transparent)] +pub struct AncestralGraphData { + marker: (), #[serde(skip)] - pub graph: GraphAncestral, + pub partition: Option, #[serde(skip)] pub gtr: Option, + #[serde(skip)] pub model_name: GtrModelName, + #[serde(skip)] + pub mask: Vec, +} + +impl AncestralGraphData { + pub fn new( + partition: Option, + gtr: Option, + model_name: GtrModelName, + mask: Vec, + ) -> Self { + Self { + marker: (), + partition, + gtr, + model_name, + mask, + } + } +} + +#[derive(Serialize)] +pub struct AncestralResult { + #[serde(skip)] + pub graph: GraphAncestral, +} + +impl std::ops::Deref for AncestralResult { + type Target = AncestralGraphData; + + fn deref(&self) -> &Self::Target { + self.graph.data() + } } diff --git a/packages/treetime/src/commands/ancestral/run.rs b/packages/treetime/src/commands/ancestral/run.rs index 524d23b62..402bce7e1 100644 --- a/packages/treetime/src/commands/ancestral/run.rs +++ b/packages/treetime/src/commands/ancestral/run.rs @@ -8,9 +8,9 @@ use crate::commands::ancestral::aa_node_data::{ }; use crate::commands::ancestral::args::TreetimeAncestralArgs; use crate::commands::ancestral::augur_node_data::write_augur_node_data_json_with_aa; -use crate::commands::ancestral::result::AncestralResult; -use crate::commands::shared::ir_projection::build_ir_with_mutations; +use crate::commands::ancestral::result::{AncestralGraphData, AncestralResult}; use crate::commands::shared::output::{CommandKind, OutputSelection}; +use crate::commands::shared::tree_output::TreeOutputAdapter; use crate::gtr::get_gtr::{GtrOutput, write_gtr_json}; use crate::make_error; use crate::partition::traits::MutationCommentProvider; @@ -20,12 +20,10 @@ use crate::seq::gap_fill::apply_gap_fill; use eyre::Report; use log::{info, warn}; use treetime_graph::node::Named; -use treetime_graph::topology_order::TopologyOrderSpec; use treetime_io::fasta::{FastaReader, FastaRecord, FastaWriter, read_many_fasta}; use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_io::nwk::nwk_read_file; -use treetime_io::tree_ir::types::TreeIrData; use treetime_utils::io::file::{create_file_or_stdout, open_stdin}; pub fn run_ancestral_reconstruction( @@ -99,6 +97,7 @@ pub fn run_ancestral_reconstruction( ), ], )?; + resolved.prepare()?; let mut output_fasta = if resolved .non_tree_outputs @@ -168,39 +167,30 @@ pub fn run_ancestral_reconstruction( None }; + let pipeline::AncestralOutputFull { output, partition } = result; + let pipeline::AncestralOutput { + graph, + gtr, + model_name, + mask, + } = output; + let mut graph = graph.map_data(AncestralGraphData::new(partition, gtr, model_name, mask)); + topology_order.apply(&mut graph)?; progress.report("Writing output", 0.9, ""); if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::AugurNodeData) { - match &result.partition { + match &graph.data().partition { Some(AncestralPartition::Fitch(partition)) => { let guard = partition.read_arc(); - write_augur_node_data_json_with_aa( - &result.output.graph, - &*guard, - &result.output.mask, - aa_node_data.as_ref(), - path, - )?; + write_augur_node_data_json_with_aa(&graph, &*guard, &graph.data().mask, aa_node_data.as_ref(), path)?; }, Some(AncestralPartition::Sparse(partition)) => { let guard = partition.read_arc(); - write_augur_node_data_json_with_aa( - &result.output.graph, - &*guard, - &result.output.mask, - aa_node_data.as_ref(), - path, - )?; + write_augur_node_data_json_with_aa(&graph, &*guard, &graph.data().mask, aa_node_data.as_ref(), path)?; }, Some(AncestralPartition::Dense(partition)) => { let guard = partition.read_arc(); - write_augur_node_data_json_with_aa( - &result.output.graph, - &*guard, - &result.output.mask, - aa_node_data.as_ref(), - path, - )?; + write_augur_node_data_json_with_aa(&graph, &*guard, &graph.data().mask, aa_node_data.as_ref(), path)?; }, None => {}, } @@ -208,9 +198,9 @@ pub fn run_ancestral_reconstruction( } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { - match result.output.gtr.as_ref() { + match graph.data().gtr.as_ref() { Some(gtr) => { - let gtr_output = GtrOutput::new(gtr, result.output.model_name); + let gtr_output = GtrOutput::new(gtr, graph.data().model_name); write_gtr_json(>r_output, path)?; }, None if ancestral_args.output_gtr.is_some() => { @@ -221,50 +211,32 @@ pub fn run_ancestral_reconstruction( } if !resolved.tree_outputs.is_empty() { - write_tree_for_partition(&result, &resolved, &topology_order)?; + write_tree_for_partition(&graph, &resolved)?; } progress.report("Done", 1.0, ""); - Ok(AncestralResult { - graph: result.output.graph, - gtr: result.output.gtr, - model_name: result.output.model_name, - }) + Ok(AncestralResult { graph }) } fn write_tree_for_partition( - result: &pipeline::AncestralOutputFull, + graph: &GraphAncestral, resolved: &crate::commands::shared::output::ResolvedOutputs, - topology_order: &TopologyOrderSpec, ) -> Result<(), Report> { - let plan = topology_order.plan(&result.output.graph)?; - let ordered = plan.ordered_graph(&result.output.graph)?; - - match &result.partition { + match &graph.data().partition { Some(AncestralPartition::Sparse(partition)) => { let guard = partition.read_arc(); - let provider = MutationCommentProvider::new(&*guard, &result.output.graph); + let provider = MutationCommentProvider::new(&*guard, graph); let providers = CommentProviders::new().with(&provider); - let data = TreeIrData { - has_mutations: true, - ..TreeIrData::default() - }; - let ir = build_ir_with_mutations(&result.output.graph, &*guard, data)?; - write_tree_outputs(&ordered, &resolved.tree_outputs, &providers, Some(&ir))?; + write_tree_outputs::(graph, &resolved.tree_outputs, &providers)?; }, Some(AncestralPartition::Dense(partition)) => { let guard = partition.read_arc(); - let provider = MutationCommentProvider::new(&*guard, &result.output.graph); + let provider = MutationCommentProvider::new(&*guard, graph); let providers = CommentProviders::new().with(&provider); - let data = TreeIrData { - has_mutations: true, - ..TreeIrData::default() - }; - let ir = build_ir_with_mutations(&result.output.graph, &*guard, data)?; - write_tree_outputs(&ordered, &resolved.tree_outputs, &providers, Some(&ir))?; + write_tree_outputs::(graph, &resolved.tree_outputs, &providers)?; }, Some(AncestralPartition::Fitch(_)) | None => { - write_tree_outputs(&ordered, &resolved.tree_outputs, &CommentProviders::new(), None)?; + write_tree_outputs::(graph, &resolved.tree_outputs, &CommentProviders::new())?; }, } Ok(()) diff --git a/packages/treetime/src/commands/clock/run.rs b/packages/treetime/src/commands/clock/run.rs index 5863e22f6..eb1f81260 100644 --- a/packages/treetime/src/commands/clock/run.rs +++ b/packages/treetime/src/commands/clock/run.rs @@ -6,8 +6,8 @@ use crate::clock::find_best_root::params::{BranchPointOptimizationParams, Optimi use crate::clock::pipeline::{self, ClockInput, ClockPipelineParams}; use crate::clock::rtt::{ClockRegressionResult, write_clock_regression_result_csv}; use crate::commands::clock::args::{BranchSplitArgs, TreetimeClockArgs}; -use crate::commands::shared::ir_projection::build_ir_topology_only; use crate::commands::shared::output::{CommandKind, OutputSelection}; +use crate::commands::shared::tree_output::TreeOutputAdapter; use crate::make_error; use crate::make_report; use eyre::{Report, WrapErr}; @@ -17,16 +17,42 @@ use treetime_graph::node::{GraphNode, Named}; use treetime_io::dates_csv::read_dates; use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::nwk_read_file; -use treetime_io::tree_ir::types::{TreeIrData, TreeIrNode}; -#[derive(Debug, serde::Serialize)] -pub struct ClockResult { +#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. +#[derive(serde::Serialize)] +#[serde(transparent)] +pub struct ClockGraphData { + marker: (), #[serde(skip)] - pub graph: GraphClock, pub clock_model: ClockModel, + #[serde(skip)] pub regression_results: Vec, } +impl ClockGraphData { + pub fn new(clock_model: ClockModel, regression_results: Vec) -> Self { + Self { + marker: (), + clock_model, + regression_results, + } + } +} + +#[derive(serde::Serialize)] +pub struct ClockResult { + #[serde(skip)] + pub graph: GraphClock, +} + +impl std::ops::Deref for ClockResult { + type Target = ClockGraphData; + + fn deref(&self) -> &Self::Target { + self.graph.data() + } +} + fn branch_split_to_params(args: &BranchSplitArgs) -> BranchPointOptimizationParams { match args.method { OptimizationMethod::Grid => BranchPointOptimizationParams::grid_with(args.grid_params.clone()), @@ -100,49 +126,38 @@ pub fn run_clock( let input = ClockInput { graph, dates }; let output = pipeline::run(¶ms, input, progress)?; + let pipeline::ClockOutput { + graph, + clock_model, + regression_results, + } = output; + let mut graph = graph.map_data(ClockGraphData::new(clock_model, regression_results)); + let topology_order = clock_args + .topology_order + .resolve_topology_order(&graph, Some(input_order))?; + topology_order.apply(&mut graph)?; + resolved.prepare()?; progress.report("Writing output", 0.8, ""); if !resolved.tree_outputs.is_empty() { - let topology_order = clock_args - .topology_order - .resolve_topology_order(&output.graph, Some(input_order))?; - let plan = topology_order.plan(&output.graph)?; - let ordered = plan.ordered_graph(&output.graph)?; - let data = TreeIrData { - has_dates: true, - has_bad_branch: true, - ..TreeIrData::default() - }; - let ir = build_ir_topology_only(&output.graph, data, |_key, node| TreeIrNode { - name: node.name.clone(), - div: Some(node.div), - date: node.time, - bad_branch: node.bad_branch || node.is_outlier, - ..TreeIrNode::default() - })?; - write_tree_outputs( - &ordered, + write_tree_outputs::( + &graph, &resolved.tree_outputs, &treetime_io::nwk::CommentProviders::new(), - Some(&ir), )?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::ClockModel) { - write_clock_model(&output.clock_model, path)?; + write_clock_model(&graph.data().clock_model, path)?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::ClockCsv) { - write_clock_regression_result_csv(&output.regression_results, path, b',')?; + write_clock_regression_result_csv(&graph.data().regression_results, path, b',')?; } progress.report("Done", 1.0, ""); - Ok(ClockResult { - graph: output.graph, - clock_model: output.clock_model, - regression_results: output.regression_results, - }) + Ok(ClockResult { graph }) } fn leaf_order(graph: &Graph) -> Result, Report> diff --git a/packages/treetime/src/commands/mugration/augur_node_data.rs b/packages/treetime/src/commands/mugration/augur_node_data.rs index 436365d48..a640ec636 100644 --- a/packages/treetime/src/commands/mugration/augur_node_data.rs +++ b/packages/treetime/src/commands/mugration/augur_node_data.rs @@ -14,8 +14,8 @@ use util_augur_node_data_json::{ }; pub fn build_augur_node_data_json(result: &MugrationResult) -> Result { - let attribute = &result.traits.attribute; - let partition = &result.partition; + let attribute = &result.graph.data().traits.attribute; + let partition = &result.graph.data().partition; let graph = &result.graph; let models = build_models(attribute, partition); @@ -80,9 +80,9 @@ fn build_models( models } -fn build_nodes( +fn build_nodes( attribute: &str, - graph: &GraphAncestral, + graph: &GraphAncestral, partition: &PartitionMarginalDiscrete, ) -> BTreeMap { let confidence_key = format!("{attribute}_confidence"); @@ -121,7 +121,7 @@ fn build_nodes( } /// Build confidence map: state -> probability, sorted descending, filtered > 0.001. -fn build_confidence_map(states: &DiscreteStates, profile: &ndarray::Array1) -> BTreeMap { +pub(crate) fn build_confidence_map(states: &DiscreteStates, profile: &ndarray::Array1) -> BTreeMap { let mut pairs: Vec<(&str, f64)> = states.iter().zip(profile.iter()).map(|(s, &p)| (s, p)).collect(); pairs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); pairs @@ -132,14 +132,14 @@ fn build_confidence_map(states: &DiscreteStates, profile: &ndarray::Array1) } /// Shannon entropy: -sum(p * ln(p + 1e-12)) over real states (excludes missing). -fn compute_entropy(profile: &ndarray::Array1) -> f64 { +pub(crate) fn compute_entropy(profile: &ndarray::Array1) -> f64 { const TINY: f64 = 1e-12; -profile.iter().map(|&p| p * (p + TINY).ln()).sum::() } -fn build_branches( +fn build_branches( attribute: &str, - graph: &GraphAncestral, + graph: &GraphAncestral, partition: &PartitionMarginalDiscrete, ) -> BTreeMap { let root_key = graph.get_exactly_one_root().ok().map(|r| r.read_arc().key()); @@ -184,8 +184,8 @@ fn build_branches( branches } -fn build_parent_trait_map( - graph: &GraphAncestral, +fn build_parent_trait_map( + graph: &GraphAncestral, partition: &PartitionMarginalDiscrete, ) -> BTreeMap> { let mut map = BTreeMap::new(); diff --git a/packages/treetime/src/commands/mugration/run.rs b/packages/treetime/src/commands/mugration/run.rs index a719c706e..638a258e0 100644 --- a/packages/treetime/src/commands/mugration/run.rs +++ b/packages/treetime/src/commands/mugration/run.rs @@ -1,7 +1,7 @@ use crate::commands::mugration::args::TreetimeMugrationArgs; use crate::commands::mugration::augur_node_data::write_augur_node_data_json; -use crate::commands::shared::ir_projection::build_ir_mugration; use crate::commands::shared::output::{CommandKind, OutputSelection}; +use crate::commands::shared::tree_output::TreeOutputAdapter; use crate::gtr::get_gtr::{GtrModelName, GtrOutput, write_gtr_json}; use crate::make_report; use crate::mugration::mugration::execute_mugration; @@ -77,7 +77,7 @@ pub fn run_mugration( progress.check_cancelled()?; progress.report("Mugration inference", 0.3, ""); - let result = execute_mugration( + let mut result = execute_mugration( graph, &traits, &mugration_args.attribute, @@ -91,34 +91,36 @@ pub fn run_mugration( mugration_args.filter_uninformative_root, )?; + let topology_order = mugration_args + .topology_order + .resolve_topology_order(&result.graph, None)?; + topology_order.apply(&mut result.graph)?; + resolved.prepare()?; + progress.report("Writing output", 0.8, ""); if !resolved.tree_outputs.is_empty() { - let provider = DiscreteCommentProvider::new(&result.partition, &result.traits.attribute); + let provider = DiscreteCommentProvider::new(&result.graph.data().partition, &result.graph.data().traits.attribute); let providers = CommentProviders::new().with(&provider); - let topology_order = mugration_args - .topology_order - .resolve_topology_order(&result.graph, None)?; - let plan = topology_order.plan(&result.graph)?; - let ordered = plan.ordered_graph(&result.graph)?; - let ir = build_ir_mugration(&result.graph, &result.partition, &result.traits.attribute)?; - write_tree_outputs(&ordered, &resolved.tree_outputs, &providers, Some(&ir))?; + write_tree_outputs::(&result.graph, &resolved.tree_outputs, &providers)?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { - let gtr_output = GtrOutput::new(result.partition.gtr(), GtrModelName::Infer) - .with_discrete_states(&result.traits.attribute, result.partition.states.iter()); + let gtr_output = GtrOutput::new(result.graph.data().partition.gtr(), GtrModelName::Infer).with_discrete_states( + &result.graph.data().traits.attribute, + result.graph.data().partition.states.iter(), + ); write_gtr_json(>r_output, path)?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::TraitsCsv) { let mut f = create_file_or_stdout(path)?; - std::io::Write::write_all(&mut f, result.traits.render_csv().as_bytes())?; + std::io::Write::write_all(&mut f, result.graph.data().traits.render_csv().as_bytes())?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::ConfidenceCsv) { let mut f = create_file_or_stdout(path)?; - std::io::Write::write_all(&mut f, result.confidence.render_csv().as_bytes())?; + std::io::Write::write_all(&mut f, result.graph.data().confidence.render_csv().as_bytes())?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::AugurNodeData) { diff --git a/packages/treetime/src/commands/optimize/augur_node_data.rs b/packages/treetime/src/commands/optimize/augur_node_data.rs index 4f4feff2d..c02bf1704 100644 --- a/packages/treetime/src/commands/optimize/augur_node_data.rs +++ b/packages/treetime/src/commands/optimize/augur_node_data.rs @@ -39,8 +39,8 @@ use util_augur_node_data_json::{ /// /// When `mutation_counts` is `Some`, `branch_length` is set to the per-edge /// mutation count instead of the ML branch length (subs/site). -pub fn build_augur_node_data_json( - graph: &GraphAncestral, +pub fn build_augur_node_data_json( + graph: &GraphAncestral, alignment: Option<&Path>, input_tree: Option<&Path>, mutation_counts: Option<&BTreeMap>, @@ -106,8 +106,8 @@ pub fn build_augur_node_data_json( }) } -pub fn write_augur_node_data_json( - graph: &GraphAncestral, +pub fn write_augur_node_data_json( + graph: &GraphAncestral, alignment: Option<&Path>, input_tree: Option<&Path>, mutation_counts: Option<&BTreeMap>, diff --git a/packages/treetime/src/commands/optimize/result.rs b/packages/treetime/src/commands/optimize/result.rs index 4751e5c8e..2b656e28c 100644 --- a/packages/treetime/src/commands/optimize/result.rs +++ b/packages/treetime/src/commands/optimize/result.rs @@ -1,8 +1,46 @@ +use crate::gtr::get_gtr::GtrModelName; +use crate::gtr::gtr::GTR; +use crate::partition::marginal_dense::PartitionMarginalDense; +use crate::partition::marginal_sparse::PartitionMarginalSparse; use crate::payload::ancestral::GraphAncestral; +use parking_lot::RwLock; use serde::Serialize; +use std::sync::Arc; -#[derive(Debug, Serialize)] +#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. +#[derive(Serialize)] +#[serde(transparent)] +pub struct OptimizeGraphData { + marker: (), + #[serde(skip)] + pub gtr: GTR, + #[serde(skip)] + pub model_name: GtrModelName, + #[serde(skip)] + pub sparse_partitions: Vec>>, + #[serde(skip)] + pub dense_partitions: Vec>>, +} + +impl OptimizeGraphData { + pub fn new( + gtr: GTR, + model_name: GtrModelName, + sparse_partitions: Vec>>, + dense_partitions: Vec>>, + ) -> Self { + Self { + marker: (), + gtr, + model_name, + sparse_partitions, + dense_partitions, + } + } +} + +#[derive(Serialize)] pub struct OptimizeResult { #[serde(skip)] - pub graph: GraphAncestral, + pub graph: GraphAncestral, } diff --git a/packages/treetime/src/commands/optimize/run.rs b/packages/treetime/src/commands/optimize/run.rs index 5923f1147..6f22e0026 100644 --- a/packages/treetime/src/commands/optimize/run.rs +++ b/packages/treetime/src/commands/optimize/run.rs @@ -1,9 +1,9 @@ use crate::alphabet::alphabet::Alphabet; use crate::commands::optimize::args::TreetimeOptimizeArgs; use crate::commands::optimize::augur_node_data::write_augur_node_data_json; -use crate::commands::optimize::result::OptimizeResult; -use crate::commands::shared::ir_projection::build_ir_with_mutations; +use crate::commands::optimize::result::{OptimizeGraphData, OptimizeResult}; use crate::commands::shared::output::{CommandKind, DivergenceUnits, OutputSelection}; +use crate::commands::shared::tree_output::TreeOutputAdapter; use crate::gtr::get_gtr::{GtrOutput, write_gtr_json}; use crate::make_error; use crate::optimize::pipeline::{self, OptimizeInput, OptimizeParams}; @@ -17,7 +17,6 @@ use treetime_io::fasta::read_many_fasta; use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_io::nwk::nwk_read_file; -use treetime_io::tree_ir::types::TreeIrData; pub fn run_optimize( args: &TreetimeOptimizeArgs, @@ -68,64 +67,67 @@ pub fn run_optimize( }; let output = pipeline::run(¶ms, input, progress)?; + let pipeline::OptimizeOutput { + graph, + gtr, + model_name, + sparse_partitions, + dense_partitions, + } = output; + let mut graph = graph.map_data(OptimizeGraphData::new( + gtr, + model_name, + sparse_partitions, + dense_partitions, + )); + let topology_order = args.topology_order.resolve_topology_order(&graph, None)?; + topology_order.apply(&mut graph)?; + resolved.prepare()?; progress.report("Writing output", 0.9, ""); if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { - let gtr_output = GtrOutput::new(&output.gtr, output.model_name); + let gtr_output = GtrOutput::new(&graph.data().gtr, graph.data().model_name); write_gtr_json(>r_output, path)?; } if !resolved.tree_outputs.is_empty() { - let topology_order = args.topology_order.resolve_topology_order(&output.graph, None)?; - let plan = topology_order.plan(&output.graph)?; - let ordered = plan.ordered_graph(&output.graph)?; - - if !output.dense_partitions.is_empty() { - let guard = output.dense_partitions[0].read_arc(); - let provider = MutationCommentProvider::new(&*guard, &output.graph); + if !graph.data().dense_partitions.is_empty() { + let guard = graph.data().dense_partitions[0].read_arc(); + let provider = MutationCommentProvider::new(&*guard, &graph); let providers = CommentProviders::new().with(&provider); - let data = TreeIrData { - has_mutations: true, - ..TreeIrData::default() - }; - let ir = build_ir_with_mutations(&output.graph, &*guard, data)?; - write_tree_outputs(&ordered, &resolved.tree_outputs, &providers, Some(&ir))?; - } else if !output.sparse_partitions.is_empty() { - let guard = output.sparse_partitions[0].read_arc(); - let provider = MutationCommentProvider::new(&*guard, &output.graph); + write_tree_outputs::(&graph, &resolved.tree_outputs, &providers)?; + } else if !graph.data().sparse_partitions.is_empty() { + let guard = graph.data().sparse_partitions[0].read_arc(); + let provider = MutationCommentProvider::new(&*guard, &graph); let providers = CommentProviders::new().with(&provider); - let data = TreeIrData { - has_mutations: true, - ..TreeIrData::default() - }; - let ir = build_ir_with_mutations(&output.graph, &*guard, data)?; - write_tree_outputs(&ordered, &resolved.tree_outputs, &providers, Some(&ir))?; + write_tree_outputs::(&graph, &resolved.tree_outputs, &providers)?; } else { - write_tree_outputs(&ordered, &resolved.tree_outputs, &CommentProviders::new(), None)?; + write_tree_outputs::(&graph, &resolved.tree_outputs, &CommentProviders::new())?; } } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::AugurNodeData) { let mutation_counts = match args.divergence_units { DivergenceUnits::Mutations => { - let partition: &dyn crate::partition::traits::PartitionBranchOps = if !output.dense_partitions.is_empty() { - &*output.dense_partitions[0].read_arc() - } else if !output.sparse_partitions.is_empty() { - &*output.sparse_partitions[0].read_arc() + let partition: &dyn crate::partition::traits::PartitionBranchOps = if !graph.data().dense_partitions.is_empty() + { + &*graph.data().dense_partitions[0].read_arc() + } else if !graph.data().sparse_partitions.is_empty() { + &*graph.data().sparse_partitions[0].read_arc() } else { return make_error!( "--divergence-units=mutations requires ancestral reconstruction but no partitions are available" ); }; - Some(compute_edge_mutation_counts(&output.graph, partition)?) + Some(compute_edge_mutation_counts(&graph, partition)?) }, DivergenceUnits::MutationsPerSite => None, }; let alignment = args.alignment.alignment.first().map(PathBuf::as_path); write_augur_node_data_json( - &output.graph, + &graph, alignment, Some(args.tree.as_path()), mutation_counts.as_ref(), @@ -135,5 +137,5 @@ pub fn run_optimize( } progress.report("Done", 1.0, ""); - Ok(OptimizeResult { graph: output.graph }) + Ok(OptimizeResult { graph }) } diff --git a/packages/treetime/src/commands/prune/result.rs b/packages/treetime/src/commands/prune/result.rs index 77f799639..1b024db9a 100644 --- a/packages/treetime/src/commands/prune/result.rs +++ b/packages/treetime/src/commands/prune/result.rs @@ -1,8 +1,33 @@ +use crate::gtr::gtr::GTR; +use crate::partition::marginal_sparse::PartitionMarginalSparse; use crate::payload::ancestral::GraphAncestral; +use parking_lot::RwLock; use serde::Serialize; +use std::sync::Arc; -#[derive(Debug, Serialize)] +#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. +#[derive(Serialize)] +#[serde(transparent)] +pub struct PruneGraphData { + marker: (), + #[serde(skip)] + pub gtr: Option, + #[serde(skip)] + pub partitions: Vec>>, +} + +impl PruneGraphData { + pub fn new(gtr: Option, partitions: Vec>>) -> Self { + Self { + marker: (), + gtr, + partitions, + } + } +} + +#[derive(Serialize)] pub struct PruneResult { #[serde(skip)] - pub graph: GraphAncestral, + pub graph: GraphAncestral, } diff --git a/packages/treetime/src/commands/prune/run.rs b/packages/treetime/src/commands/prune/run.rs index 19340b76c..65441b393 100644 --- a/packages/treetime/src/commands/prune/run.rs +++ b/packages/treetime/src/commands/prune/run.rs @@ -1,8 +1,8 @@ use crate::alphabet::alphabet::Alphabet; use crate::commands::prune::args::TreetimePruneArgs; -use crate::commands::prune::result::PruneResult; -use crate::commands::shared::ir_projection::build_ir_topology_only; +use crate::commands::prune::result::{PruneGraphData, PruneResult}; use crate::commands::shared::output::{CommandKind, OutputSelection}; +use crate::commands::shared::tree_output::TreeOutputAdapter; use crate::gtr::get_gtr::{GtrModelName, GtrOutput, write_gtr_json}; use crate::make_error; use crate::prune::pipeline::{self, PruneInput, PruneParams}; @@ -20,7 +20,6 @@ use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_io::nwk::nwk_read_file; use treetime_io::parse_delimited::{parse_delimited_file, parse_delimited_str}; -use treetime_io::tree_ir::types::{TreeIrData, TreeIrNode}; use crate::payload::ancestral::GraphAncestral; @@ -79,11 +78,16 @@ pub fn run_prune( progress.check_cancelled()?; progress.report("Pruning", 0.4, ""); let output = pipeline::run(¶ms, input)?; + let pipeline::PruneOutput { graph, gtr, partitions } = output; + let mut graph = graph.map_data(PruneGraphData::new(gtr, partitions)); + let topology_order = args.topology_order.resolve_topology_order(&graph, Some(input_order))?; + topology_order.apply(&mut graph)?; + resolved.prepare()?; progress.report("Writing output", 0.8, ""); if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { - match output.gtr.as_ref() { + match graph.data().gtr.as_ref() { Some(gtr) => { let gtr_output = GtrOutput::new(gtr, GtrModelName::JC69); write_gtr_json(>r_output, path)?; @@ -98,20 +102,11 @@ pub fn run_prune( } if !resolved.tree_outputs.is_empty() { - let topology_order = args - .topology_order - .resolve_topology_order(&output.graph, Some(input_order))?; - let plan = topology_order.plan(&output.graph)?; - let ordered = plan.ordered_graph(&output.graph)?; - let ir = build_ir_topology_only(&output.graph, TreeIrData::default(), |_key, node| TreeIrNode { - name: node.name().map(|n| n.as_ref().to_owned()), - ..TreeIrNode::default() - })?; - write_tree_outputs(&ordered, &resolved.tree_outputs, &CommentProviders::new(), Some(&ir))?; + write_tree_outputs::(&graph, &resolved.tree_outputs, &CommentProviders::new())?; } progress.report("Done", 1.0, ""); - Ok(PruneResult { graph: output.graph }) + Ok(PruneResult { graph }) } fn validate_args(args: &TreetimePruneArgs) -> Result<(), Report> { diff --git a/packages/treetime/src/commands/shared/__tests__/mod.rs b/packages/treetime/src/commands/shared/__tests__/mod.rs index d8072675b..8936213b9 100644 --- a/packages/treetime/src/commands/shared/__tests__/mod.rs +++ b/packages/treetime/src/commands/shared/__tests__/mod.rs @@ -1,2 +1,3 @@ mod test_output_args; mod test_output_resolution; +mod test_tree_output; diff --git a/packages/treetime/src/commands/shared/__tests__/test_output_resolution.rs b/packages/treetime/src/commands/shared/__tests__/test_output_resolution.rs index f0274273c..7526bdfe5 100644 --- a/packages/treetime/src/commands/shared/__tests__/test_output_resolution.rs +++ b/packages/treetime/src/commands/shared/__tests__/test_output_resolution.rs @@ -1,8 +1,9 @@ #[cfg(test)] mod tests { use crate::commands::shared::output::{CommandKind, NwkStyleArg, OutputCoreArgs, OutputSelection}; - use maplit::btreemap; + use maplit::{btreemap, btreeset}; use pretty_assertions::assert_eq; + use rstest::rstest; use std::collections::BTreeMap; use std::path::PathBuf; use tempfile::TempDir; @@ -80,6 +81,26 @@ mod tests { ); } + #[test] + fn test_resolve_rejects_duplicate_output_destinations() { + let path = PathBuf::from("same-output"); + let args = OutputCoreArgs { + output_tree_nwk: Some(path.clone()), + ..OutputCoreArgs::default() + }; + + let result = args.resolve( + CommandKind::Ancestral, + &[], + &[(OutputSelection::Gtr, Some(path.as_path()))], + ); + + assert_error!( + result, + "Output destination 'same-output' is selected more than once (Nwk(NwkWriteSpec { style: Plain }) and --output-gtr)" + ); + } + #[test] fn test_resolve_non_tree_explicit_path() { let dir = TempDir::new().unwrap(); @@ -141,65 +162,62 @@ mod tests { assert!(resolved.tree_outputs.contains_key(&TreeWriteKind::Auspice)); } - // --- All-selection expansion and producibility gating --- + // --- All-selection expansion --- - #[test] - fn test_resolve_all_expands_to_producible_set() { + #[rustfmt::skip] + #[rstest] + #[case::ancestral(CommandKind::Ancestral)] + #[case::timetree (CommandKind::Timetree)] + #[case::optimize (CommandKind::Optimize)] + #[case::mugration(CommandKind::Mugration)] + #[case::clock (CommandKind::Clock)] + #[case::prune (CommandKind::Prune)] + #[trace] + fn test_resolve_all_expands_to_complete_tree_set(#[case] command: CommandKind) { let dir = TempDir::new().unwrap(); let args = OutputCoreArgs { output_all: Some(dir.path().to_path_buf()), ..Default::default() }; - let resolved = args - .resolve( - CommandKind::Ancestral, - &[OutputSelection::All], - &[ - (OutputSelection::Gtr, None), - (OutputSelection::ReconstructedAaFasta, None), - ], - ) - .unwrap(); - - // Every tree format available to ancestral is included. - assert!(resolved.tree_outputs.contains_key(&nwk(NwkStyle::Plain))); - assert!(resolved.tree_outputs.contains_key(&nexus(NwkStyle::Plain))); - assert!(resolved.tree_outputs.contains_key(&TreeWriteKind::Phyloxml)); - assert!(resolved.tree_outputs.contains_key(&TreeWriteKind::GraphJson)); - assert!(resolved.tree_outputs.contains_key(&TreeWriteKind::Dot)); - assert!(resolved.tree_outputs.contains_key(&TreeWriteKind::MatPb)); - assert!(resolved.tree_outputs.contains_key(&TreeWriteKind::Auspice)); - // The non-default AA FASTA is part of `all`. - assert!( - resolved - .non_tree_outputs - .contains_key(&OutputSelection::ReconstructedAaFasta) - ); + let resolved = args.resolve(command, &[OutputSelection::All], &[]).unwrap(); + let actual = resolved.tree_outputs.keys().cloned().collect(); + let expected = btreeset! { + nwk(NwkStyle::Plain), + nexus(NwkStyle::Plain), + TreeWriteKind::Auspice, + TreeWriteKind::Phyloxml, + TreeWriteKind::PhyloxmlJson, + TreeWriteKind::MatPb, + TreeWriteKind::MatJson, + TreeWriteKind::GraphJson, + TreeWriteKind::Dot, + }; + assert_eq!(expected, actual); } #[test] - fn test_resolve_explicit_command_unavailable_errors() { + fn test_resolve_explicit_mat_available_on_clock() { let args = OutputCoreArgs { output_tree_mat_pb: Some(PathBuf::from("/tmp/out.pb")), ..Default::default() }; - let result = args.resolve(CommandKind::Clock, &[], &[]); - assert_error!( - result, - "Output '--output-tree-mat-pb' is not available for the clock command" + let resolved = args.resolve(CommandKind::Clock, &[], &[]).unwrap(); + assert_eq!( + &PathBuf::from("/tmp/out.pb"), + &resolved.tree_outputs[&TreeWriteKind::MatPb] ); } #[test] - fn test_resolve_explicit_auspice_unavailable_on_clock_errors() { + fn test_resolve_explicit_auspice_available_on_clock() { let args = OutputCoreArgs { output_tree_auspice: Some(PathBuf::from("/tmp/out.auspice.json")), ..Default::default() }; - let result = args.resolve(CommandKind::Clock, &[], &[]); - assert_error!( - result, - "Output '--output-tree-auspice' is not available for the clock command" + let resolved = args.resolve(CommandKind::Clock, &[], &[]).unwrap(); + assert_eq!( + &PathBuf::from("/tmp/out.auspice.json"), + &resolved.tree_outputs[&TreeWriteKind::Auspice] ); } diff --git a/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs b/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs new file mode 100644 index 000000000..666a087ce --- /dev/null +++ b/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs @@ -0,0 +1,298 @@ +#[cfg(test)] +mod tests { + use crate::commands::shared::tree_output::{ + TreeOutputAdapter, TreeOutputData, TreeOutputEdge, TreeOutputMutation, TreeOutputNode, TreeOutputTrait, + format_number, + }; + use crate::payload::ancestral::GraphAncestral; + use approx::assert_ulps_eq; + use eyre::Report; + use itertools::Itertools; + use pretty_assertions::assert_eq; + use std::collections::BTreeMap; + use treetime_graph::edge::{GraphEdge, GraphEdgeKey, HasBranchLength}; + use treetime_graph::graph::Graph; + use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; + use treetime_graph::topology_order::TopologyOrderSpec; + use treetime_io::auspice::{auspice_from_graph, auspice_write_str}; + use treetime_io::auspice_types::AuspiceTree; + use treetime_io::nwk::{EdgeToNwk, NodeToNwk, nwk_read_str}; + use treetime_io::phyloxml::{phyloxml_from_graph, phyloxml_write_str}; + use treetime_io::usher_mat::usher_from_graph; + use treetime_primitives::AsciiChar; + + #[test] + fn test_tree_output_all_schema_adapters_use_final_graph_order() -> Result<(), Report> { + let mut graph: GraphAncestral = nwk_read_str("((A,B,C)large,D)root;")?; + TopologyOrderSpec::descendant_count(false).apply(&mut graph)?; + + let auspice = auspice_from_graph::(&graph)?; + let auspice_children = auspice + .tree + .children + .iter() + .map(|child| child.name.as_str()) + .collect_vec(); + + let phyloxml = phyloxml_from_graph::(&graph)?; + let phyloxml_children = phyloxml.phylogeny[0] + .clade + .as_ref() + .expect("PhyloXML must contain a root clade") + .clade + .iter() + .map(|child| child.name.as_deref().unwrap_or_default()) + .collect_vec(); + + let usher = usher_from_graph::(&graph)?; + let usher_graph: GraphAncestral = nwk_read_str(&usher.newick)?; + let usher_children = root_child_names(&usher_graph)?; + + let expected = vec!["D", "large"]; + assert_eq!(expected, auspice_children); + assert_eq!(expected, phyloxml_children); + assert_eq!(expected, usher_children); + + Ok(()) + } + + #[test] + fn test_tree_output_adapter_preserves_supported_graph_semantics() -> Result<(), Report> { + let graph = helpers::semantic_graph()?; + + let json = auspice_write_str::(&graph)?; + let auspice: AuspiceTree = serde_json::from_str(&json)?; + let child_a = auspice.tree.children.iter().find(|child| child.name == "A").unwrap(); + let child_b = auspice.tree.children.iter().find(|child| child.name == "B").unwrap(); + assert_eq!(Some(0.5), child_a.node_attrs.div); + assert_eq!(vec!["A1T".to_owned()], child_a.branch_attrs.mutations["nuc"]); + let num_date = child_a.node_attrs.num_date.as_ref().unwrap(); + assert_ulps_eq!(2020.5, num_date.value, max_ulps = 0); + assert_eq!(Some([2020.2, 2020.8]), num_date.confidence); + assert_eq!("Yes", child_b.node_attrs.bad_branch.as_ref().unwrap().value); + assert_eq!("No", auspice.tree.node_attrs.bad_branch.as_ref().unwrap().value); + let coloring_keys = auspice + .data + .meta + .colorings + .iter() + .map(|coloring| coloring.key.as_str()) + .collect_vec(); + assert!(coloring_keys.contains(&"num_date")); + assert!(coloring_keys.contains(&"bad_branch")); + assert!(coloring_keys.contains(&"country")); + assert!(coloring_keys.contains(&"gt")); + + let xml = phyloxml_write_str::(&graph)?; + assert!(xml.contains("A")); + assert!(xml.contains("0.5")); + assert!(xml.contains("treetime:divergence")); + assert!(xml.contains("treetime:mutation")); + + let usher = usher_from_graph::(&graph)?; + let mutations = usher + .node_mutations + .iter() + .flat_map(|mutations| &mutations.mutation) + .collect_vec(); + assert_eq!(1, mutations.len()); + assert_eq!(1, mutations[0].position); + assert_eq!(0, mutations[0].ref_nuc); + assert_eq!(0, mutations[0].par_nuc); + assert_eq!(vec![3], mutations[0].mut_nuc); + + Ok(()) + } + + // Oracle: augur export_v2.format_number keeps `precision` significant figures in the + // fractional part while preserving integer digits. + #[test] + fn test_tree_output_format_number_fractional_precision() { + assert_ulps_eq!(0.123457, format_number(0.12345678, 6), max_ulps = 0); + assert_ulps_eq!(123.456789, format_number(123.456789, 6), max_ulps = 0); + assert_ulps_eq!(0.0, format_number(0.0, 6), max_ulps = 0); + assert_ulps_eq!(2020.123, format_number(2020.1234567, 3), max_ulps = 0); + } + + fn root_child_names(graph: &Graph) -> Result, Report> + where + N: GraphNode + Named, + E: GraphEdge, + D: Send + Sync, + { + let root = graph.get_exactly_one_root()?; + Ok( + graph + .children_of(&root.read_arc()) + .into_iter() + .map(|(child, _)| { + child + .read_arc() + .payload() + .read_arc() + .name() + .expect("fixture child must have a name") + .as_ref() + .to_owned() + }) + .collect(), + ) + } + + mod helpers { + use super::*; + + #[derive(Debug)] + pub struct TestNode { + name: String, + div: f64, + date: Option, + bad_branch: bool, + } + + impl GraphNode for TestNode {} + + impl Named for TestNode { + fn name(&self) -> Option> { + Some(&self.name) + } + + fn set_name(&mut self, name: Option>) { + self.name = name.map_or_else(String::new, |name| name.as_ref().to_owned()); + } + } + + impl NodeToNwk for TestNode { + fn nwk_name(&self) -> Option> { + Some(&self.name) + } + } + + impl TreeOutputNode for TestNode { + fn output_divergence(&self) -> Option { + Some(self.div) + } + + fn output_date(&self) -> Option { + self.date + } + + fn output_bad_branch(&self) -> bool { + self.bad_branch + } + } + + #[derive(Debug)] + pub struct TestEdge { + branch_length: f64, + } + + impl GraphEdge for TestEdge {} + + impl HasBranchLength for TestEdge { + fn branch_length(&self) -> Option { + Some(self.branch_length) + } + + fn set_branch_length(&mut self, branch_length: Option) { + if let Some(branch_length) = branch_length { + self.branch_length = branch_length; + } + } + } + + impl EdgeToNwk for TestEdge { + fn nwk_weight(&self) -> Option { + Some(self.branch_length) + } + } + + impl TreeOutputEdge for TestEdge {} + + #[derive(Debug, Default)] + pub struct TestData { + date_confidence: BTreeMap, + mutations: BTreeMap>, + traits: BTreeMap>, + } + + impl TreeOutputData for TestData { + fn trait_attributes(&self) -> Vec { + vec!["country".to_owned()] + } + + fn has_dates(&self) -> bool { + true + } + + fn has_bad_branch(&self) -> bool { + true + } + + fn has_mutations(&self) -> bool { + true + } + + fn root_sequences(&self, _graph: &Graph) -> Result, Report> { + Ok(BTreeMap::from([("nuc".to_owned(), "A".to_owned())])) + } + + fn date_confidence(&self, key: GraphNodeKey) -> Option<[f64; 2]> { + self.date_confidence.get(&key).copied() + } + + fn traits(&self, key: GraphNodeKey) -> BTreeMap { + self.traits.get(&key).cloned().unwrap_or_default() + } + + fn mutations( + &self, + _graph: &Graph, + key: GraphEdgeKey, + ) -> Result, Report> { + Ok(self.mutations.get(&key).cloned().unwrap_or_default()) + } + } + + pub fn semantic_graph() -> Result, Report> { + let mut graph = Graph::with_data(TestData::default()); + let root = graph.add_node(node("root", 0.0, Some(2020.0), false)); + let a = graph.add_node(node("A", 0.5, Some(2020.5), false)); + let b = graph.add_node(node("B", 1.0, None, true)); + let edge_a = graph.add_edge(root, a, TestEdge { branch_length: 0.5 })?; + graph.add_edge(root, b, TestEdge { branch_length: 1.0 })?; + graph.build()?; + + graph.data_mut().date_confidence.insert(a, [2020.2, 2020.8]); + graph.data_mut().mutations.insert( + edge_a, + vec![TreeOutputMutation { + gene: "nuc".to_owned(), + position: 0, + parent: AsciiChar::try_new(b'A')?, + child: AsciiChar::try_new(b'T')?, + }], + ); + graph.data_mut().traits.insert( + a, + BTreeMap::from([( + "country".to_owned(), + TreeOutputTrait { + value: "CH".to_owned(), + ..TreeOutputTrait::default() + }, + )]), + ); + Ok(graph) + } + + fn node(name: &str, div: f64, date: Option, bad_branch: bool) -> TestNode { + TestNode { + name: name.to_owned(), + div, + date, + bad_branch, + } + } + } +} diff --git a/packages/treetime/src/commands/shared/ir_projection.rs b/packages/treetime/src/commands/shared/ir_projection.rs deleted file mode 100644 index 913e62923..000000000 --- a/packages/treetime/src/commands/shared/ir_projection.rs +++ /dev/null @@ -1,210 +0,0 @@ -use crate::partition::discrete_states::DiscreteStates; -use crate::partition::marginal_discrete::PartitionMarginalDiscrete; -use crate::partition::traits::PartitionBranchOps; -use crate::seq::mutation::Sub; -use eyre::Report; -use std::collections::BTreeMap; -use treetime_graph::edge::{GraphEdge, HasBranchLength}; -use treetime_graph::graph::Graph; -use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; -use treetime_io::graph::TreeIrGraph; -use treetime_io::tree_ir::mutation::{NUC_GENE, TreeIrSub}; -use treetime_io::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode, TreeIrTrait}; - -pub fn sub_to_ir(sub: &Sub) -> TreeIrSub { - TreeIrSub { - gene: NUC_GENE.to_owned(), - position: sub.pos() + 1, - parent: sub.reff(), - child: sub.qry(), - } -} - -pub fn subs_to_ir(subs: &[Sub]) -> Vec { - subs.iter().map(sub_to_ir).collect() -} - -/// Build a TreeIR graph from a domain graph, extracting per-edge nucleotide -/// mutations from a partition. This is the common projection for commands that -/// have ancestral reconstruction (ancestral, optimize, timetree). -pub fn build_ir_with_mutations( - graph: &Graph, - partition: &dyn PartitionBranchOps, - data: TreeIrData, -) -> Result -where - N: GraphNode + Named, - E: GraphEdge + HasBranchLength, - D: Send + Sync, -{ - let mut ir = TreeIrGraph::with_data(data); - let mut key_map: BTreeMap = BTreeMap::new(); - - graph.iter_depth_first_preorder_forward(|node| { - let dkey = node.key; - let parent = node.parent_keys.first().copied(); - - let ir_node = TreeIrNode { - name: node.payload.name().map(|n| n.as_ref().to_owned()), - ..TreeIrNode::default() - }; - let ir_key = ir.add_node(ir_node); - key_map.insert(dkey, ir_key); - - if let Some((pkey, ekey)) = parent { - let ir_parent = key_map[&pkey]; - let branch_length = node - .parents - .first() - .and_then(|(_, edge)| edge.read_arc().branch_length()); - let subs = partition.edge_subs(graph, ekey)?; - ir.add_edge( - ir_parent, - ir_key, - TreeIrEdge { - branch_length, - mutations: subs_to_ir(&subs), - ..TreeIrEdge::default() - }, - )?; - } - - Ok(()) - })?; - - ir.build()?; - Ok(ir) -} - -/// Build a TreeIR graph from a domain graph without mutations. For commands -/// that lack ancestral reconstruction (clock, prune without alignment). -pub fn build_ir_topology_only( - graph: &Graph, - data: TreeIrData, - node_mapper: impl Fn(GraphNodeKey, &N) -> TreeIrNode, -) -> Result -where - N: GraphNode + Named, - E: GraphEdge + HasBranchLength, - D: Send + Sync, -{ - let mut ir = TreeIrGraph::with_data(data); - let mut key_map: BTreeMap = BTreeMap::new(); - - graph.iter_depth_first_preorder_forward(|node| { - let dkey = node.key; - let parent = node.parent_keys.first().copied(); - - let ir_node = node_mapper(dkey, &*node.payload); - let ir_key = ir.add_node(ir_node); - key_map.insert(dkey, ir_key); - - if let Some((pkey, _ekey)) = parent { - let ir_parent = key_map[&pkey]; - let branch_length = node - .parents - .first() - .and_then(|(_, edge)| edge.read_arc().branch_length()); - ir.add_edge( - ir_parent, - ir_key, - TreeIrEdge { - branch_length, - ..TreeIrEdge::default() - }, - )?; - } - - Ok(()) - })?; - - ir.build()?; - Ok(ir) -} - -/// Build a TreeIR graph from a mugration result with discrete trait -/// assignments and confidence. -pub fn build_ir_mugration( - graph: &Graph, - partition: &PartitionMarginalDiscrete, - attribute: &str, -) -> Result -where - N: GraphNode + Named, - E: GraphEdge + HasBranchLength, - D: Send + Sync, -{ - let data = TreeIrData { - trait_attrs: vec![attribute.to_owned()], - ..TreeIrData::default() - }; - - let mut ir = TreeIrGraph::with_data(data); - let mut key_map: BTreeMap = BTreeMap::new(); - - graph.iter_depth_first_preorder_forward(|node| { - let dkey = node.key; - let parent = node.parent_keys.first().copied(); - - let mut traits = BTreeMap::new(); - if let Some(trait_value) = partition.get_reconstructed_trait(dkey) { - let confidence = partition - .get_confidence(dkey) - .map(|profile| build_confidence_map(&partition.states, &profile)) - .unwrap_or_default(); - let entropy = partition.get_confidence(dkey).map(|profile| compute_entropy(&profile)); - traits.insert( - attribute.to_owned(), - TreeIrTrait { - value: trait_value, - confidence, - entropy, - }, - ); - } - - let ir_node = TreeIrNode { - name: node.payload.name().map(|n| n.as_ref().to_owned()), - traits, - ..TreeIrNode::default() - }; - let ir_key = ir.add_node(ir_node); - key_map.insert(dkey, ir_key); - - if let Some((pkey, _ekey)) = parent { - let ir_parent = key_map[&pkey]; - let branch_length = node - .parents - .first() - .and_then(|(_, edge)| edge.read_arc().branch_length()); - ir.add_edge( - ir_parent, - ir_key, - TreeIrEdge { - branch_length, - ..TreeIrEdge::default() - }, - )?; - } - - Ok(()) - })?; - - ir.build()?; - Ok(ir) -} - -fn build_confidence_map(states: &DiscreteStates, profile: &ndarray::Array1) -> BTreeMap { - let mut pairs: Vec<(&str, f64)> = states.iter().zip(profile.iter()).map(|(s, &p)| (s, p)).collect(); - pairs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - pairs - .into_iter() - .filter(|(_, p)| *p > 0.001) - .map(|(s, p)| (s.to_owned(), p)) - .collect() -} - -fn compute_entropy(profile: &ndarray::Array1) -> f64 { - const TINY: f64 = 1e-12; - -profile.iter().map(|&p| p * (p + TINY).ln()).sum::() -} diff --git a/packages/treetime/src/commands/shared/mod.rs b/packages/treetime/src/commands/shared/mod.rs index 2d46f6d89..87d8db01a 100644 --- a/packages/treetime/src/commands/shared/mod.rs +++ b/packages/treetime/src/commands/shared/mod.rs @@ -1,11 +1,11 @@ pub mod alignment; pub mod alphabet; pub mod gap_fill; -pub mod ir_projection; pub mod metadata; pub mod model; pub mod output; pub mod reroot; +pub mod tree_output; #[cfg(test)] mod __tests__; diff --git a/packages/treetime/src/commands/shared/output.rs b/packages/treetime/src/commands/shared/output.rs index b04046638..c26d80633 100644 --- a/packages/treetime/src/commands/shared/output.rs +++ b/packages/treetime/src/commands/shared/output.rs @@ -1,7 +1,6 @@ #[cfg(feature = "clap")] use clap::ValueHint; use eyre::{Report, WrapErr}; -use log::warn; use maplit::btreeset; use serde::{Deserialize, Serialize}; use smart_default::SmartDefault; @@ -275,7 +274,7 @@ pub enum CommandKind { impl CommandKind { /// Full set selectable on this command. pub fn all_selectable(self) -> BTreeSet { - &self.available_tree_outputs() | &self.non_tree_outputs() + &Self::available_tree_outputs() | &self.non_tree_outputs() } /// Outputs produced by `--output-all` without an explicit `--output-selection`. @@ -295,30 +294,19 @@ impl CommandKind { } #[allow(clippy::enum_glob_use)] - fn available_tree_outputs(self) -> BTreeSet { + fn available_tree_outputs() -> BTreeSet { use OutputSelection::*; - let mut tree = btreeset![Nwk, Nexus, Phyloxml, PhyloxmlJson, GraphJson, Dot]; - match self { - Self::Ancestral => { - tree.insert(Auspice); - tree.insert(MatPb); - tree.insert(MatJson); - }, - Self::Optimize => { - tree.insert(MatPb); - tree.insert(MatJson); - }, - Self::Timetree => { - tree.insert(Auspice); - tree.insert(MatPb); - tree.insert(MatJson); - }, - Self::Mugration => { - tree.insert(Auspice); - }, - Self::Clock | Self::Prune => {}, - } - tree + btreeset![ + Nwk, + Nexus, + Auspice, + Phyloxml, + PhyloxmlJson, + MatPb, + MatJson, + GraphJson, + Dot + ] } #[allow(clippy::enum_glob_use)] @@ -352,31 +340,6 @@ impl CommandKind { } } -/// Whether an output has a working writer for this command, ignoring runtime data prerequisites. -fn is_producible(command: CommandKind, sel: OutputSelection) -> bool { - if sel.is_tree() { - return command.available_tree_outputs().contains(&sel); - } - true -} - -fn unproducible_reason(command: CommandKind, _sel: OutputSelection) -> String { - format!("not available for the {} command", command.stem()) -} - -/// Per-file flags are explicit requests: a non-producible target is an error at startup rather than -/// a silent skip. -fn ensure_producible_explicit(command: CommandKind, sel: OutputSelection) -> Result<(), Report> { - if is_producible(command, sel) { - return Ok(()); - } - make_error!( - "Output '{}' is not available for the {} command", - sel.flag_name(), - command.stem() - ) -} - /// Insert a secondary filename extension before the final extension of a path. /// `my.nwk` + `.annotated` -> `my.annotated.nwk`. Empty secondary leaves the path unchanged. fn insert_secondary_ext(path: &Path, secondary: &str) -> PathBuf { @@ -579,15 +542,29 @@ pub struct ResolvedOutputs { pub non_tree_outputs: BTreeMap, } +impl ResolvedOutputs { + /// Create parent directories immediately before the command starts publishing output. + pub fn prepare(&self) -> Result<(), Report> { + for path in self.tree_outputs.values().chain(self.non_tree_outputs.values()) { + if let Some(parent) = path.parent() { + if !parent.as_os_str().is_empty() { + std::fs::create_dir_all(parent) + .wrap_err_with(|| format!("Failed to create parent directory '{}'", parent.display()))?; + } + } + } + Ok(()) + } +} + impl OutputCoreArgs { /// Resolve the three-tier output configuration into concrete file paths. /// /// `selection` is the command's `--output-selection` already converted to `OutputSelection`. /// `non_tree_fields` carries the command's per-file non-tree flag values keyed by selection. /// - /// Per-file flags are honored unconditionally and take precedence over `--output-all`; explicit - /// per-file flags targeting a non-producible output error early. Outputs requested via selection - /// (or `all`) that are not producible are warned and skipped. Runtime data prerequisites (e.g. a + /// Per-file flags are honored unconditionally and take precedence over `--output-all`. Every tree + /// format is available to every command; runtime data prerequisites for non-tree outputs (e.g. a /// fitted GTR model) are checked by each command at write time, not here. pub fn resolve( &self, @@ -610,7 +587,6 @@ impl OutputCoreArgs { let Some(Some(path)) = variant.tree_field(self) else { continue; }; - ensure_producible_explicit(command, variant)?; overridden_tree.insert(variant); if variant.is_styled_tree() { for (style, p) in expand_override_styles(path, &styles) { @@ -625,7 +601,6 @@ impl OutputCoreArgs { // Tier 3b: per-file non-tree overrides. for &(sel, ref_path) in non_tree_fields { if let Some(path) = ref_path { - ensure_producible_explicit(command, sel)?; non_tree_outputs.insert(sel, path.to_path_buf()); } } @@ -644,14 +619,6 @@ impl OutputCoreArgs { if variant.is_meta() { continue; } - if !is_producible(command, variant) { - warn!( - "Skipping output '{}': {}", - variant.flag_name(), - unproducible_reason(command, variant) - ); - continue; - } if variant.is_tree() { if overridden_tree.contains(&variant) { continue; @@ -683,18 +650,7 @@ impl OutputCoreArgs { ); } - if let Some(dir) = &self.output_all { - std::fs::create_dir_all(dir) - .wrap_err_with(|| format!("Failed to create output directory '{}'", dir.display()))?; - } - for path in tree_outputs.values().chain(non_tree_outputs.values()) { - if let Some(parent) = path.parent() { - if !parent.as_os_str().is_empty() { - std::fs::create_dir_all(parent) - .wrap_err_with(|| format!("Failed to create parent directory '{}'", parent.display()))?; - } - } - } + ensure_unique_output_paths(&tree_outputs, &non_tree_outputs)?; Ok(ResolvedOutputs { tree_outputs, @@ -719,6 +675,31 @@ impl OutputCoreArgs { } } +fn ensure_unique_output_paths( + tree_outputs: &BTreeMap, + non_tree_outputs: &BTreeMap, +) -> Result<(), Report> { + let mut destinations: BTreeMap<&Path, String> = BTreeMap::new(); + for (kind, path) in tree_outputs { + if let Some(previous) = destinations.insert(path, format!("{kind:?}")) { + return make_error!( + "Output destination '{}' is selected more than once ({previous} and {kind:?})", + path.display() + ); + } + } + for (selection, path) in non_tree_outputs { + if let Some(previous) = destinations.insert(path, selection.flag_name().to_owned()) { + return make_error!( + "Output destination '{}' is selected more than once ({previous} and {})", + path.display(), + selection.flag_name() + ); + } + } + Ok(()) +} + #[derive(Debug, Clone, SmartDefault, Serialize, Deserialize)] #[serde(default)] #[cfg_attr(feature = "clap", derive(clap::Args))] diff --git a/packages/treetime/src/commands/shared/tree_output.rs b/packages/treetime/src/commands/shared/tree_output.rs new file mode 100644 index 000000000..eb3fe856d --- /dev/null +++ b/packages/treetime/src/commands/shared/tree_output.rs @@ -0,0 +1,822 @@ +use crate::ancestral::pipeline::AncestralPartition; +use crate::clock::clock_graph::{EdgeClock, NodeClock}; +use crate::commands::ancestral::result::AncestralGraphData; +use crate::commands::clock::run::ClockGraphData; +use crate::commands::mugration::augur_node_data::{build_confidence_map, compute_entropy}; +use crate::commands::optimize::result::OptimizeGraphData; +use crate::commands::prune::result::PruneGraphData; +use crate::commands::timetree::result::TimetreeGraphData; +use crate::mugration::result::MugrationGraphData; +use crate::partition::augur::AugurNodeDataJsonAncestralPartition; +use crate::partition::traits::{BranchTopology, PartitionBranchOps}; +use crate::payload::ancestral::{EdgeAncestral, NodeAncestral}; +use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; +use crate::seq::mutation::Sub; +use eyre::Report; +use log::warn; +use serde_json::{Value, json}; +use std::collections::BTreeMap; +use treetime_graph::edge::{GraphEdge, GraphEdgeKey, HasBranchLength}; +use treetime_graph::graph::Graph; +use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; +use treetime_io::auspice::{AuspiceGraphContext, AuspiceWrite}; +use treetime_io::auspice_types::{ + AuspiceColoring, AuspiceDisplayDefaults, AuspiceNumDate, AuspiceTreeBranchAttrs, AuspiceTreeData, AuspiceTreeMeta, + AuspiceTreeNode, AuspiceTreeNodeAttr, AuspiceTreeNodeAttrs, +}; +use treetime_io::phyloxml::{ + Phyloxml, PhyloxmlClade, PhyloxmlDate, PhyloxmlFromGraph, PhyloxmlGraphContext, PhyloxmlPhylogeny, PhyloxmlProperty, +}; +use treetime_io::usher_mat::{ + UsherGraphContext, UsherMetadata, UsherMutation, UsherMutationList, UsherTreeNode, UsherWrite, +}; +use treetime_primitives::AsciiChar; +use treetime_utils::make_error; + +const APPLIES_BRANCH: &str = "parent_branch"; +const APPLIES_NODE: &str = "node"; +const BRANCH_LENGTH_UNIT: &str = "subs/site"; +const COLORING_BAD_BRANCH: &str = "bad_branch"; +const COLORING_GENOTYPE: &str = "gt"; +const COLORING_NUM_DATE: &str = "num_date"; +const DT_BOOLEAN: &str = "xsd:boolean"; +const DT_DOUBLE: &str = "xsd:double"; +const DT_STRING: &str = "xsd:string"; +const NUC_GENE: &str = "nuc"; +const REF_BAD_BRANCH: &str = "treetime:bad_branch"; +const REF_DATE_INFERRED: &str = "treetime:date_inferred"; +const REF_DIV: &str = "treetime:divergence"; +const REF_GAMMA: &str = "treetime:gamma"; +const REF_MUTATION: &str = "treetime:mutation"; +const REF_TRAIT_PREFIX: &str = "treetime:trait:"; + +pub struct TreeOutputAdapter { + has_bad_branch: bool, + warned_ambiguous: bool, + warned_non_nuc: bool, +} + +pub trait TreeOutputNode: GraphNode + Named { + fn output_divergence(&self) -> Option { + None + } + + fn output_date(&self) -> Option { + None + } + + fn output_bad_branch(&self) -> bool { + false + } +} + +pub trait TreeOutputEdge: GraphEdge + HasBranchLength { + fn output_gamma(&self) -> Option { + None + } +} + +pub trait TreeOutputData: Send + Sync +where + N: TreeOutputNode, + E: TreeOutputEdge, +{ + fn title(&self) -> Option<&str> { + None + } + + fn description(&self) -> Option<&str> { + None + } + + fn trait_attributes(&self) -> Vec { + vec![] + } + + fn has_dates(&self) -> bool { + false + } + + fn has_bad_branch(&self) -> bool { + false + } + + fn has_mutations(&self) -> bool { + false + } + + fn root_sequences(&self, _graph: &Graph) -> Result, Report> + where + Self: Sized, + { + Ok(BTreeMap::new()) + } + + fn divergence(&self, _graph: &Graph, _key: GraphNodeKey, node: &N) -> Result, Report> + where + Self: Sized, + { + Ok(node.output_divergence()) + } + + fn date_confidence(&self, _key: GraphNodeKey) -> Option<[f64; 2]> { + None + } + + fn date_is_inferred(&self, _key: GraphNodeKey) -> bool { + false + } + + fn traits(&self, _key: GraphNodeKey) -> BTreeMap { + BTreeMap::new() + } + + fn mutations(&self, _graph: &Graph, _key: GraphEdgeKey) -> Result, Report> + where + Self: Sized, + { + Ok(vec![]) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct TreeOutputMutation { + pub gene: String, + pub position: usize, + pub parent: AsciiChar, + pub child: AsciiChar, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct TreeOutputTrait { + pub value: String, + pub confidence: BTreeMap, + pub entropy: Option, +} + +impl TreeOutputNode for NodeAncestral {} + +impl TreeOutputNode for NodeClock { + fn output_divergence(&self) -> Option { + Some(self.div) + } + + fn output_date(&self) -> Option { + self.time + } + + fn output_bad_branch(&self) -> bool { + self.bad_branch || self.is_outlier + } +} + +impl TreeOutputNode for NodeTimetree { + fn output_divergence(&self) -> Option { + Some(self.div) + } + + fn output_date(&self) -> Option { + self.time + } + + fn output_bad_branch(&self) -> bool { + self.bad_branch || self.is_outlier + } +} + +impl TreeOutputEdge for EdgeAncestral {} +impl TreeOutputEdge for EdgeClock {} + +impl TreeOutputEdge for EdgeTimetree { + fn output_gamma(&self) -> Option { + Some(self.gamma) + } +} + +impl TreeOutputData for () +where + N: TreeOutputNode, + E: TreeOutputEdge, +{ +} + +impl TreeOutputData for AncestralGraphData { + fn has_mutations(&self) -> bool { + self.partition.is_some() + } + + fn root_sequences( + &self, + graph: &Graph, + ) -> Result, Report> { + self.partition.as_ref().map_or_else( + || Ok(BTreeMap::new()), + |partition| { + Ok(BTreeMap::from([( + NUC_GENE.to_owned(), + ancestral_root_sequence(partition, graph)?, + )])) + }, + ) + } + + fn mutations( + &self, + graph: &Graph, + key: GraphEdgeKey, + ) -> Result, Report> { + self + .partition + .as_ref() + .map_or_else(|| Ok(vec![]), |partition| ancestral_mutations(partition, graph, key)) + } +} + +impl TreeOutputData for OptimizeGraphData { + fn has_mutations(&self) -> bool { + !self.dense_partitions.is_empty() || !self.sparse_partitions.is_empty() + } + + fn root_sequences( + &self, + graph: &Graph, + ) -> Result, Report> { + let sequence = if let Some(partition) = self.dense_partitions.first() { + Some(partition.read_arc().root_sequence(graph)?) + } else if let Some(partition) = self.sparse_partitions.first() { + Some(partition.read_arc().root_sequence(graph)?) + } else { + None + }; + Ok( + sequence + .map(|sequence| BTreeMap::from([(NUC_GENE.to_owned(), sequence.to_string())])) + .unwrap_or_default(), + ) + } + + fn mutations( + &self, + graph: &Graph, + key: GraphEdgeKey, + ) -> Result, Report> { + if let Some(partition) = self.dense_partitions.first() { + subs_to_output(PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?) + } else if let Some(partition) = self.sparse_partitions.first() { + subs_to_output(PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?) + } else { + Ok(vec![]) + } + } +} + +impl TreeOutputData for PruneGraphData { + fn has_mutations(&self) -> bool { + !self.partitions.is_empty() + } + + fn mutations( + &self, + graph: &Graph, + key: GraphEdgeKey, + ) -> Result, Report> { + self.partitions.first().map_or_else( + || Ok(vec![]), + |partition| subs_to_output(PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?), + ) + } +} + +impl TreeOutputData for ClockGraphData { + fn title(&self) -> Option<&str> { + Some("TreeTime clock analysis") + } + + fn has_dates(&self) -> bool { + true + } + + fn has_bad_branch(&self) -> bool { + true + } +} + +impl TreeOutputData for MugrationGraphData { + fn trait_attributes(&self) -> Vec { + vec![self.traits.attribute.clone()] + } + + fn traits(&self, key: GraphNodeKey) -> BTreeMap { + let Some(value) = self.partition.get_reconstructed_trait(key) else { + return BTreeMap::new(); + }; + let profile = self.partition.get_confidence(key); + let confidence = profile + .as_ref() + .map(|profile| build_confidence_map(&self.partition.states, profile)) + .unwrap_or_default(); + let entropy = profile.as_ref().map(compute_entropy); + BTreeMap::from([( + self.traits.attribute.clone(), + TreeOutputTrait { + value, + confidence, + entropy, + }, + )]) + } +} + +impl TreeOutputData for TimetreeGraphData { + fn title(&self) -> Option<&str> { + Some("TreeTime timetree analysis") + } + + fn has_dates(&self) -> bool { + true + } + + fn has_bad_branch(&self) -> bool { + true + } + + fn has_mutations(&self) -> bool { + !self.partitions.is_empty() + } + + fn divergence( + &self, + graph: &Graph, + key: GraphNodeKey, + node: &NodeTimetree, + ) -> Result, Report> { + self.mutation_counts.as_ref().map_or(Ok(Some(node.div)), |counts| { + cumulative_mutation_count(graph, key, counts).map(Some) + }) + } + + fn date_confidence(&self, key: GraphNodeKey) -> Option<[f64; 2]> { + self + .confidence_intervals + .as_ref()? + .iter() + .find_map(|interval| (interval.key == key).then_some([interval.lower, interval.upper])) + } + + fn mutations( + &self, + graph: &Graph, + key: GraphEdgeKey, + ) -> Result, Report> { + self.partitions.first().map_or_else( + || Ok(vec![]), + |partition| subs_to_output(partition.read_arc().edge_subs(graph, key)?), + ) + } +} + +impl AuspiceWrite for TreeOutputAdapter +where + N: TreeOutputNode, + E: TreeOutputEdge, + D: TreeOutputData, +{ + fn new(graph: &Graph) -> Result { + Ok(Self { + has_bad_branch: graph.data().has_bad_branch(), + warned_ambiguous: false, + warned_non_nuc: false, + }) + } + + fn auspice_data_from_graph_data(&self, graph: &Graph) -> Result { + let data = graph.data(); + let trait_attributes = data.trait_attributes(); + let mut colorings = vec![]; + if data.has_dates() { + colorings.push(coloring(COLORING_NUM_DATE, "Date", "continuous")); + } + if data.has_bad_branch() { + colorings.push(coloring(COLORING_BAD_BRANCH, "Excluded", "categorical")); + } + for attribute in &trait_attributes { + colorings.push(coloring(attribute, attribute, "categorical")); + } + if data.has_mutations() { + colorings.push(coloring(COLORING_GENOTYPE, "Genotype", "categorical")); + } + + let color_by = trait_attributes + .first() + .cloned() + .or_else(|| data.has_bad_branch().then(|| COLORING_BAD_BRANCH.to_owned())) + .or_else(|| data.has_dates().then(|| COLORING_NUM_DATE.to_owned())); + let mut filters = trait_attributes; + if data.has_bad_branch() { + filters.push(COLORING_BAD_BRANCH.to_owned()); + } + let root_sequences = data.root_sequences(graph)?; + + Ok(AuspiceTreeData { + version: Some("v2".to_owned()), + meta: AuspiceTreeMeta { + title: data + .title() + .map(str::to_owned) + .or_else(|| Some("TreeTime analysis".to_owned())), + description: data.description().map(str::to_owned), + panels: vec!["tree".to_owned()], + colorings, + filters, + display_defaults: AuspiceDisplayDefaults { + color_by, + ..AuspiceDisplayDefaults::default() + }, + ..AuspiceTreeMeta::default() + }, + root_sequence: (!root_sequences.is_empty()).then_some(root_sequences), + other: Value::default(), + }) + } + + fn auspice_node_from_graph_components( + &mut self, + context: &AuspiceGraphContext, + ) -> Result { + let name = context.node.name().map_or_else( + || format!("node_{}", context.node_key.as_usize()), + |name| name.as_ref().to_owned(), + ); + let div = context + .graph + .data() + .divergence(context.graph, context.node_key, context.node)?; + let div = finite_number(div, 6, &name, "div")?; + let date = finite_number(context.node.output_date(), 3, &name, "date")?; + let confidence = context.graph.data().date_confidence(context.node_key); + let confidence = match confidence { + Some([lower, upper]) if lower.is_finite() && upper.is_finite() => { + Some([format_number(lower, 3), format_number(upper, 3)]) + }, + Some([lower, upper]) => return make_error!("Node '{name}' has non-finite date confidence [{lower}, {upper}]"), + None => None, + }; + let num_date = date.map(|value| AuspiceNumDate { value, confidence }); + let bad_branch = self + .has_bad_branch + .then(|| AuspiceTreeNodeAttr::new(if context.node.output_bad_branch() { "Yes" } else { "No" })); + let traits = context.graph.data().traits(context.node_key); + let mutations = match context.edge_key { + Some(edge_key) => group_mutations(&context.graph.data().mutations(context.graph, edge_key)?), + None => BTreeMap::new(), + }; + + Ok(AuspiceTreeNode { + name, + branch_attrs: AuspiceTreeBranchAttrs { + mutations, + labels: None, + other: Value::default(), + }, + node_attrs: AuspiceTreeNodeAttrs { + div, + num_date, + bad_branch, + clade_membership: None, + region: None, + country: None, + division: None, + other: build_trait_attrs(&traits), + }, + children: vec![], + other: Value::default(), + }) + } +} + +impl PhyloxmlFromGraph for TreeOutputAdapter +where + N: TreeOutputNode, + E: TreeOutputEdge, + D: TreeOutputData, +{ + fn phyloxml_data_from_graph_data(graph: &Graph) -> Result { + let data = graph.data(); + Ok(Phyloxml { + phylogeny: vec![PhyloxmlPhylogeny { + rooted: true, + rerootable: None, + branch_length_unit: Some(BRANCH_LENGTH_UNIT.to_owned()), + phylogeny_type: None, + name: data.title().map(str::to_owned), + id: None, + description: data.description().map(str::to_owned), + date: None, + confidence: vec![], + clade: None, + clade_relation: vec![], + sequence_relation: vec![], + property: vec![], + other: BTreeMap::new(), + }], + other: BTreeMap::new(), + }) + } + + fn phyloxml_node_from_graph_components(context: &PhyloxmlGraphContext) -> Result { + let mut property = vec![]; + if let Some(div) = context + .graph + .data() + .divergence(context.graph, context.node_key, context.node)? + { + property.push(make_property(REF_DIV, DT_DOUBLE, APPLIES_NODE, &div.to_string())); + } + if context.node.output_bad_branch() { + property.push(make_property(REF_BAD_BRANCH, DT_BOOLEAN, APPLIES_NODE, "true")); + } + if context.graph.data().date_is_inferred(context.node_key) { + property.push(make_property(REF_DATE_INFERRED, DT_BOOLEAN, APPLIES_NODE, "true")); + } + for (attribute, value) in context.graph.data().traits(context.node_key) { + property.push(make_property( + &format!("{REF_TRAIT_PREFIX}{attribute}"), + DT_STRING, + APPLIES_NODE, + &value.value, + )); + } + if let Some(edge) = context.edge { + if let Some(gamma) = edge.output_gamma() { + property.push(make_property(REF_GAMMA, DT_DOUBLE, APPLIES_BRANCH, &gamma.to_string())); + } + } + if let Some(edge_key) = context.edge_key { + for mutation in context.graph.data().mutations(context.graph, edge_key)? { + property.push(make_property( + REF_MUTATION, + DT_STRING, + APPLIES_BRANCH, + &format!("{}:{}", mutation.gene, mutation_string(&mutation)), + )); + } + } + let date = context.node.output_date().map(|value| { + let (minimum, maximum) = context + .graph + .data() + .date_confidence(context.node_key) + .map_or((None, None), |[lower, upper]| (Some(lower), Some(upper))); + PhyloxmlDate { + desc: None, + value: Some(value), + minimum, + maximum, + unit: Some("year".to_owned()), + } + }); + + Ok(PhyloxmlClade { + name: context.node.name().map(|name| name.as_ref().to_owned()), + branch_length_elem: context.edge.and_then(HasBranchLength::branch_length), + branch_length_attr: None, + confidence: vec![], + width: None, + color: None, + node_id: None, + taxonomy: vec![], + sequence: vec![], + events: None, + binary_characters: None, + distribution: vec![], + date, + reference: vec![], + property, + clade: vec![], + other: BTreeMap::new(), + }) + } +} + +impl UsherWrite for TreeOutputAdapter +where + N: TreeOutputNode + treetime_io::nwk::NodeToNwk, + E: TreeOutputEdge + treetime_io::nwk::EdgeToNwk, + D: TreeOutputData, +{ + fn new(graph: &Graph) -> Result { + Ok(Self { + has_bad_branch: graph.data().has_bad_branch(), + warned_ambiguous: false, + warned_non_nuc: false, + }) + } + + fn usher_node_from_graph_components( + &mut self, + context: &UsherGraphContext, + ) -> Result<(UsherTreeNode, UsherMutationList, UsherMetadata), Report> { + let reference = context + .graph + .data() + .root_sequences(context.graph)? + .remove(NUC_GENE) + .unwrap_or_default(); + let mutations = match context.edge_key { + Some(edge_key) => context.graph.data().mutations(context.graph, edge_key)?, + None => vec![], + }; + let mut encoded = vec![]; + for mutation in mutations { + if mutation.gene != NUC_GENE { + if !self.warned_non_nuc { + warn!("UShER MAT format carries only nucleotide substitutions; dropping amino-acid mutations"); + self.warned_non_nuc = true; + } + continue; + } + let (Ok(par_nuc), Ok(mut_nuc)) = (nuc_to_int(mutation.parent), nuc_to_int(mutation.child)) else { + if !self.warned_ambiguous { + warn!("UShER MAT format cannot encode non-ACGT nucleotides; skipping ambiguous substitutions"); + self.warned_ambiguous = true; + } + continue; + }; + let ref_nuc = reference + .as_bytes() + .get(mutation.position) + .copied() + .and_then(|nucleotide| AsciiChar::try_new(nucleotide).ok()) + .and_then(|nucleotide| nuc_to_int(nucleotide).ok()) + .unwrap_or(par_nuc); + encoded.push(UsherMutation { + position: i32::try_from(mutation.position + 1)?, + ref_nuc, + par_nuc, + mut_nuc: vec![mut_nuc], + chromosome: String::new(), + }); + } + + Ok(( + UsherTreeNode { + node_name: context + .node + .name() + .map(|name| name.as_ref().to_owned()) + .unwrap_or_default(), + condensed_leaves: vec![], + }, + UsherMutationList { mutation: encoded }, + UsherMetadata { + clade_annotations: vec![], + }, + )) + } +} + +fn ancestral_root_sequence(partition: &AncestralPartition, graph: &dyn BranchTopology) -> Result { + let sequence = match partition { + AncestralPartition::Fitch(partition) => partition.read_arc().root_sequence(graph)?, + AncestralPartition::Sparse(partition) => partition.read_arc().root_sequence(graph)?, + AncestralPartition::Dense(partition) => partition.read_arc().root_sequence(graph)?, + }; + Ok(sequence.to_string()) +} + +fn ancestral_mutations( + partition: &AncestralPartition, + graph: &dyn BranchTopology, + key: GraphEdgeKey, +) -> Result, Report> { + let substitutions = match partition { + AncestralPartition::Fitch(partition) => partition.read_arc().edge_subs(graph, key)?, + AncestralPartition::Sparse(partition) => PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?, + AncestralPartition::Dense(partition) => PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?, + }; + subs_to_output(substitutions) +} + +fn subs_to_output(substitutions: Vec) -> Result, Report> { + Ok( + substitutions + .into_iter() + .map(|substitution| TreeOutputMutation { + gene: NUC_GENE.to_owned(), + position: substitution.pos(), + parent: substitution.reff(), + child: substitution.qry(), + }) + .collect(), + ) +} + +fn cumulative_mutation_count( + graph: &Graph, + mut key: GraphNodeKey, + counts: &BTreeMap, +) -> Result +where + N: GraphNode, + E: GraphEdge, + D: Send + Sync, +{ + let mut count = 0; + while let Some((parent, edge)) = graph.node_parent(key)? { + count += counts.get(&edge).copied().unwrap_or_default(); + key = parent; + } + Ok(count as f64) +} + +pub(crate) fn format_number(number: f64, precision: i32) -> f64 { + if number == 0.0 || !number.is_finite() { + return number; + } + let integral = number.abs().trunc(); + let significand = if integral >= 1.0 { + integral.log10().floor() as i32 + 1 + } else { + 0 + }; + let significant_figures = (significand + precision).max(1) as usize; + format!("{number:.*e}", significant_figures - 1) + .parse() + .expect("a float formatted in scientific notation must parse back") +} + +fn finite_number(value: Option, precision: i32, node_name: &str, field: &str) -> Result, Report> { + match value { + Some(value) if value.is_finite() => Ok(Some(format_number(value, precision))), + Some(value) => make_error!("Node '{node_name}' has non-finite {field}={value}"), + None => Ok(None), + } +} + +fn coloring(key: &str, title: &str, type_: &str) -> AuspiceColoring { + AuspiceColoring { + key: key.to_owned(), + title: title.to_owned(), + type_: type_.to_owned(), + ..AuspiceColoring::default() + } +} + +fn build_trait_attrs(traits: &BTreeMap) -> Value { + let map = traits + .iter() + .map(|(attribute, value)| { + let mut fields = serde_json::Map::new(); + fields.insert("value".to_owned(), json!(value.value)); + if !value.confidence.is_empty() { + let confidence = value + .confidence + .iter() + .map(|(state, probability)| (state.clone(), format_number(*probability, 3))) + .collect::>(); + fields.insert("confidence".to_owned(), json!(confidence)); + } + if let Some(entropy) = value.entropy { + fields.insert("entropy".to_owned(), json!(format_number(entropy, 3))); + } + (attribute.clone(), Value::Object(fields)) + }) + .collect(); + Value::Object(map) +} + +fn group_mutations(mutations: &[TreeOutputMutation]) -> BTreeMap> { + let mut grouped: BTreeMap> = BTreeMap::new(); + for mutation in mutations { + grouped + .entry(mutation.gene.clone()) + .or_default() + .push(mutation_string(mutation)); + } + grouped +} + +fn mutation_string(mutation: &TreeOutputMutation) -> String { + format!("{}{}{}", mutation.parent, mutation.position + 1, mutation.child) +} + +fn make_property(ref_: &str, datatype: &str, applies_to: &str, value: &str) -> PhyloxmlProperty { + PhyloxmlProperty { + value: value.to_owned(), + ref_: ref_.to_owned(), + unit: None, + datatype: datatype.to_owned(), + applies_to: applies_to.to_owned(), + id_ref: None, + } +} + +fn nuc_to_int(nucleotide: AsciiChar) -> Result { + match char::from(nucleotide).to_ascii_uppercase() { + 'A' => Ok(0), + 'C' => Ok(1), + 'G' => Ok(2), + 'T' => Ok(3), + other => make_error!("UShER MAT cannot encode nucleotide '{other}': expected A, C, G, or T"), + } +} diff --git a/packages/treetime/src/commands/timetree/output/__tests__/mod.rs b/packages/treetime/src/commands/timetree/output/__tests__/mod.rs index f93b8d946..02fc9f9a3 100644 --- a/packages/treetime/src/commands/timetree/output/__tests__/mod.rs +++ b/packages/treetime/src/commands/timetree/output/__tests__/mod.rs @@ -2,4 +2,3 @@ mod test_augur_node_data; mod test_confidence_combine; mod test_confidence_extract; mod test_confidence_rate; -mod test_ir; diff --git a/packages/treetime/src/commands/timetree/output/__tests__/test_ir.rs b/packages/treetime/src/commands/timetree/output/__tests__/test_ir.rs deleted file mode 100644 index a6437fb31..000000000 --- a/packages/treetime/src/commands/timetree/output/__tests__/test_ir.rs +++ /dev/null @@ -1,108 +0,0 @@ -#[cfg(test)] -mod tests { - use crate::commands::timetree::output::ir::build_timetree_ir; - use approx::assert_ulps_eq; - use eyre::Report; - use maplit::btreemap; - use pretty_assertions::assert_eq; - use treetime_io::auspice::auspice_write_str; - use treetime_io::auspice_types::AuspiceTree; - use treetime_io::tree_ir::auspice::TreeIrAuspiceWriter; - - // Without per-edge mutation counts, divergence is taken directly from the node's - // stored cumulative substitutions-per-site value. - #[test] - fn test_build_timetree_ir_div_from_node_value() -> Result<(), Report> { - let graph = helpers::two_node_graph()?; - let ir = build_timetree_ir(&graph, None, None, None)?; - let tree: AuspiceTree = serde_json::from_str(&auspice_write_str::(&ir)?)?; - - assert_eq!(Some(0.0), tree.tree.node_attrs.div); - let child = &tree.tree.children[0]; - assert_eq!(Some(0.5), child.node_attrs.div); - Ok(()) - } - - // With per-edge mutation counts, divergence accumulates the counts from the root. - #[test] - fn test_build_timetree_ir_div_accumulates_mutation_counts() -> Result<(), Report> { - let (graph, edge_key) = helpers::two_node_graph_with_edge_key()?; - let counts = btreemap! { edge_key => 3_usize }; - let ir = build_timetree_ir(&graph, None, Some(&counts), None)?; - let tree: AuspiceTree = serde_json::from_str(&auspice_write_str::(&ir)?)?; - - assert_eq!(Some(0.0), tree.tree.node_attrs.div); - assert_eq!(Some(3.0), tree.tree.children[0].node_attrs.div); - Ok(()) - } - - #[test] - fn test_build_timetree_ir_maps_date_and_bad_branch() -> Result<(), Report> { - let graph = helpers::two_node_graph()?; - let ir = build_timetree_ir(&graph, None, None, None)?; - let tree: AuspiceTree = serde_json::from_str(&auspice_write_str::(&ir)?)?; - - assert_ulps_eq!( - 2020.0, - tree.tree.node_attrs.num_date.as_ref().unwrap().value, - max_ulps = 0 - ); - assert_eq!("No", tree.tree.node_attrs.bad_branch.as_ref().unwrap().value); - assert_eq!( - "Yes", - tree.tree.children[0].node_attrs.bad_branch.as_ref().unwrap().value - ); - Ok(()) - } - - #[test] - fn test_build_timetree_ir_emits_timetree_colorings() -> Result<(), Report> { - let graph = helpers::two_node_graph()?; - let ir = build_timetree_ir(&graph, None, None, None)?; - let tree: AuspiceTree = serde_json::from_str(&auspice_write_str::(&ir)?)?; - - let keys: Vec<&str> = tree.data.meta.colorings.iter().map(|c| c.key.as_str()).collect(); - assert!(keys.contains(&"num_date")); - assert!(keys.contains(&"bad_branch")); - assert_eq!(Some("bad_branch".to_owned()), tree.data.meta.display_defaults.color_by); - Ok(()) - } - - mod helpers { - use crate::partition::timetree::GraphTimetree; - use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; - use eyre::Report; - use treetime_graph::edge::{GraphEdgeKey, HasBranchLength}; - use treetime_graph::node::Named; - - pub fn node(name: &str, div: f64, time: Option, bad_branch: bool) -> NodeTimetree { - let mut node = NodeTimetree { - time, - bad_branch, - div, - ..NodeTimetree::default() - }; - node.set_name(Some(name)); - node - } - - pub fn edge(branch_length: f64) -> EdgeTimetree { - let mut edge = EdgeTimetree::default(); - edge.set_branch_length(Some(branch_length)); - edge - } - - pub fn two_node_graph() -> Result { - Ok(two_node_graph_with_edge_key()?.0) - } - - pub fn two_node_graph_with_edge_key() -> Result<(GraphTimetree, GraphEdgeKey), Report> { - let mut graph = GraphTimetree::new(); - let root = graph.add_node(node("root", 0.0, Some(2020.0), false)); - let child = graph.add_node(node("A", 0.5, Some(2020.5), true)); - let edge_key = graph.add_edge(root, child, edge(0.5))?; - graph.build()?; - Ok((graph, edge_key)) - } - } -} diff --git a/packages/treetime/src/commands/timetree/output/augur_node_data.rs b/packages/treetime/src/commands/timetree/output/augur_node_data.rs index 9e2822e2a..f034379c9 100644 --- a/packages/treetime/src/commands/timetree/output/augur_node_data.rs +++ b/packages/treetime/src/commands/timetree/output/augur_node_data.rs @@ -49,8 +49,8 @@ use util_augur_node_data_json::{ /// When `mutation_counts` is `Some`, `mutation_length` is set to the per-edge /// mutation count instead of the ML branch length (subs/site). `branch_length` /// and `clock_length` remain time-valued (years) regardless. -pub fn build_augur_node_data_json( - graph: &GraphTimetree, +pub fn build_augur_node_data_json( + graph: &GraphTimetree, clock_model: &ClockModel, confidence_intervals: Option<&[NodeConfidenceInterval]>, dates: Option<&DatesMap>, @@ -144,8 +144,8 @@ pub fn build_augur_node_data_json( }) } -pub fn write_augur_node_data_json( - graph: &GraphTimetree, +pub fn write_augur_node_data_json( + graph: &GraphTimetree, clock_model: &ClockModel, confidence_intervals: Option<&[NodeConfidenceInterval]>, dates: Option<&DatesMap>, @@ -193,7 +193,7 @@ fn build_clock(clock_model: &ClockModel) -> AugurNodeDataJsonClock { /// Inferred numeric date (`numdate`) of a node by key, used for the parent endpoint /// of `clock_length = child.numdate - parent.numdate`. -fn parent_time(graph: &GraphTimetree, parent_key: GraphNodeKey) -> Result, Report> { +fn parent_time(graph: &GraphTimetree, parent_key: GraphNodeKey) -> Result, Report> { let parent = graph .get_node(parent_key) .ok_or_else(|| make_internal_report!("Timetree node data: missing parent node {parent_key:?}"))?; diff --git a/packages/treetime/src/commands/timetree/output/ir.rs b/packages/treetime/src/commands/timetree/output/ir.rs deleted file mode 100644 index c7a0ec825..000000000 --- a/packages/treetime/src/commands/timetree/output/ir.rs +++ /dev/null @@ -1,84 +0,0 @@ -use crate::commands::shared::ir_projection::subs_to_ir; -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionBranchOps; -use crate::timetree::confidence::NodeConfidenceInterval; -use eyre::Report; -use std::collections::BTreeMap; -use treetime_graph::edge::{GraphEdgeKey, HasBranchLength}; -use treetime_graph::node::{GraphNodeKey, Named}; -use treetime_io::graph::TreeIrGraph; -use treetime_io::tree_ir::types::{TreeIrData, TreeIrEdge, TreeIrNode}; - -pub fn build_timetree_ir( - graph: &GraphTimetree, - confidence_intervals: Option<&[NodeConfidenceInterval]>, - mutation_counts: Option<&BTreeMap>, - partition: Option<&dyn PartitionBranchOps>, -) -> Result { - let ci_map: BTreeMap = confidence_intervals - .map(|cis| cis.iter().map(|ci| (ci.key, [ci.lower, ci.upper])).collect()) - .unwrap_or_default(); - - let mut ir = TreeIrGraph::with_data(TreeIrData { - title: Some("TreeTime timetree analysis".to_owned()), - has_dates: true, - has_bad_branch: true, - has_mutations: partition.is_some(), - ..TreeIrData::default() - }); - - let mut key_map: BTreeMap = BTreeMap::new(); - let mut div_map: BTreeMap = BTreeMap::new(); - - graph.iter_depth_first_preorder_forward(|node| { - let dkey = node.key; - let parent = node.parent_keys.first().copied(); - - let div = match mutation_counts { - Some(counts) => match parent { - Some((pkey, ekey)) => div_map.get(&pkey).copied().unwrap_or(0.0) + *counts.get(&ekey).unwrap_or(&0) as f64, - None => 0.0, - }, - None => node.payload.div, - }; - div_map.insert(dkey, div); - - let date = node.payload.time; - let ir_node = TreeIrNode { - name: node.payload.name().map(|n| n.as_ref().to_owned()), - div: Some(div), - date, - date_confidence: date.and_then(|_| ci_map.get(&dkey).copied()), - bad_branch: node.payload.bad_branch, - ..TreeIrNode::default() - }; - let ir_key = ir.add_node(ir_node); - key_map.insert(dkey, ir_key); - - if let Some((pkey, ekey)) = parent { - let ir_parent = key_map[&pkey]; - let branch_length = node - .parents - .first() - .and_then(|(_, edge)| edge.read_arc().branch_length()); - let mutations = match partition { - Some(p) => subs_to_ir(&p.edge_subs(graph, ekey)?), - None => vec![], - }; - ir.add_edge( - ir_parent, - ir_key, - TreeIrEdge { - branch_length, - mutations, - ..TreeIrEdge::default() - }, - )?; - } - - Ok(()) - })?; - - ir.build()?; - Ok(ir) -} diff --git a/packages/treetime/src/commands/timetree/output/mod.rs b/packages/treetime/src/commands/timetree/output/mod.rs index 152f91652..ce9f5b6e0 100644 --- a/packages/treetime/src/commands/timetree/output/mod.rs +++ b/packages/treetime/src/commands/timetree/output/mod.rs @@ -3,5 +3,4 @@ mod __tests__; pub mod augur_node_data; pub mod dates; -pub mod ir; pub mod plots; diff --git a/packages/treetime/src/commands/timetree/result.rs b/packages/treetime/src/commands/timetree/result.rs index 021e77dc0..6d2b54d32 100644 --- a/packages/treetime/src/commands/timetree/result.rs +++ b/packages/treetime/src/commands/timetree/result.rs @@ -1,12 +1,67 @@ use crate::clock::clock_model::ClockModel; +use crate::gtr::get_gtr::GtrModelName; +use crate::gtr::gtr::GTR; use crate::partition::timetree::GraphTimetree; +use crate::partition::timetree::PartitionTimetreeAllVec; use crate::timetree::confidence::NodeConfidenceInterval; use serde::Serialize; +use std::collections::BTreeMap; +use treetime_io::dates_csv::DatesMap; -#[derive(Debug, Serialize)] -pub struct TimetreeResult { +#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. +#[derive(Serialize)] +#[serde(transparent)] +pub struct TimetreeGraphData { + marker: (), #[serde(skip)] - pub graph: GraphTimetree, pub clock_model: ClockModel, + #[serde(skip)] pub confidence_intervals: Option>, + #[serde(skip)] + pub partitions: PartitionTimetreeAllVec, + #[serde(skip)] + pub dates: Option, + #[serde(skip)] + pub gtr: Option, + #[serde(skip)] + pub model_name: Option, + #[serde(skip)] + pub mutation_counts: Option>, +} + +impl TimetreeGraphData { + pub fn new( + clock_model: ClockModel, + confidence_intervals: Option>, + partitions: PartitionTimetreeAllVec, + dates: Option, + gtr: Option, + model_name: Option, + mutation_counts: Option>, + ) -> Self { + Self { + marker: (), + clock_model, + confidence_intervals, + partitions, + dates, + gtr, + model_name, + mutation_counts, + } + } +} + +#[derive(Serialize)] +pub struct TimetreeResult { + #[serde(skip)] + pub graph: GraphTimetree, +} + +impl std::ops::Deref for TimetreeResult { + type Target = TimetreeGraphData; + + fn deref(&self) -> &Self::Target { + self.graph.data() + } } diff --git a/packages/treetime/src/commands/timetree/run.rs b/packages/treetime/src/commands/timetree/run.rs index af16fc9ed..f8c9ddbcd 100644 --- a/packages/treetime/src/commands/timetree/run.rs +++ b/packages/treetime/src/commands/timetree/run.rs @@ -1,10 +1,10 @@ use crate::clock::clock_output::write_clock_model; use crate::commands::shared::output::{CommandKind, DivergenceUnits, OutputSelection}; +use crate::commands::shared::tree_output::TreeOutputAdapter; use crate::commands::timetree::args::TreetimeTimetreeArgs; use crate::commands::timetree::initialization::load_input_data; use crate::commands::timetree::output::augur_node_data::write_augur_node_data_json; -use crate::commands::timetree::output::ir::build_timetree_ir; -use crate::commands::timetree::result::TimetreeResult; +use crate::commands::timetree::result::{TimetreeGraphData, TimetreeResult}; use crate::gtr::get_gtr::{GtrOutput, write_gtr_json}; use crate::make_error; use crate::partition::traits::MutationCommentProvider; @@ -47,6 +47,7 @@ pub fn run_timetree_estimation( (OutputSelection::Tracelog, args.output_tracelog.as_deref()), ], )?; + resolved.prepare()?; let tracelog: Option> = match resolved.non_tree_outputs.get(&OutputSelection::Tracelog) { @@ -95,15 +96,49 @@ pub fn run_timetree_estimation( let output = pipeline::run(¶ms, input, tracelog, progress)?; + let mutation_counts = match args.divergence_units { + DivergenceUnits::Mutations => { + if output.partitions.is_empty() { + return make_error!( + "--divergence-units=mutations requires ancestral reconstruction; \ + incompatible with --branch-length-mode=input" + ); + } + let guard = output.partitions[0].read_arc(); + Some(compute_edge_mutation_counts(&output.graph, &*guard)?) + }, + DivergenceUnits::MutationsPerSite => None, + }; + + let pipeline::TimetreeOutput { + graph, + clock_model, + confidence_intervals, + partitions, + dates, + gtr, + model_name, + } = output; + let mut graph = graph.map_data(TimetreeGraphData::new( + clock_model, + confidence_intervals, + partitions, + dates, + gtr, + model_name, + mutation_counts, + )); + progress.report("Writing output", 0.95, ""); info!("### TreeTime: writing outputs"); let topology_order = args .topology_order - .resolve_topology_order(&output.graph, Some(input_leaf_order))?; + .resolve_topology_order(&graph, Some(input_leaf_order))?; + topology_order.apply(&mut graph)?; if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::ConfidenceTsv) { - match output.confidence_intervals.as_ref() { + match graph.data().confidence_intervals.as_ref() { Some(intervals) => { write_confidence_intervals_file(intervals, path).wrap_err("Failed to write confidence intervals")?; info!("Wrote confidence intervals to {path}", path = path.display()); @@ -119,11 +154,11 @@ pub fn run_timetree_estimation( } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::ClockModel) { - write_clock_model(&output.clock_model, path)?; + write_clock_model(&graph.data().clock_model, path)?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { - match (output.gtr.as_ref(), output.model_name) { + match (graph.data().gtr.as_ref(), graph.data().model_name) { (Some(gtr), Some(model_name)) => { let gtr_output = GtrOutput::new(gtr, model_name); write_gtr_json(>r_output, path)?; @@ -135,58 +170,27 @@ pub fn run_timetree_estimation( } } - let mutation_counts = match args.divergence_units { - DivergenceUnits::Mutations => { - if output.partitions.is_empty() { - return make_error!( - "--divergence-units=mutations requires ancestral reconstruction; \ - incompatible with --branch-length-mode=input" - ); - } - let guard = output.partitions[0].read_arc(); - Some(compute_edge_mutation_counts(&output.graph, &*guard)?) - }, - DivergenceUnits::MutationsPerSite => None, - }; - if !resolved.tree_outputs.is_empty() { - let plan = topology_order.plan(&output.graph)?; - let ordered = plan.ordered_graph(&output.graph)?; - // Build the IR from the ordered clone. The clone preserves original node and edge - // keys, keeping confidence intervals and mutation counts aligned while applying - // the same topology order to Auspice and the other TreeIR-backed formats. - if !output.partitions.is_empty() { - let guard = output.partitions[0].read_arc(); - let ir = build_timetree_ir( - &ordered, - output.confidence_intervals.as_deref(), - mutation_counts.as_ref(), - Some(&*guard), - )?; - let provider = MutationCommentProvider::new(&*guard, &output.graph); + if !graph.data().partitions.is_empty() { + let guard = graph.data().partitions[0].read_arc(); + let provider = MutationCommentProvider::new(&*guard, &graph); let providers = CommentProviders::new().with(&provider); - write_tree_outputs(&ordered, &resolved.tree_outputs, &providers, Some(&ir))?; + write_tree_outputs::(&graph, &resolved.tree_outputs, &providers)?; } else { - let ir = build_timetree_ir( - &ordered, - output.confidence_intervals.as_deref(), - mutation_counts.as_ref(), - None, - )?; - write_tree_outputs(&ordered, &resolved.tree_outputs, &CommentProviders::new(), Some(&ir))?; + write_tree_outputs::(&graph, &resolved.tree_outputs, &CommentProviders::new())?; } } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::AugurNodeData) { let alignment = args.alignment.alignment.first().map(PathBuf::as_path); write_augur_node_data_json( - &output.graph, - &output.clock_model, - output.confidence_intervals.as_deref(), - output.dates.as_ref(), + &graph, + &graph.data().clock_model, + graph.data().confidence_intervals.as_deref(), + graph.data().dates.as_ref(), alignment, args.tree.as_deref(), - mutation_counts.as_ref(), + graph.data().mutation_counts.as_ref(), path, )?; info!("Wrote augur node data JSON to {path}", path = path.display()); @@ -201,9 +205,5 @@ pub fn run_timetree_estimation( } progress.report("Done", 1.0, ""); - Ok(TimetreeResult { - graph: output.graph, - clock_model: output.clock_model, - confidence_intervals: output.confidence_intervals, - }) + Ok(TimetreeResult { graph }) } diff --git a/packages/treetime/src/mugration/result.rs b/packages/treetime/src/mugration/result.rs index a39439fbe..84cc94c3e 100644 --- a/packages/treetime/src/mugration/result.rs +++ b/packages/treetime/src/mugration/result.rs @@ -104,36 +104,55 @@ impl MugrationTraitsOutput { } } +#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. #[derive(Debug, serde::Serialize)] -pub struct MugrationResult { +#[serde(transparent)] +pub struct MugrationGraphData { + marker: (), #[serde(skip)] pub traits: MugrationTraitsOutput, #[serde(skip)] pub confidence: MugrationConfidenceOutput, - pub log_lh: f64, #[serde(skip)] - pub graph: GraphAncestral, + pub log_lh: f64, #[serde(skip)] pub partition: PartitionMarginalDiscrete, } +#[derive(Debug, serde::Serialize)] +pub struct MugrationResult { + #[serde(skip)] + pub graph: GraphAncestral, +} + +impl std::ops::Deref for MugrationResult { + type Target = MugrationGraphData; + + fn deref(&self) -> &Self::Target { + self.graph.data() + } +} + impl MugrationResult { pub fn new(graph: GraphAncestral, partition: PartitionMarginalDiscrete, attribute: &str, log_lh: f64) -> Self { let assignments = extract_trait_assignments(&graph, &partition); let traits = MugrationTraitsOutput::new(attribute, assignments); let confidence = MugrationConfidenceOutput::new(&graph, &partition); - Self { + let data = MugrationGraphData { + marker: (), traits, confidence, log_lh, - graph, partition, + }; + Self { + graph: graph.map_data(data), } } pub fn trait_assignments(&self) -> &IndexMap { - &self.traits.assignments + &self.graph.data().traits.assignments } } diff --git a/packages/treetime/src/partition/marginal_core.rs b/packages/treetime/src/partition/marginal_core.rs index dbbb09478..2475d7caa 100644 --- a/packages/treetime/src/partition/marginal_core.rs +++ b/packages/treetime/src/partition/marginal_core.rs @@ -48,9 +48,9 @@ where fn leaf_profile(&self, node_key: GraphNodeKey) -> Result; - fn backward_internal_pre(&mut self, _node: &GraphNodeBackward) {} + fn backward_internal_pre(&mut self, _node: &GraphNodeBackward) {} - fn forward_post(&mut self, _graph: &Graph, _node: &GraphNodeForward) -> Result<(), Report> { + fn forward_post(&mut self, _graph: &Graph, _node: &GraphNodeForward) -> Result<(), Report> { Ok(()) } } @@ -284,7 +284,7 @@ where pub fn marginal_process_node_backward( partition: &mut impl MarginalPartition, - node: &GraphNodeBackward, + node: &GraphNodeBackward, ) -> Result<(), Report> where N: GraphNode + Named, @@ -357,7 +357,7 @@ where pub fn marginal_process_node_forward( partition: &mut impl MarginalPartition, graph: &Graph, - node: &GraphNodeForward, + node: &GraphNodeForward, ) -> Result<(), Report> where N: GraphNode + Named, diff --git a/packages/treetime/src/partition/marginal_dense.rs b/packages/treetime/src/partition/marginal_dense.rs index 10f7c1b35..6a946fa47 100644 --- a/packages/treetime/src/partition/marginal_dense.rs +++ b/packages/treetime/src/partition/marginal_dense.rs @@ -237,7 +237,7 @@ where }) } - fn backward_internal_pre(&mut self, node: &GraphNodeBackward) { + fn backward_internal_pre(&mut self, node: &GraphNodeBackward) { let child_non_chars: Vec<&Vec<(usize, usize)>> = node .child_keys .iter() @@ -283,7 +283,7 @@ where ); } - fn forward_post(&mut self, graph: &Graph, node: &GraphNodeForward) -> Result<(), Report> { + fn forward_post(&mut self, graph: &Graph, node: &GraphNodeForward) -> Result<(), Report> { if node.is_root { let node_data = self.data.nodes.get_mut(&node.key).unwrap(); node_data.seq.variable_indel.clear(); @@ -479,7 +479,7 @@ where fn reconstruct_node_sequence( &mut self, - node: &GraphNodeForward, + node: &GraphNodeForward, include_leaves: bool, sample_mode: SampleMode, rng: &mut dyn rand::RngCore, diff --git a/packages/treetime/src/partition/marginal_sparse.rs b/packages/treetime/src/partition/marginal_sparse.rs index 2fb06cf66..1a232868d 100644 --- a/packages/treetime/src/partition/marginal_sparse.rs +++ b/packages/treetime/src/partition/marginal_sparse.rs @@ -515,7 +515,7 @@ where fn reconstruct_node_sequence( &mut self, - node: &GraphNodeForward, + node: &GraphNodeForward, include_leaves: bool, sample_mode: SampleMode, rng: &mut dyn rand::RngCore, diff --git a/packages/treetime/src/partition/timetree.rs b/packages/treetime/src/partition/timetree.rs index f485a8726..8836a7b3a 100644 --- a/packages/treetime/src/partition/timetree.rs +++ b/packages/treetime/src/partition/timetree.rs @@ -5,5 +5,5 @@ use parking_lot::RwLock; use std::sync::Arc; use treetime_graph::graph::Graph; -pub type GraphTimetree = Graph; +pub type GraphTimetree = Graph; pub type PartitionTimetreeAllVec = Vec>>>; diff --git a/packages/treetime/src/partition/traits.rs b/packages/treetime/src/partition/traits.rs index 22c7b2e1c..e48b721b9 100644 --- a/packages/treetime/src/partition/traits.rs +++ b/packages/treetime/src/partition/traits.rs @@ -176,7 +176,7 @@ where fn reconstruct_node_sequence( &mut self, - node: &GraphNodeForward, + node: &GraphNodeForward, include_leaves: bool, sample_mode: SampleMode, rng: &mut dyn rand::RngCore, @@ -292,7 +292,7 @@ where /// This allows trait objects to be used for both ancestral reconstruction and timetree inference. /// Includes `PartitionRerootOps` for reroot support with default no-op. pub trait PartitionTimetreeAll: - PartitionMarginalOps + PartitionTimetreeOps + PartitionRerootOps + HasLogLh + PartitionBranchOps + PartitionMarginalOps + PartitionTimetreeOps + PartitionRerootOps + HasLogLh where N: GraphNode + Named, E: EdgeOptimizeOps, @@ -302,7 +302,7 @@ where /// Blanket implementation: any type implementing all required traits automatically implements the combined trait impl PartitionTimetreeAll for T where - T: PartitionMarginalOps + PartitionTimetreeOps + PartitionRerootOps + HasLogLh, + T: PartitionBranchOps + PartitionMarginalOps + PartitionTimetreeOps + PartitionRerootOps + HasLogLh, N: GraphNode + Named, E: EdgeOptimizeOps, { diff --git a/packages/treetime/src/payload/ancestral.rs b/packages/treetime/src/payload/ancestral.rs index 33260ada8..8ca3695ad 100644 --- a/packages/treetime/src/payload/ancestral.rs +++ b/packages/treetime/src/payload/ancestral.rs @@ -7,7 +7,7 @@ use treetime_graph::node::{Described, GraphNode, Named}; use treetime_io::graphviz::{EdgeToGraphviz, NodeToGraphviz}; use treetime_io::nwk::{EdgeFromNwk, EdgeToNwk, NodeFromNwk, NodeToNwk}; -pub type GraphAncestral = Graph; +pub type GraphAncestral = Graph; #[derive(Clone, Default, Debug, Serialize, Deserialize)] pub struct NodeAncestral { diff --git a/packages/treetime/src/prune/pipeline.rs b/packages/treetime/src/prune/pipeline.rs index 05961cb40..e4153736f 100644 --- a/packages/treetime/src/prune/pipeline.rs +++ b/packages/treetime/src/prune/pipeline.rs @@ -33,6 +33,8 @@ pub struct PruneOutput { pub graph: GraphAncestral, #[serde(skip)] pub gtr: Option, + #[serde(skip)] + pub partitions: Vec>>, } pub fn run(params: &PruneParams, mut input: PruneInput) -> Result { @@ -84,5 +86,6 @@ pub fn run(params: &PruneParams, mut input: PruneInput) -> Result Date: Wed, 22 Jul 2026 14:32:43 +0200 Subject: [PATCH 2/4] fix(io): warn on fabricated UShER reference allele --- .../shared/__tests__/test_tree_output.rs | 39 ++++++++++++++++++- .../src/commands/shared/tree_output.rs | 19 ++++++++- 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs b/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs index 666a087ce..03095834c 100644 --- a/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs +++ b/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs @@ -104,6 +104,25 @@ mod tests { Ok(()) } + #[test] + fn test_tree_output_usher_missing_reference_uses_documented_fallback() -> Result<(), Report> { + let graph = helpers::semantic_graph_missing_root_reference()?; + + let usher = usher_from_graph::(&graph)?; + let mutations = usher + .node_mutations + .iter() + .flat_map(|mutations| &mutations.mutation) + .collect_vec(); + assert_eq!(1, mutations.len()); + let mutation = mutations[0]; + + assert_eq!(1, mutation.ref_nuc); + assert_eq!(1, mutation.par_nuc); + + Ok(()) + } + // Oracle: augur export_v2.format_number keeps `precision` significant figures in the // fractional part while preserving integer digits. #[test] @@ -213,6 +232,7 @@ mod tests { pub struct TestData { date_confidence: BTreeMap, mutations: BTreeMap>, + root_sequences: BTreeMap, traits: BTreeMap>, } @@ -234,7 +254,7 @@ mod tests { } fn root_sequences(&self, _graph: &Graph) -> Result, Report> { - Ok(BTreeMap::from([("nuc".to_owned(), "A".to_owned())])) + Ok(self.root_sequences.clone()) } fn date_confidence(&self, key: GraphNodeKey) -> Option<[f64; 2]> { @@ -254,6 +274,22 @@ mod tests { } } + pub fn semantic_graph_missing_root_reference() -> Result, Report> { + let mut graph = semantic_graph()?; + graph.data_mut().root_sequences.clear(); + let mutations = graph + .data_mut() + .mutations + .values_mut() + .next() + .expect("semantic graph must contain mutation data for one edge"); + mutations + .first_mut() + .expect("semantic graph must contain one mutation") + .parent = AsciiChar::from_byte_unchecked(b'C'); + Ok(graph) + } + pub fn semantic_graph() -> Result, Report> { let mut graph = Graph::with_data(TestData::default()); let root = graph.add_node(node("root", 0.0, Some(2020.0), false)); @@ -264,6 +300,7 @@ mod tests { graph.build()?; graph.data_mut().date_confidence.insert(a, [2020.2, 2020.8]); + graph.data_mut().root_sequences.insert("nuc".to_owned(), "A".to_owned()); graph.data_mut().mutations.insert( edge_a, vec![TreeOutputMutation { diff --git a/packages/treetime/src/commands/shared/tree_output.rs b/packages/treetime/src/commands/shared/tree_output.rs index eb3fe856d..b6dbf35cf 100644 --- a/packages/treetime/src/commands/shared/tree_output.rs +++ b/packages/treetime/src/commands/shared/tree_output.rs @@ -54,6 +54,7 @@ pub struct TreeOutputAdapter { has_bad_branch: bool, warned_ambiguous: bool, warned_non_nuc: bool, + warned_reference_fallback: bool, } pub trait TreeOutputNode: GraphNode + Named { @@ -386,6 +387,7 @@ where has_bad_branch: graph.data().has_bad_branch(), warned_ambiguous: false, warned_non_nuc: false, + warned_reference_fallback: false, }) } @@ -609,6 +611,7 @@ where has_bad_branch: graph.data().has_bad_branch(), warned_ambiguous: false, warned_non_nuc: false, + warned_reference_fallback: false, }) } @@ -647,8 +650,20 @@ where .get(mutation.position) .copied() .and_then(|nucleotide| AsciiChar::try_new(nucleotide).ok()) - .and_then(|nucleotide| nuc_to_int(nucleotide).ok()) - .unwrap_or(par_nuc); + .and_then(|nucleotide| nuc_to_int(nucleotide).ok()); + let ref_nuc = ref_nuc.unwrap_or_else(|| { + // This fallback can fabricate a global reference allele from a branch-local parent allele. + // Keep the current output behavior visible until the graph supplies a validated root sequence. + // See kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md. + if !self.warned_reference_fallback { + warn!( + "UShER MAT global reference nucleotide is unavailable at zero-based position {}; using the branch-local parent allele for ref_nuc, which can produce incorrect MAT reference alleles", + mutation.position + ); + self.warned_reference_fallback = true; + } + par_nuc + }); encoded.push(UsherMutation { position: i32::try_from(mutation.position + 1)?, ref_nuc, From 665b33cccee8ba10da09f94ccd4d798f6385a737 Mon Sep 17 00:00:00 2001 From: ivan-aksamentov Date: Wed, 22 Jul 2026 15:07:56 +0200 Subject: [PATCH 3/4] refactor(io): write tree formats from concrete graphs --- Cargo.lock | 488 +++- Cargo.toml | 1 + packages/treetime-io/src/auspice.rs | 168 +- packages/treetime-io/src/auspice_types.rs | 3 +- packages/treetime-io/src/graph.rs | 68 - packages/treetime-io/src/phyloxml.rs | 181 +- packages/treetime-io/src/usher_mat.rs | 160 +- .../treetime-utils/src/array/__tests__/mod.rs | 1 + .../src/array/__tests__/serde.rs | 26 + packages/treetime-utils/src/array/serde.rs | 10 +- packages/treetime/Cargo.toml | 1 + .../__tests__/test_dense_completeness.rs | 3 +- .../ancestral/__tests__/test_fitch_indel.rs | 4 +- packages/treetime/src/ancestral/fitch.rs | 2 +- packages/treetime/src/ancestral/pipeline.rs | 1 + .../__tests__/test_augur_node_data.rs | 1 + .../src/commands/ancestral/aa_node_data.rs | 92 +- .../treetime/src/commands/ancestral/result.rs | 12 +- .../treetime/src/commands/ancestral/run.rs | 37 +- packages/treetime/src/commands/clock/run.rs | 13 +- .../treetime/src/commands/mugration/run.rs | 7 +- .../treetime/src/commands/optimize/result.rs | 8 - .../treetime/src/commands/optimize/run.rs | 11 +- .../treetime/src/commands/prune/result.rs | 11 +- packages/treetime/src/commands/prune/run.rs | 7 +- .../__tests__/schemas/auspice/LICENSE.txt | 619 +++++ .../__tests__/schemas/auspice/README.md | 7 + .../schemas/auspice/schema-annotations.json | 95 + .../auspice/schema-auspice-config-v2.json | 329 +++ .../auspice/schema-export-root-sequence.json | 33 + .../schemas/auspice/schema-export-v2.json | 364 +++ .../shared/__tests__/test_tree_output.rs | 962 +++++-- .../treetime/src/commands/shared/output.rs | 15 - .../src/commands/shared/tree_output.rs | 2204 ++++++++++++----- .../src/commands/timetree/initialization.rs | 14 +- .../src/commands/timetree/output/dates.rs | 9 +- .../treetime/src/commands/timetree/result.rs | 11 - .../treetime/src/commands/timetree/run.rs | 9 +- packages/treetime/src/gtr/gtr.rs | 10 +- packages/treetime/src/mugration/result.rs | 18 +- .../test_initial_guess_indel_zero_bl.rs | 4 +- .../__tests__/test_initial_guess_mode.rs | 2 +- .../src/optimize/__tests__/test_no_indels.rs | 8 +- .../optimize/__tests__/test_optimize_indel.rs | 22 +- .../__tests__/test_optimize_method.rs | 2 +- .../__tests__/test_run_optimize_loop.rs | 4 +- .../topology/__tests__/test_collapse_edge.rs | 4 +- .../__tests__/test_merge_shared_mutations.rs | 6 +- .../test_prop_merge_shared_mutations.rs | 2 +- packages/treetime/src/partition/augur.rs | 15 + packages/treetime/src/partition/fitch.rs | 46 +- .../treetime/src/partition/marginal_core.rs | 3 +- .../treetime/src/partition/marginal_dense.rs | 15 +- .../src/partition/marginal_discrete.rs | 3 +- .../treetime/src/partition/marginal_sparse.rs | 19 +- packages/treetime/src/partition/timetree.rs | 213 +- packages/treetime/src/partition/traits.rs | 51 +- packages/treetime/src/payload/ancestral.rs | 17 +- .../src/prune/__tests__/test_prune.rs | 8 +- .../src/seq/__tests__/test_composition.rs | 4 +- .../treetime/src/seq/__tests__/test_indel.rs | 20 +- .../src/seq/__tests__/test_mutation.rs | 27 +- packages/treetime/src/seq/composition.rs | 2 +- packages/treetime/src/seq/indel.rs | 87 +- packages/treetime/src/seq/mutation.rs | 104 + packages/treetime/src/timetree/confidence.rs | 10 +- .../src/timetree/convergence/likelihood.rs | 12 +- .../src/timetree/convergence/optimizer.rs | 9 +- .../timetree/convergence/sequence_changes.rs | 13 +- .../test_gm_runner_marginal_dense.rs | 10 +- .../test_gm_runner_marginal_sparse.rs | 10 +- .../test_gm_runner_pre_optimize.rs | 18 +- .../test_gm_runner/test_runner_coalescent.rs | 19 +- .../optimization/__tests__/test_reroot.rs | 26 +- .../src/timetree/optimization/polytomy.rs | 11 +- .../src/timetree/optimization/reroot.rs | 10 +- packages/treetime/src/timetree/pipeline.rs | 13 +- packages/treetime/src/timetree/refinement.rs | 10 +- 78 files changed, 5092 insertions(+), 1782 deletions(-) create mode 100644 packages/treetime-utils/src/array/__tests__/serde.rs create mode 100644 packages/treetime/src/commands/shared/__tests__/schemas/auspice/LICENSE.txt create mode 100644 packages/treetime/src/commands/shared/__tests__/schemas/auspice/README.md create mode 100644 packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-annotations.json create mode 100644 packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-auspice-config-v2.json create mode 100644 packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-root-sequence.json create mode 100644 packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-v2.json diff --git a/Cargo.lock b/Cargo.lock index 489db9ce0..81bd40ccd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,6 +17,20 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.3" @@ -35,6 +49,12 @@ dependencies = [ "equator", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "android_system_properties" version = "0.1.5" @@ -542,6 +562,12 @@ dependencies = [ "generic-array", ] +[[package]] +name = "borrow-or-share" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc0b364ead1874514c8c2855ab558056ebfeb775653e7ae45ff72f28f8f3166c" + [[package]] name = "built" version = "0.7.7" @@ -1137,6 +1163,12 @@ version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef8ae57c4978a2acd8b869ce6b9ca1dfe817bff704c220209fdef2c0b75a01b9" +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "deranged" version = "0.5.3" @@ -1237,6 +1269,17 @@ dependencies = [ "windows-sys 0.61.0", ] +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "dlib" version = "0.5.2" @@ -1310,6 +1353,15 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +[[package]] +name = "email_address" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e079f19b08ca6239f47f8ba8509c11cf3ea30095831f7fed61441475edd8c449" +dependencies = [ + "serde", +] + [[package]] name = "encode_unicode" version = "1.0.0" @@ -1420,6 +1472,17 @@ dependencies = [ "once_cell", ] +[[package]] +name = "fancy-regex" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1e1dacd0d2082dfcf1351c4bdd566bbe89a2b263235a2b50058f1e130a47277" +dependencies = [ + "bit-set 0.8.0", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -1489,12 +1552,29 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ce81f49ae8a0482e4c55ea62ebbd7e5a686af544c00b9d090bba3ff9be97b3d" +[[package]] +name = "fluent-uri" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc74ac4d8359ae70623506d512209619e5cf8f347124910440dbc221714b328e" +dependencies = [ + "borrow-or-share", + "ref-cast", + "serde", +] + [[package]] name = "fnv" version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "font-kit" version = "0.14.3" @@ -1556,6 +1636,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fraction" +version = "0.15.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e076045bb43dac435333ed5f04caf35c7463631d0dae2deb2638d94dd0a5b872" +dependencies = [ + "lazy_static", + "num", +] + [[package]] name = "freetype-sys" version = "0.20.1" @@ -1701,19 +1791,32 @@ checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", "libc", - "wasi 0.11.1+wasi-snapshot-preview1", + "wasi", ] [[package]] name = "getrandom" -version = "0.3.3" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 5.3.0", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" dependencies = [ "cfg-if", "libc", - "r-efi", - "wasi 0.14.4+wasi-0.2.4", + "r-efi 6.0.0", ] [[package]] @@ -1766,6 +1869,17 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + [[package]] name = "heck" version = "0.5.0" @@ -1882,6 +1996,109 @@ dependencies = [ "cc", ] +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + [[package]] name = "image" version = "0.25.8" @@ -1935,7 +2152,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4b0f83760fb341a774ed326568e19f5a863af4a952def8c39f9ab92fd95b88e5" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.15.5", "rayon", "serde", "serde_core", @@ -2063,7 +2280,7 @@ version = "0.1.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" dependencies = [ - "getrandom 0.3.3", + "getrandom 0.3.4", "libc", ] @@ -2077,6 +2294,42 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonschema" +version = "0.48.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9baf521a926fd1e6f94af6922be9686b61c8a89a5fb7e14d6a1e21b79738d4cc" +dependencies = [ + "ahash", + "bytecount", + "data-encoding", + "email_address", + "fancy-regex", + "fraction", + "getrandom 0.3.4", + "idna", + "itoa", + "jsonschema-regex", + "num-cmp", + "num-traits", + "percent-encoding", + "referencing", + "regex", + "serde", + "serde_json", + "unicode-general-category", + "uuid-simd", +] + +[[package]] +name = "jsonschema-regex" +version = "0.48.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d668f4775619127513d9a4f190704d20a66d20ccf0efc258f8ed42548e5fd7cd" +dependencies = [ + "regex-syntax", +] + [[package]] name = "katexit" version = "0.1.5" @@ -2186,6 +2439,12 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + [[package]] name = "litrs" version = "0.4.2" @@ -2267,6 +2526,12 @@ version = "2.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" +[[package]] +name = "micromap" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a86d3146ed3995b5913c414f6664344b9617457320782e64f0bb44afd49d74" + [[package]] name = "mime" version = "0.3.17" @@ -2307,7 +2572,7 @@ checksum = "50b7e5b27aa02a74bac8c3f23f448f8d87ff11f92d3aac1a6ed369ee08cc56c1" dependencies = [ "libc", "log", - "wasi 0.11.1+wasi-snapshot-preview1", + "wasi", "windows-sys 0.61.0", ] @@ -2550,6 +2815,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-cmp" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63335b2e2c34fae2fb0aa2cecfd9f0832a1e24b3b32ecec612c3426d46dc8aaa" + [[package]] name = "num-complex" version = "0.4.6" @@ -2670,6 +2941,12 @@ dependencies = [ "num-traits", ] +[[package]] +name = "outref" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" + [[package]] name = "owo-colors" version = "4.2.2" @@ -2875,6 +3152,15 @@ dependencies = [ "portable-atomic", ] +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + [[package]] name = "powerfmt" version = "0.2.0" @@ -3110,9 +3396,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.40" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -3123,6 +3409,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + [[package]] name = "rand" version = "0.8.5" @@ -3330,6 +3622,23 @@ dependencies = [ "syn", ] +[[package]] +name = "referencing" +version = "0.48.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df843f41f0cb459649f64a8b6b10579cf454f7a7420dc702767c81a26f23d780" +dependencies = [ + "ahash", + "fluent-uri", + "getrandom 0.3.4", + "hashbrown 0.17.1", + "itoa", + "micromap", + "parking_lot", + "percent-encoding", + "serde_json", +] + [[package]] name = "regex" version = "1.11.2" @@ -3344,9 +3653,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.10" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6b9458fa0bfeeac22b5ca447c63aaf45f28439a709ccd244698632f9aa6394d6" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" dependencies = [ "aho-corasick", "memchr", @@ -3355,9 +3664,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.8.6" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "relative-path" @@ -3785,6 +4094,12 @@ dependencies = [ "windows-sys 0.61.0", ] +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + [[package]] name = "stacker" version = "0.1.21" @@ -3876,6 +4191,17 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "system-deps" version = "6.2.2" @@ -3902,7 +4228,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.3", + "getrandom 0.4.3", "once_cell", "rustix", "windows-sys 0.61.0", @@ -4041,6 +4367,16 @@ dependencies = [ "time-core", ] +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + [[package]] name = "tinytemplate" version = "1.2.1" @@ -4270,6 +4606,7 @@ dependencies = [ "indexmap", "indoc", "itertools 0.14.0", + "jsonschema", "lazy_static", "log", "maplit", @@ -4571,6 +4908,12 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" +[[package]] +name = "unicode-general-category" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b993bddc193ae5bd0d623b49ec06ac3e9312875fdae725a975c51db1cc1677f" + [[package]] name = "unicode-ident" version = "1.0.18" @@ -4607,6 +4950,12 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + [[package]] name = "utf8parse" version = "0.2.2" @@ -4667,6 +5016,16 @@ dependencies = [ "serde", ] +[[package]] +name = "uuid-simd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b082222b4f6619906941c17eb2297fff4c2fb96cb60164170522942a200bd8" +dependencies = [ + "outref", + "vsimd", +] + [[package]] name = "v_frame" version = "0.3.9" @@ -4711,6 +5070,12 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "vsimd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" + [[package]] name = "wait-timeout" version = "0.2.1" @@ -4737,10 +5102,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" [[package]] -name = "wasi" -version = "0.14.4+wasi-0.2.4" +name = "wasip2" +version = "1.0.4+wasi-0.2.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88a5f4a424faf49c3c2c344f166f0662341d470ea185e939657aaff130f0ec4a" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" dependencies = [ "wit-bindgen", ] @@ -5112,9 +5477,15 @@ dependencies = [ [[package]] name = "wit-bindgen" -version = "0.45.1" +version = "0.57.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c573471f125075647d03df72e026074b7203790d41351cd6edc96f46bcccd36" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" [[package]] name = "xml" @@ -5148,6 +5519,29 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + [[package]] name = "zerocopy" version = "0.8.26" @@ -5168,6 +5562,60 @@ dependencies = [ "syn", ] +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zstd" version = "0.13.3" diff --git a/Cargo.toml b/Cargo.toml index 831aa88a8..9e5d20527 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -100,6 +100,7 @@ indexmap = { version = "=2.11.4", features = ["rayon", "serde"] } indoc = "=2.0.6" intervallum = { version = "=1.4.4", package = "intervallum" } itertools = "=0.14.0" +jsonschema = { version = "=0.48.1", default-features = false } lazy_static = "=1.5.0" log = "=0.4.28" maplit = { version = "=1.0.2", features = [] } diff --git a/packages/treetime-io/src/auspice.rs b/packages/treetime-io/src/auspice.rs index 54ef4dec2..6e340d385 100644 --- a/packages/treetime-io/src/auspice.rs +++ b/packages/treetime-io/src/auspice.rs @@ -1,19 +1,15 @@ -use crate::auspice_types::{AuspiceTree, AuspiceTreeData, AuspiceTreeNode}; +use crate::auspice_types::{AuspiceTree, AuspiceTreeNode}; use eyre::{Report, WrapErr}; -use maplit::{btreemap, btreeset}; -use parking_lot::RwLock; use std::collections::VecDeque; use std::io::Cursor; use std::io::{Read, Write}; use std::path::Path; -use std::sync::Arc; -use treetime_graph::edge::{Edge, GraphEdge, GraphEdgeKey}; +use treetime_graph::edge::GraphEdge; use treetime_graph::graph::Graph; -use treetime_graph::node::{GraphNode, GraphNodeKey, Node}; +use treetime_graph::node::GraphNode; use treetime_utils::io::file::create_file_or_stdout; use treetime_utils::io::file::open_file_or_stdin; use treetime_utils::io::json::{JsonPretty, json_read, json_write}; -use treetime_utils::make_internal_report; pub fn auspice_read_file(filepath: impl AsRef) -> Result, Report> where @@ -49,42 +45,23 @@ where auspice_to_graph::(&tree).wrap_err("When converting Auspice v2 JSON to graph") } -pub fn auspice_write_file(filepath: impl AsRef, graph: &Graph) -> Result<(), Report> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: AuspiceWrite, -{ +pub fn auspice_write_file(filepath: impl AsRef, tree: &AuspiceTree) -> Result<(), Report> { let filepath = filepath.as_ref(); let mut f = create_file_or_stdout(filepath)?; - auspice_write::(&mut f, graph) + auspice_write(&mut f, tree) .wrap_err_with(|| format!("When writing Auspice v2 JSON file '{}'", filepath.display()))?; writeln!(f)?; Ok(()) } -pub fn auspice_write_str(graph: &Graph) -> Result -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: AuspiceWrite, -{ +pub fn auspice_write_str(tree: &AuspiceTree) -> Result { let mut buf = Vec::new(); - auspice_write::(&mut buf, graph).wrap_err("When writing Auspice v2 JSON string")?; + auspice_write(&mut buf, tree).wrap_err("When writing Auspice v2 JSON string")?; Ok(String::from_utf8(buf)?) } -pub fn auspice_write(writer: &mut impl Write, graph: &Graph) -> Result<(), Report> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: AuspiceWrite, -{ - let tree = auspice_from_graph::(graph)?; - json_write(writer, &tree, JsonPretty(true)).wrap_err("When writing Auspice v2 JSON") +pub fn auspice_write(writer: &mut impl Write, tree: &AuspiceTree) -> Result<(), Report> { + json_write(writer, tree, JsonPretty(true)).wrap_err("When writing Auspice v2 JSON") } pub struct AuspiceTreeContext<'a> { @@ -144,130 +121,3 @@ where graph.build()?; Ok(graph) } - -pub struct AuspiceGraphContext<'a, N, E, D> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, -{ - pub node_key: GraphNodeKey, - pub node: &'a N, - pub parent_key: Option, - pub parent: Option<&'a N>, - pub edge_key: Option, - pub edge: Option<&'a E>, - pub graph: &'a Graph, -} - -pub trait AuspiceWrite: Sized -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, -{ - fn new(graph: &Graph) -> Result; - - fn auspice_data_from_graph_data(&self, graph: &Graph) -> Result; - - fn auspice_node_from_graph_components( - &mut self, - context: &AuspiceGraphContext, - ) -> Result; -} - -/// Convert graph to Auspice v2 JSON -pub fn auspice_from_graph(graph: &Graph) -> Result -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: AuspiceWrite, -{ - let mut converter = C::new(graph)?; - auspice_from_graph_with(&mut converter, graph) -} - -/// Convert graph to Auspice v2 JSON using a pre-constructed converter. -/// -/// Use this when the converter needs initialization beyond what -/// `AuspiceWrite::new(graph)` provides (e.g. extra data not on the graph). -pub fn auspice_from_graph_with(converter: &mut C, graph: &Graph) -> Result -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: AuspiceWrite, -{ - let root = graph.get_exactly_one_root().wrap_err("When writing Auspice v2 JSON")?; - - // Pre-order iteration to construct the nodes from graph nodes and edges - let mut node_map = { - let mut node_map = btreemap! {}; - let mut queue = VecDeque::from([(Arc::clone(&root), None, None)]); - while let Some((current_node, current_parent, current_edge)) = queue.pop_front() { - let node_key = current_node.read_arc().key(); - let node = &*current_node.read_arc().payload().read_arc(); - let parent_key = current_parent - .as_ref() - .map(|parent: &Arc>>| parent.read_arc().key()); - let parent = current_parent - .as_ref() - .map(|parent: &Arc>>| parent.read_arc().payload().read_arc()); - let parent = parent.as_deref(); - let edge_key = current_edge - .as_ref() - .map(|edge: &Arc>>| edge.read_arc().key()); - let edge = current_edge - .as_ref() - .map(|edge: &Arc>>| edge.read_arc().payload().read_arc()); - let edge = edge.as_deref(); - let context = AuspiceGraphContext { - node_key, - node, - parent_key, - parent, - edge_key, - edge, - graph, - }; - let current_tree_node = converter.auspice_node_from_graph_components(&context)?; - for (child, edge) in graph.children_of(¤t_node.read_arc()) { - queue.push_back((child, Some(Arc::clone(¤t_node)), Some(edge))); - } - node_map.insert(current_node.read_arc().key(), current_tree_node); - } - node_map - }; - - // Post-order traversal to populate .children array - let mut visited = btreeset! {}; - let mut stack = vec![Arc::clone(&root)]; - while let Some(node) = stack.pop() { - if visited.contains(&node.read_arc().key()) { - let mut children_to_add = vec![]; - for (child, _) in graph.children_of(&node.read_arc()) { - let child_key = child.read_arc().key(); - if let Some(child_node) = node_map.remove(&child_key) { - children_to_add.push(child_node); - } - } - if let Some(node) = node_map.get_mut(&node.read_arc().key()) { - node.children.extend(children_to_add); - } - } else { - visited.insert(node.read_arc().key()); - stack.push(Arc::clone(&node)); - for (child, _) in graph.children_of(&node.read_arc()) { - stack.push(child); - } - } - } - - let data = converter.auspice_data_from_graph_data(graph)?; - let root_key = root.read_arc().key(); - let tree = node_map - .remove(&root_key) - .ok_or_else(|| make_internal_report!("Root node {root_key} not found in node_map"))?; - Ok(AuspiceTree { data, tree }) -} diff --git a/packages/treetime-io/src/auspice_types.rs b/packages/treetime-io/src/auspice_types.rs index a491b6c2c..356017012 100644 --- a/packages/treetime-io/src/auspice_types.rs +++ b/packages/treetime-io/src/auspice_types.rs @@ -235,7 +235,8 @@ pub struct AuspiceGenomeAnnotationCds { #[derive(Clone, Debug, Serialize, Deserialize)] pub struct AuspiceGenomeAnnotations { - pub nuc: AuspiceGenomeAnnotationNuc, + #[serde(skip_serializing_if = "Option::is_none")] + pub nuc: Option, #[serde(flatten)] pub cdses: BTreeMap, diff --git a/packages/treetime-io/src/graph.rs b/packages/treetime-io/src/graph.rs index 6922c0f93..39d6bb2ad 100644 --- a/packages/treetime-io/src/graph.rs +++ b/packages/treetime-io/src/graph.rs @@ -1,17 +1,3 @@ -use crate::auspice::{AuspiceWrite, auspice_write_file}; -use crate::graphviz::{EdgeToGraphviz, NodeToGraphviz, graphviz_write_file}; -use crate::nex::{NexWriteOptions, nex_write_file_with}; -use crate::nwk::{CommentProviders, EdgeToNwk, NodeToNwk, NwkWriteOptions, nwk_write_file_with}; -use crate::phyloxml::{PhyloxmlFromGraph, PhyloxmlJsonOptions, phyloxml_json_write_file, phyloxml_write_file}; -use crate::usher_mat::{UsherMatJsonOptions, UsherWrite, usher_mat_json_write_file, usher_mat_pb_write_file}; -use eyre::Report; -use serde::Serialize; -use std::collections::BTreeMap; -use std::path::PathBuf; -use treetime_graph::edge::{GraphEdge, HasBranchLength}; -use treetime_graph::graph::Graph; -use treetime_graph::node::{GraphNode, Named}; -use treetime_utils::io::json::{JsonPretty, json_write_file}; use util_newick::NwkStyle; /// Dispatch tag for tree output formats. Used as keys in the resolved output map. @@ -42,57 +28,3 @@ impl TreeWriteKind { pub struct NwkWriteSpec { pub style: NwkStyle, } - -/// Write every requested tree format from one authoritative graph. -pub fn write_tree_outputs( - graph: &Graph, - outputs: &BTreeMap, - providers: &CommentProviders, -) -> Result<(), Report> -where - N: GraphNode + Named + NodeToNwk + NodeToGraphviz + Serialize, - E: GraphEdge + HasBranchLength + EdgeToNwk + EdgeToGraphviz + Serialize, - D: Send + Sync + Serialize, - C: AuspiceWrite + PhyloxmlFromGraph + UsherWrite, -{ - for (kind, path) in outputs { - match kind { - TreeWriteKind::Nwk(spec) => { - let options = NwkWriteOptions { - style: spec.style, - ..NwkWriteOptions::default() - }; - nwk_write_file_with(path, graph, &options, providers)?; - }, - TreeWriteKind::Nexus(spec) => { - let options = NexWriteOptions { - style: spec.style, - ..NexWriteOptions::default() - }; - nex_write_file_with(path, graph, &options, providers)?; - }, - TreeWriteKind::GraphJson => { - json_write_file(path, graph, JsonPretty(true))?; - }, - TreeWriteKind::Dot => { - graphviz_write_file(path, graph)?; - }, - TreeWriteKind::Auspice => { - auspice_write_file::(path, graph)?; - }, - TreeWriteKind::Phyloxml => { - phyloxml_write_file::(path, graph)?; - }, - TreeWriteKind::PhyloxmlJson => { - phyloxml_json_write_file::(path, graph, &PhyloxmlJsonOptions::default())?; - }, - TreeWriteKind::MatPb => { - usher_mat_pb_write_file::(path, graph)?; - }, - TreeWriteKind::MatJson => { - usher_mat_json_write_file::(path, graph, &UsherMatJsonOptions::default())?; - }, - } - } - Ok(()) -} diff --git a/packages/treetime-io/src/phyloxml.rs b/packages/treetime-io/src/phyloxml.rs index 21e5e8362..63c5c07f5 100644 --- a/packages/treetime-io/src/phyloxml.rs +++ b/packages/treetime-io/src/phyloxml.rs @@ -1,12 +1,9 @@ use eyre::{Report, WrapErr}; -use maplit::{btreemap, btreeset}; -use parking_lot::RwLock; use smart_default::SmartDefault; use std::collections::VecDeque; use std::io::{Cursor, Read, Write}; use std::path::Path; -use std::sync::Arc; -use treetime_graph::edge::{Edge, GraphEdge}; +use treetime_graph::edge::GraphEdge; use treetime_graph::graph::Graph; use treetime_graph::node::GraphNode; use treetime_utils::io::file::create_file_or_stdout; @@ -14,7 +11,7 @@ use treetime_utils::io::file::open_file_or_stdin; use treetime_utils::io::json::{ JsonPretty, json_read, json_read_file, json_read_str, json_write, json_write_file, json_write_str, }; -use treetime_utils::{make_internal_error, make_internal_report}; +use treetime_utils::make_internal_error; pub use util_phyloxml::{ Phyloxml, PhyloxmlAccession, PhyloxmlAnnotation, PhyloxmlBinaryCharacterList, PhyloxmlBinaryCharacters, PhyloxmlBranchColor, PhyloxmlClade, PhyloxmlCladeRelation, PhyloxmlConfidence, PhyloxmlDate, PhyloxmlDistribution, @@ -91,89 +88,45 @@ where phyloxml_to_graph::(&pxml) } -pub fn phyloxml_write_file(filepath: impl AsRef, graph: &Graph) -> Result<(), Report> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: PhyloxmlFromGraph, -{ +pub fn phyloxml_write_file(filepath: impl AsRef, phyloxml: &Phyloxml) -> Result<(), Report> { let filepath = filepath.as_ref(); let mut f = create_file_or_stdout(filepath)?; - phyloxml_write::(&mut f, graph) - .wrap_err_with(|| format!("When writing PhyloXML file '{}'", filepath.display()))?; + phyloxml_write(&mut f, phyloxml).wrap_err_with(|| format!("When writing PhyloXML file '{}'", filepath.display()))?; writeln!(f)?; Ok(()) } -pub fn phyloxml_write_str(graph: &Graph) -> Result -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: PhyloxmlFromGraph, -{ +pub fn phyloxml_write_str(phyloxml: &Phyloxml) -> Result { let mut buf = Vec::new(); - phyloxml_write::(&mut buf, graph).wrap_err("When writing PhyloXML string")?; + phyloxml_write(&mut buf, phyloxml).wrap_err("When writing PhyloXML string")?; String::from_utf8(buf).wrap_err("PhyloXML output is not valid UTF-8") } -pub fn phyloxml_write(writer: &mut impl Write, graph: &Graph) -> Result<(), Report> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: PhyloxmlFromGraph, -{ - let pxml = phyloxml_from_graph::(graph)?; - util_phyloxml::phyloxml_write(writer, &pxml).wrap_err("When writing PhyloXML") +pub fn phyloxml_write(writer: &mut impl Write, phyloxml: &Phyloxml) -> Result<(), Report> { + util_phyloxml::phyloxml_write(writer, phyloxml).wrap_err("When writing PhyloXML") } -pub fn phyloxml_json_write_file( +pub fn phyloxml_json_write_file( filepath: impl AsRef, - graph: &Graph, + phyloxml: &Phyloxml, options: &PhyloxmlJsonOptions, -) -> Result<(), Report> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: PhyloxmlFromGraph, -{ +) -> Result<(), Report> { let filepath = filepath.as_ref(); - let pxml = phyloxml_from_graph::(graph)?; - json_write_file(filepath, &pxml, JsonPretty(options.pretty)) + json_write_file(filepath, phyloxml, JsonPretty(options.pretty)) .wrap_err_with(|| format!("When writing PhyloXML JSON file: '{}'", filepath.display()))?; Ok(()) } -pub fn phyloxml_json_write_str( - graph: &Graph, - options: &PhyloxmlJsonOptions, -) -> Result -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: PhyloxmlFromGraph, -{ - let pxml = phyloxml_from_graph::(graph)?; - json_write_str(&pxml, JsonPretty(options.pretty)).wrap_err("When writing PhyloXML JSON string") +pub fn phyloxml_json_write_str(phyloxml: &Phyloxml, options: &PhyloxmlJsonOptions) -> Result { + json_write_str(phyloxml, JsonPretty(options.pretty)).wrap_err("When writing PhyloXML JSON string") } -pub fn phyloxml_json_write( +pub fn phyloxml_json_write( writer: &mut impl Write, - graph: &Graph, + phyloxml: &Phyloxml, options: &PhyloxmlJsonOptions, -) -> Result<(), Report> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: PhyloxmlFromGraph, -{ - let pxml = phyloxml_from_graph::(graph)?; - json_write(writer, &pxml, JsonPretty(options.pretty)).wrap_err("When writing PhyloXML JSON") +) -> Result<(), Report> { + json_write(writer, phyloxml, JsonPretty(options.pretty)).wrap_err("When writing PhyloXML JSON") } pub trait PhyloxmlDataToGraphData: Sized { @@ -239,101 +192,3 @@ where Ok(graph) } - -pub struct PhyloxmlGraphContext<'a, N, E, D> -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, -{ - pub node_key: treetime_graph::node::GraphNodeKey, - pub node: &'a N, - pub edge_key: Option, - pub edge: Option<&'a E>, - pub graph: &'a Graph, -} - -/// Describes conversion to Phyloxml tree node data when writing to PhyloXML -pub trait PhyloxmlFromGraph -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, -{ - fn phyloxml_data_from_graph_data(graph: &Graph) -> Result; - - fn phyloxml_node_from_graph_components(context: &PhyloxmlGraphContext) -> Result; -} - -/// Convert graph to PhyloXML -pub fn phyloxml_from_graph(graph: &Graph) -> Result -where - N: GraphNode, - E: GraphEdge, - D: Sync + Send, - C: PhyloxmlFromGraph, -{ - let root = graph.get_exactly_one_root()?; - - // Pre-order iteration to construct the nodes from graph nodes and edges - let mut node_map = { - let mut node_map = btreemap! {}; - let mut queue = VecDeque::from([(Arc::clone(&root), None)]); - while let Some((current_node, current_edge)) = queue.pop_front() { - let current_node = current_node.read_arc(); - - let node = &*current_node.payload().read_arc(); - let edge_key = current_edge - .as_ref() - .map(|edge: &Arc>>| edge.read_arc().key()); - let edge = current_edge - .as_ref() - .map(|edge: &Arc>>| edge.read_arc().payload().read_arc()); - let edge = edge.as_deref(); - let current_tree_node = C::phyloxml_node_from_graph_components(&PhyloxmlGraphContext { - node_key: current_node.key(), - node, - edge_key, - edge, - graph, - })?; - for (child, edge) in graph.children_of(¤t_node) { - queue.push_back((child, Some(edge))); - } - node_map.insert(current_node.key(), current_tree_node); - } - node_map - }; - - // Post-order traversal to populate .children array - let mut visited = btreeset! {}; - let mut stack = vec![Arc::clone(&root)]; - while let Some(node) = stack.pop() { - if visited.contains(&node.read_arc().key()) { - let mut children_to_add = vec![]; - for (child, _) in graph.children_of(&node.read_arc()) { - let child_key = child.read_arc().key(); - if let Some(child_node) = node_map.remove(&child_key) { - children_to_add.push(child_node); - } - } - if let Some(node) = node_map.get_mut(&node.read_arc().key()) { - node.clade.extend(children_to_add); - } - } else { - visited.insert(node.read_arc().key()); - stack.push(Arc::clone(&node)); - for (child, _) in graph.children_of(&node.read_arc()) { - stack.push(child); - } - } - } - - let mut data = C::phyloxml_data_from_graph_data(graph)?; - let root_key = root.read_arc().key(); - let clade = node_map - .remove(&root_key) - .ok_or_else(|| make_internal_report!("Root node not found in node_map during PhyloXML export"))?; - data.phylogeny[0].clade = Some(clade); - Ok(data) -} diff --git a/packages/treetime-io/src/usher_mat.rs b/packages/treetime-io/src/usher_mat.rs index 63754d1e0..1c260c62d 100644 --- a/packages/treetime-io/src/usher_mat.rs +++ b/packages/treetime-io/src/usher_mat.rs @@ -1,12 +1,12 @@ -use crate::nwk::{EdgeFromNwk, EdgeToNwk, NodeFromNwk, NodeToNwk, NwkWriteOptions, nwk_read_str, nwk_write_str}; -use bytes::Buf; +use crate::nwk::{EdgeFromNwk, NodeFromNwk, nwk_read_str}; +use bytes::{Buf, BytesMut}; use eyre::{Report, WrapErr}; use smart_default::SmartDefault; use std::io::{Read, Write}; use std::path::Path; -use treetime_graph::edge::{GraphEdge, GraphEdgeKey, HasBranchLength}; +use treetime_graph::edge::{GraphEdge, HasBranchLength}; use treetime_graph::graph::Graph; -use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; +use treetime_graph::node::{GraphNode, Named}; use treetime_utils::io::file::create_file_or_stdout; use treetime_utils::io::file::open_file_or_stdin; use treetime_utils::io::json::{ @@ -85,92 +85,46 @@ where usher_to_graph::(&tree) } -pub fn usher_mat_pb_write_file(filepath: impl AsRef, graph: &Graph) -> Result<(), Report> -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, - C: UsherWrite, -{ +pub fn usher_mat_pb_write_file(filepath: impl AsRef, tree: &UsherTree) -> Result<(), Report> { let filepath = filepath.as_ref(); let mut f = create_file_or_stdout(filepath)?; - usher_mat_pb_write::(&mut f, graph) + usher_mat_pb_write(&mut f, tree) .wrap_err_with(|| format!("When writing Usher MAT protobuf file '{}'", filepath.display()))?; writeln!(f)?; Ok(()) } -pub fn usher_mat_pb_write_bytes(graph: &Graph) -> Result<(), Report> -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, - C: UsherWrite, -{ - let tree = usher_from_graph::(graph)?; - let mut buf = Vec::new(); - util_usher_mat::usher_mat_pb_write_bytes(&mut buf, &tree).wrap_err("When writing Usher MAT protobuf bytes") +pub fn usher_mat_pb_write_bytes(tree: &UsherTree) -> Result, Report> { + let mut buf = BytesMut::new(); + util_usher_mat::usher_mat_pb_write_bytes(&mut buf, tree).wrap_err("When writing Usher MAT protobuf bytes")?; + Ok(buf.to_vec()) } -pub fn usher_mat_pb_write(writer: &mut impl Write, graph: &Graph) -> Result<(), Report> -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, - C: UsherWrite, -{ - let tree = usher_from_graph::(graph)?; - util_usher_mat::usher_mat_pb_write(writer, &tree).wrap_err("When writing Usher MAT protobuf") +pub fn usher_mat_pb_write(writer: &mut impl Write, tree: &UsherTree) -> Result<(), Report> { + util_usher_mat::usher_mat_pb_write(writer, tree).wrap_err("When writing Usher MAT protobuf") } -pub fn usher_mat_json_write_file( +pub fn usher_mat_json_write_file( filepath: impl AsRef, - graph: &Graph, - + tree: &UsherTree, options: &UsherMatJsonOptions, -) -> Result<(), Report> -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, - C: UsherWrite, -{ +) -> Result<(), Report> { let filepath = filepath.as_ref(); - let tree = usher_from_graph::(graph)?; - json_write_file(filepath, &tree, JsonPretty(options.pretty)) + json_write_file(filepath, tree, JsonPretty(options.pretty)) .wrap_err_with(|| format!("When writing Usher MAT JSON file: '{}'", filepath.display()))?; Ok(()) } -pub fn usher_mat_json_write_str( - graph: &Graph, - - options: &UsherMatJsonOptions, -) -> Result -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, - C: UsherWrite, -{ - let tree = usher_from_graph::(graph)?; - json_write_str(&tree, JsonPretty(options.pretty)).wrap_err("When writing Usher MAT JSON string") +pub fn usher_mat_json_write_str(tree: &UsherTree, options: &UsherMatJsonOptions) -> Result { + json_write_str(tree, JsonPretty(options.pretty)).wrap_err("When writing Usher MAT JSON string") } -pub fn usher_mat_json_write( +pub fn usher_mat_json_write( writer: &mut impl Write, - graph: &Graph, - + tree: &UsherTree, options: &UsherMatJsonOptions, -) -> Result<(), Report> -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, - C: UsherWrite, -{ - let tree = usher_from_graph::(graph)?; - json_write(writer, &tree, JsonPretty(options.pretty)).wrap_err("When writing Usher MAT JSON") +) -> Result<(), Report> { + json_write(writer, tree, JsonPretty(options.pretty)).wrap_err("When writing Usher MAT JSON") } #[derive(SmartDefault)] @@ -260,73 +214,3 @@ where Ok(graph) } - -pub struct UsherGraphContext<'a, N, E, D> -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, -{ - pub node_key: GraphNodeKey, - pub node: &'a N, - pub edge_key: Option, - pub edge: Option<&'a E>, - pub graph: &'a Graph, -} - -pub trait UsherWrite: Sized -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, -{ - fn new(graph: &Graph) -> Result; - - fn usher_node_from_graph_components( - &mut self, - context: &UsherGraphContext, - ) -> Result<(UsherTreeNode, UsherMutationList, UsherMetadata), Report>; -} - -/// Convert graph to Usher MAT protobuf -pub fn usher_from_graph(graph: &Graph) -> Result -where - N: GraphNode + NodeToNwk, - E: GraphEdge + EdgeToNwk, - D: Sync + Send, - C: UsherWrite, -{ - let mut node_mutations = vec![]; - let mut condensed_nodes = vec![]; - let mut metadata = vec![]; - - let mut converter = C::new(graph)?; - graph.iter_depth_first_preorder_forward(|node| { - let node_key = node.key; - let edge = node.parents.first().map(|(_, edge)| edge.read_arc()); - let edge_key = node.parent_keys.first().map(|(_, edge_key)| *edge_key); - let edge = edge.as_deref(); - let node = &node.payload; - - let (node, mutations, meta) = converter.usher_node_from_graph_components(&UsherGraphContext { - node_key, - node, - edge_key, - edge, - graph, - })?; - - node_mutations.push(mutations); - condensed_nodes.push(node); - metadata.push(meta); - Ok(()) - })?; - - let newick = nwk_write_str(graph, &NwkWriteOptions::default())?; - Ok(UsherTree { - newick, - node_mutations, - condensed_nodes, - metadata, - }) -} diff --git a/packages/treetime-utils/src/array/__tests__/mod.rs b/packages/treetime-utils/src/array/__tests__/mod.rs index 246d2d7da..07c2a90e2 100644 --- a/packages/treetime-utils/src/array/__tests__/mod.rs +++ b/packages/treetime-utils/src/array/__tests__/mod.rs @@ -1 +1,2 @@ mod ndarray; +mod serde; diff --git a/packages/treetime-utils/src/array/__tests__/serde.rs b/packages/treetime-utils/src/array/__tests__/serde.rs new file mode 100644 index 000000000..ae9b2a4c7 --- /dev/null +++ b/packages/treetime-utils/src/array/__tests__/serde.rs @@ -0,0 +1,26 @@ +#[cfg(test)] +mod tests { + use crate::array::serde::array2_as_vec; + use crate::io::json::{JsonPretty, json_write_str}; + use ndarray::{Array2, array}; + use pretty_assertions::assert_eq; + use serde::Serialize; + + #[test] + fn serde_serializes_non_contiguous_array2() { + let value = Matrix { + values: array![[1, 2], [3, 4]].reversed_axes(), + }; + let expected = r#"{"values":[[1,3],[2,4]]}"#; + + let actual = json_write_str(&value, JsonPretty(false)).unwrap(); + + assert_eq!(expected, actual); + } + + #[derive(Serialize)] + struct Matrix { + #[serde(serialize_with = "array2_as_vec")] + values: Array2, + } +} diff --git a/packages/treetime-utils/src/array/serde.rs b/packages/treetime-utils/src/array/serde.rs index 0bd8b7cf6..4e5854121 100644 --- a/packages/treetime-utils/src/array/serde.rs +++ b/packages/treetime-utils/src/array/serde.rs @@ -29,7 +29,7 @@ where T: Serialize, S: Serializer, { - array.as_slice().unwrap().serialize(serializer) + serializer.collect_seq(array.iter()) } /// Deserialize Array1 from a simple JSON array @@ -56,7 +56,11 @@ where T: Serialize, S: Serializer, { - let rows: Vec<&[T]> = array.rows().into_iter().map(|row| row.to_slice().unwrap()).collect(); + let rows = array + .rows() + .into_iter() + .map(|row| row.into_iter().collect::>()) + .collect::>(); rows.serialize(serializer) } @@ -88,7 +92,7 @@ where S: Serializer, { match array { - Some(arr) => arr.as_slice().unwrap().serialize(serializer), + Some(arr) => serializer.collect_seq(arr.iter()), None => serializer.serialize_none(), } } diff --git a/packages/treetime/Cargo.toml b/packages/treetime/Cargo.toml index f8432ba57..1e663481a 100644 --- a/packages/treetime/Cargo.toml +++ b/packages/treetime/Cargo.toml @@ -64,6 +64,7 @@ bio = { workspace = true } criterion = { workspace = true } ctor = { workspace = true } indoc = { workspace = true } +jsonschema = { workspace = true } lazy_static = { workspace = true } petgraph = { workspace = true } pretty_assertions = { workspace = true } diff --git a/packages/treetime/src/ancestral/__tests__/test_dense_completeness.rs b/packages/treetime/src/ancestral/__tests__/test_dense_completeness.rs index 60a08fb5f..d617ee9be 100644 --- a/packages/treetime/src/ancestral/__tests__/test_dense_completeness.rs +++ b/packages/treetime/src/ancestral/__tests__/test_dense_completeness.rs @@ -10,6 +10,7 @@ mod tests { use crate::partition::traits::PartitionOptimizeOps; use crate::payload::ancestral::GraphAncestral; use crate::seq::alignment::get_common_length; + use crate::seq::indel::InDel; use eyre::Report; use parking_lot::RwLock; @@ -188,7 +189,7 @@ ACGTACGTAC ); // Verify indel directions exist (at least one deletion) - let has_deletion = p.data.edges.values().any(|e| e.indels.iter().any(|i| i.deletion)); + let has_deletion = p.data.edges.values().any(|e| e.indels.iter().any(InDel::is_deletion)); assert!( has_deletion, "At least one deletion should be detected from gap-bearing leaves" diff --git a/packages/treetime/src/ancestral/__tests__/test_fitch_indel.rs b/packages/treetime/src/ancestral/__tests__/test_fitch_indel.rs index 7835ad64a..24a2b402f 100644 --- a/packages/treetime/src/ancestral/__tests__/test_fitch_indel.rs +++ b/packages/treetime/src/ancestral/__tests__/test_fitch_indel.rs @@ -138,7 +138,7 @@ mod tests { let (indels, new_gaps) = resolve_indels_forward(&variable_indel, &node_gaps, &[], &[], &parent_seq, &node_seq); assert_eq!(indels.len(), 1); - assert!(indels[0].deletion); + assert!(indels[0].is_deletion()); assert_eq!(indels[0].range, (2, 5)); assert_eq!(indels[0].seq, Seq::try_from_str("GTA").unwrap()); assert_eq!(new_gaps, vec![(2, 5)]); @@ -154,7 +154,7 @@ mod tests { let (indels, new_gaps) = resolve_indels_forward(&variable_indel, &[], &[], &parent_gaps, &parent_seq, &node_seq); assert_eq!(indels.len(), 1); - assert!(!indels[0].deletion); + assert!(!indels[0].is_deletion()); assert_eq!(indels[0].range, (2, 5)); assert_eq!(indels[0].seq, Seq::try_from_str("GTA").unwrap()); assert!(new_gaps.is_empty()); diff --git a/packages/treetime/src/ancestral/fitch.rs b/packages/treetime/src/ancestral/fitch.rs index 11e862597..aab0f2d31 100644 --- a/packages/treetime/src/ancestral/fitch.rs +++ b/packages/treetime/src/ancestral/fitch.rs @@ -384,7 +384,7 @@ fn run_fitch_reconstruction( } for indel in &edge_part.indels { - if indel.deletion { + if indel.is_deletion() { sequence[indel.range.0..indel.range.1].fill(alphabet.gap()); } else { sequence[indel.range.0..indel.range.1].copy_from_slice(&indel.seq); diff --git a/packages/treetime/src/ancestral/pipeline.rs b/packages/treetime/src/ancestral/pipeline.rs index cda5a6322..763af9708 100644 --- a/packages/treetime/src/ancestral/pipeline.rs +++ b/packages/treetime/src/ancestral/pipeline.rs @@ -44,6 +44,7 @@ pub struct AncestralInput { pub sequences: Vec, } +#[derive(Serialize)] pub enum AncestralPartition { Fitch(Arc>), Sparse(Arc>), diff --git a/packages/treetime/src/commands/ancestral/__tests__/test_augur_node_data.rs b/packages/treetime/src/commands/ancestral/__tests__/test_augur_node_data.rs index 3526628a8..c53323ce8 100644 --- a/packages/treetime/src/commands/ancestral/__tests__/test_augur_node_data.rs +++ b/packages/treetime/src/commands/ancestral/__tests__/test_augur_node_data.rs @@ -310,6 +310,7 @@ mod tests { name_to_key["B"] => vec![], name_to_key["root"] => vec![], }, + node_mutations: btreemap! {}, }, Some(AugurNodeDataJsonAnnotationEntry { start: Some(1), diff --git a/packages/treetime/src/commands/ancestral/aa_node_data.rs b/packages/treetime/src/commands/ancestral/aa_node_data.rs index 46a156732..0150ae958 100644 --- a/packages/treetime/src/commands/ancestral/aa_node_data.rs +++ b/packages/treetime/src/commands/ancestral/aa_node_data.rs @@ -4,9 +4,10 @@ use crate::make_error; use crate::partition::augur::AugurNodeDataJsonAncestralPartition; use crate::partition::traits::BranchTopology; use crate::payload::ancestral::GraphAncestral; -use crate::seq::mutation::Sub; +use crate::seq::mutation::{Mutation, MutationEvent, MutationTrack, Sub}; use eyre::Report; use itertools::Itertools; +use serde::Serialize; use serde_json::json; use std::collections::{BTreeMap, BTreeSet}; use std::path::{Path, PathBuf}; @@ -16,7 +17,7 @@ use treetime_io::gff::{GffCdsFeature, read_gff3_cds_features_filtered}; use treetime_primitives::{AsciiChar, Seq}; use util_augur_node_data_json::{AugurNodeDataJsonAnnotationEntry, AugurNodeDataJsonAnnotationSegment}; -#[derive(Clone, Debug, Default, PartialEq, Eq)] +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)] pub struct AaNodeData { pub annotations: BTreeMap, pub reference: BTreeMap, @@ -24,6 +25,7 @@ pub struct AaNodeData { // shared tree as the nucleotide partition, so per-node results join by node identity (key) rather // than by reconstructing identity from a synthesized node name across independent graphs. pub node_aa_muts: BTreeMap>>, + pub node_aa_mutations: BTreeMap>>, pub root_aa_sequences: BTreeMap, } @@ -41,14 +43,22 @@ impl AaNodeData { .or_default() .insert(cds.to_owned(), muts); } + for (node_key, mutations) in cds_data.node_mutations { + self + .node_aa_mutations + .entry(node_key) + .or_default() + .insert(cds.to_owned(), mutations); + } } } -#[derive(Clone, Debug, Default, PartialEq, Eq)] +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)] pub struct AaCdsNodeData { pub reference: String, pub root_sequence: String, pub node_muts: BTreeMap>, + pub node_mutations: BTreeMap>, } pub fn validate_aa_args( @@ -233,6 +243,7 @@ pub fn collect_aa_cds_node_data( } let mut node_muts = BTreeMap::new(); + let mut node_mutations = BTreeMap::new(); for node in graph.get_nodes() { let node_guard = node.read_arc(); let node_key = node_guard.key(); @@ -241,31 +252,42 @@ pub fn collect_aa_cds_node_data( .name() .map_or_else(|| format!("node_{}", node_key.0), |n| n.as_ref().to_owned()); - let muts = if node_key == root_key { + let mutations = if node_key == root_key { diff_sequences(&reference, &inferred_root, partition.ambiguous_char())? + .into_iter() + .map(MutationEvent::Substitution) + .collect() } else { let (_parent_key, edge_key) = graph .node_parent(node_key)? .ok_or_else(|| eyre::eyre!("Non-root node '{node_name}' has no parent while collecting AA node data"))?; - partition + let substitutions = partition .edge_subs(graph, edge_key)? .into_iter() .sorted_by_key(Sub::pos) - .map(|sub| sub.to_string()) - .collect() + .map(MutationEvent::Substitution) + .map(Ok); + let indels = partition + .edge_indels(edge_key) + .into_iter() + .map(|indel| Mutation::indel(MutationTrack::AminoAcid(cds.to_owned()), &indel).map(|mutation| mutation.event)); + substitutions.chain(indels).collect::, Report>>()? }; + let muts = mutations.iter().flat_map(mutation_event_strings).collect(); node_muts.insert(node_key, muts); + node_mutations.insert(node_key, mutations); } Ok(AaCdsNodeData { reference: reference.as_str().to_owned(), root_sequence: inferred_root.as_str().to_owned(), node_muts, + node_mutations, }) } -fn diff_sequences(reference: &Seq, query: &Seq, unknown: AsciiChar) -> Result, Report> { +fn diff_sequences(reference: &Seq, query: &Seq, unknown: AsciiChar) -> Result, Report> { if reference.len() != query.len() { return make_error!( "Cannot diff sequences with lengths {} and {}", @@ -274,15 +296,31 @@ fn diff_sequences(reference: &Seq, query: &Seq, unknown: AsciiChar) -> Result Vec { + match event { + MutationEvent::Substitution(substitution) => vec![substitution.to_string()], + MutationEvent::Insertion(segment) => segment + .sequence .iter() - .zip(query.iter()) .enumerate() - .filter(|(_pos, (reff, qry))| reff != qry && is_reportable_sub(**reff, **qry, unknown)) - .map(|(pos, (reff, qry))| format!("{}{}{}", char::from(*reff), pos + 1, char::from(*qry))) + .map(|(offset, state)| format!("-{}{state}", segment.range.0 + offset + 1)) .collect(), - ) + MutationEvent::Deletion(segment) => segment + .sequence + .iter() + .enumerate() + .map(|(offset, state)| format!("{state}{}-", segment.range.0 + offset + 1)) + .collect(), + } } fn is_reportable_sub(reff: AsciiChar, qry: AsciiChar, unknown: AsciiChar) -> bool { @@ -410,7 +448,20 @@ mod tests { let actual = diff_sequences(&reference, &query, AsciiChar::from_byte_unchecked(b'X')).unwrap(); - let expected = vec![o!("C2D"), o!("D3Q")]; + let expected = vec![ + Sub::new( + AsciiChar::from_byte_unchecked(b'C'), + 1_usize, + AsciiChar::from_byte_unchecked(b'D'), + ) + .unwrap(), + Sub::new( + AsciiChar::from_byte_unchecked(b'D'), + 2_usize, + AsciiChar::from_byte_unchecked(b'Q'), + ) + .unwrap(), + ]; assert_eq!(expected, actual); } @@ -438,6 +489,13 @@ mod tests { name_to_key["B"] => vec![], name_to_key["root"] => vec![o!("A2C")], }, + node_mutations: btreemap! { + name_to_key["A"] => vec![], + name_to_key["B"] => vec![], + name_to_key["root"] => vec![MutationEvent::Substitution( + Sub::new(AsciiChar::from_byte_unchecked(b'A'), 1_usize, AsciiChar::from_byte_unchecked(b'C')).unwrap() + )], + }, }; assert_eq!(expected, actual); } @@ -504,6 +562,10 @@ mod tests { Ok(vec![]) } + fn edge_indels(&self, _edge_key: GraphEdgeKey) -> Vec { + vec![] + } + fn ambiguous_char(&self) -> AsciiChar { self.unknown } diff --git a/packages/treetime/src/commands/ancestral/result.rs b/packages/treetime/src/commands/ancestral/result.rs index b0989751c..3610578c4 100644 --- a/packages/treetime/src/commands/ancestral/result.rs +++ b/packages/treetime/src/commands/ancestral/result.rs @@ -1,22 +1,17 @@ use crate::ancestral::pipeline::AncestralPartition; +use crate::commands::ancestral::aa_node_data::AaNodeData; use crate::gtr::get_gtr::GtrModelName; use crate::gtr::gtr::GTR; use crate::payload::ancestral::GraphAncestral; use serde::Serialize; -#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. #[derive(Serialize)] -#[serde(transparent)] pub struct AncestralGraphData { - marker: (), - #[serde(skip)] pub partition: Option, - #[serde(skip)] pub gtr: Option, - #[serde(skip)] pub model_name: GtrModelName, - #[serde(skip)] pub mask: Vec, + pub aa_node_data: Option, } impl AncestralGraphData { @@ -25,13 +20,14 @@ impl AncestralGraphData { gtr: Option, model_name: GtrModelName, mask: Vec, + aa_node_data: Option, ) -> Self { Self { - marker: (), partition, gtr, model_name, mask, + aa_node_data, } } } diff --git a/packages/treetime/src/commands/ancestral/run.rs b/packages/treetime/src/commands/ancestral/run.rs index 402bce7e1..096f5dc64 100644 --- a/packages/treetime/src/commands/ancestral/run.rs +++ b/packages/treetime/src/commands/ancestral/run.rs @@ -10,7 +10,7 @@ use crate::commands::ancestral::args::TreetimeAncestralArgs; use crate::commands::ancestral::augur_node_data::write_augur_node_data_json_with_aa; use crate::commands::ancestral::result::{AncestralGraphData, AncestralResult}; use crate::commands::shared::output::{CommandKind, OutputSelection}; -use crate::commands::shared::tree_output::TreeOutputAdapter; +use crate::commands::shared::tree_output::write_ancestral_tree_outputs; use crate::gtr::get_gtr::{GtrOutput, write_gtr_json}; use crate::make_error; use crate::partition::traits::MutationCommentProvider; @@ -21,7 +21,6 @@ use eyre::Report; use log::{info, warn}; use treetime_graph::node::Named; use treetime_io::fasta::{FastaReader, FastaRecord, FastaWriter, read_many_fasta}; -use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_io::nwk::nwk_read_file; use treetime_utils::io::file::{create_file_or_stdout, open_stdin}; @@ -97,8 +96,6 @@ pub fn run_ancestral_reconstruction( ), ], )?; - resolved.prepare()?; - let mut output_fasta = if resolved .non_tree_outputs .contains_key(&OutputSelection::ReconstructedNucFasta) @@ -174,7 +171,7 @@ pub fn run_ancestral_reconstruction( model_name, mask, } = output; - let mut graph = graph.map_data(AncestralGraphData::new(partition, gtr, model_name, mask)); + let mut graph = graph.map_data(AncestralGraphData::new(partition, gtr, model_name, mask, aa_node_data)); topology_order.apply(&mut graph)?; progress.report("Writing output", 0.9, ""); @@ -182,15 +179,33 @@ pub fn run_ancestral_reconstruction( match &graph.data().partition { Some(AncestralPartition::Fitch(partition)) => { let guard = partition.read_arc(); - write_augur_node_data_json_with_aa(&graph, &*guard, &graph.data().mask, aa_node_data.as_ref(), path)?; + write_augur_node_data_json_with_aa( + &graph, + &*guard, + &graph.data().mask, + graph.data().aa_node_data.as_ref(), + path, + )?; }, Some(AncestralPartition::Sparse(partition)) => { let guard = partition.read_arc(); - write_augur_node_data_json_with_aa(&graph, &*guard, &graph.data().mask, aa_node_data.as_ref(), path)?; + write_augur_node_data_json_with_aa( + &graph, + &*guard, + &graph.data().mask, + graph.data().aa_node_data.as_ref(), + path, + )?; }, Some(AncestralPartition::Dense(partition)) => { let guard = partition.read_arc(); - write_augur_node_data_json_with_aa(&graph, &*guard, &graph.data().mask, aa_node_data.as_ref(), path)?; + write_augur_node_data_json_with_aa( + &graph, + &*guard, + &graph.data().mask, + graph.data().aa_node_data.as_ref(), + path, + )?; }, None => {}, } @@ -227,16 +242,16 @@ fn write_tree_for_partition( let guard = partition.read_arc(); let provider = MutationCommentProvider::new(&*guard, graph); let providers = CommentProviders::new().with(&provider); - write_tree_outputs::(graph, &resolved.tree_outputs, &providers)?; + write_ancestral_tree_outputs(graph, &resolved.tree_outputs, &providers)?; }, Some(AncestralPartition::Dense(partition)) => { let guard = partition.read_arc(); let provider = MutationCommentProvider::new(&*guard, graph); let providers = CommentProviders::new().with(&provider); - write_tree_outputs::(graph, &resolved.tree_outputs, &providers)?; + write_ancestral_tree_outputs(graph, &resolved.tree_outputs, &providers)?; }, Some(AncestralPartition::Fitch(_)) | None => { - write_tree_outputs::(graph, &resolved.tree_outputs, &CommentProviders::new())?; + write_ancestral_tree_outputs(graph, &resolved.tree_outputs, &CommentProviders::new())?; }, } Ok(()) diff --git a/packages/treetime/src/commands/clock/run.rs b/packages/treetime/src/commands/clock/run.rs index eb1f81260..ff287be89 100644 --- a/packages/treetime/src/commands/clock/run.rs +++ b/packages/treetime/src/commands/clock/run.rs @@ -7,7 +7,7 @@ use crate::clock::pipeline::{self, ClockInput, ClockPipelineParams}; use crate::clock::rtt::{ClockRegressionResult, write_clock_regression_result_csv}; use crate::commands::clock::args::{BranchSplitArgs, TreetimeClockArgs}; use crate::commands::shared::output::{CommandKind, OutputSelection}; -use crate::commands::shared::tree_output::TreeOutputAdapter; +use crate::commands::shared::tree_output::write_clock_tree_outputs; use crate::make_error; use crate::make_report; use eyre::{Report, WrapErr}; @@ -15,24 +15,17 @@ use treetime_graph::edge::GraphEdge; use treetime_graph::graph::Graph; use treetime_graph::node::{GraphNode, Named}; use treetime_io::dates_csv::read_dates; -use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::nwk_read_file; -#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. #[derive(serde::Serialize)] -#[serde(transparent)] pub struct ClockGraphData { - marker: (), - #[serde(skip)] pub clock_model: ClockModel, - #[serde(skip)] pub regression_results: Vec, } impl ClockGraphData { pub fn new(clock_model: ClockModel, regression_results: Vec) -> Self { Self { - marker: (), clock_model, regression_results, } @@ -136,12 +129,10 @@ pub fn run_clock( .topology_order .resolve_topology_order(&graph, Some(input_order))?; topology_order.apply(&mut graph)?; - resolved.prepare()?; - progress.report("Writing output", 0.8, ""); if !resolved.tree_outputs.is_empty() { - write_tree_outputs::( + write_clock_tree_outputs( &graph, &resolved.tree_outputs, &treetime_io::nwk::CommentProviders::new(), diff --git a/packages/treetime/src/commands/mugration/run.rs b/packages/treetime/src/commands/mugration/run.rs index 638a258e0..9f8c3e886 100644 --- a/packages/treetime/src/commands/mugration/run.rs +++ b/packages/treetime/src/commands/mugration/run.rs @@ -1,7 +1,7 @@ use crate::commands::mugration::args::TreetimeMugrationArgs; use crate::commands::mugration::augur_node_data::write_augur_node_data_json; use crate::commands::shared::output::{CommandKind, OutputSelection}; -use crate::commands::shared::tree_output::TreeOutputAdapter; +use crate::commands::shared::tree_output::write_mugration_tree_outputs; use crate::gtr::get_gtr::{GtrModelName, GtrOutput, write_gtr_json}; use crate::make_report; use crate::mugration::mugration::execute_mugration; @@ -13,7 +13,6 @@ use eyre::Report; use log::info; use std::collections::BTreeMap; use treetime_io::discrete_states_csv::read_discrete_attrs; -use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_io::nwk::nwk_read_file; use treetime_utils::io::file::create_file_or_stdout; @@ -95,14 +94,12 @@ pub fn run_mugration( .topology_order .resolve_topology_order(&result.graph, None)?; topology_order.apply(&mut result.graph)?; - resolved.prepare()?; - progress.report("Writing output", 0.8, ""); if !resolved.tree_outputs.is_empty() { let provider = DiscreteCommentProvider::new(&result.graph.data().partition, &result.graph.data().traits.attribute); let providers = CommentProviders::new().with(&provider); - write_tree_outputs::(&result.graph, &resolved.tree_outputs, &providers)?; + write_mugration_tree_outputs(&result.graph, &resolved.tree_outputs, &providers)?; } if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { diff --git a/packages/treetime/src/commands/optimize/result.rs b/packages/treetime/src/commands/optimize/result.rs index 2b656e28c..783fdc021 100644 --- a/packages/treetime/src/commands/optimize/result.rs +++ b/packages/treetime/src/commands/optimize/result.rs @@ -7,18 +7,11 @@ use parking_lot::RwLock; use serde::Serialize; use std::sync::Arc; -#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. #[derive(Serialize)] -#[serde(transparent)] pub struct OptimizeGraphData { - marker: (), - #[serde(skip)] pub gtr: GTR, - #[serde(skip)] pub model_name: GtrModelName, - #[serde(skip)] pub sparse_partitions: Vec>>, - #[serde(skip)] pub dense_partitions: Vec>>, } @@ -30,7 +23,6 @@ impl OptimizeGraphData { dense_partitions: Vec>>, ) -> Self { Self { - marker: (), gtr, model_name, sparse_partitions, diff --git a/packages/treetime/src/commands/optimize/run.rs b/packages/treetime/src/commands/optimize/run.rs index 6f22e0026..eacc7c2e4 100644 --- a/packages/treetime/src/commands/optimize/run.rs +++ b/packages/treetime/src/commands/optimize/run.rs @@ -3,7 +3,7 @@ use crate::commands::optimize::args::TreetimeOptimizeArgs; use crate::commands::optimize::augur_node_data::write_augur_node_data_json; use crate::commands::optimize::result::{OptimizeGraphData, OptimizeResult}; use crate::commands::shared::output::{CommandKind, DivergenceUnits, OutputSelection}; -use crate::commands::shared::tree_output::TreeOutputAdapter; +use crate::commands::shared::tree_output::write_optimize_tree_outputs; use crate::gtr::get_gtr::{GtrOutput, write_gtr_json}; use crate::make_error; use crate::optimize::pipeline::{self, OptimizeInput, OptimizeParams}; @@ -14,7 +14,6 @@ use eyre::Report; use log::info; use std::path::PathBuf; use treetime_io::fasta::read_many_fasta; -use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_io::nwk::nwk_read_file; @@ -82,8 +81,6 @@ pub fn run_optimize( )); let topology_order = args.topology_order.resolve_topology_order(&graph, None)?; topology_order.apply(&mut graph)?; - resolved.prepare()?; - progress.report("Writing output", 0.9, ""); if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { @@ -96,14 +93,14 @@ pub fn run_optimize( let guard = graph.data().dense_partitions[0].read_arc(); let provider = MutationCommentProvider::new(&*guard, &graph); let providers = CommentProviders::new().with(&provider); - write_tree_outputs::(&graph, &resolved.tree_outputs, &providers)?; + write_optimize_tree_outputs(&graph, &resolved.tree_outputs, &providers)?; } else if !graph.data().sparse_partitions.is_empty() { let guard = graph.data().sparse_partitions[0].read_arc(); let provider = MutationCommentProvider::new(&*guard, &graph); let providers = CommentProviders::new().with(&provider); - write_tree_outputs::(&graph, &resolved.tree_outputs, &providers)?; + write_optimize_tree_outputs(&graph, &resolved.tree_outputs, &providers)?; } else { - write_tree_outputs::(&graph, &resolved.tree_outputs, &CommentProviders::new())?; + write_optimize_tree_outputs(&graph, &resolved.tree_outputs, &CommentProviders::new())?; } } diff --git a/packages/treetime/src/commands/prune/result.rs b/packages/treetime/src/commands/prune/result.rs index 1b024db9a..a14f3782c 100644 --- a/packages/treetime/src/commands/prune/result.rs +++ b/packages/treetime/src/commands/prune/result.rs @@ -5,24 +5,15 @@ use parking_lot::RwLock; use serde::Serialize; use std::sync::Arc; -#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. #[derive(Serialize)] -#[serde(transparent)] pub struct PruneGraphData { - marker: (), - #[serde(skip)] pub gtr: Option, - #[serde(skip)] pub partitions: Vec>>, } impl PruneGraphData { pub fn new(gtr: Option, partitions: Vec>>) -> Self { - Self { - marker: (), - gtr, - partitions, - } + Self { gtr, partitions } } } diff --git a/packages/treetime/src/commands/prune/run.rs b/packages/treetime/src/commands/prune/run.rs index 65441b393..30d3ff1db 100644 --- a/packages/treetime/src/commands/prune/run.rs +++ b/packages/treetime/src/commands/prune/run.rs @@ -2,7 +2,7 @@ use crate::alphabet::alphabet::Alphabet; use crate::commands::prune::args::TreetimePruneArgs; use crate::commands::prune::result::{PruneGraphData, PruneResult}; use crate::commands::shared::output::{CommandKind, OutputSelection}; -use crate::commands::shared::tree_output::TreeOutputAdapter; +use crate::commands::shared::tree_output::write_prune_tree_outputs; use crate::gtr::get_gtr::{GtrModelName, GtrOutput, write_gtr_json}; use crate::make_error; use crate::prune::pipeline::{self, PruneInput, PruneParams}; @@ -16,7 +16,6 @@ use treetime_graph::edge::GraphEdge; use treetime_graph::graph::Graph; use treetime_graph::node::{GraphNode, Named}; use treetime_io::fasta::read_many_fasta; -use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_io::nwk::nwk_read_file; use treetime_io::parse_delimited::{parse_delimited_file, parse_delimited_str}; @@ -82,8 +81,6 @@ pub fn run_prune( let mut graph = graph.map_data(PruneGraphData::new(gtr, partitions)); let topology_order = args.topology_order.resolve_topology_order(&graph, Some(input_order))?; topology_order.apply(&mut graph)?; - resolved.prepare()?; - progress.report("Writing output", 0.8, ""); if let Some(path) = resolved.non_tree_outputs.get(&OutputSelection::Gtr) { @@ -102,7 +99,7 @@ pub fn run_prune( } if !resolved.tree_outputs.is_empty() { - write_tree_outputs::(&graph, &resolved.tree_outputs, &CommentProviders::new())?; + write_prune_tree_outputs(&graph, &resolved.tree_outputs, &CommentProviders::new())?; } progress.report("Done", 1.0, ""); diff --git a/packages/treetime/src/commands/shared/__tests__/schemas/auspice/LICENSE.txt b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/LICENSE.txt new file mode 100644 index 000000000..4ec8c3f73 --- /dev/null +++ b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/LICENSE.txt @@ -0,0 +1,619 @@ + GNU AFFERO GENERAL PUBLIC LICENSE + Version 3, 19 November 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU Affero General Public License is a free, copyleft license for +software and other kinds of works, specifically designed to ensure +cooperation with the community in the case of network server software. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +our General Public Licenses are intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + Developers that use our General Public Licenses protect your rights +with two steps: (1) assert copyright on the software, and (2) offer +you this License which gives you legal permission to copy, distribute +and/or modify the software. + + A secondary benefit of defending all users' freedom is that +improvements made in alternate versions of the program, if they +receive widespread use, become available for other developers to +incorporate. Many developers of free software are heartened and +encouraged by the resulting cooperation. However, in the case of +software used on network servers, this result may fail to come about. +The GNU General Public License permits making a modified version and +letting the public access it on a server without ever releasing its +source code to the public. + + The GNU Affero General Public License is designed specifically to +ensure that, in such cases, the modified source code becomes available +to the community. It requires the operator of a network server to +provide the source code of the modified version running there to the +users of that server. Therefore, public use of a modified version, on +a publicly accessible server, gives the public access to the source +code of the modified version. + + An older license, called the Affero General Public License and +published by Affero, was designed to accomplish similar goals. This is +a different license, not a version of the Affero GPL, but Affero has +released a new version of the Affero GPL which permits relicensing under +this license. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU Affero General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Remote Network Interaction; Use with the GNU General Public License. + + Notwithstanding any other provision of this License, if you modify the +Program, your modified version must prominently offer all users +interacting with it remotely through a computer network (if your version +supports such interaction) an opportunity to receive the Corresponding +Source of your version by providing access to the Corresponding Source +from a network server at no charge, through some standard or customary +means of facilitating copying of software. This Corresponding Source +shall include the Corresponding Source for any work covered by version 3 +of the GNU General Public License that is incorporated pursuant to the +following paragraph. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the work with which it is combined will remain governed by version +3 of the GNU General Public License. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU Affero General Public License from time to time. Such new versions +will be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU Affero General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU Affero General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU Affero General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS diff --git a/packages/treetime/src/commands/shared/__tests__/schemas/auspice/README.md b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/README.md new file mode 100644 index 000000000..bc9c3b634 --- /dev/null +++ b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/README.md @@ -0,0 +1,7 @@ +# Auspice v2 schemas + +These test fixtures are copied from Nextstrain Augur 34.0.0 at commit `d8e38736037ba9474a809f9a5a63bc2b279d2407`: + + + +They are used unchanged to validate every command's Auspice v2 output model. The files in this directory retain Augur's AGPL-3.0 license; see `LICENSE.txt`. diff --git a/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-annotations.json b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-annotations.json new file mode 100644 index 000000000..1c0bcf431 --- /dev/null +++ b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-annotations.json @@ -0,0 +1,95 @@ +{ + "type" : "object", + "$schema": "http://json-schema.org/draft-06/schema#", + "$id": "https://nextstrain.org/schemas/augur/annotations", + "title": "Schema for the 'annotations' property (node-data JSON) or the 'genome_annotations' property (auspice JSON)", + "properties": { + "nuc": { + "description": "Genome/gene/segment/sequence coordinates in one-based nuc(leotide) coordinates", + "type": "object", + "allOf": [{ "$ref": "#/$defs/startend" }], + "properties": { + "start": { + "enum": [1], + "$comment": "nuc must begin at 1" + }, + "strand": { + "type": "string", + "enum":["+"], + "description": "Strand is optional for nuc, as it should be +ve for all genomes (-ve strand genomes are reverse complemented)", + "$comment": "Auspice will not proceed if the JSON has strand='-'" + } + }, + "additionalProperties": true, + "$comment": "All other properties are unused by Auspice." + } + }, + "required": [], + "minProperties": 1, + "additionalProperties": false, + "patternProperties": { + "^(?!nuc$)[a-zA-Z0-9*_.()-]+$": { + "$comment": "Each object here defines a single CDS", + "type": "object", + "oneOf": [{ "$ref": "#/$defs/startend" }, { "$ref": "#/$defs/segments" }], + "additionalProperties": true, + "required": ["strand"], + "properties": { + "gene": { + "type": "string", + "description": "The name of the gene the CDS is from. Optional.", + "$comment": "Shown in on-hover infobox & influences default CDS colors" + }, + "strand": { + "description": "Strand of the CDS", + "type": "string", + "enum": ["-", "+"] + }, + "color": { + "type": "string", + "description": "A CSS color or a color hex code. Optional." + }, + "display_name": { + "type": "string", + "$comment": "Shown in the on-hover info box" + }, + "description": { + "type": "string", + "$comment": "Shown in the on-hover info box" + } + } + } + }, + "$defs": { + "startend": { + "type": "object", + "required": ["start", "end"], + "properties": { + "start": { + "type": "integer", + "minimum": 1, + "description": "Start position (one-based, following GFF format)" + }, + "end": { + "type": "integer", + "minimum": 2, + "description": "End position (one-based, following GFF format). This value _must_ be greater than the start." + } + } + }, + "segments": { + "type": "object", + "required": ["segments"], + "properties": { + "segments": { + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "allOf": [{ "$ref": "#/$defs/startend" }] + } + } + } + } + } +} diff --git a/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-auspice-config-v2.json b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-auspice-config-v2.json new file mode 100644 index 000000000..1a5878cfc --- /dev/null +++ b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-auspice-config-v2.json @@ -0,0 +1,329 @@ +{ + "$schema": "http://json-schema.org/draft-06/schema#", + "$id": "https://nextstrain.org/schemas/auspice/config/v2", + "type": "object", + "title": "Auspice config file to be supplied to `augur export v2`", + "$comment": "This schema includes deprecated-but-handled-by-augur-export-v1 properties, but their schema definitions are somewhat incomplete", + "additionalProperties": false, + "required": [], + "properties" : { + "title": { + "description": "Title to be displayed in auspice", + "type" : "string" + }, + "colorings": { + "description": "Traits available as color-by options", + "type": "array", + "items": { + "type": "object", + "description": "Each object here is an individual coloring, which will populate the sidebar dropdown in auspice", + "additionalProperties": false, + "required": ["key"], + "properties": { + "key": { + "description": "They key used to access the value of this coloring on each node", + "type": "string", + "not": {"const": "none"} + }, + "title": { + "description": "Text to be displayed in the \"color by\" dropdown and legends", + "$comment": "string is parsed unchanged by Auspice", + "type": "string" + }, + "type": { + "description": "Defines how the color scale should be constructed", + "$comment": "[augur export v2] will (try to) infer the type if this is not present", + "type": "string", + "enum": ["continuous", "temporal", "ordinal", "categorical", "boolean"] + }, + "scale": { + "description": "Provided mapping between trait values & hex values. For continuous scales at least 2 items must be specified", + "$comment": "[auspice export v2] preferentially uses this over colors TSV", + "type": "array", + "items": { + "type": "array", + "items": [ + { + "type": ["string", "number"], + "description": "For categorical/ordinal scales, this is the (string) value of the trait to associate with the colour. For continuous scales this is the (numeric) value to associate with the colour, and interpolation will be used to span the domain" + }, + {"type": "string", "description": "color hex value", "pattern": "^#[0-9A-Fa-f]{6}$"} + ] + } + }, + "legend": { + "description": "Specify the entries displayed in the legend. This can be used to restrict the entries in the legend for display without otherwise affecting the data viz", + "type": "array", + "items": { + "type": "object", + "required": ["value"], + "properties": { + "value": { + "description": "value to associate with this legend entry. Used to determine colour. For non-continuous scales this also determines the matching between legend items and data.", + "type": ["string", "number"], + "$comment": "Continuous scales must use a numeric value. Other scales can use either." + }, + "display": { + "description": "Label to display in the legend. Optional - `value` will be used if this is not provided.", + "type": ["string", "number"] + }, + "bounds": { + "description": "(for continuous scales only) provide the lower & upper bounds to match data to this legend entry. Bounds from different legend entries must not overlap. Matching is (a, b] - exclusive of the lower bound, inclusive of the upper.", + "type": "array", + "items": [ + {"type": "number", "description": "lower bound"}, + {"type": "number", "description": "upper bound"} + ] + } + } + } + } + } + } + }, + "color_options": { + "description": "DEPRECATED v1 syntax for defining colorings", + "deprecated": true, + "type": "object" + }, + "geo_resolutions": { + "description": "Traits to be interpreted as 'geo resolution' options -- i.e. associated with lat/longs & made points on the map", + "$comment": "Note that array entries can be different structures & you can mix & match", + "type": "array", + "uniqueItems": true, + "minItems": 1, + "items": { + "oneOf": [ + { + "type": "object", + "description": "An indiviual geo resolution", + "additionalProperties": false, + "required": ["key"], + "properties": { + "key": { + "type": "string", + "description": "Trait key - must be specified on nodes (e.g. 'country')" + }, + "title": { + "type": "string", + "description": "The title to display in the geo resolution dropdown. Optional -- if not provided then `key` will be used." + } + } + }, + { + "type": "string", + "description": "An indiviual geo resolution key" + } + ] + } + }, + "geo": { + "description": "DEPRECATED v1 syntax for defining geo_resolutions", + "deprecated": true, + "type": "array" + }, + "maintainers": { + "description": "Who maintains this dataset?", + "$comment": "order similar to a publication", + "type": "array", + "uniqueItems": true, + "minItems": 1, + "items": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "url": {"type": "string"} + }, + "required": ["name"] + } + }, + "maintainer": { + "description": "DEPRECATED v1 syntax for defining maintainers (but you could only define one!)", + "deprecated": true, + "type": "array" + }, + "build_url": { + "description": "URL with instructions to reproduce build, usually expected to be a GitHub repo URL", + "$comment": "Auspice displays this at the top of the page as part of a byline", + "type": "string" + }, + "build_avatar": { + "description": "The custom avatar image URL allows users to include a logo, even if it is hosted outside of GitHub.", + "$comment": "Auspice displays this at the top of the page as part of a byline", + "type": "string" + }, + "filters": { + "description": "These appear as filters in the footer of Auspice (which populates the displayed values based upon the tree)", + "$comment": "These values must be present as keys on a tree node -> trait", + "type": "array", + "uniqueItems": true, + "items": {"type": "string"} + }, + "display_defaults": { + "description": "Set the defaults for certain display options in Auspice. All are optional.", + "$comment": "Anything able to be encoded in the auspice URL should eventually be an option here, so this will expand over time", + "type": "object", + "additionalProperties": false, + "properties": { + "map_triplicate": { + "type": "boolean" + }, + "geo_resolution": { + "$comment": "The value here must be present in geo_resolutions (see above)", + "type": "string" + }, + "color_by": { + "$comment": "The value here must be present in the colorings (see above)", + "type": "string" + }, + "distance_measure": { + "type": "string", + "enum": ["num_date", "div"] + }, + "layout": { + "type": "string", + "enum": ["rect", "radial", "unrooted", "clock"] + }, + "branch_label": { + "description": "What branch label should be displayed by default, or 'none' to hide labels by default.", + "$comment": "Should be a key present in the per-node branch_attrs.labels object of the exported JSON; pattern is from the schema for that object", + "type": "string", + "pattern": "^(none|[a-zA-Z0-9]+)$" + }, + "label": { + "description": "branch label that tree starts zoomed to, expressed as :", + "type": "string", + "pattern": "^[a-zA-Z0-9]+:[^:]+$" + }, + "tip_label": { + "description": "What tip label should be displayed by default, or 'none' to hide labels by default.", + "$comment": "Should be a key present in (at least some) node_attrs", + "type": "string" + }, + "stream_label": { + "description": "Start with streamtrees enabled using this branch label", + "type": "string" + }, + "transmission_lines": { + "$comment": "Transmission lines depend on the geo_resolution being defined for internal nodes", + "type": "boolean" + }, + "language": { + "type": "string", + "minLength": 1, + "description": "A BCP 47 language tag specifying the default language in which to display Auspice's interface (if supported)" + }, + "sidebar": { + "type": "string", + "enum": ["open", "closed"] + }, + "panels": { + "type": "array", + "description": "Panels which start toggled on (default is for all available to be shown)", + "minItems": 1, + "items": { + "type": "string", + "enum": ["tree", "map", "frequencies", "entropy"] + } + } + } + }, + "defaults": { + "description": "DEPRECATED v1 syntax for defining auspice view defaults", + "deprecated": true, + "type": "object" + }, + "stream_labels": { + "description": "Branch labels which are may be used as streamtree defining labels (otherwise all available labels will be used). Use an empty array to prevent streamtrees from being used.", + "type": "array", + "minItems": 0, + "items": { + "type": "string" + } + }, + "updated": { + "description": "DEPRECATED v1 (or older) syntax for defining when the build was updated", + "$comment": "unused in augur v6", + "deprecated": true, + "type": "string" + }, + "panels": { + "description": "The panels available for display", + "$comment": "The frequencies & measurements panel will only be available if defined here (and if their sidecar files are available)", + "type": "array", + "minItems": 1, + "items": { + "type": "string", + "enum": ["tree", "map", "frequencies", "entropy", "measurements"] + } + }, + "vaccine_choices": { + "type": "object", + "description": "UNUSED v1 syntax for defining vaccine choices", + "$comment": "This is unused in `augur export v2` which gets vaccine info vis a node-data JSON file. It remains in the schema so that v1 config files can be used by `augur export v2`" + }, + "data_provenance": { + "description": "Specify provenance of data included in this analysis", + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "description": "An individual data source", + "additionalProperties": false, + "required": ["name"], + "properties": { + "name": { + "description": "Name of the data source", + "type": "string" + }, + "url": { + "description": "URL to use in link to data source", + "type": "string" + } + } + } + }, + "sharing": { + "type": "object", + "description": "Control which assets Auspice enables for download. By default all are enabled (true). Note: Auspice will overwrite some of these for GISAID datasets to comply with data sharing agreements.", + "additionalProperties": false, + "properties": { + "dataset_json": { + "type": "boolean", + "description": "Enable/disable JSON dataset downloading" + }, + "metadata_tsv": { + "type": "boolean", + "description": "Enable/disable downloading metadata TSVs" + }, + "authors": { + "type": "boolean", + "description": "Enable/disable downloading authorship TSVs (where the dataset supports it)" + }, + "trees": { + "type": "boolean", + "description": "Enable/disable downloading Newick / Nexus trees" + }, + "entropy": { + "type": "boolean", + "description": "Enable/disable downloading entropy panel data (where the dataset supports it)" + }, + "screenshot": { + "type": "boolean", + "description": "Enable/disable downloading a SVG screenshot" + } + } + }, + "metadata_columns": { + "description": "Metadata TSV columns to export in addition to columns provided as colorings.", + "$comment": "These columns will not be used as coloring options in Auspice but will be visible in the tree.", + "type": "array", + "uniqueItems": true, + "items": {"type": "string", "not": {"const": "none"}} + }, + "extensions": { + "description": "Data to be passed through to the the resulting dataset JSON", + "$comment": "Any type is accepted" + } + } +} diff --git a/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-root-sequence.json b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-root-sequence.json new file mode 100644 index 000000000..3bf82cbb0 --- /dev/null +++ b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-root-sequence.json @@ -0,0 +1,33 @@ +{ + "$schema": "http://json-schema.org/draft-06/schema#", + "$id": "https://nextstrain.org/schemas/dataset/root-sequence", + "title": "Nextstrain root-sequence sidecar for datasets", + "description": "Typically produced by Augur and consumed by Auspice. Applicable to the `--root-sequence` output of `augur export v2` as well as the `--output-sequence` option of `augur export v1`.", + "oneOf": [ + { + "$comment": "This is sort of weird, but `augur export v1` can explicitly produce an empty object.", + "description": "An empty object", + "type": "object", + "properties": {}, + "additionalProperties": false + }, + { + "description": "An object containing at least a \"nuc\" key and optionally additional keys for genome annotations (e.g. genes)", + "type": "object", + "required": ["nuc"], + "properties": { + "nuc": { + "description": "Nucleotide sequence of whole genome (from the output of `augur ancestral`)", + "type": "string" + } + }, + "patternProperties": { + "^[a-zA-Z0-9*_-]+$": { + "$comment": "This pattern is the same pattern used in the corresponding parts of schema-export-v2.json.", + "description": "Amino acid sequence of genome annotation (e.g. gene) identified by this key (from the output of `augur translate`)", + "type": "string" + } + } + } + ] +} diff --git a/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-v2.json b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-v2.json new file mode 100644 index 000000000..734ac2140 --- /dev/null +++ b/packages/treetime/src/commands/shared/__tests__/schemas/auspice/schema-export-v2.json @@ -0,0 +1,364 @@ +{ + "$schema": "http://json-schema.org/draft-06/schema#", + "$id": "https://nextstrain.org/schemas/dataset/v2", + "type": "object", + "title": "Nextstrain dataset v2", + "description": "A file representing an Auspice dataset. Typically produced by Augur (`augur export v2`) and consumed by Auspice.", + "additionalProperties": false, + "required": ["version", "meta", "tree"], + "properties": { + "version" : { + "description": "Major schema version", + "const": "v2" + }, + "meta": { + "type": "object", + "$comment": "Metadata associated with phylogeny", + "additionalProperties": false, + "required": ["updated", "panels"], + "properties" : { + "title" : { + "description": "Auspice displays this at the top of the page", + "type" : "string" + }, + "updated" : { + "description": "Auspice displays this in the byline and footer.", + "type" : "string", + "pattern": "^[0-9X]{4}-[0-9X]{2}-[0-9X]{2}$" + }, + "build_url" : { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/build_url" + }, + "build_avatar" : { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/build_avatar" + }, + "description" : { + "description": "Auspice displays this currently in the footer.", + "$comment": "Generally a description of the phylogeny and/or acknowledgements in Markdown format.", + "type": "string" + }, + "warning" : { + "description": "Text in Markdown format to be displayed by Auspice as a warning banner under the byline.", + "type": "string" + }, + "maintainers": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/maintainers" + }, + "genome_annotations": { + "$ref": "https://nextstrain.org/schemas/augur/annotations" + }, + "filters": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/filters" + }, + "panels": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/panels" + }, + "extensions": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/extensions" + }, + "geo_resolutions": { + "description": "The available options for the geographic resolution dropdown, and their lat/long information", + "type": "array", + "minItems": 1, + "items": { + "type": "object", + "description": "Each object here is an individual geo resolution", + "additionalProperties": false, + "required": ["key", "demes"], + "properties": { + "key": { + "type": "string", + "description": "Trait key - must be specified on nodes (e.g. 'country')" + }, + "title": { + "type": "string", + "description": "The title to display in the geo resolution dropdown. Optional -- if not provided then `key` will be used." + }, + "demes": { + "type": "object", + "description": "Mapping from deme (trait values) to lat/long", + "$comment": "Each value defined across the tree needs to be present here, else Auspice cannot display the deme appropriately", + "patternProperties": { + "^[a-z_]+$": { + "type": "object", + "additionalProperties": false, + "properties": { + "latitude": { + "type": "number", + "minimum": -90, + "maximum": 90 + }, + "longitude": { + "type": "number", + "minimum": -180, + "maximum": 180 + } + } + } + } + } + } + } + }, + "colorings": { + "type": "array", + "items": { + "allOf": [ + {"$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/colorings/items"} + ], + "required": ["key", "type"] + } + }, + "display_defaults": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/display_defaults" + }, + "stream_labels": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/stream_labels" + }, + "data_provenance": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/data_provenance" + }, + "sharing": { + "$ref": "https://nextstrain.org/schemas/auspice/config/v2#/properties/sharing" + } + } + }, + "tree": { + "description": "One or more phylogenies using a nested JSON structure", + "oneOf": [ + {"$ref": "#/$defs/tree"}, + { + "type": "array", + "minItems": 1, + "items": {"$ref": "#/$defs/tree"} + } + ] + }, + "root_sequence": { + "$ref": "https://nextstrain.org/schemas/dataset/root-sequence" + } + }, + "$defs": { + "tree": { + "type" : "object", + "$comment": "The phylogeny in a nested JSON structure", + "additionalProperties": false, + "required": ["name", "node_attrs"], + "properties": { + "name": { + "description": "Strain name. Must be unique. No spaces", + "type": "string" + }, + "node_attrs": { + "description": "attributes associated with the node (sequence, date, location) as opposed to changes from one node to another.", + "type": "object", + "anyOf": [ + {"required": ["div"]}, + {"required": ["num_date"]} + ], + "properties": { + "div": { + "description": "Node (phylogenetic) divergence", + "$comment": "Cumulative (root = 0)", + "type": "number" + }, + "num_date": { + "$comment": "Sample date in decimal format (e.g. 2012.1234)", + "$comment": "This is the only date information. We no longer have string dates.", + "type": "object", + "required": ["value"], + "properties": { + "value": {"type": "number"}, + "confidence": { + "description": "Confidence of the node date", + "type": "array", + "items": [ + {"type": "number"}, + {"type": "number"} + ] + }, + "inferred": { + "type": "boolean", + "description": "[terminal nodes only] was the 'value' inferred or known?" + }, + "raw_value": { + "type": "string", + "description": "[terminal nodes only, and only if inferred=true] the known (ambiguous) date string" + } + } + }, + "vaccine": { + "description": "Vaccine information", + "properties": { + "serum": { + "description": "strain used to raise sera (for ???)", + "$comment": "Currently in the flu & dengue trees", + "type": "boolean" + }, + "selection_date": { + "description": "Vaccine selection date", + "$comment": "this is currently in the metadata JSON", + "type": "string", + "pattern": "^[0-9X]{4}-[0-9X]{2}-[0-9X]{2}$" + }, + "start_date": { + "description": "Vaccine usage start date", + "type": "string", + "pattern": "^[0-9X]{4}-[0-9X]{2}-[0-9X]{2}$" + }, + "end_date": { + "description": "When the vaccine was stopped", + "$comment": "if vaccine still in use, don't include this property", + "type": "string", + "pattern": "^[0-9X]{4}-[0-9X]{2}-[0-9X]{2}$" + } + } + }, + "hidden": { + "$comment": "Instruct auspice to hide the branches from this node to it's children", + "type": "string", + "enum": ["always", "timetree", "divtree"] + }, + "url": { + "description": "URL of the sequence (usually https://www.ncbi.nlm.nih.gov/nuccore/...)", + "$comment": "terminal nodes only", + "type": "string", + "pattern": "^https?://.+$" + }, + "author": { + "description": "Author information (terminal nodes only)", + "type": "object", + "required": ["value"], + "properties": { + "value": { + "description": "unique value for this publication. Displayed as-is by auspice.", + "type": "string" + }, + "title": { + "description": "Publication title", + "type": "string" + }, + "journal": { + "description": "Journal title (including year, if applicable)", + "type": "string" + }, + "paper_url": { + "description": "URL link to paper (if available)", + "type": "string", + "pattern": "^https?://.+$" + } + } + }, + "accession": { + "description": "Sequence accession number", + "$comment": "terminal nodes only", + "type": "string", + "pattern": "^[0-9A-Za-z-_.]+$" + } + }, + "patternProperties": { + "(?!div|num_date|vaccine|hidden|url|author|accession)(^.+$)": { + "description": "coloring / geo resolution information attached to this node", + "$comment": "property name must match the `key` property provided as a coloring / geo_resolution", + "$comment": "JSON schema defined properties must also validate agaisnt patternProperties if they match. Hence the complex regex.", + "type": "object", + "required": ["value"], + "properties": { + "value": { + "type": ["string", "number", "boolean"] + }, + "url": { + "description": "URL for the property value.", + "$comment": "If a url is specified, then the value is rendered as a link in the on-click panel.", + "type": "string", + "pattern": "^https?://.+$" + }, + "confidence": { + "description": "Confidence of the trait date", + "$comment": "Should we use different keys for the two structures here?", + "oneOf": [ + { + "type": "array", + "$comment": "Only used if the value is numeric", + "items": [ + {"type": "number"}, + {"type": "number"} + ] + }, + { + "type": "object", + "$comment": "Only used if the value is a string", + "$comment": "alternative possibilities & their confidence values", + "patternProperties": { + "^.+$": { + "type": "number", + "minimum": 0.0, + "maximum": 1.0 + } + } + } + ] + }, + "entropy": { + "$comment": "Auspice uses this to control opacity of the color-by", + "type": "number" + } + } + }, + "^none$": false + } + }, + "branch_attrs": { + "description": "attributes associated with the branch from the parent node to this node, such as branch lengths, mutations, support values", + "type": "object", + "properties": { + "labels": { + "description": "Node labels", + "$comment": "Auspice scans this to generate the branch labels dropdown", + "patternProperties": { + "^[a-zA-Z0-9]+$": { + "$comment": "e.g. clade->3c3a", + "$comment": "string is parsed unchanged by Auspice", + "type": "string" + }, + "^none$": false + } + }, + "mutations": { + "description": "Mutations on the branch leading to this node. 1-based numbering (same as genome_annotations)", + "type": "object", + "additionalProperties": false, + "properties": { + "nuc": { + "description": "nucleotide mutations", + "type": "array", + "items": { + "type": "string", + "pattern": "^[ATCGNYRWSKMDVHB-][0-9]+[ATCGNYRWSKMDVHB-]$" + } + } + }, + "patternProperties": { + "^(?!nuc$)[a-zA-Z0-9*_.()-]+$": { + "description": "Amino acid mutations for this CDS", + "$comment": "properties must exist in the meta.JSON -> annotation object", + "type": "array", + "items": { + "pattern": "^[A-Z*-][0-9]+[A-Z*-]$" + } + } + } + } + } + }, + "children": { + "description": "Child nodes. Recursive structure. Terminal nodes do not have this property.", + "$comment": "Polytomies (more than 2 items) allowed, as are nodes with a single child.", + "type": "array", + "minItems": 1, + "items": {"$ref": "#/$defs/tree"} + } + } + } + } +} diff --git a/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs b/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs index 03095834c..761d8ae0c 100644 --- a/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs +++ b/packages/treetime/src/commands/shared/__tests__/test_tree_output.rs @@ -1,124 +1,345 @@ #[cfg(test)] mod tests { + use crate::alphabet::alphabet::{Alphabet, AlphabetName}; + use crate::ancestral::pipeline::AncestralPartition; + use crate::commands::ancestral::aa_node_data::AaNodeData; + use crate::commands::ancestral::result::AncestralGraphData; use crate::commands::shared::tree_output::{ - TreeOutputAdapter, TreeOutputData, TreeOutputEdge, TreeOutputMutation, TreeOutputNode, TreeOutputTrait, - format_number, + ancestral_to_auspice, ancestral_to_mat, ancestral_to_phyloxml, clock_to_auspice, clock_to_mat, clock_to_phyloxml, + format_number, mat_mutation, mugration_to_auspice, mugration_to_mat, mugration_to_phyloxml, optimize_to_auspice, + optimize_to_mat, optimize_to_phyloxml, prune_to_auspice, prune_to_mat, prune_to_phyloxml, timetree_to_auspice, + timetree_to_mat, timetree_to_phyloxml, write_ancestral_tree_outputs, }; + use crate::gtr::get_gtr::GtrModelName; + use crate::partition::fitch::PartitionFitch; + use crate::partition::sparse::{SparseEdgePartition, SparseNodePartition}; + use crate::partition::traits::BranchTopology; use crate::payload::ancestral::GraphAncestral; + use crate::seq::indel::InDel; + use crate::seq::mutation::{Mutation, MutationEvent, MutationTrack, Sub}; use approx::assert_ulps_eq; use eyre::Report; - use itertools::Itertools; + use maplit::btreemap; + use parking_lot::RwLock; use pretty_assertions::assert_eq; - use std::collections::BTreeMap; - use treetime_graph::edge::{GraphEdge, GraphEdgeKey, HasBranchLength}; - use treetime_graph::graph::Graph; - use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; - use treetime_graph::topology_order::TopologyOrderSpec; - use treetime_io::auspice::{auspice_from_graph, auspice_write_str}; - use treetime_io::auspice_types::AuspiceTree; - use treetime_io::nwk::{EdgeToNwk, NodeToNwk, nwk_read_str}; - use treetime_io::phyloxml::{phyloxml_from_graph, phyloxml_write_str}; - use treetime_io::usher_mat::usher_from_graph; - use treetime_primitives::AsciiChar; + use rstest::rstest; + use serde_json::Value; + use std::sync::Arc; + use tempfile::TempDir; + use treetime_graph::node::{GraphNodeKey, Named}; + use treetime_io::graph::TreeWriteKind; + use treetime_io::nwk::{CommentProviders, NwkStyle, nwk_read_str}; + use treetime_primitives::{AsciiChar, Seq}; + use treetime_utils::io::json::{JsonPretty, json_read_file, json_read_str, json_write_str}; #[test] - fn test_tree_output_all_schema_adapters_use_final_graph_order() -> Result<(), Report> { - let mut graph: GraphAncestral = nwk_read_str("((A,B,C)large,D)root;")?; - TopologyOrderSpec::descendant_count(false).apply(&mut graph)?; - - let auspice = auspice_from_graph::(&graph)?; - let auspice_children = auspice - .tree - .children + fn test_tree_output_ancestral_models_preserve_semantics() -> Result<(), Report> { + let graph = helpers::ancestral_graph(helpers::Mutations::NucleotideSubstitution)?; + + let auspice = ancestral_to_auspice(&graph, "2026-07-19")?; + let child = helpers::auspice_child(&auspice, "A"); + assert_eq!(Some("2026-07-19"), auspice.data.meta.updated.as_deref()); + assert_eq!(vec!["tree".to_owned()], auspice.data.meta.panels); + assert_eq!(Some(0.5), child.node_attrs.div); + assert_eq!(vec!["A1T".to_owned()], child.branch_attrs.mutations["nuc"]); + + let phyloxml = ancestral_to_phyloxml(&graph)?; + let child = helpers::phyloxml_child(&phyloxml, "A"); + assert_eq!(Some(0.5), child.branch_length_elem); + assert_eq!(Some(0.9), child.confidence.first().map(|confidence| confidence.value)); + assert!( + child + .property + .iter() + .any(|property| { property.ref_ == "treetime:mutation" && property.value == "nuc:sub:A1T" }) + ); + assert_eq!( + Some("TCG"), + child + .sequence + .iter() + .find(|sequence| sequence.name.as_deref() == Some("nuc")) + .and_then(|sequence| sequence.mol_seq.as_ref()) + .map(|sequence| sequence.sequence.as_str()) + ); + + let mat = ancestral_to_mat(&graph)?; + let mutation = mat + .node_mutations .iter() - .map(|child| child.name.as_str()) - .collect_vec(); + .flat_map(|mutations| &mutations.mutation) + .next() + .expect("fixture must contain one MAT mutation"); + assert_eq!(1, mutation.position); + assert_eq!(0, mutation.ref_nuc); + assert_eq!(0, mutation.par_nuc); + assert_eq!(vec![3], mutation.mut_nuc); - let phyloxml = phyloxml_from_graph::(&graph)?; - let phyloxml_children = phyloxml.phylogeny[0] - .clade - .as_ref() - .expect("PhyloXML must contain a root clade") - .clade + Ok(()) + } + + #[test] + fn test_tree_output_phyloxml_encodes_aa_track_and_grouped_indel() -> Result<(), Report> { + let graph = helpers::ancestral_graph(helpers::Mutations::IndelAndAminoAcid)?; + + let phyloxml = ancestral_to_phyloxml(&graph)?; + let child = helpers::phyloxml_child(&phyloxml, "A"); + let properties = child + .property .iter() - .map(|child| child.name.as_deref().unwrap_or_default()) - .collect_vec(); + .map(|property| property.value.as_str()) + .collect::>(); + assert!(properties.contains(&"nuc:del:2-3:CG")); + assert!(properties.contains(&"aa:S%2F1%3Aweird:sub:A2T")); + + let graph = helpers::ancestral_graph(helpers::Mutations::Indel)?; + let auspice = ancestral_to_auspice(&graph, "2026-07-19")?; + let child = helpers::auspice_child(&auspice, "A"); + assert_eq!( + vec!["C2-".to_owned(), "G3-".to_owned()], + child.branch_attrs.mutations["nuc"] + ); + + let graph = helpers::ancestral_graph(helpers::Mutations::AminoAcid)?; + let auspice = ancestral_to_auspice(&graph, "2026-07-19")?; + let child = helpers::auspice_child(&auspice, "A"); + assert_eq!(vec!["A2T".to_owned()], child.branch_attrs.mutations["S"]); + let annotations = auspice + .data + .meta + .genome_annotations + .as_ref() + .expect("AA fixture must have genome annotations"); + assert!(annotations.nuc.is_some()); + assert!(annotations.cdses.contains_key("S")); + + Ok(()) + } - let usher = usher_from_graph::(&graph)?; - let usher_graph: GraphAncestral = nwk_read_str(&usher.newick)?; - let usher_children = root_child_names(&usher_graph)?; + #[test] + fn test_tree_output_mat_rejects_unsupported_events() -> Result<(), Report> { + let graph = helpers::ancestral_graph(helpers::Mutations::Indel)?; + let error = ancestral_to_mat(&graph).expect_err("MAT must reject indels"); + assert!(error.to_string().contains("insertion or deletion")); - let expected = vec!["D", "large"]; - assert_eq!(expected, auspice_children); - assert_eq!(expected, phyloxml_children); - assert_eq!(expected, usher_children); + let graph = helpers::ancestral_graph(helpers::Mutations::AminoAcid)?; + let error = ancestral_to_mat(&graph).expect_err("MAT must reject amino-acid mutations"); + assert!(error.to_string().contains("amino-acid mutation")); Ok(()) } #[test] - fn test_tree_output_adapter_preserves_supported_graph_semantics() -> Result<(), Report> { - let graph = helpers::semantic_graph()?; - - let json = auspice_write_str::(&graph)?; - let auspice: AuspiceTree = serde_json::from_str(&json)?; - let child_a = auspice.tree.children.iter().find(|child| child.name == "A").unwrap(); - let child_b = auspice.tree.children.iter().find(|child| child.name == "B").unwrap(); - assert_eq!(Some(0.5), child_a.node_attrs.div); - assert_eq!(vec!["A1T".to_owned()], child_a.branch_attrs.mutations["nuc"]); - let num_date = child_a.node_attrs.num_date.as_ref().unwrap(); - assert_ulps_eq!(2020.5, num_date.value, max_ulps = 0); - assert_eq!(Some([2020.2, 2020.8]), num_date.confidence); - assert_eq!("Yes", child_b.node_attrs.bad_branch.as_ref().unwrap().value); - assert_eq!("No", auspice.tree.node_attrs.bad_branch.as_ref().unwrap().value); - let coloring_keys = auspice - .data - .meta - .colorings - .iter() - .map(|coloring| coloring.key.as_str()) - .collect_vec(); - assert!(coloring_keys.contains(&"num_date")); - assert!(coloring_keys.contains(&"bad_branch")); - assert!(coloring_keys.contains(&"country")); - assert!(coloring_keys.contains(&"gt")); - - let xml = phyloxml_write_str::(&graph)?; - assert!(xml.contains("A")); - assert!(xml.contains("0.5")); - assert!(xml.contains("treetime:divergence")); - assert!(xml.contains("treetime:mutation")); - - let usher = usher_from_graph::(&graph)?; - let mutations = usher - .node_mutations - .iter() - .flat_map(|mutations| &mutations.mutation) - .collect_vec(); - assert_eq!(1, mutations.len()); - assert_eq!(1, mutations[0].position); - assert_eq!(0, mutations[0].ref_nuc); - assert_eq!(0, mutations[0].par_nuc); - assert_eq!(vec![3], mutations[0].mut_nuc); + fn test_tree_output_mat_uses_one_global_reference_for_recurrent_mutations() -> Result<(), Report> { + let first = Mutation::substitution( + MutationTrack::Nucleotide, + Sub::new(helpers::c(b'A'), 0_usize, helpers::c(b'T'))?, + ); + let recurrent = Mutation::substitution( + MutationTrack::Nucleotide, + Sub::new(helpers::c(b'T'), 0_usize, helpers::c(b'C'))?, + ); + + let first = mat_mutation(&first, Some("A"), "inner")?; + let recurrent = mat_mutation(&recurrent, Some("A"), "leaf")?; + assert_eq!((0, 0, vec![3]), (first.ref_nuc, first.par_nuc, first.mut_nuc)); + assert_eq!( + (0, 3, vec![1]), + (recurrent.ref_nuc, recurrent.par_nuc, recurrent.mut_nuc) + ); + Ok(()) + } + #[test] + fn test_tree_output_mat_rejects_missing_reference() -> Result<(), Report> { + let mutation = Mutation::substitution( + MutationTrack::Nucleotide, + Sub::new(helpers::c(b'A'), 0_usize, helpers::c(b'T'))?, + ); + let error = mat_mutation(&mutation, None, "A").expect_err("MAT must require a global reference"); + assert!(error.to_string().contains("requires a root nucleotide reference")); Ok(()) } #[test] - fn test_tree_output_usher_missing_reference_uses_documented_fallback() -> Result<(), Report> { - let graph = helpers::semantic_graph_missing_root_reference()?; + fn test_tree_output_mat_rejects_reference_lookup_out_of_range() -> Result<(), Report> { + let mutation = Mutation::substitution( + MutationTrack::Nucleotide, + Sub::new(helpers::c(b'A'), 1_usize, helpers::c(b'T'))?, + ); + let error = mat_mutation(&mutation, Some("A"), "A").expect_err("MAT must check the reference length"); + assert!(error.to_string().contains("outside the root nucleotide reference")); + Ok(()) + } - let usher = usher_from_graph::(&graph)?; - let mutations = usher - .node_mutations - .iter() - .flat_map(|mutations| &mutations.mutation) - .collect_vec(); - assert_eq!(1, mutations.len()); - let mutation = mutations[0]; + #[test] + fn test_tree_output_mat_rejects_coordinate_above_i32() -> Result<(), Report> { + let position = usize::try_from(i32::MAX)?; + let mutation = Mutation::substitution( + MutationTrack::Nucleotide, + Sub::new(helpers::c(b'A'), position, helpers::c(b'T'))?, + ); + let error = mat_mutation(&mutation, Some("A"), "A").expect_err("MAT must check its coordinate range"); + assert!(error.to_string().contains("exceeds the UShER MAT i32 coordinate range")); + Ok(()) + } - assert_eq!(1, mutation.ref_nuc); - assert_eq!(1, mutation.par_nuc); + #[rustfmt::skip] + #[rstest] + #[case::root_reference(("N", b'A', b'T'), "root reference nucleotide 'N'")] + #[case::parent( ("A", b'N', b'T'), "parent nucleotide 'N'")] + #[case::child( ("A", b'A', b'N'), "child nucleotide 'N'")] + #[trace] + fn test_tree_output_mat_rejects_noncanonical_nucleotide( + #[case] (reference, parent, child): (&str, u8, u8), + #[case] expected: &str, + ) -> Result<(), Report> { + let mutation = Mutation::substitution( + MutationTrack::Nucleotide, + Sub::new(helpers::c(parent), 0_usize, helpers::c(child))?, + ); + let error = mat_mutation(&mutation, Some(reference), "A").expect_err("MAT must accept only A, C, G, or T"); + assert!(error.to_string().contains(expected)); + Ok(()) + } + + #[test] + fn test_tree_output_conversion_failure_does_not_create_target_and_keeps_prior_file() -> Result<(), Report> { + let graph = helpers::ancestral_graph(helpers::Mutations::Indel)?; + let dir = TempDir::new()?; + let nwk_path = dir.path().join("tree.nwk"); + let mat_path = dir.path().join("tree.mat.json"); + let outputs = btreemap! { + TreeWriteKind::nwk(NwkStyle::Plain) => nwk_path.clone(), + TreeWriteKind::MatJson => mat_path.clone(), + }; + + let error = + write_ancestral_tree_outputs(&graph, &outputs, &CommentProviders::new()).expect_err("MAT conversion must fail"); + assert!(error.to_string().contains("insertion or deletion")); + assert!(nwk_path.is_file()); + assert!(!mat_path.exists()); + + Ok(()) + } + + #[test] + fn test_tree_output_graph_json_dumps_concrete_graph_data() -> Result<(), Report> { + let graph = helpers::ancestral_graph(helpers::Mutations::None)?; + let dir = TempDir::new()?; + let path = dir.path().join("graph.json"); + let outputs = btreemap! { TreeWriteKind::GraphJson => path.clone() }; + + write_ancestral_tree_outputs(&graph, &outputs, &CommentProviders::new())?; + let actual: Value = json_read_file(&path)?; + assert_eq!(Value::String("JC69".to_owned()), actual["data"]["model_name"]); + assert_eq!(Value::Array(vec![Value::Bool(false); 3]), actual["data"]["mask"]); + assert!(actual["data"]["partition"]["Fitch"].is_object()); + assert_eq!(Value::from(0.0), helpers::edge_branch_length(&actual, "B")); + + Ok(()) + } + + #[test] + fn test_tree_output_mutation_free_mat_needs_no_reference() -> Result<(), Report> { + let graph = helpers::ancestral_graph_without_partition()?; + let mat = ancestral_to_mat(&graph)?; + assert!(mat.node_mutations.iter().all(|mutations| mutations.mutation.is_empty())); + Ok(()) + } + + // Oracle: Nextstrain Augur's dataset v2 schema at + // d8e38736037ba9474a809f9a5a63bc2b279d2407. + #[test] + fn test_tree_output_all_auspice_models_match_augur_v2_schema() -> Result<(), Report> { + let documents = helpers::all_auspice_documents()?; + let validator = helpers::auspice_validator()?; + + for (command, document) in ["ancestral", "optimize", "prune", "clock", "mugration", "timetree"] + .into_iter() + .zip(&documents) + { + let errors = validator + .iter_errors(document) + .map(|error| error.to_string()) + .collect::>() + .join("\n"); + assert!(errors.is_empty(), "{command} Auspice schema errors:\n{errors}"); + } + + let mut malformed = documents[0].clone(); + malformed["meta"] + .as_object_mut() + .expect("Auspice meta must be an object") + .remove("updated"); + assert!(!validator.is_valid(&malformed)); + + Ok(()) + } + + #[test] + fn test_tree_output_auspice_rejects_node_without_divergence_or_date() -> Result<(), Report> { + let error = helpers::optimize_auspice_without_required_node_data() + .expect_err("Auspice must reject a node with neither divergence nor numerical date"); + assert!(error.to_string().contains("requires divergence or numerical date")); + Ok(()) + } + + #[test] + fn test_tree_output_auspice_rejects_invalid_amino_acid_track_name() -> Result<(), Report> { + let graph = helpers::ancestral_graph(helpers::Mutations::IndelAndAminoAcid)?; + let error = ancestral_to_auspice(&graph, "2026-07-19") + .expect_err("Auspice must reject an amino-acid track outside its schema grammar"); + assert!(error.to_string().contains("cannot represent amino-acid mutation track")); + Ok(()) + } + + #[test] + fn test_tree_output_all_phyloxml_models_have_one_rooted_phylogeny() -> Result<(), Report> { + let documents = helpers::all_phyloxml_documents()?; + assert_eq!(6, documents.len()); + assert!(documents.iter().all(|document| { + document.phylogeny.len() == 1 && document.phylogeny[0].rooted && document.phylogeny[0].clade.is_some() + })); + Ok(()) + } + + #[test] + fn test_tree_output_all_mat_models_preserve_embedded_newick_lengths() -> Result<(), Report> { + let documents = helpers::all_mat_documents()?; + assert_eq!(6, documents.len()); + assert!(documents.iter().all(|document| { + document.node_mutations.len() == 4 + && document + .node_mutations + .iter() + .all(|mutations| mutations.mutation.is_empty()) + })); + + for (command, document) in ["ancestral", "optimize", "prune", "clock", "mugration", "timetree"] + .into_iter() + .zip(documents) + { + let graph: GraphAncestral = nwk_read_str(&document.newick)?; + assert_eq!( + None, + helpers::branch_length(&graph, "A")?, + "{command}: {}", + document.newick + ); + assert_eq!( + Some(0.0), + helpers::branch_length(&graph, "B")?, + "{command}: {}", + document.newick + ); + assert_eq!( + Some(0.5), + helpers::branch_length(&graph, "C")?, + "{command}: {}", + document.newick + ); + } Ok(()) } @@ -133,202 +354,447 @@ mod tests { assert_ulps_eq!(2020.123, format_number(2020.1234567, 3), max_ulps = 0); } - fn root_child_names(graph: &Graph) -> Result, Report> - where - N: GraphNode + Named, - E: GraphEdge, - D: Send + Sync, - { - let root = graph.get_exactly_one_root()?; - Ok( - graph - .children_of(&root.read_arc()) - .into_iter() - .map(|(child, _)| { - child - .read_arc() - .payload() - .read_arc() - .name() - .expect("fixture child must have a name") - .as_ref() - .to_owned() - }) - .collect(), - ) - } - mod helpers { use super::*; + use crate::clock::clock_graph::GraphClock; + use crate::clock::clock_model::ClockModel; + use crate::commands::clock::run::ClockGraphData; + use crate::commands::optimize::result::OptimizeGraphData; + use crate::commands::prune::result::PruneGraphData; + use crate::commands::timetree::result::TimetreeGraphData; + use crate::gtr::get_gtr::{JC69Params, jc69}; + use crate::gtr::gtr::{GTR, GTRParams}; + use crate::mugration::result::{MugrationGraphData, MugrationResult}; + use crate::partition::dense::{DenseNodePartition, DenseSeqDistribution, DenseSeqInfo}; + use crate::partition::discrete_states::DiscreteStates; + use crate::partition::marginal_discrete::PartitionMarginalDiscrete; + use crate::partition::timetree::GraphTimetree; + use crate::payload::clock_set::ClockSet; + use jsonschema::{Retrieve, Uri, Validator}; + use ndarray::array; + use serde::Serialize; + use std::collections::BTreeMap; + use std::error::Error as StdError; + use std::io; + use treetime_graph::edge::{GraphEdge, HasBranchLength}; + use treetime_graph::graph::Graph; + use treetime_graph::node::GraphNode; + use treetime_io::auspice_types::{AuspiceTree, AuspiceTreeNode}; + use treetime_io::phyloxml::{Phyloxml, PhyloxmlClade}; + use treetime_io::usher_mat::UsherTree; + use util_augur_node_data_json::AugurNodeDataJsonAnnotationEntry; + + const AUSPICE_SCHEMA: &str = include_str!("schemas/auspice/schema-export-v2.json"); + const AUSPICE_CONFIG_SCHEMA: &str = include_str!("schemas/auspice/schema-auspice-config-v2.json"); + const ANNOTATIONS_SCHEMA: &str = include_str!("schemas/auspice/schema-annotations.json"); + const ROOT_SEQUENCE_SCHEMA: &str = include_str!("schemas/auspice/schema-export-root-sequence.json"); + const MODEL_TREE: &str = "(A:0.1,B:0,C:0.5)root;"; + + #[derive(Clone, Copy)] + pub enum Mutations { + None, + NucleotideSubstitution, + Indel, + AminoAcid, + IndelAndAminoAcid, + } - #[derive(Debug)] - pub struct TestNode { - name: String, - div: f64, - date: Option, - bad_branch: bool, + pub fn ancestral_graph(mutations: Mutations) -> Result, Report> { + let graph: GraphAncestral = nwk_read_str("(A:0.5,B:0)root;")?; + let root_key = node_key(&graph, "root"); + let a_key = node_key(&graph, "A"); + let b_key = node_key(&graph, "B"); + graph + .get_node(a_key) + .unwrap() + .write_arc() + .payload() + .write_arc() + .confidence = Some(0.9); + let a_edge = graph.node_parent(a_key)?.unwrap().1; + let b_edge = graph.node_parent(b_key)?.unwrap().1; + let alphabet = Alphabet::new(AlphabetName::Nuc)?; + let root_sequence = Seq::try_from_str("ACG")?; + let a_sequence = Seq::try_from_str("TCG")?; + let b_sequence = root_sequence.clone(); + + let include_substitution = matches!(mutations, Mutations::NucleotideSubstitution); + let include_indel = matches!(mutations, Mutations::Indel | Mutations::IndelAndAminoAcid); + let include_aa = matches!(mutations, Mutations::AminoAcid | Mutations::IndelAndAminoAcid); + let mut a_edge_data = if include_substitution { + SparseEdgePartition::with_fitch_subs(vec![Sub::new(c(b'A'), 0_usize, c(b'T'))?]) + } else { + SparseEdgePartition::default() + }; + if include_indel { + a_edge_data.indels = vec![InDel::del((1, 3), Seq::try_from_str("CG")?)?]; + } + let partition = PartitionFitch { + index: 0, + alphabet: alphabet.clone(), + length: 3, + nodes: btreemap! { + root_key => SparseNodePartition::new(&root_sequence, &alphabet)?, + a_key => SparseNodePartition::new(&a_sequence, &alphabet)?, + b_key => SparseNodePartition::new(&b_sequence, &alphabet)?, + }, + edges: btreemap! { + a_edge => a_edge_data, + b_edge => SparseEdgePartition::default(), + }, + }; + + let aa_node_data = include_aa.then(|| { + let mut aa = AaNodeData::default(); + let track = if matches!(mutations, Mutations::IndelAndAminoAcid) { + "S/1:weird" + } else { + "S" + }; + aa.annotations.insert( + "S".to_owned(), + AugurNodeDataJsonAnnotationEntry { + start: Some(1), + end: Some(3), + strand: Some("+".to_owned()), + entry_type: Some("CDS".to_owned()), + ..AugurNodeDataJsonAnnotationEntry::default() + }, + ); + aa.root_aa_sequences.insert(track.to_owned(), "AA".to_owned()); + aa.node_aa_mutations.insert( + a_key, + btreemap! { + track.to_owned() => vec![MutationEvent::Substitution( + Sub::new(c(b'A'), 1_usize, c(b'T')).unwrap(), + )], + }, + ); + aa + }); + let data = AncestralGraphData::new( + Some(AncestralPartition::Fitch(Arc::new(RwLock::new(partition)))), + None, + GtrModelName::JC69, + vec![false; 3], + aa_node_data, + ); + Ok(graph.map_data(data)) } - impl GraphNode for TestNode {} + pub fn ancestral_graph_without_partition() -> Result, Report> { + let graph: GraphAncestral = nwk_read_str(MODEL_TREE)?; + Ok(graph.map_data(AncestralGraphData::new(None, None, GtrModelName::JC69, vec![], None))) + } - impl Named for TestNode { - fn name(&self) -> Option> { - Some(&self.name) - } + pub fn all_auspice_documents() -> Result, Report> { + let ancestral = ancestral_to_auspice(&ancestral_graph(Mutations::NucleotideSubstitution)?, "2026-07-19")?; + let optimize = optimize_to_auspice(&optimize_graph()?, "2026-07-19")?; + let prune = prune_to_auspice(&prune_graph()?, "2026-07-19")?; + let clock = clock_to_auspice(&clock_graph()?, "2026-07-19")?; + let mugration = mugration_to_auspice(&mugration_graph()?, "2026-07-19")?; + let timetree = timetree_to_auspice(&timetree_graph()?, "2026-07-19")?; + + [ancestral, optimize, prune, clock, mugration, timetree] + .iter() + .map(json_value) + .collect() + } - fn set_name(&mut self, name: Option>) { - self.name = name.map_or_else(String::new, |name| name.as_ref().to_owned()); - } + pub fn all_phyloxml_documents() -> Result, Report> { + Ok(vec![ + ancestral_to_phyloxml(&ancestral_graph_without_partition()?)?, + optimize_to_phyloxml(&optimize_graph()?)?, + prune_to_phyloxml(&prune_graph()?)?, + clock_to_phyloxml(&clock_graph()?)?, + mugration_to_phyloxml(&mugration_graph()?)?, + timetree_to_phyloxml(&timetree_graph()?)?, + ]) } - impl NodeToNwk for TestNode { - fn nwk_name(&self) -> Option> { - Some(&self.name) - } + pub fn all_mat_documents() -> Result, Report> { + let ancestral = ancestral_graph_without_partition()?; + set_mat_branch_lengths(&ancestral)?; + let optimize = optimize_graph()?; + set_mat_branch_lengths(&optimize)?; + let prune = prune_graph()?; + set_mat_branch_lengths(&prune)?; + let clock = clock_graph()?; + set_mat_branch_lengths(&clock)?; + let mugration = mugration_graph()?; + set_mat_branch_lengths(&mugration)?; + let timetree = timetree_graph()?; + set_timetree_mat_branch_lengths(&timetree)?; + + Ok(vec![ + ancestral_to_mat(&ancestral)?, + optimize_to_mat(&optimize)?, + prune_to_mat(&prune)?, + clock_to_mat(&clock)?, + mugration_to_mat(&mugration)?, + timetree_to_mat(&timetree)?, + ]) } - impl TreeOutputNode for TestNode { - fn output_divergence(&self) -> Option { - Some(self.div) - } + pub fn auspice_validator() -> Result { + let schema = json_read_str(AUSPICE_SCHEMA)?; + Ok( + jsonschema::draft6::options() + .with_retriever(AuspiceSchemaRetriever::new()?) + .build(&schema)?, + ) + } - fn output_date(&self) -> Option { - self.date - } + pub fn optimize_auspice_without_required_node_data() -> Result { + let graph = optimize_graph()?; + set_branch_length(&graph, "A", None)?; + optimize_to_auspice(&graph, "2026-07-19") + } - fn output_bad_branch(&self) -> bool { - self.bad_branch - } + pub fn auspice_child<'a>(tree: &'a AuspiceTree, name: &str) -> &'a AuspiceTreeNode { + tree + .tree + .children + .iter() + .find(|child| child.name == name) + .expect("fixture child must exist") } - #[derive(Debug)] - pub struct TestEdge { - branch_length: f64, + pub fn phyloxml_child<'a>(tree: &'a Phyloxml, name: &str) -> &'a PhyloxmlClade { + tree.phylogeny[0] + .clade + .as_ref() + .expect("PhyloXML fixture must have a root") + .clade + .iter() + .find(|child| child.name.as_deref() == Some(name)) + .expect("fixture child must exist") } - impl GraphEdge for TestEdge {} + pub fn edge_branch_length(graph: &Value, target_name: &str) -> Value { + let target_key = graph["nodes"] + .as_array() + .unwrap() + .iter() + .position(|node| node["data"]["name"] == target_name) + .unwrap(); + graph["edges"] + .as_array() + .unwrap() + .iter() + .find(|edge| edge["target"] == target_key) + .unwrap()["data"]["branch_length"] + .clone() + } - impl HasBranchLength for TestEdge { - fn branch_length(&self) -> Option { - Some(self.branch_length) - } + pub fn branch_length(graph: &GraphAncestral, target_name: &str) -> Result, Report> { + let node_key = node_key(graph, target_name); + let edge_key = graph + .node_parent(node_key)? + .expect("fixture target must not be the root") + .1; + Ok( + graph + .get_edge(edge_key) + .expect("fixture edge must exist") + .read_arc() + .payload() + .read_arc() + .branch_length(), + ) + } - fn set_branch_length(&mut self, branch_length: Option) { - if let Some(branch_length) = branch_length { - self.branch_length = branch_length; - } - } + fn node_key(graph: &GraphAncestral, name: &str) -> GraphNodeKey { + graph + .get_nodes() + .into_iter() + .find_map(|node| { + let node = node.read_arc(); + let payload = node.payload().read_arc(); + (payload.name().as_ref().map(AsRef::as_ref) == Some(name)).then(|| node.key()) + }) + .expect("fixture node must exist") } - impl EdgeToNwk for TestEdge { - fn nwk_weight(&self) -> Option { - Some(self.branch_length) - } + pub fn c(value: u8) -> AsciiChar { + AsciiChar::from_byte_unchecked(value) } - impl TreeOutputEdge for TestEdge {} + fn fixed_clock_model() -> Result { + ClockModel::with_fixed_rate(&ClockSet::leaf_contribution(Some(2020.0)), 1.0) + } - #[derive(Debug, Default)] - pub struct TestData { - date_confidence: BTreeMap, - mutations: BTreeMap>, - root_sequences: BTreeMap, - traits: BTreeMap>, + fn optimize_graph() -> Result, Report> { + let graph: GraphAncestral = nwk_read_str(MODEL_TREE)?; + Ok(graph.map_data(OptimizeGraphData::new( + jc69(JC69Params::default())?, + GtrModelName::JC69, + vec![], + vec![], + ))) } - impl TreeOutputData for TestData { - fn trait_attributes(&self) -> Vec { - vec!["country".to_owned()] - } + fn prune_graph() -> Result, Report> { + let graph: GraphAncestral = nwk_read_str(MODEL_TREE)?; + Ok(graph.map_data(PruneGraphData::new(Some(jc69(JC69Params::default())?), vec![]))) + } - fn has_dates(&self) -> bool { - true + fn clock_graph() -> Result, Report> { + let graph: GraphClock = nwk_read_str(MODEL_TREE)?; + for (index, node) in graph.get_nodes().into_iter().enumerate() { + let node = node.write_arc(); + let mut payload = node.payload().write_arc(); + payload.div = index as f64 / 2.0; + payload.time = Some(2020.0 + index as f64); } + Ok(graph.map_data(ClockGraphData::new(fixed_clock_model()?, vec![]))) + } - fn has_bad_branch(&self) -> bool { - true - } + fn mugration_graph() -> Result, Report> { + let graph: GraphAncestral = nwk_read_str(MODEL_TREE)?; + let states = DiscreteStates::from_values(["CH", "US"].into_iter(), "?"); + let gtr = GTR::new(GTRParams { + n_states: 2, + mu: 1.0, + W: None, + pi: array![0.5, 0.5], + })?; + let mut partition = PartitionMarginalDiscrete::new(gtr, states, 1e-8, false); + partition.data.nodes = graph + .get_nodes() + .into_iter() + .enumerate() + .map(|(index, node)| { + let key = node.read_arc().key(); + let profile = if index % 2 == 0 { + array![[1.0, 0.0]] + } else { + array![[0.0, 1.0]] + }; + ( + key, + DenseNodePartition { + seq: DenseSeqInfo::default(), + profile: DenseSeqDistribution::new(profile, 0.0), + }, + ) + }) + .collect(); + Ok(MugrationResult::new(graph, partition, "country", 0.0).graph) + } - fn has_mutations(&self) -> bool { - true + fn timetree_graph() -> Result, Report> { + let graph: GraphTimetree = nwk_read_str(MODEL_TREE)?; + for (index, node) in graph.get_nodes().into_iter().enumerate() { + let node = node.write_arc(); + let mut payload = node.payload().write_arc(); + payload.div = index as f64 / 2.0; + payload.time = Some(2020.0 + index as f64); } + Ok(graph.map_data(TimetreeGraphData::new( + fixed_clock_model()?, + None, + vec![], + None, + None, + None, + None, + ))) + } - fn root_sequences(&self, _graph: &Graph) -> Result, Report> { - Ok(self.root_sequences.clone()) - } + fn set_mat_branch_lengths(graph: &Graph) -> Result<(), Report> + where + N: GraphNode + Named, + E: GraphEdge + HasBranchLength, + D: Send + Sync, + { + set_branch_length(graph, "A", None)?; + set_branch_length(graph, "B", Some(0.0))?; + set_branch_length(graph, "C", Some(0.5)) + } - fn date_confidence(&self, key: GraphNodeKey) -> Option<[f64; 2]> { - self.date_confidence.get(&key).copied() - } + fn set_branch_length(graph: &Graph, name: &str, length: Option) -> Result<(), Report> + where + N: GraphNode + Named, + E: GraphEdge + HasBranchLength, + D: Send + Sync, + { + let key = graph + .get_nodes() + .into_iter() + .find_map(|node| { + let node = node.read_arc(); + let payload = node.payload().read_arc(); + (payload.name().as_ref().map(AsRef::as_ref) == Some(name)).then(|| node.key()) + }) + .expect("fixture node must exist"); + let edge_key = graph.node_parent(key)?.expect("fixture node must have a parent").1; + graph + .get_edge(edge_key) + .expect("fixture edge must exist") + .write_arc() + .payload() + .write_arc() + .set_branch_length(length); + Ok(()) + } - fn traits(&self, key: GraphNodeKey) -> BTreeMap { - self.traits.get(&key).cloned().unwrap_or_default() + fn set_timetree_mat_branch_lengths(graph: &GraphTimetree) -> Result<(), Report> { + for (name, length) in [("A", None), ("B", Some(0.0)), ("C", Some(0.5))] { + let key = graph + .get_nodes() + .into_iter() + .find_map(|node| { + let node = node.read_arc(); + let payload = node.payload().read_arc(); + (payload.name().as_ref().map(AsRef::as_ref) == Some(name)).then(|| node.key()) + }) + .expect("fixture node must exist"); + let edge_key = graph.node_parent(key)?.expect("fixture node must have a parent").1; + graph + .get_edge(edge_key) + .expect("fixture edge must exist") + .write_arc() + .payload() + .write_arc() + .time_length = length; } + Ok(()) + } - fn mutations( - &self, - _graph: &Graph, - key: GraphEdgeKey, - ) -> Result, Report> { - Ok(self.mutations.get(&key).cloned().unwrap_or_default()) - } + fn json_value(value: &impl Serialize) -> Result { + json_write_str(value, JsonPretty(false)).and_then(|json| json_read_str(&json)) } - pub fn semantic_graph_missing_root_reference() -> Result, Report> { - let mut graph = semantic_graph()?; - graph.data_mut().root_sequences.clear(); - let mutations = graph - .data_mut() - .mutations - .values_mut() - .next() - .expect("semantic graph must contain mutation data for one edge"); - mutations - .first_mut() - .expect("semantic graph must contain one mutation") - .parent = AsciiChar::from_byte_unchecked(b'C'); - Ok(graph) + struct AuspiceSchemaRetriever { + schemas: BTreeMap, } - pub fn semantic_graph() -> Result, Report> { - let mut graph = Graph::with_data(TestData::default()); - let root = graph.add_node(node("root", 0.0, Some(2020.0), false)); - let a = graph.add_node(node("A", 0.5, Some(2020.5), false)); - let b = graph.add_node(node("B", 1.0, None, true)); - let edge_a = graph.add_edge(root, a, TestEdge { branch_length: 0.5 })?; - graph.add_edge(root, b, TestEdge { branch_length: 1.0 })?; - graph.build()?; - - graph.data_mut().date_confidence.insert(a, [2020.2, 2020.8]); - graph.data_mut().root_sequences.insert("nuc".to_owned(), "A".to_owned()); - graph.data_mut().mutations.insert( - edge_a, - vec![TreeOutputMutation { - gene: "nuc".to_owned(), - position: 0, - parent: AsciiChar::try_new(b'A')?, - child: AsciiChar::try_new(b'T')?, - }], - ); - graph.data_mut().traits.insert( - a, - BTreeMap::from([( - "country".to_owned(), - TreeOutputTrait { - value: "CH".to_owned(), - ..TreeOutputTrait::default() - }, - )]), - ); - Ok(graph) + impl AuspiceSchemaRetriever { + fn new() -> Result { + Ok(Self { + schemas: [AUSPICE_CONFIG_SCHEMA, ANNOTATIONS_SCHEMA, ROOT_SEQUENCE_SCHEMA] + .into_iter() + .map(|schema| { + let schema: Value = json_read_str(schema)?; + let id = schema["$id"] + .as_str() + .ok_or_else(|| io::Error::new(io::ErrorKind::InvalidData, "Vendored schema has no $id"))? + .to_owned(); + Ok((id, schema)) + }) + .collect::>()?, + }) + } } - fn node(name: &str, div: f64, date: Option, bad_branch: bool) -> TestNode { - TestNode { - name: name.to_owned(), - div, - date, - bad_branch, + impl Retrieve for AuspiceSchemaRetriever { + fn retrieve(&self, uri: &Uri) -> Result> { + self + .schemas + .get(uri.as_str()) + .cloned() + .ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, format!("Schema not found: {uri}")).into()) } } } diff --git a/packages/treetime/src/commands/shared/output.rs b/packages/treetime/src/commands/shared/output.rs index c26d80633..c15423aaa 100644 --- a/packages/treetime/src/commands/shared/output.rs +++ b/packages/treetime/src/commands/shared/output.rs @@ -542,21 +542,6 @@ pub struct ResolvedOutputs { pub non_tree_outputs: BTreeMap, } -impl ResolvedOutputs { - /// Create parent directories immediately before the command starts publishing output. - pub fn prepare(&self) -> Result<(), Report> { - for path in self.tree_outputs.values().chain(self.non_tree_outputs.values()) { - if let Some(parent) = path.parent() { - if !parent.as_os_str().is_empty() { - std::fs::create_dir_all(parent) - .wrap_err_with(|| format!("Failed to create parent directory '{}'", parent.display()))?; - } - } - } - Ok(()) - } -} - impl OutputCoreArgs { /// Resolve the three-tier output configuration into concrete file paths. /// diff --git a/packages/treetime/src/commands/shared/tree_output.rs b/packages/treetime/src/commands/shared/tree_output.rs index b6dbf35cf..3980698bc 100644 --- a/packages/treetime/src/commands/shared/tree_output.rs +++ b/packages/treetime/src/commands/shared/tree_output.rs @@ -1,5 +1,5 @@ use crate::ancestral::pipeline::AncestralPartition; -use crate::clock::clock_graph::{EdgeClock, NodeClock}; +use crate::clock::clock_graph::GraphClock; use crate::commands::ancestral::result::AncestralGraphData; use crate::commands::clock::run::ClockGraphData; use crate::commands::mugration::augur_node_data::{build_confidence_map, compute_entropy}; @@ -7,31 +7,45 @@ use crate::commands::optimize::result::OptimizeGraphData; use crate::commands::prune::result::PruneGraphData; use crate::commands::timetree::result::TimetreeGraphData; use crate::mugration::result::MugrationGraphData; -use crate::partition::augur::AugurNodeDataJsonAncestralPartition; use crate::partition::traits::{BranchTopology, PartitionBranchOps}; -use crate::payload::ancestral::{EdgeAncestral, NodeAncestral}; +use crate::payload::ancestral::{EdgeAncestral, GraphAncestral, NodeAncestral}; use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; -use crate::seq::mutation::Sub; -use eyre::Report; -use log::warn; +use crate::seq::mutation::{Mutation, MutationEvent, MutationTrack, mutation_event_strings}; +use chrono::Utc; +use eyre::{Report, WrapErr}; +use maplit::{btreemap, btreeset}; +use parking_lot::RwLock; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; +use serde::Serialize; use serde_json::{Value, json}; -use std::collections::BTreeMap; -use treetime_graph::edge::{GraphEdge, GraphEdgeKey, HasBranchLength}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::path::PathBuf; +use std::sync::Arc; +use treetime_graph::edge::{Edge, GraphEdge, GraphEdgeKey, HasBranchLength}; use treetime_graph::graph::Graph; -use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; -use treetime_io::auspice::{AuspiceGraphContext, AuspiceWrite}; +use treetime_graph::node::{GraphNode, GraphNodeKey, Named, Node}; +use treetime_io::auspice::auspice_write_file; use treetime_io::auspice_types::{ - AuspiceColoring, AuspiceDisplayDefaults, AuspiceNumDate, AuspiceTreeBranchAttrs, AuspiceTreeData, AuspiceTreeMeta, - AuspiceTreeNode, AuspiceTreeNodeAttr, AuspiceTreeNodeAttrs, + AuspiceColoring, AuspiceDisplayDefaults, AuspiceGenomeAnnotationCds, AuspiceGenomeAnnotationNuc, + AuspiceGenomeAnnotations, AuspiceNumDate, AuspiceTree, AuspiceTreeBranchAttrs, AuspiceTreeBranchAttrsLabels, + AuspiceTreeData, AuspiceTreeMeta, AuspiceTreeNode, AuspiceTreeNodeAttr, AuspiceTreeNodeAttrs, Segments, StartEnd, }; +use treetime_io::graph::TreeWriteKind; +use treetime_io::graphviz::{EdgeToGraphviz, NodeToGraphviz, graphviz_write_file}; +use treetime_io::nex::{NexWriteOptions, nex_write_file_with}; +use treetime_io::nwk::{CommentProviders, EdgeToNwk, NodeToNwk, NwkWriteOptions, nwk_write_file_with, nwk_write_str}; use treetime_io::phyloxml::{ - Phyloxml, PhyloxmlClade, PhyloxmlDate, PhyloxmlFromGraph, PhyloxmlGraphContext, PhyloxmlPhylogeny, PhyloxmlProperty, + Phyloxml, PhyloxmlClade, PhyloxmlConfidence, PhyloxmlDate, PhyloxmlJsonOptions, PhyloxmlMolSeq, PhyloxmlPhylogeny, + PhyloxmlProperty, PhyloxmlSequence, phyloxml_json_write_file, phyloxml_write_file, }; use treetime_io::usher_mat::{ - UsherGraphContext, UsherMetadata, UsherMutation, UsherMutationList, UsherTreeNode, UsherWrite, + UsherMatJsonOptions, UsherMetadata, UsherMutation, UsherMutationList, UsherTree, UsherTreeNode, + usher_mat_json_write_file, usher_mat_pb_write_file, }; use treetime_primitives::AsciiChar; -use treetime_utils::make_error; +use treetime_utils::io::json::{JsonPretty, json_write_file}; +use treetime_utils::{make_error, make_internal_report, make_report}; +use util_augur_node_data_json::AugurNodeDataJsonAnnotationEntry; const APPLIES_BRANCH: &str = "parent_branch"; const APPLIES_NODE: &str = "node"; @@ -42,706 +56,1671 @@ const COLORING_NUM_DATE: &str = "num_date"; const DT_BOOLEAN: &str = "xsd:boolean"; const DT_DOUBLE: &str = "xsd:double"; const DT_STRING: &str = "xsd:string"; -const NUC_GENE: &str = "nuc"; +const NUC_TRACK: &str = "nuc"; const REF_BAD_BRANCH: &str = "treetime:bad_branch"; const REF_DATE_INFERRED: &str = "treetime:date_inferred"; const REF_DIV: &str = "treetime:divergence"; const REF_GAMMA: &str = "treetime:gamma"; const REF_MUTATION: &str = "treetime:mutation"; +const REF_TRAIT_CONFIDENCE_PREFIX: &str = "treetime:trait_confidence:"; +const REF_TRAIT_ENTROPY_PREFIX: &str = "treetime:trait_entropy:"; const REF_TRAIT_PREFIX: &str = "treetime:trait:"; +const REF_TRAIT_TRANSITION_PREFIX: &str = "treetime:trait_transition:"; +const TYPE_INPUT_BRANCH_SUPPORT: &str = "treetime:input_branch_support"; +const PROPERTY_TOKEN_ENCODE_SET: &AsciiSet = &NON_ALPHANUMERIC.remove(b'-').remove(b'.').remove(b'_').remove(b'~'); + +pub fn write_ancestral_tree_outputs( + graph: &GraphAncestral, + outputs: &BTreeMap, + providers: &CommentProviders, +) -> Result<(), Report> { + let updated = generation_date(); + write_tree_outputs( + graph, + outputs, + providers, + "ancestral", + || ancestral_to_auspice(graph, &updated), + || ancestral_to_phyloxml(graph), + || ancestral_to_mat(graph), + ) +} -pub struct TreeOutputAdapter { - has_bad_branch: bool, - warned_ambiguous: bool, - warned_non_nuc: bool, - warned_reference_fallback: bool, +pub fn write_optimize_tree_outputs( + graph: &GraphAncestral, + outputs: &BTreeMap, + providers: &CommentProviders, +) -> Result<(), Report> { + let updated = generation_date(); + write_tree_outputs( + graph, + outputs, + providers, + "optimize", + || optimize_to_auspice(graph, &updated), + || optimize_to_phyloxml(graph), + || optimize_to_mat(graph), + ) } -pub trait TreeOutputNode: GraphNode + Named { - fn output_divergence(&self) -> Option { - None - } +pub fn write_prune_tree_outputs( + graph: &GraphAncestral, + outputs: &BTreeMap, + providers: &CommentProviders, +) -> Result<(), Report> { + let updated = generation_date(); + write_tree_outputs( + graph, + outputs, + providers, + "prune", + || prune_to_auspice(graph, &updated), + || prune_to_phyloxml(graph), + || prune_to_mat(graph), + ) +} - fn output_date(&self) -> Option { - None - } +pub fn write_clock_tree_outputs( + graph: &GraphClock, + outputs: &BTreeMap, + providers: &CommentProviders, +) -> Result<(), Report> { + let updated = generation_date(); + write_tree_outputs( + graph, + outputs, + providers, + "clock", + || clock_to_auspice(graph, &updated), + || clock_to_phyloxml(graph), + || clock_to_mat(graph), + ) +} - fn output_bad_branch(&self) -> bool { - false - } +pub fn write_mugration_tree_outputs( + graph: &GraphAncestral, + outputs: &BTreeMap, + providers: &CommentProviders, +) -> Result<(), Report> { + let updated = generation_date(); + write_tree_outputs( + graph, + outputs, + providers, + "mugration", + || mugration_to_auspice(graph, &updated), + || mugration_to_phyloxml(graph), + || mugration_to_mat(graph), + ) } -pub trait TreeOutputEdge: GraphEdge + HasBranchLength { - fn output_gamma(&self) -> Option { - None - } +pub fn write_timetree_tree_outputs( + graph: &Graph, + outputs: &BTreeMap, + providers: &CommentProviders, +) -> Result<(), Report> { + let updated = generation_date(); + write_tree_outputs( + graph, + outputs, + providers, + "timetree", + || timetree_to_auspice(graph, &updated), + || timetree_to_phyloxml(graph), + || timetree_to_mat(graph), + ) } -pub trait TreeOutputData: Send + Sync +fn write_tree_outputs( + graph: &Graph, + outputs: &BTreeMap, + providers: &CommentProviders, + command: &str, + to_auspice: A, + to_phyloxml: P, + to_mat: M, +) -> Result<(), Report> where - N: TreeOutputNode, - E: TreeOutputEdge, + N: GraphNode + Named + NodeToNwk + NodeToGraphviz + Serialize, + E: GraphEdge + HasBranchLength + EdgeToNwk + EdgeToGraphviz + Serialize, + D: Send + Sync + Serialize, + A: Fn() -> Result, + P: Fn() -> Result, + M: Fn() -> Result, { - fn title(&self) -> Option<&str> { - None + for (kind, path) in outputs { + match kind { + TreeWriteKind::Nwk(spec) => { + graph + .get_exactly_one_root() + .wrap_err_with(|| format!("When converting {command} graph to Newick"))?; + nwk_write_file_with( + path, + graph, + &NwkWriteOptions { + style: spec.style, + ..NwkWriteOptions::default() + }, + providers, + )?; + }, + TreeWriteKind::Nexus(spec) => { + graph + .get_exactly_one_root() + .wrap_err_with(|| format!("When converting {command} graph to Nexus"))?; + nex_write_file_with( + path, + graph, + &NexWriteOptions { + style: spec.style, + ..NexWriteOptions::default() + }, + providers, + )?; + }, + TreeWriteKind::Auspice => auspice_write_file(path, &to_auspice()?)?, + TreeWriteKind::Phyloxml => phyloxml_write_file(path, &to_phyloxml()?)?, + TreeWriteKind::PhyloxmlJson => { + phyloxml_json_write_file(path, &to_phyloxml()?, &PhyloxmlJsonOptions::default())?; + }, + TreeWriteKind::MatPb => usher_mat_pb_write_file(path, &to_mat()?)?, + TreeWriteKind::MatJson => { + usher_mat_json_write_file(path, &to_mat()?, &UsherMatJsonOptions::default())?; + }, + TreeWriteKind::GraphJson => { + // Graph JSON is an unstable debug dump of the concrete internal structs. It has no portable schema or input contract. + json_write_file(path, graph, JsonPretty(true))?; + }, + TreeWriteKind::Dot => graphviz_write_file(path, graph)?, + } } + Ok(()) +} - fn description(&self) -> Option<&str> { - None - } +pub(crate) fn ancestral_to_auspice( + graph: &GraphAncestral, + updated: &str, +) -> Result { + let root_sequences = ancestral_root_sequences(graph)?; + let genome_annotations = ancestral_genome_annotations(graph, &root_sequences)?; + let data = auspice_data( + "TreeTime ancestral analysis", + updated, + vec![], + vec![], + None, + genome_annotations, + Some(root_sequences), + !ancestral_all_mutations(graph)?.is_empty(), + ); + auspice_from_graph(graph, data, |context| { + let mutations = ancestral_node_mutations(graph, context.node_key, context.edge_key)?; + ancestral_auspice_node(context, mutations, None, None) + }) +} - fn trait_attributes(&self) -> Vec { - vec![] - } +pub(crate) fn optimize_to_auspice( + graph: &GraphAncestral, + updated: &str, +) -> Result { + let root_sequences = optimize_root_sequences(graph)?; + let data = auspice_data( + "TreeTime optimize analysis", + updated, + vec![], + vec![], + None, + None, + Some(root_sequences), + optimize_has_mutations(graph), + ); + auspice_from_graph(graph, data, |context| { + ancestral_auspice_node(context, optimize_mutations(graph, context.edge_key)?, None, None) + }) +} - fn has_dates(&self) -> bool { - false - } +pub(crate) fn prune_to_auspice(graph: &GraphAncestral, updated: &str) -> Result { + let root_sequences = prune_root_sequences(graph)?; + let data = auspice_data( + "TreeTime prune analysis", + updated, + vec![], + vec![], + None, + None, + Some(root_sequences), + !graph.data().partitions.is_empty(), + ); + auspice_from_graph(graph, data, |context| { + ancestral_auspice_node(context, prune_mutations(graph, context.edge_key)?, None, None) + }) +} - fn has_bad_branch(&self) -> bool { - false - } +pub(crate) fn clock_to_auspice(graph: &GraphClock, updated: &str) -> Result { + let data = auspice_data( + "TreeTime clock analysis", + updated, + vec![ + coloring(COLORING_NUM_DATE, "Date", "continuous"), + coloring(COLORING_BAD_BRANCH, "Excluded", "categorical"), + ], + vec![COLORING_BAD_BRANCH.to_owned()], + Some(COLORING_BAD_BRANCH.to_owned()), + None, + None, + false, + ); + auspice_from_graph(graph, data, |context| { + let name = node_name(context.node_key, context.node); + Ok(auspice_node( + name.clone(), + finite_number(Some(context.node.div), 6, "clock", &name, "div")?, + finite_number(context.node.time, 3, "clock", &name, "date")?, + None, + Some(context.node.bad_branch || context.node.is_outlier), + BTreeMap::new(), + BTreeMap::new(), + None, + )) + }) +} - fn has_mutations(&self) -> bool { - false - } +pub(crate) fn mugration_to_auspice( + graph: &GraphAncestral, + updated: &str, +) -> Result { + let attribute = &graph.data().traits.attribute; + let data = auspice_data( + "TreeTime mugration analysis", + updated, + vec![coloring(attribute, attribute, "categorical")], + vec![attribute.clone()], + Some(attribute.clone()), + None, + None, + false, + ); + auspice_from_graph(graph, data, |context| { + let name = node_name(context.node_key, context.node); + let traits = mugration_traits(graph, context.node_key, &name)?; + Ok(auspice_node( + name.clone(), + finite_number( + cumulative_branch_length(graph, context.node_key)?, + 6, + "mugration", + &name, + "div", + )?, + None, + None, + None, + traits, + BTreeMap::new(), + mugration_transition_label(graph, context.node_key)?, + )) + }) +} - fn root_sequences(&self, _graph: &Graph) -> Result, Report> - where - Self: Sized, - { - Ok(BTreeMap::new()) - } +pub(crate) fn timetree_to_auspice( + graph: &Graph, + updated: &str, +) -> Result { + let root_sequences = timetree_root_sequences(graph)?; + let data = auspice_data( + "TreeTime timetree analysis", + updated, + vec![ + coloring(COLORING_NUM_DATE, "Date", "continuous"), + coloring(COLORING_BAD_BRANCH, "Excluded", "categorical"), + ], + vec![COLORING_BAD_BRANCH.to_owned()], + Some(COLORING_BAD_BRANCH.to_owned()), + None, + Some(root_sequences), + !graph.data().partitions.is_empty(), + ); + auspice_from_graph(graph, data, |context| { + let name = node_name(context.node_key, context.node); + let div = timetree_divergence(graph, context.node_key, context.node)?; + let confidence = timetree_date_confidence(graph, context.node_key, &name)?; + Ok(auspice_node( + name.clone(), + finite_number(Some(div), 6, "timetree", &name, "div")?, + finite_number(context.node.time, 3, "timetree", &name, "date")?, + confidence, + Some(context.node.bad_branch || context.node.is_outlier), + BTreeMap::new(), + group_mutations(timetree_mutations(graph, context.edge_key)?)?, + None, + )) + }) +} - fn divergence(&self, _graph: &Graph, _key: GraphNodeKey, node: &N) -> Result, Report> - where - Self: Sized, - { - Ok(node.output_divergence()) +fn ancestral_auspice_node( + context: &GraphNodeContext, + mutations: Vec, + date: Option, + bad_branch: Option, +) -> Result { + let name = node_name(context.node_key, context.node); + let div = cumulative_branch_length(context.graph, context.node_key)?; + let mut other = serde_json::Map::new(); + if let Some(confidence) = context.node.confidence { + ensure_finite(confidence, "tree output", &name, "input branch support")?; + other.insert("confidence".to_owned(), json!({ "value": confidence })); } + let mut node = auspice_node( + name.clone(), + finite_number(div, 6, "tree output", &name, "div")?, + finite_number(date, 3, "tree output", &name, "date")?, + None, + bad_branch, + BTreeMap::new(), + group_mutations(mutations)?, + None, + ); + if let Value::Object(target) = &mut node.node_attrs.other { + target.extend(other); + } + Ok(node) +} - fn date_confidence(&self, _key: GraphNodeKey) -> Option<[f64; 2]> { - None +fn auspice_data( + title: &str, + updated: &str, + mut colorings: Vec, + filters: Vec, + color_by: Option, + genome_annotations: Option, + root_sequences: Option>, + has_mutations: bool, +) -> AuspiceTreeData { + if has_mutations { + colorings.push(coloring(COLORING_GENOTYPE, "Genotype", "categorical")); + } + AuspiceTreeData { + version: Some("v2".to_owned()), + meta: AuspiceTreeMeta { + title: Some(title.to_owned()), + updated: Some(updated.to_owned()), + panels: vec!["tree".to_owned()], + genome_annotations, + colorings, + filters, + display_defaults: AuspiceDisplayDefaults { + color_by, + ..AuspiceDisplayDefaults::default() + }, + ..AuspiceTreeMeta::default() + }, + root_sequence: root_sequences.filter(|sequences| !sequences.is_empty()), + other: Value::default(), } +} - fn date_is_inferred(&self, _key: GraphNodeKey) -> bool { - false +fn auspice_node( + name: String, + div: Option, + date: Option, + date_confidence: Option<[f64; 2]>, + bad_branch: Option, + traits: BTreeMap, + mutations: BTreeMap>, + labels: Option, +) -> AuspiceTreeNode { + AuspiceTreeNode { + name, + branch_attrs: AuspiceTreeBranchAttrs { + mutations, + labels, + other: Value::default(), + }, + node_attrs: AuspiceTreeNodeAttrs { + div, + num_date: date.map(|value| AuspiceNumDate { + value, + confidence: date_confidence, + }), + bad_branch: bad_branch.map(|bad| AuspiceTreeNodeAttr::new(if bad { "Yes" } else { "No" })), + clade_membership: None, + region: None, + country: None, + division: None, + other: build_trait_attrs(traits), + }, + children: vec![], + other: Value::default(), } +} - fn traits(&self, _key: GraphNodeKey) -> BTreeMap { - BTreeMap::new() +fn auspice_from_graph( + graph: &Graph, + data: AuspiceTreeData, + mut convert: F, +) -> Result +where + N: GraphNode, + E: GraphEdge, + D: Send + Sync, + F: FnMut(&GraphNodeContext) -> Result, +{ + let root = graph + .get_exactly_one_root() + .wrap_err("When converting graph to Auspice v2 JSON")?; + let mut node_map = btreemap! {}; + let mut queue = VecDeque::from([(Arc::clone(&root), None)]); + while let Some((current_node, current_edge)) = queue.pop_front() { + let node_key = current_node.read_arc().key(); + let node = current_node.read_arc().payload().read_arc(); + let edge_key = current_edge + .as_ref() + .map(|edge: &Arc>>| edge.read_arc().key()); + let edge = current_edge + .as_ref() + .map(|edge: &Arc>>| edge.read_arc().payload().read_arc()); + let converted = convert(&GraphNodeContext { + node_key, + node: &node, + edge_key, + edge: edge.as_deref(), + graph, + })?; + if converted.node_attrs.div.is_none() && converted.node_attrs.num_date.is_none() { + return make_error!( + "Auspice v2 node '{}' requires divergence or numerical date data", + converted.name + ); + } + node_map.insert(node_key, converted); + for (child, edge) in graph.children_of(¤t_node.read_arc()) { + queue.push_back((child, Some(edge))); + } } + attach_auspice_children(graph, &root, &mut node_map)?; + let root_key = root.read_arc().key(); + let tree = node_map + .remove(&root_key) + .ok_or_else(|| make_internal_report!("Auspice root node {root_key} was not converted"))?; + Ok(AuspiceTree { data, tree }) +} - fn mutations(&self, _graph: &Graph, _key: GraphEdgeKey) -> Result, Report> - where - Self: Sized, - { - Ok(vec![]) +fn attach_auspice_children( + graph: &Graph, + root: &Arc>>, + node_map: &mut BTreeMap, +) -> Result<(), Report> +where + N: GraphNode, + E: GraphEdge, + D: Send + Sync, +{ + let mut visited = btreeset! {}; + let mut stack = vec![Arc::clone(root)]; + while let Some(node) = stack.pop() { + let key = node.read_arc().key(); + if visited.contains(&key) { + let children = graph + .children_of(&node.read_arc()) + .into_iter() + .map(|(child, _)| { + let child_key = child.read_arc().key(); + node_map + .remove(&child_key) + .ok_or_else(|| make_internal_report!("Auspice child node {child_key} was not converted")) + }) + .collect::, _>>()?; + node_map + .get_mut(&key) + .ok_or_else(|| make_internal_report!("Auspice parent node {key} was not converted"))? + .children = children; + } else { + visited.insert(key); + stack.push(Arc::clone(&node)); + stack.extend(graph.children_of(&node.read_arc()).into_iter().map(|(child, _)| child)); + } } + Ok(()) } -#[derive(Clone, Debug, Eq, PartialEq)] -pub struct TreeOutputMutation { - pub gene: String, - pub position: usize, - pub parent: AsciiChar, - pub child: AsciiChar, +struct GraphNodeContext<'a, N, E, D> +where + N: GraphNode, + E: GraphEdge, + D: Send + Sync, +{ + node_key: GraphNodeKey, + node: &'a N, + edge_key: Option, + edge: Option<&'a E>, + graph: &'a Graph, } #[derive(Clone, Debug, Default, PartialEq)] -pub struct TreeOutputTrait { - pub value: String, - pub confidence: BTreeMap, - pub entropy: Option, +struct TraitValue { + value: String, + confidence: BTreeMap, + entropy: Option, } -impl TreeOutputNode for NodeAncestral {} +pub(crate) fn ancestral_to_phyloxml(graph: &GraphAncestral) -> Result { + phyloxml_from_graph(graph, "TreeTime ancestral analysis", |context| { + let mutations = ancestral_node_mutations(graph, context.node_key, context.edge_key)?; + ancestral_phyloxml_clade(context, mutations, &ancestral_node_sequences(graph, context.node_key)) + }) +} -impl TreeOutputNode for NodeClock { - fn output_divergence(&self) -> Option { - Some(self.div) - } +pub(crate) fn optimize_to_phyloxml(graph: &GraphAncestral) -> Result { + phyloxml_from_graph(graph, "TreeTime optimize analysis", |context| { + ancestral_phyloxml_clade( + context, + optimize_mutations(graph, context.edge_key)?, + &optimize_node_sequences(graph, context.node_key), + ) + }) +} - fn output_date(&self) -> Option { - self.time - } +pub(crate) fn prune_to_phyloxml(graph: &GraphAncestral) -> Result { + phyloxml_from_graph(graph, "TreeTime prune analysis", |context| { + ancestral_phyloxml_clade( + context, + prune_mutations(graph, context.edge_key)?, + &prune_node_sequences(graph, context.node_key), + ) + }) +} - fn output_bad_branch(&self) -> bool { - self.bad_branch || self.is_outlier - } +pub(crate) fn clock_to_phyloxml(graph: &GraphClock) -> Result { + phyloxml_from_graph(graph, "TreeTime clock analysis", |context| { + let name = node_name(context.node_key, context.node); + ensure_optional_finite(context.node.time, "clock", &name, "date")?; + ensure_finite(context.node.div, "clock", &name, "divergence")?; + let property = vec![ + property(REF_DIV, DT_DOUBLE, APPLIES_NODE, &context.node.div.to_string()), + property( + REF_BAD_BRANCH, + DT_BOOLEAN, + APPLIES_NODE, + if context.node.bad_branch || context.node.is_outlier { + "true" + } else { + "false" + }, + ), + ]; + Ok(PhyloxmlClade { + name: context.node.name.clone(), + branch_length_elem: context.edge.and_then(HasBranchLength::branch_length), + branch_length_attr: None, + confidence: vec![], + width: None, + color: None, + node_id: None, + taxonomy: vec![], + sequence: vec![], + events: None, + binary_characters: None, + distribution: vec![], + date: context.node.time.map(phyloxml_date), + reference: vec![], + property, + clade: vec![], + other: BTreeMap::new(), + }) + }) } -impl TreeOutputNode for NodeTimetree { - fn output_divergence(&self) -> Option { - Some(self.div) - } +pub(crate) fn mugration_to_phyloxml(graph: &GraphAncestral) -> Result { + phyloxml_from_graph(graph, "TreeTime mugration analysis", |context| { + let name = node_name(context.node_key, context.node); + let div = cumulative_branch_length(graph, context.node_key)?; + ensure_optional_finite(div, "mugration", &name, "divergence")?; + let traits = mugration_traits(graph, context.node_key, &name)?; + let mut properties = div + .map(|div| vec![property(REF_DIV, DT_DOUBLE, APPLIES_NODE, &div.to_string())]) + .unwrap_or_default(); + properties.extend(trait_properties(&traits, &name)?); + if let Some((parent, child)) = mugration_transition(graph, context.node_key)? { + let attribute = encode_property_token(&graph.data().traits.attribute); + properties.push(property( + &format!("{REF_TRAIT_TRANSITION_PREFIX}{attribute}"), + DT_STRING, + APPLIES_BRANCH, + &format!("{}:{}", encode_property_token(&parent), encode_property_token(&child)), + )); + } + let mut clade = empty_phyloxml_clade( + context.node.name.clone(), + context.edge.and_then(HasBranchLength::branch_length), + ); + clade.confidence = input_branch_confidence(context.node.confidence, "mugration", &name)?; + clade.property = properties; + Ok(clade) + }) +} - fn output_date(&self) -> Option { - self.time +pub(crate) fn timetree_to_phyloxml( + graph: &Graph, +) -> Result { + phyloxml_from_graph(graph, "TreeTime timetree analysis", |context| { + let name = node_name(context.node_key, context.node); + let divergence = timetree_divergence(graph, context.node_key, context.node)?; + ensure_finite(divergence, "timetree", &name, "divergence")?; + ensure_optional_finite(context.node.time, "timetree", &name, "date")?; + ensure_finite(context.edge.map_or(1.0, |edge| edge.gamma), "timetree", &name, "gamma")?; + let mut properties = vec![ + property(REF_DIV, DT_DOUBLE, APPLIES_NODE, &divergence.to_string()), + property( + REF_BAD_BRANCH, + DT_BOOLEAN, + APPLIES_NODE, + if context.node.bad_branch || context.node.is_outlier { + "true" + } else { + "false" + }, + ), + ]; + if timetree_date_is_inferred(graph, context.node_key, context.node) == Some(true) { + properties.push(property(REF_DATE_INFERRED, DT_BOOLEAN, APPLIES_NODE, "true")); + } + if let Some(edge) = context.edge { + properties.push(property(REF_GAMMA, DT_DOUBLE, APPLIES_BRANCH, &edge.gamma.to_string())); + } + for mutation in timetree_mutations(graph, context.edge_key)? { + properties.push(mutation_property(&mutation)?); + } + let date = context.node.time.map(|value| { + let confidence = graph + .data() + .confidence_intervals + .as_ref() + .and_then(|intervals| intervals.iter().find(|interval| interval.key == context.node_key)); + PhyloxmlDate { + desc: None, + value: Some(value), + minimum: confidence.map(|interval| interval.lower), + maximum: confidence.map(|interval| interval.upper), + unit: Some("year".to_owned()), + } + }); + let mut clade = empty_phyloxml_clade( + context.node.base.name.clone(), + context.edge.and_then(HasBranchLength::branch_length), + ); + clade.confidence = input_branch_confidence(context.node.base.confidence, "timetree", &name)?; + clade.date = date; + clade.property = properties; + clade.sequence = phyloxml_sequences(&timetree_node_sequences(graph, context.node_key)); + Ok(clade) + }) +} + +fn ancestral_phyloxml_clade( + context: &GraphNodeContext, + mutations: Vec, + sequences: &BTreeMap, +) -> Result { + let name = node_name(context.node_key, context.node); + let divergence = cumulative_branch_length(context.graph, context.node_key)?; + ensure_optional_finite(divergence, "tree output", &name, "divergence")?; + let mut properties = divergence + .map(|divergence| vec![property(REF_DIV, DT_DOUBLE, APPLIES_NODE, &divergence.to_string())]) + .unwrap_or_default(); + properties.extend( + mutations + .into_iter() + .map(|mutation| mutation_property(&mutation)) + .collect::, _>>()?, + ); + let mut clade = empty_phyloxml_clade( + context.node.name.clone(), + context.edge.and_then(HasBranchLength::branch_length), + ); + clade.confidence = input_branch_confidence(context.node.confidence, "tree output", &name)?; + clade.property = properties; + clade.sequence = phyloxml_sequences(sequences); + Ok(clade) +} + +fn phyloxml_from_graph(graph: &Graph, title: &str, mut convert: F) -> Result +where + N: GraphNode, + E: GraphEdge, + D: Send + Sync, + F: FnMut(&GraphNodeContext) -> Result, +{ + let root = graph + .get_exactly_one_root() + .wrap_err("When converting graph to PhyloXML")?; + let mut node_map = btreemap! {}; + let mut queue = VecDeque::from([(Arc::clone(&root), None)]); + while let Some((current_node, current_edge)) = queue.pop_front() { + let node_key = current_node.read_arc().key(); + let node = current_node.read_arc().payload().read_arc(); + let edge_key = current_edge + .as_ref() + .map(|edge: &Arc>>| edge.read_arc().key()); + let edge = current_edge + .as_ref() + .map(|edge: &Arc>>| edge.read_arc().payload().read_arc()); + node_map.insert( + node_key, + convert(&GraphNodeContext { + node_key, + node: &node, + edge_key, + edge: edge.as_deref(), + graph, + })?, + ); + for (child, edge) in graph.children_of(¤t_node.read_arc()) { + queue.push_back((child, Some(edge))); + } } + attach_phyloxml_children(graph, &root, &mut node_map)?; + let root_key = root.read_arc().key(); + let clade = node_map + .remove(&root_key) + .ok_or_else(|| make_internal_report!("PhyloXML root node {root_key} was not converted"))?; + Ok(Phyloxml { + phylogeny: vec![PhyloxmlPhylogeny { + rooted: true, + rerootable: None, + branch_length_unit: Some(BRANCH_LENGTH_UNIT.to_owned()), + phylogeny_type: None, + name: Some(title.to_owned()), + id: None, + description: None, + date: None, + confidence: vec![], + clade: Some(clade), + clade_relation: vec![], + sequence_relation: vec![], + property: vec![], + other: BTreeMap::new(), + }], + other: BTreeMap::new(), + }) +} - fn output_bad_branch(&self) -> bool { - self.bad_branch || self.is_outlier +fn attach_phyloxml_children( + graph: &Graph, + root: &Arc>>, + node_map: &mut BTreeMap, +) -> Result<(), Report> +where + N: GraphNode, + E: GraphEdge, + D: Send + Sync, +{ + let mut visited = BTreeSet::new(); + let mut stack = vec![Arc::clone(root)]; + while let Some(node) = stack.pop() { + let key = node.read_arc().key(); + if visited.contains(&key) { + let children = graph + .children_of(&node.read_arc()) + .into_iter() + .map(|(child, _)| { + let child_key = child.read_arc().key(); + node_map + .remove(&child_key) + .ok_or_else(|| make_internal_report!("PhyloXML child node {child_key} was not converted")) + }) + .collect::, _>>()?; + node_map + .get_mut(&key) + .ok_or_else(|| make_internal_report!("PhyloXML parent node {key} was not converted"))? + .clade = children; + } else { + visited.insert(key); + stack.push(Arc::clone(&node)); + stack.extend(graph.children_of(&node.read_arc()).into_iter().map(|(child, _)| child)); + } } + Ok(()) } -impl TreeOutputEdge for EdgeAncestral {} -impl TreeOutputEdge for EdgeClock {} +pub(crate) fn ancestral_to_mat(graph: &GraphAncestral) -> Result { + let reference = ancestral_root_sequences(graph)?.remove(NUC_TRACK); + mat_from_graph(graph, reference.as_deref(), |node_key, edge_key| { + ancestral_node_mutations(graph, node_key, Some(edge_key)) + }) +} -impl TreeOutputEdge for EdgeTimetree { - fn output_gamma(&self) -> Option { - Some(self.gamma) - } +pub(crate) fn optimize_to_mat(graph: &GraphAncestral) -> Result { + let reference = optimize_root_sequences(graph)?.remove(NUC_TRACK); + mat_from_graph(graph, reference.as_deref(), |_node_key, edge_key| { + optimize_mutations(graph, Some(edge_key)) + }) +} + +pub(crate) fn prune_to_mat(graph: &GraphAncestral) -> Result { + let reference = prune_root_sequences(graph)?.remove(NUC_TRACK); + mat_from_graph(graph, reference.as_deref(), |_node_key, edge_key| { + prune_mutations(graph, Some(edge_key)) + }) +} + +pub(crate) fn clock_to_mat(graph: &GraphClock) -> Result { + mutation_free_mat(graph) +} + +pub(crate) fn mugration_to_mat(graph: &GraphAncestral) -> Result { + mutation_free_mat(graph) +} + +pub(crate) fn timetree_to_mat( + graph: &Graph, +) -> Result { + let reference = timetree_root_sequences(graph)?.remove(NUC_TRACK); + mat_from_graph(graph, reference.as_deref(), |_node_key, edge_key| { + timetree_mutations(graph, Some(edge_key)) + }) } -impl TreeOutputData for () +fn mutation_free_mat(graph: &Graph) -> Result where - N: TreeOutputNode, - E: TreeOutputEdge, + N: GraphNode + Named + NodeToNwk, + E: GraphEdge + EdgeToNwk, + D: Send + Sync, { + mat_from_graph(graph, None, |_node_key, _edge_key| Ok(vec![])) } -impl TreeOutputData for AncestralGraphData { - fn has_mutations(&self) -> bool { - self.partition.is_some() - } +fn mat_from_graph( + graph: &Graph, + reference: Option<&str>, + mut edge_mutations: F, +) -> Result +where + N: GraphNode + Named + NodeToNwk, + E: GraphEdge + EdgeToNwk, + D: Send + Sync, + F: FnMut(GraphNodeKey, GraphEdgeKey) -> Result, Report>, +{ + graph + .get_exactly_one_root() + .wrap_err("When converting graph to UShER MAT")?; + let mut node_mutations = vec![]; + let mut condensed_nodes = vec![]; + let mut metadata = vec![]; + graph.iter_depth_first_preorder_forward(|node| { + let name = node + .payload + .name() + .map(|name| name.as_ref().to_owned()) + .unwrap_or_default(); + let mutations = node + .parent_keys + .first() + .map(|(_, edge_key)| edge_mutations(node.key, *edge_key)) + .transpose()? + .unwrap_or_default(); + let mutations = mutations + .iter() + .map(|mutation| mat_mutation(mutation, reference, &name)) + .collect::, _>>()?; + node_mutations.push(UsherMutationList { mutation: mutations }); + condensed_nodes.push(UsherTreeNode { + node_name: name, + condensed_leaves: vec![], + }); + metadata.push(UsherMetadata { + clade_annotations: vec![], + }); + Ok(()) + })?; + Ok(UsherTree { + newick: nwk_write_str(graph, &NwkWriteOptions::default())?, + node_mutations, + condensed_nodes, + metadata, + }) +} - fn root_sequences( - &self, - graph: &Graph, - ) -> Result, Report> { - self.partition.as_ref().map_or_else( - || Ok(BTreeMap::new()), - |partition| { - Ok(BTreeMap::from([( - NUC_GENE.to_owned(), - ancestral_root_sequence(partition, graph)?, - )])) - }, +pub(crate) fn mat_mutation( + mutation: &Mutation, + reference: Option<&str>, + node_name: &str, +) -> Result { + if mutation.track != MutationTrack::Nucleotide { + return make_error!("Node '{node_name}' has an amino-acid mutation that UShER MAT cannot represent"); + } + let MutationEvent::Substitution(substitution) = &mutation.event else { + return make_error!("Node '{node_name}' has an insertion or deletion that UShER MAT cannot represent"); + }; + let reference = reference.ok_or_else(|| { + eyre::eyre!("Node '{node_name}' has nucleotide mutations, but UShER MAT requires a root nucleotide reference") + })?; + let position = substitution + .pos() + .checked_add(1) + .ok_or_else(|| eyre::eyre!("Node '{node_name}' mutation coordinate overflow"))?; + let position = i32::try_from(position).wrap_err_with(|| { + format!("Node '{node_name}' mutation position {position} exceeds the UShER MAT i32 coordinate range") + })?; + let reference_state = reference.as_bytes().get(substitution.pos()).copied().ok_or_else(|| { + eyre::eyre!( + "Node '{node_name}' mutation position {} is outside the root nucleotide reference of length {}", + substitution.pos() + 1, + reference.len() ) + })?; + let reference_state = AsciiChar::try_new(reference_state)?; + Ok(UsherMutation { + position, + ref_nuc: mat_nucleotide(reference_state, node_name, "root reference")?, + par_nuc: mat_nucleotide(substitution.reff(), node_name, "parent")?, + mut_nuc: vec![mat_nucleotide(substitution.qry(), node_name, "child")?], + chromosome: String::new(), + }) +} + +fn mat_nucleotide(nucleotide: AsciiChar, node_name: &str, role: &str) -> Result { + match char::from(nucleotide).to_ascii_uppercase() { + 'A' => Ok(0), + 'C' => Ok(1), + 'G' => Ok(2), + 'T' => Ok(3), + state => { + make_error!("Node '{node_name}' has {role} nucleotide '{state}', but UShER MAT accepts only A, C, G, or T") + }, } +} - fn mutations( - &self, - graph: &Graph, - key: GraphEdgeKey, - ) -> Result, Report> { - self - .partition - .as_ref() - .map_or_else(|| Ok(vec![]), |partition| ancestral_mutations(partition, graph, key)) +fn ancestral_root_sequences(graph: &GraphAncestral) -> Result, Report> { + let mut sequences = BTreeMap::new(); + if let Some(partition) = graph.data().partition.as_ref() { + let sequence = match partition { + AncestralPartition::Fitch(partition) => partition.read_arc().root_sequence(graph)?, + AncestralPartition::Sparse(partition) => partition.read_arc().root_sequence(graph)?, + AncestralPartition::Dense(partition) => partition.read_arc().root_sequence(graph)?, + }; + sequences.insert(NUC_TRACK.to_owned(), sequence.to_string()); + } + if let Some(aa) = graph.data().aa_node_data.as_ref() { + sequences.extend(aa.root_aa_sequences.clone()); } + Ok(sequences) } -impl TreeOutputData for OptimizeGraphData { - fn has_mutations(&self) -> bool { - !self.dense_partitions.is_empty() || !self.sparse_partitions.is_empty() +fn ancestral_node_sequences( + graph: &GraphAncestral, + node_key: GraphNodeKey, +) -> BTreeMap { + let mut sequences = graph + .data() + .partition + .as_ref() + .map(|partition| { + let sequence = match partition { + AncestralPartition::Fitch(partition) => partition.read_arc().node_sequence(node_key), + AncestralPartition::Sparse(partition) => partition.read_arc().node_sequence(node_key), + AncestralPartition::Dense(partition) => partition.read_arc().node_sequence(node_key), + }; + btreemap! { NUC_TRACK.to_owned() => sequence.to_string() } + }) + .unwrap_or_default(); + if graph.is_root(node_key) + && let Some(aa) = graph.data().aa_node_data.as_ref() + { + sequences.extend(aa.root_aa_sequences.clone()); } + sequences +} - fn root_sequences( - &self, - graph: &Graph, - ) -> Result, Report> { - let sequence = if let Some(partition) = self.dense_partitions.first() { - Some(partition.read_arc().root_sequence(graph)?) - } else if let Some(partition) = self.sparse_partitions.first() { - Some(partition.read_arc().root_sequence(graph)?) - } else { - None - }; - Ok( - sequence - .map(|sequence| BTreeMap::from([(NUC_GENE.to_owned(), sequence.to_string())])) - .unwrap_or_default(), - ) +fn ancestral_genome_annotations( + graph: &GraphAncestral, + root_sequences: &BTreeMap, +) -> Result, Report> { + let nuc = root_sequences + .get(NUC_TRACK) + .map(|sequence| -> Result<_, Report> { + Ok(AuspiceGenomeAnnotationNuc { + start: 1, + end: isize::try_from(sequence.len()).wrap_err("Nucleotide sequence length does not fit Auspice coordinates")?, + strand: Some("+".to_owned()), + r#type: Some("source".to_owned()), + other: Value::default(), + }) + }) + .transpose()?; + let cdses = graph + .data() + .aa_node_data + .as_ref() + .map(|data| { + data + .annotations + .iter() + .map(|(name, annotation)| Ok((name.clone(), auspice_cds_annotation(name, annotation)?))) + .collect::, Report>>() + }) + .transpose()? + .unwrap_or_default(); + if nuc.is_none() && cdses.is_empty() { + return Ok(None); } + Ok(Some(AuspiceGenomeAnnotations { + nuc, + cdses, + other: Value::default(), + })) +} - fn mutations( - &self, - graph: &Graph, - key: GraphEdgeKey, - ) -> Result, Report> { - if let Some(partition) = self.dense_partitions.first() { - subs_to_output(PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?) - } else if let Some(partition) = self.sparse_partitions.first() { - subs_to_output(PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?) - } else { - Ok(vec![]) +fn auspice_cds_annotation( + name: &str, + annotation: &AugurNodeDataJsonAnnotationEntry, +) -> Result { + let strand = annotation + .strand + .clone() + .ok_or_else(|| make_report!("CDS annotation '{name}' has no strand for Auspice output"))?; + let other = Value::Object(annotation.other.clone().into_iter().collect()); + let segments = if let Some(segments) = annotation.segments.as_ref() { + if segments.is_empty() { + return make_error!("CDS annotation '{name}' has no segments for Auspice output"); } - } + Segments::MultipleSegments { + segments: segments + .iter() + .map(|segment| { + Ok(StartEnd { + start: auspice_coordinate(segment.start, name, "segment start")?, + end: auspice_coordinate(segment.end, name, "segment end")?, + other: Value::Object(segment.other.clone().into_iter().collect()), + }) + }) + .collect::, Report>>()?, + other, + } + } else { + Segments::OneSegment(StartEnd { + start: auspice_coordinate( + annotation + .start + .ok_or_else(|| make_report!("CDS annotation '{name}' has no start for Auspice output"))?, + name, + "start", + )?, + end: auspice_coordinate( + annotation + .end + .ok_or_else(|| make_report!("CDS annotation '{name}' has no end for Auspice output"))?, + name, + "end", + )?, + other, + }) + }; + Ok(AuspiceGenomeAnnotationCds { + r#type: annotation.entry_type.clone(), + gene: None, + color: None, + display_name: None, + description: None, + strand: Some(strand), + segments, + }) } -impl TreeOutputData for PruneGraphData { - fn has_mutations(&self) -> bool { - !self.partitions.is_empty() - } +fn auspice_coordinate(coordinate: i64, name: &str, field: &str) -> Result { + isize::try_from(coordinate) + .wrap_err_with(|| format!("CDS annotation '{name}' {field} does not fit Auspice coordinates")) +} - fn mutations( - &self, - graph: &Graph, - key: GraphEdgeKey, - ) -> Result, Report> { - self.partitions.first().map_or_else( - || Ok(vec![]), - |partition| subs_to_output(PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?), - ) - } +fn optimize_root_sequences(graph: &GraphAncestral) -> Result, Report> { + let sequence = if let Some(partition) = graph.data().dense_partitions.first() { + Some(partition.read_arc().root_sequence(graph)?) + } else if let Some(partition) = graph.data().sparse_partitions.first() { + Some(partition.read_arc().root_sequence(graph)?) + } else { + None + }; + Ok( + sequence + .map(|sequence| btreemap! { NUC_TRACK.to_owned() => sequence.to_string() }) + .unwrap_or_default(), + ) } -impl TreeOutputData for ClockGraphData { - fn title(&self) -> Option<&str> { - Some("TreeTime clock analysis") - } +fn optimize_node_sequences( + graph: &GraphAncestral, + node_key: GraphNodeKey, +) -> BTreeMap { + let sequence = if let Some(partition) = graph.data().dense_partitions.first() { + Some(partition.read_arc().node_sequence(node_key)) + } else { + graph + .data() + .sparse_partitions + .first() + .map(|partition| partition.read_arc().node_sequence(node_key)) + }; + sequence + .map(|sequence| btreemap! { NUC_TRACK.to_owned() => sequence.to_string() }) + .unwrap_or_default() +} - fn has_dates(&self) -> bool { - true - } +fn prune_root_sequences(graph: &GraphAncestral) -> Result, Report> { + graph.data().partitions.first().map_or_else( + || Ok(BTreeMap::new()), + |partition| { + Ok(btreemap! { + NUC_TRACK.to_owned() => partition.read_arc().root_sequence(graph)?.to_string(), + }) + }, + ) +} - fn has_bad_branch(&self) -> bool { - true - } +fn prune_node_sequences(graph: &GraphAncestral, node_key: GraphNodeKey) -> BTreeMap { + graph.data().partitions.first().map_or_else(BTreeMap::new, |partition| { + btreemap! { + NUC_TRACK.to_owned() => partition.read_arc().node_sequence(node_key).to_string(), + } + }) } -impl TreeOutputData for MugrationGraphData { - fn trait_attributes(&self) -> Vec { - vec![self.traits.attribute.clone()] - } +fn timetree_root_sequences( + graph: &Graph, +) -> Result, Report> { + graph.data().partitions.first().map_or_else( + || Ok(BTreeMap::new()), + |partition| { + Ok(btreemap! { + NUC_TRACK.to_owned() => partition.read_arc().root_sequence(graph)?.to_string(), + }) + }, + ) +} - fn traits(&self, key: GraphNodeKey) -> BTreeMap { - let Some(value) = self.partition.get_reconstructed_trait(key) else { - return BTreeMap::new(); - }; - let profile = self.partition.get_confidence(key); - let confidence = profile - .as_ref() - .map(|profile| build_confidence_map(&self.partition.states, profile)) - .unwrap_or_default(); - let entropy = profile.as_ref().map(compute_entropy); - BTreeMap::from([( - self.traits.attribute.clone(), - TreeOutputTrait { - value, - confidence, - entropy, +fn timetree_node_sequences( + graph: &Graph, + node_key: GraphNodeKey, +) -> BTreeMap { + graph.data().partitions.first().map_or_else(BTreeMap::new, |partition| { + btreemap! { + NUC_TRACK.to_owned() => partition.read_arc().node_sequence(node_key).to_string(), + } + }) +} + +fn ancestral_edge_mutations( + graph: &GraphAncestral, + edge_key: GraphEdgeKey, +) -> Result, Report> { + graph.data().partition.as_ref().map_or_else( + || Ok(vec![]), + |partition| match partition { + AncestralPartition::Fitch(partition) => { + partition + .read_arc() + .edge_mutations(graph, edge_key, MutationTrack::Nucleotide) }, - )]) - } + AncestralPartition::Sparse(partition) => { + partition + .read_arc() + .edge_mutations(graph, edge_key, MutationTrack::Nucleotide) + }, + AncestralPartition::Dense(partition) => { + partition + .read_arc() + .edge_mutations(graph, edge_key, MutationTrack::Nucleotide) + }, + }, + ) } -impl TreeOutputData for TimetreeGraphData { - fn title(&self) -> Option<&str> { - Some("TreeTime timetree analysis") +fn ancestral_node_mutations( + graph: &GraphAncestral, + node_key: GraphNodeKey, + edge_key: Option, +) -> Result, Report> { + let mut mutations = edge_key + .map(|edge_key| ancestral_edge_mutations(graph, edge_key)) + .transpose()? + .unwrap_or_default(); + if let Some(aa) = graph.data().aa_node_data.as_ref() + && let Some(tracks) = aa.node_aa_mutations.get(&node_key) + { + mutations.extend(tracks.iter().flat_map(|(track, events)| { + events.iter().cloned().map(|event| Mutation { + track: MutationTrack::AminoAcid(track.clone()), + event, + }) + })); } + Ok(mutations) +} - fn has_dates(&self) -> bool { - true - } +fn ancestral_all_mutations(graph: &GraphAncestral) -> Result, Report> { + graph + .get_nodes() + .into_iter() + .map(|node| { + let node = node.read_arc(); + ancestral_node_mutations(graph, node.key(), node.inbound().first().copied()) + }) + .collect::, _>>() + .map(|mutations| mutations.into_iter().flatten().collect()) +} - fn has_bad_branch(&self) -> bool { - true - } +fn optimize_has_mutations(graph: &GraphAncestral) -> bool { + !graph.data().dense_partitions.is_empty() || !graph.data().sparse_partitions.is_empty() +} - fn has_mutations(&self) -> bool { - !self.partitions.is_empty() +fn optimize_mutations( + graph: &GraphAncestral, + edge_key: Option, +) -> Result, Report> { + let Some(edge_key) = edge_key else { + return Ok(vec![]); + }; + if let Some(partition) = graph.data().dense_partitions.first() { + partition + .read_arc() + .edge_mutations(graph, edge_key, MutationTrack::Nucleotide) + } else if let Some(partition) = graph.data().sparse_partitions.first() { + partition + .read_arc() + .edge_mutations(graph, edge_key, MutationTrack::Nucleotide) + } else { + Ok(vec![]) } +} - fn divergence( - &self, - graph: &Graph, - key: GraphNodeKey, - node: &NodeTimetree, - ) -> Result, Report> { - self.mutation_counts.as_ref().map_or(Ok(Some(node.div)), |counts| { - cumulative_mutation_count(graph, key, counts).map(Some) - }) +fn prune_mutations( + graph: &GraphAncestral, + edge_key: Option, +) -> Result, Report> { + match (graph.data().partitions.first(), edge_key) { + (Some(partition), Some(edge_key)) => { + partition + .read_arc() + .edge_mutations(graph, edge_key, MutationTrack::Nucleotide) + }, + _ => Ok(vec![]), } +} - fn date_confidence(&self, key: GraphNodeKey) -> Option<[f64; 2]> { - self - .confidence_intervals - .as_ref()? - .iter() - .find_map(|interval| (interval.key == key).then_some([interval.lower, interval.upper])) +fn timetree_mutations( + graph: &Graph, + edge_key: Option, +) -> Result, Report> { + match (graph.data().partitions.first(), edge_key) { + (Some(partition), Some(edge_key)) => { + partition + .read_arc() + .edge_mutations(graph, edge_key, MutationTrack::Nucleotide) + }, + _ => Ok(vec![]), } +} - fn mutations( - &self, - graph: &Graph, - key: GraphEdgeKey, - ) -> Result, Report> { - self.partitions.first().map_or_else( - || Ok(vec![]), - |partition| subs_to_output(partition.read_arc().edge_subs(graph, key)?), - ) +fn mugration_traits( + graph: &GraphAncestral, + node_key: GraphNodeKey, + node_name: &str, +) -> Result, Report> { + let Some(value) = graph.data().partition.get_reconstructed_trait(node_key) else { + return Ok(BTreeMap::new()); + }; + let profile = graph.data().partition.get_confidence(node_key); + if let Some(profile) = profile.as_ref() { + for (state, probability) in graph.data().partition.states.iter().zip(profile) { + ensure_finite( + *probability, + "mugration", + node_name, + &format!("trait state '{state}' probability"), + )?; + } + } + let confidence = profile + .as_ref() + .map(|profile| build_confidence_map(&graph.data().partition.states, profile)) + .unwrap_or_default(); + let entropy = profile.as_ref().map(compute_entropy); + if let Some(entropy) = entropy { + ensure_finite(entropy, "mugration", node_name, "trait entropy")?; } + Ok(btreemap! { + graph.data().traits.attribute.clone() => TraitValue { value, confidence, entropy }, + }) } -impl AuspiceWrite for TreeOutputAdapter -where - N: TreeOutputNode, - E: TreeOutputEdge, - D: TreeOutputData, -{ - fn new(graph: &Graph) -> Result { - Ok(Self { - has_bad_branch: graph.data().has_bad_branch(), - warned_ambiguous: false, - warned_non_nuc: false, - warned_reference_fallback: false, - }) - } +fn mugration_transition( + graph: &GraphAncestral, + node_key: GraphNodeKey, +) -> Result, Report> { + let Some((parent_key, _edge_key)) = graph.node_parent(node_key)? else { + return Ok(None); + }; + let parent = graph.data().partition.get_reconstructed_trait(parent_key); + let child = graph.data().partition.get_reconstructed_trait(node_key); + Ok(match (parent, child) { + (Some(parent), Some(child)) if parent != child => Some((parent, child)), + _ => None, + }) +} - fn auspice_data_from_graph_data(&self, graph: &Graph) -> Result { - let data = graph.data(); - let trait_attributes = data.trait_attributes(); - let mut colorings = vec![]; - if data.has_dates() { - colorings.push(coloring(COLORING_NUM_DATE, "Date", "continuous")); - } - if data.has_bad_branch() { - colorings.push(coloring(COLORING_BAD_BRANCH, "Excluded", "categorical")); - } - for attribute in &trait_attributes { - colorings.push(coloring(attribute, attribute, "categorical")); - } - if data.has_mutations() { - colorings.push(coloring(COLORING_GENOTYPE, "Genotype", "categorical")); - } +fn mugration_transition_label( + graph: &GraphAncestral, + node_key: GraphNodeKey, +) -> Result, Report> { + let Some((parent, child)) = mugration_transition(graph, node_key)? else { + return Ok(None); + }; + Ok(Some(AuspiceTreeBranchAttrsLabels { + aa: None, + clade: None, + other: json!({ graph.data().traits.attribute.clone(): format!("{parent} → {child}") }), + })) +} - let color_by = trait_attributes - .first() - .cloned() - .or_else(|| data.has_bad_branch().then(|| COLORING_BAD_BRANCH.to_owned())) - .or_else(|| data.has_dates().then(|| COLORING_NUM_DATE.to_owned())); - let mut filters = trait_attributes; - if data.has_bad_branch() { - filters.push(COLORING_BAD_BRANCH.to_owned()); +fn timetree_divergence( + graph: &Graph, + node_key: GraphNodeKey, + node: &NodeTimetree, +) -> Result { + graph.data().mutation_counts.as_ref().map_or(Ok(node.div), |counts| { + let mut key = node_key; + let mut count = 0; + while let Some((parent, edge)) = graph.node_parent(key)? { + count += counts.get(&edge).copied().unwrap_or_default(); + key = parent; } - let root_sequences = data.root_sequences(graph)?; - - Ok(AuspiceTreeData { - version: Some("v2".to_owned()), - meta: AuspiceTreeMeta { - title: data - .title() - .map(str::to_owned) - .or_else(|| Some("TreeTime analysis".to_owned())), - description: data.description().map(str::to_owned), - panels: vec!["tree".to_owned()], - colorings, - filters, - display_defaults: AuspiceDisplayDefaults { - color_by, - ..AuspiceDisplayDefaults::default() - }, - ..AuspiceTreeMeta::default() - }, - root_sequence: (!root_sequences.is_empty()).then_some(root_sequences), - other: Value::default(), - }) + Ok(count as f64) + }) +} + +fn timetree_date_confidence( + graph: &Graph, + node_key: GraphNodeKey, + node_name: &str, +) -> Result, Report> { + let confidence = graph + .data() + .confidence_intervals + .as_ref() + .and_then(|intervals| intervals.iter().find(|interval| interval.key == node_key)) + .map(|interval| [interval.lower, interval.upper]); + if let Some([lower, upper]) = confidence { + ensure_finite(lower, "timetree", node_name, "date confidence lower bound")?; + ensure_finite(upper, "timetree", node_name, "date confidence upper bound")?; + Ok(Some([format_number(lower, 3), format_number(upper, 3)])) + } else { + Ok(None) } +} - fn auspice_node_from_graph_components( - &mut self, - context: &AuspiceGraphContext, - ) -> Result { - let name = context.node.name().map_or_else( - || format!("node_{}", context.node_key.as_usize()), - |name| name.as_ref().to_owned(), - ); - let div = context - .graph - .data() - .divergence(context.graph, context.node_key, context.node)?; - let div = finite_number(div, 6, &name, "div")?; - let date = finite_number(context.node.output_date(), 3, &name, "date")?; - let confidence = context.graph.data().date_confidence(context.node_key); - let confidence = match confidence { - Some([lower, upper]) if lower.is_finite() && upper.is_finite() => { - Some([format_number(lower, 3), format_number(upper, 3)]) - }, - Some([lower, upper]) => return make_error!("Node '{name}' has non-finite date confidence [{lower}, {upper}]"), - None => None, - }; - let num_date = date.map(|value| AuspiceNumDate { value, confidence }); - let bad_branch = self - .has_bad_branch - .then(|| AuspiceTreeNodeAttr::new(if context.node.output_bad_branch() { "Yes" } else { "No" })); - let traits = context.graph.data().traits(context.node_key); - let mutations = match context.edge_key { - Some(edge_key) => group_mutations(&context.graph.data().mutations(context.graph, edge_key)?), - None => BTreeMap::new(), - }; +fn timetree_date_is_inferred( + graph: &Graph, + node_key: GraphNodeKey, + node: &NodeTimetree, +) -> Option { + let dates = graph.data().dates.as_ref()?; + let name = node.base.name.as_ref(); + Some( + name.and_then(|name| dates.get(name)).and_then(Option::as_ref).is_none() + && node.time.is_some() + && graph.get_node(node_key).is_some(), + ) +} - Ok(AuspiceTreeNode { - name, - branch_attrs: AuspiceTreeBranchAttrs { - mutations, - labels: None, - other: Value::default(), - }, - node_attrs: AuspiceTreeNodeAttrs { - div, - num_date, - bad_branch, - clade_membership: None, - region: None, - country: None, - division: None, - other: build_trait_attrs(&traits), +fn group_mutations(mutations: Vec) -> Result>, Report> { + let mut grouped = BTreeMap::new(); + for mutation in mutations { + let track = match mutation.track { + MutationTrack::Nucleotide => NUC_TRACK.to_owned(), + MutationTrack::AminoAcid(track) => { + if track.is_empty() + || !track + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'*' | b'_' | b'.' | b'(' | b')' | b'-')) + { + return make_error!("Auspice v2 cannot represent amino-acid mutation track '{track}'"); + } + track }, - children: vec![], - other: Value::default(), - }) + }; + grouped + .entry(track) + .or_insert_with(Vec::new) + .extend(mutation_event_strings(&mutation.event)?); } + Ok(grouped) } -impl PhyloxmlFromGraph for TreeOutputAdapter -where - N: TreeOutputNode, - E: TreeOutputEdge, - D: TreeOutputData, -{ - fn phyloxml_data_from_graph_data(graph: &Graph) -> Result { - let data = graph.data(); - Ok(Phyloxml { - phylogeny: vec![PhyloxmlPhylogeny { - rooted: true, - rerootable: None, - branch_length_unit: Some(BRANCH_LENGTH_UNIT.to_owned()), - phylogeny_type: None, - name: data.title().map(str::to_owned), - id: None, - description: data.description().map(str::to_owned), - date: None, - confidence: vec![], - clade: None, - clade_relation: vec![], - sequence_relation: vec![], - property: vec![], - other: BTreeMap::new(), - }], - other: BTreeMap::new(), - }) - } +fn mutation_property(mutation: &Mutation) -> Result { + let track = match &mutation.track { + MutationTrack::Nucleotide => NUC_TRACK.to_owned(), + MutationTrack::AminoAcid(track) => format!("aa:{}", encode_property_token(track)), + }; + let value = match &mutation.event { + MutationEvent::Substitution(substitution) => { + let position = substitution + .pos() + .checked_add(1) + .ok_or_else(|| eyre::eyre!("Mutation coordinate overflow at {}", substitution.pos()))?; + format!("{track}:sub:{}{position}{}", substitution.reff(), substitution.qry()) + }, + MutationEvent::Insertion(segment) => { + let (start, end) = segment.one_based_inclusive_range()?; + format!("{track}:ins:{start}-{end}:{}", segment.sequence) + }, + MutationEvent::Deletion(segment) => { + let (start, end) = segment.one_based_inclusive_range()?; + format!("{track}:del:{start}-{end}:{}", segment.sequence) + }, + }; + Ok(property(REF_MUTATION, DT_STRING, APPLIES_BRANCH, &value)) +} - fn phyloxml_node_from_graph_components(context: &PhyloxmlGraphContext) -> Result { - let mut property = vec![]; - if let Some(div) = context - .graph - .data() - .divergence(context.graph, context.node_key, context.node)? - { - property.push(make_property(REF_DIV, DT_DOUBLE, APPLIES_NODE, &div.to_string())); - } - if context.node.output_bad_branch() { - property.push(make_property(REF_BAD_BRANCH, DT_BOOLEAN, APPLIES_NODE, "true")); - } - if context.graph.data().date_is_inferred(context.node_key) { - property.push(make_property(REF_DATE_INFERRED, DT_BOOLEAN, APPLIES_NODE, "true")); - } - for (attribute, value) in context.graph.data().traits(context.node_key) { - property.push(make_property( - &format!("{REF_TRAIT_PREFIX}{attribute}"), - DT_STRING, +fn trait_properties(traits: &BTreeMap, node_name: &str) -> Result, Report> { + let mut properties = vec![]; + for (attribute, value) in traits { + let attribute = encode_property_token(attribute); + properties.push(property( + &format!("{REF_TRAIT_PREFIX}{attribute}"), + DT_STRING, + APPLIES_NODE, + &value.value, + )); + for (state, probability) in &value.confidence { + ensure_finite( + *probability, + "mugration", + node_name, + &format!("trait confidence '{state}'"), + )?; + properties.push(property( + &format!( + "{REF_TRAIT_CONFIDENCE_PREFIX}{attribute}:{}", + encode_property_token(state) + ), + DT_DOUBLE, APPLIES_NODE, - &value.value, + &probability.to_string(), )); } - if let Some(edge) = context.edge { - if let Some(gamma) = edge.output_gamma() { - property.push(make_property(REF_GAMMA, DT_DOUBLE, APPLIES_BRANCH, &gamma.to_string())); - } - } - if let Some(edge_key) = context.edge_key { - for mutation in context.graph.data().mutations(context.graph, edge_key)? { - property.push(make_property( - REF_MUTATION, - DT_STRING, - APPLIES_BRANCH, - &format!("{}:{}", mutation.gene, mutation_string(&mutation)), - )); - } + if let Some(entropy) = value.entropy { + ensure_finite(entropy, "mugration", node_name, "trait entropy")?; + properties.push(property( + &format!("{REF_TRAIT_ENTROPY_PREFIX}{attribute}"), + DT_DOUBLE, + APPLIES_NODE, + &entropy.to_string(), + )); } - let date = context.node.output_date().map(|value| { - let (minimum, maximum) = context - .graph - .data() - .date_confidence(context.node_key) - .map_or((None, None), |[lower, upper]| (Some(lower), Some(upper))); - PhyloxmlDate { - desc: None, - value: Some(value), - minimum, - maximum, - unit: Some("year".to_owned()), - } - }); + } + Ok(properties) +} - Ok(PhyloxmlClade { - name: context.node.name().map(|name| name.as_ref().to_owned()), - branch_length_elem: context.edge.and_then(HasBranchLength::branch_length), - branch_length_attr: None, - confidence: vec![], - width: None, - color: None, - node_id: None, - taxonomy: vec![], - sequence: vec![], - events: None, - binary_characters: None, - distribution: vec![], - date, - reference: vec![], - property, - clade: vec![], - other: BTreeMap::new(), - }) +fn empty_phyloxml_clade(name: Option, branch_length: Option) -> PhyloxmlClade { + PhyloxmlClade { + name, + branch_length_elem: branch_length, + branch_length_attr: None, + confidence: vec![], + width: None, + color: None, + node_id: None, + taxonomy: vec![], + sequence: vec![], + events: None, + binary_characters: None, + distribution: vec![], + date: None, + reference: vec![], + property: vec![], + clade: vec![], + other: BTreeMap::new(), } } -impl UsherWrite for TreeOutputAdapter -where - N: TreeOutputNode + treetime_io::nwk::NodeToNwk, - E: TreeOutputEdge + treetime_io::nwk::EdgeToNwk, - D: TreeOutputData, -{ - fn new(graph: &Graph) -> Result { - Ok(Self { - has_bad_branch: graph.data().has_bad_branch(), - warned_ambiguous: false, - warned_non_nuc: false, - warned_reference_fallback: false, +fn input_branch_confidence( + confidence: Option, + command: &str, + node_name: &str, +) -> Result, Report> { + confidence.map_or_else( + || Ok(vec![]), + |value| { + ensure_finite(value, command, node_name, "input branch support")?; + Ok(vec![PhyloxmlConfidence { + value, + type_: TYPE_INPUT_BRANCH_SUPPORT.to_owned(), + }]) + }, + ) +} + +fn phyloxml_sequences(sequences: &BTreeMap) -> Vec { + sequences + .iter() + .map(|(track, sequence)| PhyloxmlSequence { + symbol: None, + accession: None, + name: Some(track.clone()), + location: None, + mol_seq: Some(PhyloxmlMolSeq { + sequence: sequence.clone(), + is_aligned: Some(true), + }), + uri: None, + annotation: vec![], + domain_architecture: None, + other: BTreeMap::new(), }) - } + .collect() +} - fn usher_node_from_graph_components( - &mut self, - context: &UsherGraphContext, - ) -> Result<(UsherTreeNode, UsherMutationList, UsherMetadata), Report> { - let reference = context - .graph - .data() - .root_sequences(context.graph)? - .remove(NUC_GENE) - .unwrap_or_default(); - let mutations = match context.edge_key { - Some(edge_key) => context.graph.data().mutations(context.graph, edge_key)?, - None => vec![], - }; - let mut encoded = vec![]; - for mutation in mutations { - if mutation.gene != NUC_GENE { - if !self.warned_non_nuc { - warn!("UShER MAT format carries only nucleotide substitutions; dropping amino-acid mutations"); - self.warned_non_nuc = true; - } - continue; - } - let (Ok(par_nuc), Ok(mut_nuc)) = (nuc_to_int(mutation.parent), nuc_to_int(mutation.child)) else { - if !self.warned_ambiguous { - warn!("UShER MAT format cannot encode non-ACGT nucleotides; skipping ambiguous substitutions"); - self.warned_ambiguous = true; - } - continue; - }; - let ref_nuc = reference - .as_bytes() - .get(mutation.position) - .copied() - .and_then(|nucleotide| AsciiChar::try_new(nucleotide).ok()) - .and_then(|nucleotide| nuc_to_int(nucleotide).ok()); - let ref_nuc = ref_nuc.unwrap_or_else(|| { - // This fallback can fabricate a global reference allele from a branch-local parent allele. - // Keep the current output behavior visible until the graph supplies a validated root sequence. - // See kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md. - if !self.warned_reference_fallback { - warn!( - "UShER MAT global reference nucleotide is unavailable at zero-based position {}; using the branch-local parent allele for ref_nuc, which can produce incorrect MAT reference alleles", - mutation.position - ); - self.warned_reference_fallback = true; - } - par_nuc - }); - encoded.push(UsherMutation { - position: i32::try_from(mutation.position + 1)?, - ref_nuc, - par_nuc, - mut_nuc: vec![mut_nuc], - chromosome: String::new(), - }); - } +fn phyloxml_date(value: f64) -> PhyloxmlDate { + PhyloxmlDate { + desc: None, + value: Some(value), + minimum: None, + maximum: None, + unit: Some("year".to_owned()), + } +} - Ok(( - UsherTreeNode { - node_name: context - .node - .name() - .map(|name| name.as_ref().to_owned()) - .unwrap_or_default(), - condensed_leaves: vec![], - }, - UsherMutationList { mutation: encoded }, - UsherMetadata { - clade_annotations: vec![], - }, - )) +fn property(ref_: &str, datatype: &str, applies_to: &str, value: &str) -> PhyloxmlProperty { + PhyloxmlProperty { + value: value.to_owned(), + ref_: ref_.to_owned(), + unit: None, + datatype: datatype.to_owned(), + applies_to: applies_to.to_owned(), + id_ref: None, } } -fn ancestral_root_sequence(partition: &AncestralPartition, graph: &dyn BranchTopology) -> Result { - let sequence = match partition { - AncestralPartition::Fitch(partition) => partition.read_arc().root_sequence(graph)?, - AncestralPartition::Sparse(partition) => partition.read_arc().root_sequence(graph)?, - AncestralPartition::Dense(partition) => partition.read_arc().root_sequence(graph)?, - }; - Ok(sequence.to_string()) -} - -fn ancestral_mutations( - partition: &AncestralPartition, - graph: &dyn BranchTopology, - key: GraphEdgeKey, -) -> Result, Report> { - let substitutions = match partition { - AncestralPartition::Fitch(partition) => partition.read_arc().edge_subs(graph, key)?, - AncestralPartition::Sparse(partition) => PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?, - AncestralPartition::Dense(partition) => PartitionBranchOps::edge_subs(&*partition.read_arc(), graph, key)?, - }; - subs_to_output(substitutions) +fn encode_property_token(value: &str) -> String { + utf8_percent_encode(value, PROPERTY_TOKEN_ENCODE_SET).to_string() } -fn subs_to_output(substitutions: Vec) -> Result, Report> { - Ok( - substitutions +fn build_trait_attrs(traits: BTreeMap) -> Value { + Value::Object( + traits .into_iter() - .map(|substitution| TreeOutputMutation { - gene: NUC_GENE.to_owned(), - position: substitution.pos(), - parent: substitution.reff(), - child: substitution.qry(), + .map(|(attribute, value)| { + let mut fields = serde_json::Map::new(); + fields.insert("value".to_owned(), json!(value.value)); + if !value.confidence.is_empty() { + fields.insert( + "confidence".to_owned(), + json!( + value + .confidence + .iter() + .map(|(state, probability)| (state.clone(), format_number(*probability, 3))) + .collect::>() + ), + ); + } + if let Some(entropy) = value.entropy { + fields.insert("entropy".to_owned(), json!(format_number(entropy, 3))); + } + (attribute, Value::Object(fields)) }) .collect(), ) } -fn cumulative_mutation_count( - graph: &Graph, - mut key: GraphNodeKey, - counts: &BTreeMap, -) -> Result +fn cumulative_branch_length(graph: &Graph, mut key: GraphNodeKey) -> Result, Report> where N: GraphNode, - E: GraphEdge, + E: GraphEdge + HasBranchLength, D: Send + Sync, { - let mut count = 0; - while let Some((parent, edge)) = graph.node_parent(key)? { - count += counts.get(&edge).copied().unwrap_or_default(); + let mut total = 0.0; + while let Some((parent, edge_key)) = graph.node_parent(key)? { + let edge = graph + .get_edge(edge_key) + .ok_or_else(|| make_internal_report!("Edge {edge_key} disappeared while summing branch lengths"))?; + let Some(length) = edge.read_arc().payload().read_arc().branch_length() else { + return Ok(None); + }; + total += length; key = parent; } - Ok(count as f64) + Ok(Some(total)) +} + +fn node_name(key: GraphNodeKey, node: &N) -> String { + node + .name() + .map_or_else(|| format!("node_{}", key.as_usize()), |name| name.as_ref().to_owned()) +} + +fn ensure_optional_finite(value: Option, command: &str, node_name: &str, field: &str) -> Result<(), Report> { + if let Some(value) = value { + ensure_finite(value, command, node_name, field)?; + } + Ok(()) +} + +fn ensure_finite(value: f64, command: &str, node_name: &str, field: &str) -> Result<(), Report> { + if value.is_finite() { + Ok(()) + } else { + make_error!("{command} node '{node_name}' has non-finite {field}={value}") + } +} + +fn finite_number( + value: Option, + precision: i32, + command: &str, + node_name: &str, + field: &str, +) -> Result, Report> { + ensure_optional_finite(value, command, node_name, field)?; + Ok(value.map(|value| format_number(value, precision))) } pub(crate) fn format_number(number: f64, precision: i32) -> f64 { @@ -760,14 +1739,6 @@ pub(crate) fn format_number(number: f64, precision: i32) -> f64 { .expect("a float formatted in scientific notation must parse back") } -fn finite_number(value: Option, precision: i32, node_name: &str, field: &str) -> Result, Report> { - match value { - Some(value) if value.is_finite() => Ok(Some(format_number(value, precision))), - Some(value) => make_error!("Node '{node_name}' has non-finite {field}={value}"), - None => Ok(None), - } -} - fn coloring(key: &str, title: &str, type_: &str) -> AuspiceColoring { AuspiceColoring { key: key.to_owned(), @@ -777,61 +1748,6 @@ fn coloring(key: &str, title: &str, type_: &str) -> AuspiceColoring { } } -fn build_trait_attrs(traits: &BTreeMap) -> Value { - let map = traits - .iter() - .map(|(attribute, value)| { - let mut fields = serde_json::Map::new(); - fields.insert("value".to_owned(), json!(value.value)); - if !value.confidence.is_empty() { - let confidence = value - .confidence - .iter() - .map(|(state, probability)| (state.clone(), format_number(*probability, 3))) - .collect::>(); - fields.insert("confidence".to_owned(), json!(confidence)); - } - if let Some(entropy) = value.entropy { - fields.insert("entropy".to_owned(), json!(format_number(entropy, 3))); - } - (attribute.clone(), Value::Object(fields)) - }) - .collect(); - Value::Object(map) -} - -fn group_mutations(mutations: &[TreeOutputMutation]) -> BTreeMap> { - let mut grouped: BTreeMap> = BTreeMap::new(); - for mutation in mutations { - grouped - .entry(mutation.gene.clone()) - .or_default() - .push(mutation_string(mutation)); - } - grouped -} - -fn mutation_string(mutation: &TreeOutputMutation) -> String { - format!("{}{}{}", mutation.parent, mutation.position + 1, mutation.child) -} - -fn make_property(ref_: &str, datatype: &str, applies_to: &str, value: &str) -> PhyloxmlProperty { - PhyloxmlProperty { - value: value.to_owned(), - ref_: ref_.to_owned(), - unit: None, - datatype: datatype.to_owned(), - applies_to: applies_to.to_owned(), - id_ref: None, - } -} - -fn nuc_to_int(nucleotide: AsciiChar) -> Result { - match char::from(nucleotide).to_ascii_uppercase() { - 'A' => Ok(0), - 'C' => Ok(1), - 'G' => Ok(2), - 'T' => Ok(3), - other => make_error!("UShER MAT cannot encode nucleotide '{other}': expected A, C, G, or T"), - } +fn generation_date() -> String { + Utc::now().format("%Y-%m-%d").to_string() } diff --git a/packages/treetime/src/commands/timetree/initialization.rs b/packages/treetime/src/commands/timetree/initialization.rs index 0170c3c05..3eba23629 100644 --- a/packages/treetime/src/commands/timetree/initialization.rs +++ b/packages/treetime/src/commands/timetree/initialization.rs @@ -9,10 +9,7 @@ use crate::make_report; use crate::optimize::params::BranchLengthMode; use crate::partition::algo::infer_dense::infer_dense; use crate::partition::marginal_dense::PartitionMarginalDense; -use crate::partition::timetree::{GraphTimetree, PartitionTimetreeAllVec}; -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::partition::timetree::{GraphTimetree, PartitionTimetree, PartitionTimetreeAllVec}; use crate::seq::alignment::get_common_length; use crate::seq::gap_fill::apply_gap_fill; use crate::timetree::utils::initialize_node_divergences; @@ -126,8 +123,7 @@ pub fn initialize_partitions( log_gtr(>r, model_name); let partition = fitch.into_marginal_sparse(gtr, graph)?; - let sparse_partition: Arc>> = - Arc::new(RwLock::new(partition)); + let sparse_partition = Arc::new(RwLock::new(PartitionTimetree::Sparse(partition))); Ok(vec![sparse_partition]) } else if model_name == GtrModelName::Infer { let aln_data = aln.ok_or_else(|| make_report!("Alignment required for dense GTR inference"))?; @@ -136,8 +132,7 @@ pub fn initialize_partitions( log_gtr(>r, model_name); let partition = fitch.into_marginal_dense(gtr); - let dense_partition: Arc>> = - Arc::new(RwLock::new(partition)); + let dense_partition = Arc::new(RwLock::new(PartitionTimetree::Dense(partition))); Ok(vec![dense_partition]) } else { info!("GTR model: {model_name}"); @@ -145,8 +140,7 @@ pub fn initialize_partitions( log_gtr(>r, model_name); let partition = PartitionMarginalDense::new(0, gtr, alphabet, length); - let dense_partition: Arc>> = - Arc::new(RwLock::new(partition)); + let dense_partition = Arc::new(RwLock::new(PartitionTimetree::Dense(partition))); Ok(vec![dense_partition]) } } diff --git a/packages/treetime/src/commands/timetree/output/dates.rs b/packages/treetime/src/commands/timetree/output/dates.rs index 7b93cfab8..82d8b0051 100644 --- a/packages/treetime/src/commands/timetree/output/dates.rs +++ b/packages/treetime/src/commands/timetree/output/dates.rs @@ -1,11 +1,6 @@ -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::partition::timetree::{GraphTimetree, PartitionTimetreeRef}; use eyre::Report; -use parking_lot::RwLock; use std::path::Path; -use std::sync::Arc; /// Write inferred node dates to file. /// @@ -14,7 +9,7 @@ use std::sync::Arc; /// How: TSV file with node_name, numdate, time_before_present. pub fn write_node_dates( _graph: &GraphTimetree, - _partitions: &[Arc>>], + _partitions: &[PartitionTimetreeRef], _out_base: &Path, ) -> Result<(), Report> { todo!("Write node dates to TSV file") diff --git a/packages/treetime/src/commands/timetree/result.rs b/packages/treetime/src/commands/timetree/result.rs index 6d2b54d32..35aea324c 100644 --- a/packages/treetime/src/commands/timetree/result.rs +++ b/packages/treetime/src/commands/timetree/result.rs @@ -8,24 +8,14 @@ use serde::Serialize; use std::collections::BTreeMap; use treetime_io::dates_csv::DatesMap; -#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. #[derive(Serialize)] -#[serde(transparent)] pub struct TimetreeGraphData { - marker: (), - #[serde(skip)] pub clock_model: ClockModel, - #[serde(skip)] pub confidence_intervals: Option>, - #[serde(skip)] pub partitions: PartitionTimetreeAllVec, - #[serde(skip)] pub dates: Option, - #[serde(skip)] pub gtr: Option, - #[serde(skip)] pub model_name: Option, - #[serde(skip)] pub mutation_counts: Option>, } @@ -40,7 +30,6 @@ impl TimetreeGraphData { mutation_counts: Option>, ) -> Self { Self { - marker: (), clock_model, confidence_intervals, partitions, diff --git a/packages/treetime/src/commands/timetree/run.rs b/packages/treetime/src/commands/timetree/run.rs index f8c9ddbcd..f40660506 100644 --- a/packages/treetime/src/commands/timetree/run.rs +++ b/packages/treetime/src/commands/timetree/run.rs @@ -1,6 +1,6 @@ use crate::clock::clock_output::write_clock_model; use crate::commands::shared::output::{CommandKind, DivergenceUnits, OutputSelection}; -use crate::commands::shared::tree_output::TreeOutputAdapter; +use crate::commands::shared::tree_output::write_timetree_tree_outputs; use crate::commands::timetree::args::TreetimeTimetreeArgs; use crate::commands::timetree::initialization::load_input_data; use crate::commands::timetree::output::augur_node_data::write_augur_node_data_json; @@ -14,7 +14,6 @@ use crate::timetree::pipeline::{self, TimetreeInput, TimetreeParams}; use eyre::{Report, WrapErr}; use log::{info, warn}; use std::path::PathBuf; -use treetime_io::graph::write_tree_outputs; use treetime_io::nwk::CommentProviders; use treetime_utils::io::file::create_file_or_stdout; @@ -47,8 +46,6 @@ pub fn run_timetree_estimation( (OutputSelection::Tracelog, args.output_tracelog.as_deref()), ], )?; - resolved.prepare()?; - let tracelog: Option> = match resolved.non_tree_outputs.get(&OutputSelection::Tracelog) { Some(path) => Some(Box::new(create_file_or_stdout(path)?)), @@ -175,9 +172,9 @@ pub fn run_timetree_estimation( let guard = graph.data().partitions[0].read_arc(); let provider = MutationCommentProvider::new(&*guard, &graph); let providers = CommentProviders::new().with(&provider); - write_tree_outputs::(&graph, &resolved.tree_outputs, &providers)?; + write_timetree_tree_outputs(&graph, &resolved.tree_outputs, &providers)?; } else { - write_tree_outputs::(&graph, &resolved.tree_outputs, &CommentProviders::new())?; + write_timetree_tree_outputs(&graph, &resolved.tree_outputs, &CommentProviders::new())?; } } diff --git a/packages/treetime/src/gtr/gtr.rs b/packages/treetime/src/gtr/gtr.rs index 500aae17e..9de302b39 100644 --- a/packages/treetime/src/gtr/gtr.rs +++ b/packages/treetime/src/gtr/gtr.rs @@ -4,7 +4,9 @@ use ndarray::prelude::*; use ndarray_linalg::Eigh; use ndarray_linalg::UPLO::Lower; use num_traits::abs; +use serde::Serialize; use treetime_utils::array::ndarray::{clamp_min, outer}; +use treetime_utils::array::serde::{array1_as_vec, array2_as_vec, option_array1_as_vec}; /// Compute the average substitution rate for normalization. /// @@ -164,7 +166,7 @@ pub struct GTRParams { /// /// See `eig_single_site()` for details on the symmetrization trick that ensures /// numerical stability. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct GTR { pub debug: bool, /// Average substitution rate before normalization, used for scaling. @@ -172,14 +174,19 @@ pub struct GTR { /// Overall substitution rate (incorporates average_rate after normalization). pub mu: f64, /// Symmetric exchangeability matrix (zero diagonal, normalized by average_rate). + #[serde(serialize_with = "array2_as_vec")] pub W: Array2, /// Equilibrium frequencies (sum to 1). + #[serde(serialize_with = "array1_as_vec")] pub pi: Array1, /// Eigenvalues of the rate matrix (all <= 0, exactly one zero). + #[serde(serialize_with = "array1_as_vec")] pub eigvals: Array1, /// Right transformation matrix for P(t) = v * exp(Lambda * t) * v_inv. + #[serde(serialize_with = "array2_as_vec")] pub v: Array2, /// Left transformation matrix (inverse of v, adjusted for non-symmetric Q). + #[serde(serialize_with = "array2_as_vec")] pub v_inv: Array2, /// Per-site rate multipliers for among-site rate variation. /// @@ -190,6 +197,7 @@ pub struct GTR { /// /// The effective rate at site `a` is `mu * site_rates[a]`. The matrix exponential /// becomes `P_a(t) = V * diag(exp(eigvals * mu * site_rates[a] * t)) * V_inv`. + #[serde(serialize_with = "option_array1_as_vec")] pub site_rates: Option>, /// Whether the one-dimensional branch-length likelihood $L(t)$ is guaranteed /// unimodal on $(0, \infty)$ for this model. diff --git a/packages/treetime/src/mugration/result.rs b/packages/treetime/src/mugration/result.rs index 84cc94c3e..6a3622a03 100644 --- a/packages/treetime/src/mugration/result.rs +++ b/packages/treetime/src/mugration/result.rs @@ -3,19 +3,21 @@ use crate::payload::ancestral::GraphAncestral; use indexmap::IndexMap; use itertools::Itertools; use ndarray::Array1; +use serde::Serialize; use std::collections::BTreeMap; use std::fmt::Write; use treetime_graph::node::Named; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct ConfidenceRow { /// Node name. pub node: String, /// Probability for each state (in state order). + #[serde(serialize_with = "treetime_utils::array::serde::array1_as_vec")] pub profile: Array1, } /// Structured confidence output for all nodes. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct MugrationConfidenceOutput { /// State names in order (column headers). pub states: Vec, @@ -75,7 +77,7 @@ impl MugrationConfidenceOutput { } /// Structured trait output for all nodes. -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct MugrationTraitsOutput { /// Attribute name (column header). pub attribute: String, @@ -104,18 +106,11 @@ impl MugrationTraitsOutput { } } -#[allow(clippy::manual_non_exhaustive, clippy::partial_pub_fields)] // The private unit field preserves Graph JSON's `data: null` shape. -#[derive(Debug, serde::Serialize)] -#[serde(transparent)] +#[derive(Debug, Serialize)] pub struct MugrationGraphData { - marker: (), - #[serde(skip)] pub traits: MugrationTraitsOutput, - #[serde(skip)] pub confidence: MugrationConfidenceOutput, - #[serde(skip)] pub log_lh: f64, - #[serde(skip)] pub partition: PartitionMarginalDiscrete, } @@ -140,7 +135,6 @@ impl MugrationResult { let confidence = MugrationConfidenceOutput::new(&graph, &partition); let data = MugrationGraphData { - marker: (), traits, confidence, log_lh, diff --git a/packages/treetime/src/optimize/__tests__/test_initial_guess_indel_zero_bl.rs b/packages/treetime/src/optimize/__tests__/test_initial_guess_indel_zero_bl.rs index 62f03f7d0..492cc631a 100644 --- a/packages/treetime/src/optimize/__tests__/test_initial_guess_indel_zero_bl.rs +++ b/packages/treetime/src/optimize/__tests__/test_initial_guess_indel_zero_bl.rs @@ -55,8 +55,8 @@ mod tests { let edge_entry = partition.data.edges.entry(edge_key).or_insert(edge_data); edge_entry.indels.push(InDel { range: (4, 7), - seq: Seq::default(), - deletion: true, + seq: Seq::try_from_str("ACG")?, + kind: crate::seq::indel::InDelKind::Deletion, }); } diff --git a/packages/treetime/src/optimize/__tests__/test_initial_guess_mode.rs b/packages/treetime/src/optimize/__tests__/test_initial_guess_mode.rs index 00fddffad..aff8360b2 100644 --- a/packages/treetime/src/optimize/__tests__/test_initial_guess_mode.rs +++ b/packages/treetime/src/optimize/__tests__/test_initial_guess_mode.rs @@ -356,7 +356,7 @@ pub mod tests { let edge_key = graph.get_edges()[0].read_arc().key(); for partition in partitions { let mut partition = partition.write_arc(); - partition.data.edges.get_mut(&edge_key).unwrap().indels = vec![InDel::del((4, 7), Seq::try_from_str("ACG")?)]; + partition.data.edges.get_mut(&edge_key).unwrap().indels = vec![InDel::del((4, 7), Seq::try_from_str("ACG")?)?]; } Ok(()) } diff --git a/packages/treetime/src/optimize/__tests__/test_no_indels.rs b/packages/treetime/src/optimize/__tests__/test_no_indels.rs index c1973c50a..e8e6a6585 100644 --- a/packages/treetime/src/optimize/__tests__/test_no_indels.rs +++ b/packages/treetime/src/optimize/__tests__/test_no_indels.rs @@ -44,7 +44,7 @@ mod tests { .edges .get_mut(&first_edge_key) .unwrap() - .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; let first_edge_key_without = graph_without.get_edges()[0].read_arc().key(); graph_without.get_edges()[0] @@ -57,7 +57,7 @@ mod tests { .edges .get_mut(&first_edge_key_without) .unwrap() - .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; let result_with = run_optimize_loop( &mut graph_with, @@ -109,7 +109,7 @@ mod tests { .edges .get_mut(&first_edge_key) .unwrap() - .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; update_marginal(&graph, &sparse_partitions)?; @@ -203,7 +203,7 @@ mod tests { let graph_with_indel: GraphAncestral = nwk_read_str(TREE_NEWICK)?; let (dense_with_indel, sparse_with_indel, partitions_with_indel) = setup_identical_partitions(&graph_with_indel)?; - let indels = vec![InDel::del((0, 2), Seq::try_from_str("AC")?)]; + let indels = vec![InDel::del((0, 2), Seq::try_from_str("AC")?)?]; inject_indels_on_first_edge( &graph_with_indel, &dense_with_indel, diff --git a/packages/treetime/src/optimize/__tests__/test_optimize_indel.rs b/packages/treetime/src/optimize/__tests__/test_optimize_indel.rs index 0f2b8e491..6f3f4cfeb 100644 --- a/packages/treetime/src/optimize/__tests__/test_optimize_indel.rs +++ b/packages/treetime/src/optimize/__tests__/test_optimize_indel.rs @@ -115,7 +115,7 @@ pub mod tests { let aln = simple_alignment()?; let (dense_partitions, sparse_partitions, mixed_partitions) = setup_partitions(&graph, &aln)?; - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); let rate = estimate_indel_rate(&graph, &mixed_partitions); @@ -137,7 +137,7 @@ pub mod tests { let aln = simple_alignment()?; let (dense_partitions, sparse_partitions, mixed_partitions) = setup_partitions(&graph, &aln)?; - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; let first_edge_key = inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); graph.get_edges()[0] .write_arc() @@ -171,7 +171,7 @@ pub mod tests { let aln = simple_alignment()?; let (dense_partitions, sparse_partitions, mixed_partitions) = setup_partitions(&graph, &aln)?; - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); graph.get_edges()[0] .write_arc() @@ -213,7 +213,7 @@ pub mod tests { let (dense_partitions_high, sparse_partitions_high, mixed_partitions_high) = setup_identical_partitions(&graph_high)?; - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; inject_indels_on_first_edge(&graph_low, &dense_partitions_low, &sparse_partitions_low, &indels); inject_indels_on_first_edge(&graph_high, &dense_partitions_high, &sparse_partitions_high, &indels); @@ -254,7 +254,7 @@ pub mod tests { let graph: GraphAncestral = nwk_read_str(TREE_NEWICK)?; let (dense_partitions, sparse_partitions, mixed_partitions) = setup_identical_partitions(&graph)?; - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); initial_guess_mixed(&graph, &mixed_partitions, true, false)?; @@ -286,7 +286,7 @@ pub mod tests { } // Inject indels on the first edge (after zeroing, so indel_rate starts at 0) - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); // indel_rate is 0 at this point (all BL = 0), but initial_guess should bootstrap @@ -324,8 +324,8 @@ pub mod tests { // Inject indels AFTER setup (which includes update_marginal) to avoid being wiped // by the backward pass that recreates DenseEdgePartition from scratch. let indels = vec![ - InDel::del((0, 3), Seq::try_from_str("ACG")?), - InDel::del((5, 8), Seq::try_from_str("ACG")?), + InDel::del((0, 3), Seq::try_from_str("ACG")?)?, + InDel::del((5, 8), Seq::try_from_str("ACG")?)?, ]; inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); @@ -361,7 +361,7 @@ pub mod tests { .write_arc() .set_branch_length(Some(-0.1)); if has_indels { - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); } @@ -400,7 +400,7 @@ pub mod tests { .edges .get_mut(&first_edge_key) .unwrap() - .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; initial_guess_mixed(&graph, &mixed_partitions, true, false)?; @@ -646,7 +646,7 @@ pub mod tests { let graph: GraphAncestral = nwk_read_str(TREE_NEWICK)?; let (dense_partitions, sparse_partitions, mixed_partitions) = setup_identical_partitions(&graph)?; - let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + let indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; let _first_edge_key = inject_indels_on_first_edge(&graph, &dense_partitions, &sparse_partitions, &indels); // Set a very small initial branch length to test clamping diff --git a/packages/treetime/src/optimize/__tests__/test_optimize_method.rs b/packages/treetime/src/optimize/__tests__/test_optimize_method.rs index 2ad9b7fea..6f6114f8e 100644 --- a/packages/treetime/src/optimize/__tests__/test_optimize_method.rs +++ b/packages/treetime/src/optimize/__tests__/test_optimize_method.rs @@ -979,7 +979,7 @@ mod tests { let first_edge_key = graph.get_edges()[0].read_arc().key(); let indels: Vec = (0..n_indels) - .map(|i| InDel::del((i * 3, i * 3 + 3), Seq::try_from_str("ACG").unwrap())) + .map(|i| InDel::del((i * 3, i * 3 + 3), Seq::try_from_str("ACG").unwrap()).unwrap()) .collect(); for p in &dense_partitions { diff --git a/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs b/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs index b405f1126..587459eb3 100644 --- a/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs +++ b/packages/treetime/src/optimize/__tests__/test_run_optimize_loop.rs @@ -128,7 +128,7 @@ mod tests { .edges .get_mut(&first_edge_key) .unwrap() - .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; let sparse_lh = update_marginal(&graph, &sparse_partitions)?; let dense_lh = update_marginal(&graph, &dense_partitions)?; @@ -256,7 +256,7 @@ mod tests { .edges .get_mut(&first_edge_key) .unwrap() - .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)]; + .indels = vec![InDel::del((0, 3), Seq::try_from_str("ACG")?)?]; let initial_sparse_lh = update_marginal(&graph, &sparse_partitions)?; let initial_dense_lh = update_marginal(&graph, &dense_partitions)?; diff --git a/packages/treetime/src/optimize/topology/__tests__/test_collapse_edge.rs b/packages/treetime/src/optimize/topology/__tests__/test_collapse_edge.rs index 84376dc2b..73a87fe16 100644 --- a/packages/treetime/src/optimize/topology/__tests__/test_collapse_edge.rs +++ b/packages/treetime/src/optimize/topology/__tests__/test_collapse_edge.rs @@ -227,8 +227,8 @@ mod tests { let mut partition = make_sparse_partition(100)?; populate_test_nodes(&mut partition, &graph); - let collapsed_indel = InDel::ins((0, 3), [c(b'A'), c(b'C'), c(b'G')].as_slice()); - let child_a_indel = InDel::del((10, 12), [c(b'T'), c(b'T')].as_slice()); + let collapsed_indel = InDel::ins((0, 3), [c(b'A'), c(b'C'), c(b'G')].as_slice()).unwrap(); + let child_a_indel = InDel::del((10, 12), [c(b'T'), c(b'T')].as_slice()).unwrap(); partition.edges.insert( ri_key, diff --git a/packages/treetime/src/optimize/topology/__tests__/test_merge_shared_mutations.rs b/packages/treetime/src/optimize/topology/__tests__/test_merge_shared_mutations.rs index 16e84081a..cb4d1db62 100644 --- a/packages/treetime/src/optimize/topology/__tests__/test_merge_shared_mutations.rs +++ b/packages/treetime/src/optimize/topology/__tests__/test_merge_shared_mutations.rs @@ -667,7 +667,7 @@ mod tests { let mut p = partition.write_arc(); let edge_a = find_edge_key(&graph, "root", "A").expect("edge root->A"); p.edges.get_mut(&edge_a).expect("partition edge A").indels = - vec![InDel::del((10, 13), Seq::try_from_str("GTA").unwrap())]; + vec![InDel::del((10, 13), Seq::try_from_str("GTA").unwrap()).unwrap()]; } let partitions = vec![partition]; @@ -770,7 +770,7 @@ mod tests { ], )?; - let shared_indel = InDel::del((5, 8), Seq::try_from_str("GTA").unwrap()); + let shared_indel = InDel::del((5, 8), Seq::try_from_str("GTA").unwrap()).unwrap(); { let mut p = partition.write_arc(); let edge_a = find_edge_key(&graph, "root", "A").expect("edge root->A"); @@ -800,7 +800,7 @@ mod tests { ], )?; - let shared_indel = InDel::del((5, 8), Seq::try_from_str("GTA").unwrap()); + let shared_indel = InDel::del((5, 8), Seq::try_from_str("GTA").unwrap()).unwrap(); { let mut p = partition.write_arc(); let edge_a = find_edge_key(&graph, "root", "A").expect("edge root->A"); diff --git a/packages/treetime/src/optimize/topology/__tests__/test_prop_merge_shared_mutations.rs b/packages/treetime/src/optimize/topology/__tests__/test_prop_merge_shared_mutations.rs index 3f1d8e946..0023756a0 100644 --- a/packages/treetime/src/optimize/topology/__tests__/test_prop_merge_shared_mutations.rs +++ b/packages/treetime/src/optimize/topology/__tests__/test_prop_merge_shared_mutations.rs @@ -326,7 +326,7 @@ mod tests { ], )?; - let shared_indel = InDel::del((10, 13), Seq::try_from_str("GTA").unwrap()); + let shared_indel = InDel::del((10, 13), Seq::try_from_str("GTA").unwrap()).unwrap(); { let mut p = partition.write_arc(); let edge_c = find_edge_key(&graph, "root", "C").expect("edge root->C"); diff --git a/packages/treetime/src/partition/augur.rs b/packages/treetime/src/partition/augur.rs index 5b2a027bc..1c955eb97 100644 --- a/packages/treetime/src/partition/augur.rs +++ b/packages/treetime/src/partition/augur.rs @@ -3,6 +3,7 @@ use crate::partition::marginal_dense::PartitionMarginalDense; use crate::partition::marginal_sparse::PartitionMarginalSparse; use crate::partition::traits::{BranchTopology, PartitionBranchOps, PartitionMarginalOps}; use crate::payload::ancestral::{EdgeAncestral, NodeAncestral}; +use crate::seq::indel::InDel; use crate::seq::mutation::Sub; use eyre::Report; use treetime_graph::edge::GraphEdgeKey; @@ -30,6 +31,8 @@ pub trait AugurNodeDataJsonAncestralPartition { /// Substitutions on the parent edge of one node (parent -> child). fn edge_subs(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result, Report>; + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec; + /// Ambiguous (unknown) character of the partition alphabet, used to fill /// masked positions in per-node output sequences. fn ambiguous_char(&self) -> AsciiChar; @@ -53,6 +56,10 @@ impl AugurNodeDataJsonAncestralPartition for PartitionFitch { Ok(self.edges[&edge_key].fitch_subs().to_vec()) } + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec { + self.edges[&edge_key].indels.clone() + } + fn ambiguous_char(&self) -> AsciiChar { self.alphabet.unknown() } @@ -71,6 +78,10 @@ impl AugurNodeDataJsonAncestralPartition for PartitionMarginalSparse { PartitionBranchOps::edge_subs(self, graph, edge_key) } + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec { + PartitionBranchOps::edge_indels(self, edge_key) + } + fn ambiguous_char(&self) -> AsciiChar { self.alphabet.unknown() } @@ -91,6 +102,10 @@ impl AugurNodeDataJsonAncestralPartition for PartitionMarginalDense { PartitionBranchOps::edge_subs(self, graph, edge_key) } + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec { + PartitionBranchOps::edge_indels(self, edge_key) + } + fn ambiguous_char(&self) -> AsciiChar { self.alphabet.unknown() } diff --git a/packages/treetime/src/partition/fitch.rs b/packages/treetime/src/partition/fitch.rs index 250fb12d2..7cd90af07 100644 --- a/packages/treetime/src/partition/fitch.rs +++ b/packages/treetime/src/partition/fitch.rs @@ -3,15 +3,16 @@ use crate::gtr::gtr::GTR; use crate::partition::marginal_dense::PartitionMarginalDense; use crate::partition::marginal_sparse::PartitionMarginalSparse; use crate::partition::sparse::{SparseEdgePartition, SparseNodePartition}; -use crate::partition::traits::PartitionCompressed; +use crate::partition::traits::{BranchTopology, PartitionBranchOps, PartitionCompressed}; use eyre::Report; +use serde::Serialize; use std::collections::BTreeMap; use treetime_graph::edge::{GraphEdge, GraphEdgeKey}; use treetime_graph::graph::Graph; use treetime_graph::node::{GraphNode, GraphNodeKey}; -use treetime_primitives::seq; +use treetime_primitives::{Seq, seq}; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct PartitionFitch { pub index: usize, pub alphabet: Alphabet, @@ -88,3 +89,42 @@ impl PartitionCompressed for PartitionFitch { (&mut self.nodes, &mut self.edges) } } + +impl PartitionBranchOps for PartitionFitch { + fn sequence_length(&self) -> usize { + self.length + } + + fn edge_subs( + &self, + _graph: &dyn BranchTopology, + edge_key: GraphEdgeKey, + ) -> Result, Report> { + Ok(self.edges[&edge_key].fitch_subs().to_vec()) + } + + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec { + self.edges[&edge_key].indels.clone() + } + + fn root_sequence(&self, graph: &dyn BranchTopology) -> Result { + Ok(self.nodes[&graph.root_key()?].seq.sequence.clone()) + } + + fn node_sequence(&self, node_key: GraphNodeKey) -> Seq { + self.nodes[&node_key].seq.sequence.clone() + } + + fn edge_effective_length(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result { + let (parent_key, child_key) = graph.edge_endpoints(edge_key)?; + Ok( + self.nodes[&parent_key] + .seq + .sequence + .iter() + .zip(&self.nodes[&child_key].seq.sequence) + .filter(|(parent, child)| self.alphabet.is_canonical(**parent) && self.alphabet.is_canonical(**child)) + .count(), + ) + } +} diff --git a/packages/treetime/src/partition/marginal_core.rs b/packages/treetime/src/partition/marginal_core.rs index 2475d7caa..3a8ff22f0 100644 --- a/packages/treetime/src/partition/marginal_core.rs +++ b/packages/treetime/src/partition/marginal_core.rs @@ -9,6 +9,7 @@ use eyre::Report; use itertools::{Itertools, izip}; use ndarray::prelude::*; use rayon::prelude::*; +use serde::Serialize; use std::collections::BTreeMap; use treetime_graph::edge::{EdgeOptimizeOps, GraphEdge, GraphEdgeKey, HasBranchLength}; use treetime_graph::graph::Graph; @@ -18,7 +19,7 @@ use treetime_utils::array::ndarray::argmax_first; use treetime_utils::array::softmax_with_log_norm::softmax_with_log_norm; use treetime_utils::collections::container::get_exactly_one; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct MarginalData { pub gtr: GTR, pub nodes: BTreeMap, diff --git a/packages/treetime/src/partition/marginal_dense.rs b/packages/treetime/src/partition/marginal_dense.rs index 6a946fa47..9b83c3acf 100644 --- a/packages/treetime/src/partition/marginal_dense.rs +++ b/packages/treetime/src/partition/marginal_dense.rs @@ -19,6 +19,7 @@ use crate::seq::mutation::Sub; use eyre::Report; use itertools::{Itertools, izip}; use maplit::btreemap; +use serde::Serialize; use std::collections::{BTreeMap, BTreeSet}; use treetime_graph::edge::{EdgeOptimizeOps, GraphEdgeKey}; use treetime_graph::graph::Graph; @@ -31,7 +32,7 @@ use treetime_utils::collections::container::get_exactly_one; use treetime_utils::interval::range::range_contains; use treetime_utils::interval::range_union::range_union; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct PartitionMarginalDense { pub data: MarginalData, pub index: usize, @@ -158,6 +159,18 @@ impl PartitionBranchOps for PartitionMarginalDense { Ok(subs) } + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec { + self.data.edges[&edge_key].indels.clone() + } + + fn root_sequence(&self, graph: &dyn BranchTopology) -> Result { + Ok(assign_sequence(&self.data.nodes[&graph.root_key()?], &self.alphabet)) + } + + fn node_sequence(&self, node_key: GraphNodeKey) -> Seq { + assign_sequence(&self.data.nodes[&node_key], &self.alphabet) + } + fn edge_effective_length(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result { let (parent_key, child_key) = graph.edge_endpoints(edge_key)?; let parent_non_char = &self.data.nodes[&parent_key].seq.non_char; diff --git a/packages/treetime/src/partition/marginal_discrete.rs b/packages/treetime/src/partition/marginal_discrete.rs index f831989f7..9b4779bae 100644 --- a/packages/treetime/src/partition/marginal_discrete.rs +++ b/packages/treetime/src/partition/marginal_discrete.rs @@ -14,6 +14,7 @@ use itertools::Itertools; use log::warn; use maplit::btreemap; use ndarray::{Array1, Array2}; +use serde::Serialize; use std::collections::BTreeMap; use treetime_graph::edge::EdgeOptimizeOps; use treetime_graph::graph::Graph; @@ -21,7 +22,7 @@ use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; use treetime_io::nwk::NodeCommentProvider; use treetime_utils::array::ndarray::argmax_first; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct PartitionMarginalDiscrete { pub data: MarginalData, pub states: DiscreteStates, diff --git a/packages/treetime/src/partition/marginal_sparse.rs b/packages/treetime/src/partition/marginal_sparse.rs index 1a232868d..b8e34be0e 100644 --- a/packages/treetime/src/partition/marginal_sparse.rs +++ b/packages/treetime/src/partition/marginal_sparse.rs @@ -15,6 +15,7 @@ use crate::seq::mutation::Sub; use crate::{make_error, make_internal_report}; use eyre::Report; use ndarray::{Array1, Array2}; +use serde::Serialize; use std::collections::{BTreeMap, BTreeSet}; use treetime_graph::edge::{EdgeOptimizeOps, GraphEdgeKey}; use treetime_graph::graph::Graph; @@ -27,7 +28,7 @@ use treetime_utils::array::ndarray::argmax_first; use treetime_utils::collections::container::get_exactly_one; use treetime_utils::interval::range_union::range_union; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, Serialize)] pub struct PartitionMarginalSparse { pub index: usize, pub gtr: GTR, @@ -74,7 +75,7 @@ pub(crate) fn reconstruct_map_seq_sampled( seq[m.pos()] = m.qry(); } for indel in &edge.indels { - if indel.deletion { + if indel.is_deletion() { seq[indel.range.0..indel.range.1].fill(alphabet.gap()); } else { seq[indel.range.0..indel.range.1].copy_from_slice(&indel.seq); @@ -373,7 +374,7 @@ impl PartitionMarginalSparse { } for indel in &edge.indels { if indel.range.0 < seq.len() && indel.range.1 <= seq.len() { - if indel.deletion { + if indel.is_deletion() { seq[indel.range.0..indel.range.1].copy_from_slice(&indel.seq); } else { seq[indel.range.0..indel.range.1].fill(alphabet.gap()); @@ -418,6 +419,18 @@ impl PartitionBranchOps for PartitionMarginalSparse { } } + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec { + self.edges[&edge_key].indels.clone() + } + + fn root_sequence(&self, _graph: &dyn BranchTopology) -> Result { + Ok(self.root_sequence.clone()) + } + + fn node_sequence(&self, node_key: GraphNodeKey) -> Seq { + self.nodes[&node_key].seq.sequence.clone() + } + fn edge_effective_length(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result { let (parent_key, child_key) = graph.edge_endpoints(edge_key)?; let parent_non_char = &self.nodes[&parent_key].seq.non_char; diff --git a/packages/treetime/src/partition/timetree.rs b/packages/treetime/src/partition/timetree.rs index 8836a7b3a..e31bd8344 100644 --- a/packages/treetime/src/partition/timetree.rs +++ b/packages/treetime/src/partition/timetree.rs @@ -1,9 +1,214 @@ -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::ancestral::sample::SampleMode; +use crate::gtr::gtr::GTR; +use crate::partition::marginal_dense::PartitionMarginalDense; +use crate::partition::marginal_sparse::PartitionMarginalSparse; +use crate::partition::optimization_contribution::OptimizationContribution; +use crate::partition::traits::{ + BranchTopology, HasGtr, HasLogLh, PartitionBranchOps, PartitionMarginalOps, PartitionMarginalPasses, + PartitionOptimizeOps, PartitionRerootOps, PartitionTimetreeOps, +}; +use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; +use crate::seq::indel::InDel; +use crate::seq::mutation::Sub; +use eyre::Report; use parking_lot::RwLock; +use serde::Serialize; use std::sync::Arc; +use treetime_graph::edge::{EdgeOptimizeOps, GraphEdgeKey}; use treetime_graph::graph::Graph; +use treetime_graph::graph_traverse::GraphNodeForward; +use treetime_graph::node::{GraphNode, GraphNodeKey, Named}; +use treetime_graph::reroot::RerootChanges; +use treetime_io::fasta::FastaRecord; +use treetime_primitives::Seq; pub type GraphTimetree = Graph; -pub type PartitionTimetreeAllVec = Vec>>>; +pub type PartitionTimetreeRef = Arc>; +pub type PartitionTimetreeAllVec = Vec; + +#[derive(Debug, Serialize)] +pub enum PartitionTimetree { + Dense(PartitionMarginalDense), + Sparse(PartitionMarginalSparse), +} + +impl HasGtr for PartitionTimetree { + fn gtr(&self) -> >R { + match self { + Self::Dense(partition) => partition.gtr(), + Self::Sparse(partition) => partition.gtr(), + } + } + + fn gtr_mut(&mut self) -> &mut GTR { + match self { + Self::Dense(partition) => partition.gtr_mut(), + Self::Sparse(partition) => partition.gtr_mut(), + } + } + + fn sequence_length(&self) -> usize { + match self { + Self::Dense(partition) => HasGtr::sequence_length(partition), + Self::Sparse(partition) => HasGtr::sequence_length(partition), + } + } +} + +impl HasLogLh for PartitionTimetree { + fn get_log_lh(&self, node_key: GraphNodeKey) -> f64 { + match self { + Self::Dense(partition) => partition.get_log_lh(node_key), + Self::Sparse(partition) => partition.get_log_lh(node_key), + } + } + + fn reset_node_log_likelihoods(&mut self) { + match self { + Self::Dense(partition) => partition.reset_node_log_likelihoods(), + Self::Sparse(partition) => partition.reset_node_log_likelihoods(), + } + } +} + +impl PartitionBranchOps for PartitionTimetree { + fn sequence_length(&self) -> usize { + match self { + Self::Dense(partition) => PartitionBranchOps::sequence_length(partition), + Self::Sparse(partition) => PartitionBranchOps::sequence_length(partition), + } + } + + fn edge_subs(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result, Report> { + match self { + Self::Dense(partition) => partition.edge_subs(graph, edge_key), + Self::Sparse(partition) => partition.edge_subs(graph, edge_key), + } + } + + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec { + match self { + Self::Dense(partition) => partition.edge_indels(edge_key), + Self::Sparse(partition) => partition.edge_indels(edge_key), + } + } + + fn root_sequence(&self, graph: &dyn BranchTopology) -> Result { + match self { + Self::Dense(partition) => partition.root_sequence(graph), + Self::Sparse(partition) => partition.root_sequence(graph), + } + } + + fn node_sequence(&self, node_key: GraphNodeKey) -> Seq { + match self { + Self::Dense(partition) => partition.node_sequence(node_key), + Self::Sparse(partition) => partition.node_sequence(node_key), + } + } + + fn edge_effective_length(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result { + match self { + Self::Dense(partition) => partition.edge_effective_length(graph, edge_key), + Self::Sparse(partition) => partition.edge_effective_length(graph, edge_key), + } + } +} + +impl PartitionOptimizeOps for PartitionTimetree { + fn create_edge_contribution(&self, edge_key: GraphEdgeKey) -> Result { + match self { + Self::Dense(partition) => partition.create_edge_contribution(edge_key), + Self::Sparse(partition) => partition.create_edge_contribution(edge_key), + } + } + + fn edge_indel_count(&self, edge_key: GraphEdgeKey) -> usize { + match self { + Self::Dense(partition) => partition.edge_indel_count(edge_key), + Self::Sparse(partition) => partition.edge_indel_count(edge_key), + } + } +} + +impl PartitionRerootOps for PartitionTimetree { + fn apply_reroot(&mut self, changes: &RerootChanges) -> Result<(), Report> { + match self { + Self::Dense(partition) => partition.apply_reroot(changes), + Self::Sparse(partition) => partition.apply_reroot(changes), + } + } +} + +impl PartitionTimetreeOps for PartitionTimetree +where + N: GraphNode + Named, + E: EdgeOptimizeOps, +{ + fn reconcile_topology(&mut self, graph: &Graph) { + match self { + Self::Dense(partition) => partition.reconcile_topology(graph), + Self::Sparse(partition) => partition.reconcile_topology(graph), + } + } +} + +impl PartitionMarginalPasses for PartitionTimetree { + fn process_backward_pass(&mut self, graph: &GraphTimetree) -> Result<(), Report> { + match self { + Self::Dense(partition) => partition.process_backward_pass(graph), + Self::Sparse(partition) => partition.process_backward_pass(graph), + } + } + + fn process_forward_pass(&mut self, graph: &GraphTimetree) -> Result<(), Report> { + match self { + Self::Dense(partition) => partition.process_forward_pass(graph), + Self::Sparse(partition) => partition.process_forward_pass(graph), + } + } + + fn get_sequence_length(&self) -> usize { + match self { + Self::Dense(partition) => partition.get_sequence_length(), + Self::Sparse(partition) => partition.get_sequence_length(), + } + } +} + +impl PartitionMarginalOps for PartitionTimetree { + fn attach_sequences(&mut self, graph: &GraphTimetree, aln: &[FastaRecord]) -> Result<(), Report> { + match self { + Self::Dense(partition) => partition.attach_sequences(graph, aln), + Self::Sparse(partition) => partition.attach_sequences(graph, aln), + } + } + + fn extract_ancestral_sequence(&self, node_key: GraphNodeKey) -> Seq { + match self { + Self::Dense(partition) => { + >::extract_ancestral_sequence( + partition, node_key, + ) + }, + Self::Sparse(partition) => { + >::extract_ancestral_sequence( + partition, node_key, + ) + }, + } + } + + fn reconstruct_node_sequence( + &mut self, + node: &GraphNodeForward, + include_leaves: bool, + sample_mode: SampleMode, + rng: &mut dyn rand::RngCore, + ) -> Option { + match self { + Self::Dense(partition) => partition.reconstruct_node_sequence(node, include_leaves, sample_mode, rng), + Self::Sparse(partition) => partition.reconstruct_node_sequence(node, include_leaves, sample_mode, rng), + } + } +} diff --git a/packages/treetime/src/partition/traits.rs b/packages/treetime/src/partition/traits.rs index e48b721b9..f0cc492cd 100644 --- a/packages/treetime/src/partition/traits.rs +++ b/packages/treetime/src/partition/traits.rs @@ -6,7 +6,8 @@ use crate::make_internal_error; use crate::make_internal_report; use crate::partition::optimization_contribution::OptimizationContribution; use crate::partition::sparse::{SparseEdgePartition, SparseNodePartition}; -use crate::seq::mutation::Sub; +use crate::seq::indel::InDel; +use crate::seq::mutation::{Mutation, MutationEvent, MutationTrack, Sub, mutation_event_strings}; use eyre::Report; use itertools::Itertools; use maplit::btreemap; @@ -121,6 +122,34 @@ pub trait PartitionBranchOps: Send + Sync { /// and gap positions are excluded in both cases. fn edge_subs(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result, Report>; + /// Return grouped aligned insertions and deletions for one edge. + fn edge_indels(&self, edge_key: GraphEdgeKey) -> Vec; + + /// Return the reconstructed root sequence represented by this partition. + fn root_sequence(&self, graph: &dyn BranchTopology) -> Result; + + /// Return the reconstructed sequence for one node. + fn node_sequence(&self, node_key: GraphNodeKey) -> Seq; + + fn edge_mutations( + &self, + graph: &dyn BranchTopology, + edge_key: GraphEdgeKey, + track: MutationTrack, + ) -> Result, Report> { + self + .edge_subs(graph, edge_key)? + .into_iter() + .map(|substitution| Ok(Mutation::substitution(track.clone(), substitution))) + .chain( + self + .edge_indels(edge_key) + .iter() + .map(|indel| Mutation::indel(track.clone(), indel)), + ) + .collect() + } + /// Return the number of alignment positions where both parent and child /// have canonical (non-gap, non-ambiguous) states for one edge. fn edge_effective_length(&self, graph: &dyn BranchTopology, edge_key: GraphEdgeKey) -> Result; @@ -142,13 +171,25 @@ impl NodeCommentProvider for MutationCommentProvider<'_> { let Some((_parent_key, edge_key)) = self.graph.node_parent(key)? else { return Ok(BTreeMap::new()); }; - let mut subs = self.partition.edge_subs(self.graph, edge_key)?; - if subs.is_empty() { + let mut mutations = self + .partition + .edge_mutations(self.graph, edge_key, MutationTrack::Nucleotide)?; + if mutations.is_empty() { return Ok(BTreeMap::new()); } - subs.sort_by_key(|s| s.pos()); + mutations.sort_by_key(|mutation| match &mutation.event { + MutationEvent::Substitution(substitution) => substitution.pos(), + MutationEvent::Insertion(segment) | MutationEvent::Deletion(segment) => segment.range.0, + }); + let mutations = mutations + .iter() + .map(|mutation| mutation_event_strings(&mutation.event)) + .collect::, _>>()? + .into_iter() + .flatten() + .join(","); Ok(btreemap! { - "mutations".to_owned() => subs.iter().join(","), + "mutations".to_owned() => mutations, }) } } diff --git a/packages/treetime/src/payload/ancestral.rs b/packages/treetime/src/payload/ancestral.rs index 8ca3695ad..b00338df6 100644 --- a/packages/treetime/src/payload/ancestral.rs +++ b/packages/treetime/src/payload/ancestral.rs @@ -108,13 +108,14 @@ mod tests { use crate::partition::sparse::{SparseEdgePartition, SparseNodePartition}; use crate::partition::traits::MutationCommentProvider; use crate::payload::ancestral::GraphAncestral; + use crate::seq::indel::InDel; use crate::seq::mutation::Sub; use eyre::Report; use maplit::btreemap; use pretty_assertions::assert_eq; use treetime_graph::node::GraphNodeKey; use treetime_io::nwk::{NodeCommentProvider, nwk_read_str}; - use treetime_primitives::AsciiChar; + use treetime_primitives::{AsciiChar, Seq}; fn c(b: u8) -> AsciiChar { AsciiChar::from_byte_unchecked(b) @@ -126,7 +127,7 @@ mod tests { edge_subs: &[(usize, Vec)], ) -> Result { let alphabet = Alphabet::default(); - let mut ref_seq: treetime_primitives::Seq = std::iter::repeat_with(|| c(b'A')).take(length).collect(); + let mut ref_seq: Seq = std::iter::repeat_with(|| c(b'A')).take(length).collect(); for (_, subs) in edge_subs { for s in subs { if s.pos() < length { @@ -170,9 +171,9 @@ mod tests { } #[test] - fn test_mutation_comment_provider_formats_1_based_positions() -> Result<(), Report> { + fn test_mutation_comment_provider_formats_1_based_substitutions_and_indels() -> Result<(), Report> { let graph: GraphAncestral = nwk_read_str("(A:0.1)root;")?; - let partition = make_test_partition( + let mut partition = make_test_partition( &graph, 100, &[( @@ -183,9 +184,15 @@ mod tests { ], )], )?; + let edge_key = graph.get_edges()[0].read_arc().key(); + partition + .edges + .get_mut(&edge_key) + .expect("fixture edge partition must exist") + .indels = vec![InDel::del((1, 3), Seq::try_from_str("CG")?)?]; let provider = MutationCommentProvider::new(&partition, &graph); let comments = provider.node_comments(leaf_key(&graph))?; - assert_eq!(comments.get("mutations").map(String::as_str), Some("A1T,G6C")); + assert_eq!(comments.get("mutations").map(String::as_str), Some("A1T,C2-,G3-,G6C")); Ok(()) } diff --git a/packages/treetime/src/prune/__tests__/test_prune.rs b/packages/treetime/src/prune/__tests__/test_prune.rs index 4263caea9..bdccf7169 100644 --- a/packages/treetime/src/prune/__tests__/test_prune.rs +++ b/packages/treetime/src/prune/__tests__/test_prune.rs @@ -1282,8 +1282,8 @@ mod tests { root_sequence: seq![], }; - let parent_indel = InDel::del((10, 15), [c(b'A'), c(b'C'), c(b'G'), c(b'T'), c(b'A')].as_slice()); - let child_indel = InDel::ins((20, 23), [c(b'G'), c(b'G'), c(b'C')].as_slice()); + let parent_indel = InDel::del((10, 15), [c(b'A'), c(b'C'), c(b'G'), c(b'T'), c(b'A')].as_slice())?; + let child_indel = InDel::ins((20, 23), [c(b'G'), c(b'G'), c(b'C')].as_slice())?; partition.edges.insert( root_internal_edge_key, @@ -1320,9 +1320,9 @@ mod tests { // Parent indel first, then child indel assert_eq!(edge_partition.indels.len(), 2); assert_eq!(edge_partition.indels[0].range, (10, 15)); - assert!(edge_partition.indels[0].deletion); + assert!(edge_partition.indels[0].is_deletion()); assert_eq!(edge_partition.indels[1].range, (20, 23)); - assert!(!edge_partition.indels[1].deletion); + assert!(!edge_partition.indels[1].is_deletion()); } else if target_name.as_deref() == Some("B") { let edge_partition = &partition.edges[&edge_key]; // Only parent indel (child had none) diff --git a/packages/treetime/src/seq/__tests__/test_composition.rs b/packages/treetime/src/seq/__tests__/test_composition.rs index 692862bdd..750928cfe 100644 --- a/packages/treetime/src/seq/__tests__/test_composition.rs +++ b/packages/treetime/src/seq/__tests__/test_composition.rs @@ -110,7 +110,7 @@ mod tests { #[test] fn test_composition_add_deletion() -> Result<(), Report> { let mut actual = Composition::with_seq_str("AAAGCTTACGGGGTCAAGTCC", chars("ACGT-"), c(b'-'))?; - let indel = InDel::del((1, 5), Seq::try_from_str("AAGC")?); + let indel = InDel::del((1, 5), Seq::try_from_str("AAGC")?)?; actual.add_indel(&indel); let expected = Composition::from_counts( btreemap! { c(b'-') => 4, c(b'A') => 4, c(b'C') => 4, c(b'G') => 5, c(b'T') => 4}, @@ -123,7 +123,7 @@ mod tests { #[test] fn test_composition_add_insertion() -> Result<(), Report> { let mut actual = Composition::with_seq_str("AAAGCTTACGGGGTCAAGTCC", chars("ACGT-"), c(b'-'))?; - let indel = InDel::ins((3, 6), Seq::try_from_str("ATC")?); + let indel = InDel::ins((3, 6), Seq::try_from_str("ATC")?)?; actual.add_indel(&indel); let expected = Composition::from_counts( btreemap! { c(b'-') => 0, c(b'A') => 7, c(b'C') => 6, c(b'G') => 6, c(b'T') => 5}, diff --git a/packages/treetime/src/seq/__tests__/test_indel.rs b/packages/treetime/src/seq/__tests__/test_indel.rs index b03db011e..911affbcb 100644 --- a/packages/treetime/src/seq/__tests__/test_indel.rs +++ b/packages/treetime/src/seq/__tests__/test_indel.rs @@ -6,11 +6,11 @@ mod tests { use treetime_primitives::Seq; fn del(start: usize, end: usize, seq: &str) -> InDel { - InDel::del((start, end), Seq::try_from_str(seq).unwrap()) + InDel::del((start, end), Seq::try_from_str(seq).unwrap()).unwrap() } fn ins(start: usize, end: usize, seq: &str) -> InDel { - InDel::ins((start, end), Seq::try_from_str(seq).unwrap()) + InDel::ins((start, end), Seq::try_from_str(seq).unwrap()).unwrap() } #[test] @@ -230,4 +230,20 @@ mod tests { sort_indels(&mut indels); assert_eq!(indels, vec![del(1, 3, "CG"), del(3, 5, "GT"), del(5, 7, "AC")]); } + + #[rustfmt::skip] + #[rstest] + #[case::reversed((3, 2), "A", "non-empty and ordered")] + #[case::empty( (2, 2), "", "non-empty and ordered")] + #[case::short( (2, 4), "A", "sequence has length 1")] + #[trace] + fn test_indel_constructor_rejects_invalid_event( + #[case] range: (usize, usize), + #[case] sequence: &str, + #[case] message: &str, + ) { + let sequence = Seq::try_from_str(sequence).unwrap(); + let error = InDel::del(range, sequence).expect_err("invalid indel must be rejected"); + assert!(error.to_string().contains(message)); + } } diff --git a/packages/treetime/src/seq/__tests__/test_mutation.rs b/packages/treetime/src/seq/__tests__/test_mutation.rs index 6cd589b11..cfd0adade 100644 --- a/packages/treetime/src/seq/__tests__/test_mutation.rs +++ b/packages/treetime/src/seq/__tests__/test_mutation.rs @@ -1,9 +1,11 @@ #[cfg(test)] mod tests { - use crate::seq::mutation::{Sub, compose_substitutions}; + use crate::seq::mutation::{AlignedMutation, MutationEvent, Sub, compose_substitutions, mutation_event_strings}; use eyre::Report; use pretty_assertions::assert_eq; + use proptest::prelude::*; use treetime_primitives::AsciiChar; + use treetime_primitives::seq; #[test] fn test_mutation_compose_substitutions_both_empty() -> Result<(), Report> { @@ -133,6 +135,29 @@ mod tests { Ok(()) } + #[test] + fn test_mutation_event_strings_expand_aligned_deletion() -> Result<(), Report> { + // Oracle: TreeTime v0 spells gap transitions as one-based per-position substitutions. + let event = MutationEvent::Deletion(AlignedMutation::new((1, 3), seq![helpers::c(b'C'), helpers::c(b'G')])?); + let actual = mutation_event_strings(&event)?; + let expected = vec!["C2-".to_owned(), "G3-".to_owned()]; + assert_eq!(expected, actual); + Ok(()) + } + + proptest! { + #[test] + fn test_prop_mutation_event_strings_preserve_range(start in 0_usize..10_000, length in 1_usize..32) { + let sequence = seq![helpers::c(b'A'); length]; + let event = MutationEvent::Insertion(AlignedMutation::new((start, start + length), sequence).unwrap()); + let actual = mutation_event_strings(&event).unwrap(); + let expected = (start + 1..=start + length) + .map(|position| format!("-{position}A")) + .collect::>(); + prop_assert_eq!(expected, actual); + } + } + mod helpers { use super::*; diff --git a/packages/treetime/src/seq/composition.rs b/packages/treetime/src/seq/composition.rs index 92058e30d..a9f00692f 100644 --- a/packages/treetime/src/seq/composition.rs +++ b/packages/treetime/src/seq/composition.rs @@ -84,7 +84,7 @@ impl Composition { /// Reflect sequence indel in the composition counts pub fn add_indel(&mut self, indel: &InDel) { - let adjust_by = if indel.deletion { -1 } else { 1 }; + let adjust_by = if indel.is_deletion() { -1 } else { 1 }; for nuc in &indel.seq { self.adjust_count(*nuc, adjust_by); self.adjust_count(self.gap, -adjust_by); diff --git a/packages/treetime/src/seq/indel.rs b/packages/treetime/src/seq/indel.rs index 07e44f874..f55b8d148 100644 --- a/packages/treetime/src/seq/indel.rs +++ b/packages/treetime/src/seq/indel.rs @@ -1,3 +1,4 @@ +use eyre::Report; use itertools::Itertools; use serde::{Deserialize, Serialize}; use std::collections::BTreeSet; @@ -6,39 +7,65 @@ use treetime_primitives::Seq; use treetime_utils::interval::range_difference::range_difference; use treetime_utils::interval::range_intersection::{range_intersection, range_intersection_iter}; use treetime_utils::interval::range_union::range_union_iter; +use treetime_utils::make_error; #[derive(Clone, Debug, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)] pub struct InDel { pub range: (usize, usize), pub seq: Seq, - pub deletion: bool, // deletion if True, insertion if False + pub kind: InDelKind, } impl InDel { - pub fn del(range: (usize, usize), seq: impl Into) -> Self { - Self::new(range, seq, true) + pub fn del(range: (usize, usize), seq: impl Into) -> Result { + Self::new(range, seq, InDelKind::Deletion) } - pub fn ins(range: (usize, usize), seq: impl Into) -> Self { - Self::new(range, seq, false) + pub fn ins(range: (usize, usize), seq: impl Into) -> Result { + Self::new(range, seq, InDelKind::Insertion) } - pub fn new(range: (usize, usize), seq: impl Into, deletion: bool) -> Self { + pub fn new(range: (usize, usize), seq: impl Into, kind: InDelKind) -> Result { let seq = seq.into(); - assert!(range.0 <= range.1); - assert_eq!(seq.len(), range.1 - range.0); - Self { range, seq, deletion } + let Some(length) = range.1.checked_sub(range.0).filter(|length| *length > 0) else { + return make_error!( + "Indel range must be non-empty and ordered, got {}..{}", + range.0, + range.1 + ); + }; + if seq.len() != length { + return make_error!( + "Indel range {}..{} has length {length}, but its sequence has length {}", + range.0, + range.1, + seq.len() + ); + } + Ok(Self { range, seq, kind }) } - /// Invert the indel direction by toggling deletion flag pub fn invert(&mut self) { - self.deletion = !self.deletion; + self.kind = match self.kind { + InDelKind::Insertion => InDelKind::Deletion, + InDelKind::Deletion => InDelKind::Insertion, + }; } + + pub fn is_deletion(&self) -> bool { + self.kind == InDelKind::Deletion + } +} + +#[derive(Clone, Copy, Debug, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)] +pub enum InDelKind { + Insertion, + Deletion, } impl fmt::Display for InDel { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - let delta_str = if self.deletion { + let delta_str = if self.is_deletion() { format!("{} -> {}", &self.seq.as_str(), "-".repeat(self.seq.len())) } else { format!("{} -> {}", "-".repeat(self.seq.len()), &self.seq.as_str()) @@ -95,7 +122,7 @@ pub fn compose_indels(parent_indels: &[InDel], child_indels: &[InDel]) -> Vec Vec { + match (p.kind, c.kind) { + (InDelKind::Deletion, InDelKind::Deletion) => { // Cases 1, 6: overlapping deletions - use parent's seq for overlap region let merged_start = p.range.0.min(c.range.0); let merged_end = p.range.1.max(c.range.1); @@ -153,13 +184,17 @@ pub fn compose_indels(parent_indels: &[InDel], child_indels: &[InDel]) -> Vec { + (InDelKind::Insertion, InDelKind::Insertion) => { // Case 3: overlapping insertions - child's seq wins result.push(c.clone()); }, - (true, false) | (false, true) => { + (InDelKind::Deletion, InDelKind::Insertion) | (InDelKind::Insertion, InDelKind::Deletion) => { if !(p.range == c.range && p.seq == c.seq) { // Emit in sorted order by range.0 to maintain output sortedness if p.range.0 <= c.range.0 { @@ -194,7 +229,7 @@ fn merge_adjacent_deletions(indels: Vec) -> Vec { #[allow(clippy::suspicious_operation_groupings)] let should_merge = merged .last() - .is_some_and(|prev: &InDel| prev.deletion && indel.deletion && prev.range.1 == indel.range.0); + .is_some_and(|prev: &InDel| prev.is_deletion() && indel.is_deletion() && prev.range.1 == indel.range.0); if should_merge { let prev = merged.last_mut().unwrap(); prev.range.1 = indel.range.1; @@ -349,7 +384,11 @@ pub fn resolve_indels_forward( last.range.1 = hi; last.seq.extend(parent_sequence[lo..hi].iter().copied()); } else { - deletions.push(InDel::del((lo, hi), &parent_sequence[lo..hi])); + deletions.push(InDel { + range: (lo, hi), + seq: (&parent_sequence[lo..hi]).into(), + kind: InDelKind::Deletion, + }); } if let Some(last) = new_node_gaps.last_mut().filter(|l| l.1 == lo) { last.1 = hi; @@ -362,7 +401,11 @@ pub fn resolve_indels_forward( last.range.1 = hi; last.seq.extend(node_sequence[lo..hi].iter().copied()); } else { - insertions.push(InDel::ins((lo, hi), &node_sequence[lo..hi])); + insertions.push(InDel { + range: (lo, hi), + seq: (&node_sequence[lo..hi]).into(), + kind: InDelKind::Insertion, + }); } } else if in_parent { // Parent has gap; node is gapped, non_char, or variable: inherit parent gap. diff --git a/packages/treetime/src/seq/mutation.rs b/packages/treetime/src/seq/mutation.rs index 61b5c7c1c..3fe4804f8 100644 --- a/packages/treetime/src/seq/mutation.rs +++ b/packages/treetime/src/seq/mutation.rs @@ -1,4 +1,5 @@ use crate::alphabet::alphabet::Alphabet; +use crate::seq::indel::{InDel, InDelKind}; use crate::{make_error, make_internal_error}; use eyre::{Report, WrapErr}; use getset::CopyGetters; @@ -9,12 +10,115 @@ use std::fmt; use std::str::FromStr; use std::sync::LazyLock; use treetime_primitives::AsciiChar; +use treetime_primitives::Seq; use treetime_utils::error::to_eyre_error; static NUC_MUT_RE: LazyLock = LazyLock::new(|| { Regex::new(r"^(?P[A-Z])(?P\d{1,10})(?P[A-Z])$").expect("NUC_MUT_RE regex compilation") }); +#[derive(Clone, Debug, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)] +pub struct Mutation { + pub track: MutationTrack, + pub event: MutationEvent, +} + +impl Mutation { + pub fn substitution(track: MutationTrack, substitution: Sub) -> Self { + Self { + track, + event: MutationEvent::Substitution(substitution), + } + } + + pub fn indel(track: MutationTrack, indel: &InDel) -> Result { + let segment = AlignedMutation::new(indel.range, indel.seq.clone())?; + let event = match indel.kind { + InDelKind::Insertion => MutationEvent::Insertion(segment), + InDelKind::Deletion => MutationEvent::Deletion(segment), + }; + Ok(Self { track, event }) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)] +pub enum MutationTrack { + Nucleotide, + AminoAcid(String), +} + +#[derive(Clone, Debug, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)] +pub enum MutationEvent { + Substitution(Sub), + Insertion(AlignedMutation), + Deletion(AlignedMutation), +} + +pub fn mutation_event_strings(event: &MutationEvent) -> Result, Report> { + match event { + MutationEvent::Substitution(substitution) => Ok(vec![substitution.to_string()]), + MutationEvent::Insertion(segment) => segment + .sequence + .iter() + .enumerate() + .map(|(offset, state)| mutation_position(segment.range.0, offset).map(|position| format!("-{position}{state}"))) + .collect(), + MutationEvent::Deletion(segment) => segment + .sequence + .iter() + .enumerate() + .map(|(offset, state)| mutation_position(segment.range.0, offset).map(|position| format!("{state}{position}-"))) + .collect(), + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq)] +pub struct AlignedMutation { + pub range: (usize, usize), + pub sequence: Seq, +} + +impl AlignedMutation { + pub fn new(range: (usize, usize), sequence: Seq) -> Result { + if range.0 >= range.1 { + return make_error!( + "Aligned mutation range must be non-empty and ordered, but found {}..{}", + range.0, + range.1 + ); + } + let range_length = range + .1 + .checked_sub(range.0) + .ok_or_else(|| eyre::eyre!("Aligned mutation range underflow for {}..{}", range.0, range.1))?; + if sequence.len() != range_length { + return make_error!( + "Aligned mutation range {}..{} has length {range_length}, but its sequence has length {}", + range.0, + range.1, + sequence.len() + ); + } + Ok(Self { range, sequence }) + } + + pub fn one_based_inclusive_range(&self) -> Result<(usize, usize), Report> { + let start = self + .range + .0 + .checked_add(1) + .ok_or_else(|| eyre::eyre!("Mutation start coordinate overflow at {}", self.range.0))?; + Ok((start, self.range.1)) + } +} + +fn mutation_position(start: usize, offset: usize) -> Result { + start + .checked_add(offset) + .and_then(|position| position.checked_add(1)) + .ok_or_else(|| eyre::eyre!("Mutation coordinate overflow at start {start} and offset {offset}")) +} + #[derive(Clone, Debug, Serialize, Deserialize, Ord, PartialOrd, Eq, PartialEq, CopyGetters)] #[getset(get_copy = "pub")] pub struct Sub { diff --git a/packages/treetime/src/timetree/confidence.rs b/packages/treetime/src/timetree/confidence.rs index 9d8d0769b..23e696d3b 100644 --- a/packages/treetime/src/timetree/confidence.rs +++ b/packages/treetime/src/timetree/confidence.rs @@ -1,23 +1,18 @@ use crate::clock::clock_model::{ClockModel, ClockModelStats}; use crate::make_error; -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::partition::timetree::{GraphTimetree, PartitionTimetreeRef}; use crate::payload::traits::TimetreeNode; use crate::timetree::inference::runner::run_timetree; use eyre::{Report, WrapErr}; use itertools::Itertools; use log::{info, warn}; use ordered_float::OrderedFloat; -use parking_lot::RwLock; use serde::Serialize; use statrs::function::erf::erf_inv; use std::collections::BTreeMap; use std::f64::consts::SQRT_2; use std::io::Write; use std::path::Path; -use std::sync::Arc; use treetime_distribution::Distribution; use treetime_graph::edge::GraphEdgeKey; use treetime_graph::node::{GraphNodeKey, Named, TimeConstraint}; @@ -57,7 +52,7 @@ const CI_UPPER_QUANTILE: f64 = 1.0 - (1.0 - CI_FRACTION) * 0.5; // 0.95 /// caller can proceed with further passes (e.g., final marginal reconstruction). pub fn compute_rate_susceptibility( graph: &mut GraphTimetree, - partitions: &[Arc>>], + partitions: &[PartitionTimetreeRef], clock_model: &ClockModel, coalescent_tc: Option<&Distribution>, rate_std: f64, @@ -238,7 +233,6 @@ pub fn extract_confidence_intervals(graph: &GraphTimetree) -> Vec>>], -) -> Option { +pub fn compute_sequence_likelihood(graph: &GraphTimetree, partitions: &[PartitionTimetreeRef]) -> Option { if partitions.is_empty() { return None; } diff --git a/packages/treetime/src/timetree/convergence/optimizer.rs b/packages/treetime/src/timetree/convergence/optimizer.rs index 458b361f8..859f3d12d 100644 --- a/packages/treetime/src/timetree/convergence/optimizer.rs +++ b/packages/treetime/src/timetree/convergence/optimizer.rs @@ -1,16 +1,11 @@ -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::partition::timetree::{GraphTimetree, PartitionTimetreeRef}; use crate::timetree::convergence::likelihood::{ compute_coalescent_likelihood, compute_positional_likelihood, compute_sequence_likelihood, }; use crate::timetree::convergence::metrics::ConvergenceMetrics; use eyre::Report; use log::info; -use parking_lot::RwLock; use std::io::Write; -use std::sync::Arc; use treetime_distribution::Distribution; use treetime_io::csv::CsvStructWriter; @@ -55,7 +50,7 @@ impl TimetreeOptimizer { n_diff: usize, n_resolved: usize, graph: &GraphTimetree, - partitions: &[Arc>>], + partitions: &[PartitionTimetreeRef], coalescent_tc: Option<&Distribution>, ) -> Result<(), Report> { let lh_seq = compute_sequence_likelihood(graph, partitions); diff --git a/packages/treetime/src/timetree/convergence/sequence_changes.rs b/packages/treetime/src/timetree/convergence/sequence_changes.rs index 6aabe246e..29bf8441f 100644 --- a/packages/treetime/src/timetree/convergence/sequence_changes.rs +++ b/packages/treetime/src/timetree/convergence/sequence_changes.rs @@ -1,11 +1,7 @@ -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::partition::timetree::{GraphTimetree, PartitionTimetreeRef}; +use crate::partition::traits::PartitionMarginalOps; use log::debug; -use parking_lot::RwLock; use std::collections::BTreeMap; -use std::sync::Arc; use treetime_graph::node::GraphNodeKey; use treetime_primitives::Seq; @@ -47,10 +43,7 @@ pub fn count_sequence_changes(previous: &AncestralStateSnapshot, current: &Ances } /// Snapshot current ancestral sequences for all internal nodes across partitions. -pub fn capture_ancestral_states( - graph: &GraphTimetree, - partitions: &[Arc>>], -) -> AncestralStateSnapshot { +pub fn capture_ancestral_states(graph: &GraphTimetree, partitions: &[PartitionTimetreeRef]) -> AncestralStateSnapshot { if partitions.is_empty() { return vec![]; } diff --git a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_dense.rs b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_dense.rs index 80b7f29c0..13036e338 100644 --- a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_dense.rs +++ b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_dense.rs @@ -9,9 +9,7 @@ mod tests { use crate::clock::find_best_root::params::BranchPointOptimizationParams; use crate::gtr::get_gtr::{JC69Params, jc69}; use crate::partition::marginal_dense::PartitionMarginalDense; - use crate::partition::timetree::GraphTimetree; - use crate::partition::traits::PartitionTimetreeAll; - use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; + use crate::partition::timetree::{GraphTimetree, PartitionTimetree, PartitionTimetreeAllVec}; use crate::timetree::inference::runner::run_timetree; use crate::timetree::utils::{ extract_node_times, initialize_clock_totals_from_time_distributions, initialize_node_divergences, @@ -50,14 +48,14 @@ mod tests { load_date_constraints(&dates, &graph)?; let aln = load_alignment_for_dataset(dataset)?; - let dense_partition = Arc::new(RwLock::new(PartitionMarginalDense::new( + let dense_partition = Arc::new(RwLock::new(PartitionTimetree::Dense(PartitionMarginalDense::new( 0, jc69(JC69Params::default())?, ALPHABET.clone(), case.sequence_length(), - ))); + )))); - let partitions: Vec>>> = vec![dense_partition]; + let partitions: PartitionTimetreeAllVec = vec![dense_partition]; initialize_marginal(&graph, &partitions, &aln)?; initialize_node_divergences(&graph)?; diff --git a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_sparse.rs b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_sparse.rs index df6ed7b49..a3745c580 100644 --- a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_sparse.rs +++ b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_marginal_sparse.rs @@ -10,14 +10,12 @@ mod tests { use crate::clock::date_constraints::load_date_constraints; use crate::clock::find_best_root::params::BranchPointOptimizationParams; use crate::gtr::get_gtr::{JC69Params, jc69}; - use crate::partition::traits::PartitionTimetreeAll; use crate::timetree::inference::runner::run_timetree; use crate::timetree::utils::{ extract_node_times, initialize_clock_totals_from_time_distributions, initialize_node_divergences, }; - use crate::partition::timetree::GraphTimetree; - use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; + use crate::partition::timetree::{GraphTimetree, PartitionTimetree, PartitionTimetreeAllVec}; use eyre::Report; use parking_lot::RwLock; @@ -57,9 +55,11 @@ mod tests { let aln = load_alignment_for_dataset(dataset)?; let fitch = create_fitch_partition(&graph, 0, ALPHABET.clone(), &aln)?; - let sparse_partition = Arc::new(RwLock::new(fitch.into_marginal_sparse(jc69(JC69Params::default())?, &graph)?)); + let sparse_partition = Arc::new(RwLock::new(PartitionTimetree::Sparse( + fitch.into_marginal_sparse(jc69(JC69Params::default())?, &graph)?, + ))); - let partitions: Vec>>> = vec![sparse_partition]; + let partitions: PartitionTimetreeAllVec = vec![sparse_partition]; initialize_marginal(&graph, &partitions, &aln)?; initialize_node_divergences(&graph)?; diff --git a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_pre_optimize.rs b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_pre_optimize.rs index 4a82af193..d9b2ca2ca 100644 --- a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_pre_optimize.rs +++ b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_gm_runner_pre_optimize.rs @@ -13,14 +13,12 @@ mod tests { use crate::optimize::dispatch::run_optimize_mixed; use crate::optimize::params::BranchOptMethod; use crate::partition::traits::PartitionOptimizeOps; - use crate::partition::traits::PartitionTimetreeAll; use crate::timetree::inference::runner::run_timetree; use crate::timetree::utils::{ extract_node_times, initialize_clock_totals_from_time_distributions, initialize_node_divergences, }; - use crate::partition::timetree::GraphTimetree; - use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; + use crate::partition::timetree::{GraphTimetree, PartitionTimetree, PartitionTimetreeAllVec}; use eyre::Report; use itertools::Itertools; @@ -54,10 +52,11 @@ mod tests { let graph: GraphTimetree = nwk_read_str(case.rerooted_tree_nwk())?; let aln = load_alignment_for_dataset(dataset)?; let fitch = create_fitch_partition(&graph, 0, ALPHABET.clone(), &aln)?; - let sparse_partition = Arc::new(RwLock::new(fitch.into_marginal_sparse(jc69(JC69Params::default())?, &graph)?)); + let sparse_partition = Arc::new(RwLock::new(PartitionTimetree::Sparse( + fitch.into_marginal_sparse(jc69(JC69Params::default())?, &graph)?, + ))); - let partitions: Vec>>> = - vec![sparse_partition]; + let partitions: PartitionTimetreeAllVec = vec![sparse_partition]; initialize_marginal(&graph, &partitions, &aln)?; let before = extract_branch_lengths(&graph); @@ -105,10 +104,11 @@ mod tests { let aln = load_alignment_for_dataset(dataset)?; let fitch = create_fitch_partition(&graph, 0, ALPHABET.clone(), &aln)?; - let sparse_partition = Arc::new(RwLock::new(fitch.into_marginal_sparse(jc69(JC69Params::default())?, &graph)?)); + let sparse_partition = Arc::new(RwLock::new(PartitionTimetree::Sparse( + fitch.into_marginal_sparse(jc69(JC69Params::default())?, &graph)?, + ))); - let partitions: Vec>>> = - vec![sparse_partition]; + let partitions: PartitionTimetreeAllVec = vec![sparse_partition]; initialize_marginal(&graph, &partitions, &aln)?; initialize_node_divergences(&graph)?; diff --git a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_runner_coalescent.rs b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_runner_coalescent.rs index 21226499e..4ad51b924 100644 --- a/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_runner_coalescent.rs +++ b/packages/treetime/src/timetree/inference/__tests__/test_gm_runner/test_runner_coalescent.rs @@ -10,9 +10,7 @@ mod tests { use crate::clock::find_best_root::params::BranchPointOptimizationParams; use crate::gtr::get_gtr::{JC69Params, jc69}; use crate::partition::marginal_dense::PartitionMarginalDense; - use crate::partition::timetree::GraphTimetree; - use crate::partition::traits::PartitionTimetreeAll; - use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; + use crate::partition::timetree::{GraphTimetree, PartitionTimetree, PartitionTimetreeAllVec}; use crate::timetree::inference::runner::run_timetree; use crate::timetree::utils::{ extract_node_times, initialize_clock_totals_from_time_distributions, initialize_node_divergences, @@ -64,27 +62,20 @@ mod tests { fn build_timetree_setup( dataset: &str, case: &super::super::test_gm_runner_support::support::DatasetOutputs, - ) -> Result< - ( - GraphTimetree, - Vec>>>, - ClockModel, - ), - Report, - > { + ) -> Result<(GraphTimetree, PartitionTimetreeAllVec, ClockModel), Report> { let mut graph: GraphTimetree = nwk_read_str(case.rerooted_tree_nwk())?; let dates = load_dates_for_dataset(dataset)?; load_date_constraints(&dates, &graph)?; let aln = load_alignment_for_dataset(dataset)?; - let dense_partition = Arc::new(RwLock::new(PartitionMarginalDense::new( + let dense_partition = Arc::new(RwLock::new(PartitionTimetree::Dense(PartitionMarginalDense::new( 0, jc69(JC69Params::default())?, ALPHABET.clone(), case.sequence_length(), - ))); + )))); - let partitions: Vec>>> = vec![dense_partition]; + let partitions: PartitionTimetreeAllVec = vec![dense_partition]; initialize_marginal(&graph, &partitions, &aln)?; initialize_node_divergences(&graph)?; diff --git a/packages/treetime/src/timetree/optimization/__tests__/test_reroot.rs b/packages/treetime/src/timetree/optimization/__tests__/test_reroot.rs index 2bd36819c..8d5819a20 100644 --- a/packages/treetime/src/timetree/optimization/__tests__/test_reroot.rs +++ b/packages/treetime/src/timetree/optimization/__tests__/test_reroot.rs @@ -10,10 +10,8 @@ mod tests { use crate::o; use crate::partition::marginal_sparse::PartitionMarginalSparse; use crate::partition::sparse::{SparseEdgePartition, SparseNodePartition, SparseSeqDistribution}; - use crate::partition::timetree::GraphTimetree; - use crate::partition::traits::{PartitionRerootOps, PartitionTimetreeAll}; - use crate::payload::timetree::EdgeTimetree; - use crate::payload::timetree::NodeTimetree; + use crate::partition::timetree::{GraphTimetree, PartitionTimetree}; + use crate::partition::traits::PartitionRerootOps; use crate::pretty_assert_ulps_eq; use crate::seq::indel::InDel; use crate::seq::mutation::Sub; @@ -88,14 +86,15 @@ mod tests { })?; let fitch = create_fitch_partition(&graph, 0, alphabet, &aln)?; - let sparse_partition = Arc::new(RwLock::new(fitch.into_marginal_sparse(gtr, &graph)?)); + let sparse_partition = Arc::new(RwLock::new(PartitionTimetree::Sparse( + fitch.into_marginal_sparse(gtr, &graph)?, + ))); let clock_params = ClockParams::default(); clock_regression_backward(&graph, &clock_params, None)?; clock_regression_forward(&graph, &clock_params, None)?; - let partition: Arc>> = sparse_partition; - let partitions = vec![partition]; + let partitions = vec![sparse_partition]; // Record initial state let initial_leaf_count = graph.get_leaves().len(); @@ -192,7 +191,7 @@ mod tests { AsciiChar::from_byte_unchecked(b'A'), AsciiChar::from_byte_unchecked(b'C') ], - ); // deletion=true + )?; // deletion=true let mut sparse_partition = PartitionMarginalSparse { index: 0, @@ -233,7 +232,7 @@ mod tests { // Verify indel is inverted let indel_after = &edge_data.indels[0]; - assert_eq!(indel_after.deletion, false, "Indel deletion flag should be toggled"); + assert!(!indel_after.is_deletion(), "Indel direction should be toggled"); // Verify msg_from_child is cleared pretty_assert_ulps_eq!(edge_data.msg_from_child.log_lh, 0.0, max_ulps = 5); @@ -359,7 +358,7 @@ mod tests { .ok_or_else(|| make_report!("Edge to A not found"))?; let root_seq = Seq::try_from_slice(b"ACGTACGT")?; - let indel = InDel::del((2, 4), seq![c(b'G'), c(b'T')]); + let indel = InDel::del((2, 4), seq![c(b'G'), c(b'T')])?; let mut sparse_partition = PartitionMarginalSparse { index: 0, @@ -484,14 +483,15 @@ mod tests { })?; let fitch = create_fitch_partition(&graph, 0, alphabet, &aln)?; - let sparse_partition = Arc::new(RwLock::new(fitch.into_marginal_sparse(gtr, &graph)?)); + let sparse_partition = Arc::new(RwLock::new(PartitionTimetree::Sparse( + fitch.into_marginal_sparse(gtr, &graph)?, + ))); let clock_params = ClockParams::default(); clock_regression_backward(&graph, &clock_params, None)?; clock_regression_forward(&graph, &clock_params, None)?; - let partition: Arc>> = sparse_partition; - let partitions = vec![partition]; + let partitions = vec![sparse_partition]; // Record initial state let initial_leaf_count = graph.get_leaves().len(); diff --git a/packages/treetime/src/timetree/optimization/polytomy.rs b/packages/treetime/src/timetree/optimization/polytomy.rs index 740567b43..fb4772a0f 100644 --- a/packages/treetime/src/timetree/optimization/polytomy.rs +++ b/packages/treetime/src/timetree/optimization/polytomy.rs @@ -1,14 +1,11 @@ use crate::optimize::topology::polytomy_nodes::find_polytomy_nodes; -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionTimetreeAll; +use crate::partition::timetree::{GraphTimetree, PartitionTimetreeRef}; use crate::payload::clock_set::ClockSet; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; use argmin::core::{CostFunction, Executor}; use argmin::solver::brent::BrentOpt; use eyre::Report; use log::debug; -use parking_lot::RwLock; use std::collections::BTreeMap; use std::sync::Arc; use treetime_distribution::Distribution; @@ -32,7 +29,7 @@ pub const DEFAULT_RESOLUTION_THRESHOLD: f64 = 0.05; /// `clock_rate` = substitutions per site per time unit. pub fn resolve_polytomies( graph: &mut GraphTimetree, - partitions: &[Arc>>], + partitions: &[PartitionTimetreeRef], zero_branch_slope: f64, clock_rate: f64, ) -> Result { @@ -55,7 +52,7 @@ pub fn resolve_polytomies( /// - `merge_compressed`: if true, also merge compressed children after stretched pub fn resolve_polytomies_with_options( graph: &mut GraphTimetree, - partitions: &[Arc>>], + partitions: &[PartitionTimetreeRef], resolution_threshold: f64, zero_branch_slope: f64, clock_rate: f64, diff --git a/packages/treetime/src/timetree/optimization/reroot.rs b/packages/treetime/src/timetree/optimization/reroot.rs index 876a29f50..0462138ba 100644 --- a/packages/treetime/src/timetree/optimization/reroot.rs +++ b/packages/treetime/src/timetree/optimization/reroot.rs @@ -3,14 +3,10 @@ use crate::clock::clock_model::ClockModel; use crate::clock::clock_regression::{ClockParams, estimate_clock_model_with_reroot_policy}; use crate::clock::find_best_root::params::{BranchPointOptimizationParams, RerootSpec}; use crate::clock::reroot::RerootParams; -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::partition::timetree::{GraphTimetree, PartitionTimetreeRef}; +use crate::partition::traits::PartitionRerootOps; use eyre::{Report, WrapErr}; use log::info; -use parking_lot::RwLock; -use std::sync::Arc; use treetime_graph::reroot::RerootChanges; /// Reroot tree for optimal temporal signal and update partition state. @@ -19,7 +15,7 @@ use treetime_graph::reroot::RerootChanges; /// with bundled topology changes (edge split, edge merge, inverted edges). pub fn reroot_tree( graph: &mut GraphTimetree, - partitions: &[Arc>>], + partitions: &[PartitionTimetreeRef], clock_params: &ClockParams, clock_rate: Option, branch_params: &BranchPointOptimizationParams, diff --git a/packages/treetime/src/timetree/pipeline.rs b/packages/treetime/src/timetree/pipeline.rs index 1df3da71a..faf43e927 100644 --- a/packages/treetime/src/timetree/pipeline.rs +++ b/packages/treetime/src/timetree/pipeline.rs @@ -15,9 +15,8 @@ use crate::optimize::dispatch::{run_optimize_mixed, run_optimize_mixed_inner}; use crate::optimize::iteration::{apply_damping, save_branch_lengths}; use crate::optimize::params::{BranchLengthMode, BranchOptMethod}; use crate::partition::create::{MarginalPartition, create_marginal_partition}; -use crate::partition::timetree::{GraphTimetree, PartitionTimetreeAllVec}; -use crate::partition::traits::{HasGtr, PartitionTimetreeAll}; -use crate::payload::timetree::{EdgeTimetree, NodeTimetree}; +use crate::partition::timetree::{GraphTimetree, PartitionTimetree, PartitionTimetreeAllVec, PartitionTimetreeRef}; +use crate::partition::traits::HasGtr; use crate::progress::ProgressSink; use crate::timetree::confidence::{ NodeConfidenceInterval, compute_rate_susceptibility, determine_rate_std, extract_confidence_intervals, @@ -481,9 +480,9 @@ fn initialize_partitions_from_params( MarginalPartition::Dense(p) => p.gtr().clone(), }; - let partition: Arc>> = match created.partition { - MarginalPartition::Sparse(p) => Arc::new(RwLock::new(p)), - MarginalPartition::Dense(p) => Arc::new(RwLock::new(p)), + let partition = match created.partition { + MarginalPartition::Sparse(p) => Arc::new(RwLock::new(PartitionTimetree::Sparse(p))), + MarginalPartition::Dense(p) => Arc::new(RwLock::new(PartitionTimetree::Dense(p))), }; Ok(PartitionInitResult { @@ -495,7 +494,7 @@ fn initialize_partitions_from_params( fn optimize_branch_lengths_pre_step( graph: &GraphTimetree, - partitions: &[Arc>>], + partitions: &[PartitionTimetreeRef], no_indels: bool, ) -> Result<(), Report> { let old_branch_lengths = save_branch_lengths(graph); diff --git a/packages/treetime/src/timetree/refinement.rs b/packages/treetime/src/timetree/refinement.rs index 966dc838a..850abffc1 100644 --- a/packages/treetime/src/timetree/refinement.rs +++ b/packages/treetime/src/timetree/refinement.rs @@ -2,18 +2,14 @@ use crate::ancestral::marginal::update_marginal; use crate::clock::clock_model::ClockModel; use crate::clock::clock_regression::{ClockParams, estimate_clock_model_with_reroot}; use crate::clock::find_best_root::params::BranchPointOptimizationParams; -use crate::partition::timetree::GraphTimetree; -use crate::partition::traits::PartitionTimetreeAll; -use crate::payload::timetree::EdgeTimetree; -use crate::payload::timetree::NodeTimetree; +use crate::partition::timetree::{GraphTimetree, PartitionTimetreeRef}; +use crate::partition::traits::{PartitionMarginalPasses, PartitionTimetreeOps}; use crate::timetree::convergence::sequence_changes::{capture_ancestral_states, count_sequence_changes}; use crate::timetree::inference::runner::run_timetree; use crate::timetree::optimization::polytomy::{prepare_tree_after_topology_change, resolve_polytomies}; use crate::timetree::optimization::relaxed_clock::apply_relaxed_clock; use eyre::{Report, WrapErr}; use log::info; -use parking_lot::RwLock; -use std::sync::Arc; use treetime_distribution::Distribution; use treetime_graph::assign_node_names::assign_node_names; @@ -27,7 +23,7 @@ pub struct RefinementParams { pub fn run_refinement_iteration( params: &RefinementParams, graph: &mut GraphTimetree, - partitions: &[Arc>>], + partitions: &[PartitionTimetreeRef], clock_model: &mut ClockModel, clock_params: &ClockParams, branch_params: &BranchPointOptimizationParams, From e46b717ad614b1083a281afb8129f71bd87b497f Mon Sep 17 00:00:00 2001 From: ivan-aksamentov Date: Wed, 22 Jul 2026 16:14:03 +0200 Subject: [PATCH 4/4] docs(kb): update kb --- kb/features/ancestral.md | 2 +- kb/features/io.md | 6 +- kb/features/timetree.md | 4 +- ...-io-auspice-v2-required-updated-missing.md | 18 ------ .../H-io-usher-ref-nuc-uses-parent-allele.md | 25 -------- ...tion-and-format-projection-inconsistent.md | 14 ++--- ...ice-entropy-perturbs-shannon-definition.md | 6 +- ...spice-trait-confidence-silently-coerced.md | 9 ++- ...M-io-phyloxml-booleans-silently-coerced.md | 10 ++-- ...o-tree-backed-output-order-inconsistent.md | 21 ------- ...-io-usher-mat-mutation-loss-is-implicit.md | 51 ----------------- ...sher-missing-branch-length-becomes-zero.md | 25 ++++++++ ...M-io-usher-zero-branch-length-collapsed.md | 19 ------- ...ee-output-inference-metadata-incomplete.md | 17 +++--- .../N-ancestral-auspice-json-not-produced.md | 7 ++- ...-formatting-missing-augur-golden-master.md | 11 ++-- ...ml-mutation-property-contract-undecided.md | 3 +- .../N-io-tree-ir-architecture-unapproved.md | 57 ------------------- ...ation-confidence-rows-copied-for-output.md | 23 ++++++++ kb/issues/N-test-coverage-gaps.md | 2 +- kb/issues/N-timetree-dead-cli-flags.md | 1 - kb/issues/N-timetree-missing-output-files.md | 4 -- ...etree-node-data-date-string-fp-boundary.md | 4 -- kb/proposals/output-format-selection.md | 9 ++- kb/proposals/tree-format-model-embedding.md | 4 +- kb/proposals/unified-input-format-support.md | 4 +- .../io-parse-auspice-traits-fallibly.md | 3 + .../io-parse-phyloxml-booleans-fallibly.md | 3 + ...e-and-validate-usher-reference-sequence.md | 17 ------ ...o-preserve-usher-missing-branch-lengths.md | 20 +++++++ .../io-preserve-usher-zero-branch-lengths.md | 19 ------- ...ore-and-schema-validate-auspice-updated.md | 21 ------- kb/tickets/io-unify-tree-output-ordering.md | 21 ------- .../mugration-borrow-confidence-rows.md | 22 +++++++ ...dardize-mutations-and-format-projection.md | 4 +- ...complete-tree-output-inference-metadata.md | 3 + .../auspice-nucleotide-strand-invalid.md | 2 +- 37 files changed, 155 insertions(+), 336 deletions(-) delete mode 100644 kb/issues/H-io-auspice-v2-required-updated-missing.md delete mode 100644 kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md delete mode 100644 kb/issues/M-io-tree-backed-output-order-inconsistent.md delete mode 100644 kb/issues/M-io-usher-mat-mutation-loss-is-implicit.md create mode 100644 kb/issues/M-io-usher-missing-branch-length-becomes-zero.md delete mode 100644 kb/issues/M-io-usher-zero-branch-length-collapsed.md delete mode 100644 kb/issues/N-io-tree-ir-architecture-unapproved.md create mode 100644 kb/issues/N-mugration-confidence-rows-copied-for-output.md delete mode 100644 kb/tickets/io-populate-and-validate-usher-reference-sequence.md create mode 100644 kb/tickets/io-preserve-usher-missing-branch-lengths.md delete mode 100644 kb/tickets/io-preserve-usher-zero-branch-lengths.md delete mode 100644 kb/tickets/io-restore-and-schema-validate-auspice-updated.md delete mode 100644 kb/tickets/io-unify-tree-output-ordering.md create mode 100644 kb/tickets/mugration-borrow-confidence-rows.md diff --git a/kb/features/ancestral.md b/kb/features/ancestral.md index dbc9318e5..1f3a19981 100644 --- a/kb/features/ancestral.md +++ b/kb/features/ancestral.md @@ -65,7 +65,7 @@ - [x] `--reconstruct-tip-states` (overwrite ambiguous tips, controls leaf emission) - [x] `--model` relevant only on marginal paths (parsimony bypasses GTR) - [ ] `--keep-overhangs` (parsed but not wired; gap handling not implemented) -- [ ] `--zero-based` indexing (parsed but not wired; [kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md](../issues/M-core-mutation-representation-and-format-projection-inconsistent.md)) +- [ ] `--zero-based` indexing (parsed but not wired: [kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md](../issues/M-core-mutation-representation-and-format-projection-inconsistent.md)) - [ ] `--report-ambiguous` (parsed but not wired) - [x] `--seed` for reproducibility - [ ] `--gtr-params` custom GTR parameters (parsed but not wired) diff --git a/kb/features/io.md b/kb/features/io.md index 632fad03d..de598faef 100644 --- a/kb/features/io.md +++ b/kb/features/io.md @@ -24,9 +24,9 @@ - [x] CSV (clock regression, confidence intervals) - [x] SVG/PNG charts (clock regression) - [x] Graphviz DOT -- [/] Output topology ordering (direct formats honor the plan; TreeIR-backed formats can project the unordered graph: [kb/issues/M-io-tree-backed-output-order-inconsistent.md](../issues/M-io-tree-backed-output-order-inconsistent.md)) +- [x] Output topology ordering - [ ] VCF output (v0 writes .vcf for VCF inputs: [kb/issues/M-io-vcf-input-output-unimplemented.md](../issues/M-io-vcf-input-output-unimplemented.md)) -- [/] Auspice JSON (substitutions implemented; required `meta.updated` is absent and inference metadata remains incomplete: [kb/issues/H-io-auspice-v2-required-updated-missing.md](../issues/H-io-auspice-v2-required-updated-missing.md), [kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md](../issues/M-timetree-tree-output-inference-metadata-incomplete.md)) +- [/] Auspice v2 JSON (schema-validated for all tree-writing commands; entropy perturbs the Shannon definition and inference metadata is incomplete: [kb/issues/M-io-auspice-entropy-perturbs-shannon-definition.md](../issues/M-io-auspice-entropy-perturbs-shannon-definition.md), [kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md](../issues/M-timetree-tree-output-inference-metadata-incomplete.md)) - [ ] Skyline TSV/plot - [ ] Substitution rates TSV - [ ] Outliers TSV @@ -36,7 +36,7 @@ ## v1-Only Formats - [x] PhyloXML -- [/] UShER MAT (partial - reference nucleotides can use the parent allele, and ancestral parsimony cannot supply TreeIR; [kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md](../issues/H-io-usher-ref-nuc-uses-parent-allele.md), [kb/issues/N-ancestral-auspice-json-not-produced.md](../issues/N-ancestral-auspice-json-not-produced.md)) +- [/] UShER MAT (output validates global reference nucleotides; input converts missing branch lengths to zero: [kb/issues/M-io-usher-missing-branch-length-becomes-zero.md](../issues/M-io-usher-missing-branch-length-becomes-zero.md)) - [x] YAML serialization - [x] Compressed FASTA output - [x] Streaming readers/writers with automatic decompression diff --git a/kb/features/timetree.md b/kb/features/timetree.md index 219a2912b..147ed2e6d 100644 --- a/kb/features/timetree.md +++ b/kb/features/timetree.md @@ -119,12 +119,12 @@ - [x] Timetree Newick - [x] Timetree Nexus -- [/] Tree-format topology ordering ([kb/issues/M-io-tree-backed-output-order-inconsistent.md](../issues/M-io-tree-backed-output-order-inconsistent.md)) +- [x] Tree-format topology ordering - [x] Clock model JSON with `timetree.*` basename - [x] Confidence TSV - [ ] Node dates TSV (`write_node_dates()` is `todo!()` - [kb/issues/N-timetree-node-dates-output-unimplemented.md](../issues/N-timetree-node-dates-output-unimplemented.md)) - [ ] Substitution rates TSV (v0 writes `substitution_rates.tsv` when `--relax` is used) -- [/] Auspice JSON (v1 writes substitutions, `num_date`, `div`, and `bad_branch`; required `meta.updated` and remaining inference metadata are incomplete: [kb/issues/H-io-auspice-v2-required-updated-missing.md](../issues/H-io-auspice-v2-required-updated-missing.md), [kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md](../issues/M-timetree-tree-output-inference-metadata-incomplete.md)) +- [/] Auspice v2 JSON (schema-valid substitutions, dates, divergence, outlier state, sequences, and genome annotations; entropy perturbs the Shannon definition and inference metadata is incomplete: [kb/issues/M-io-auspice-entropy-perturbs-shannon-definition.md](../issues/M-io-auspice-entropy-perturbs-shannon-definition.md), [kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md](../issues/M-timetree-tree-output-inference-metadata-incomplete.md)) - [ ] Outliers TSV (v0 writes `outliers.tsv`) - [ ] Tracelog run (v0 `tracelog_run()` with detailed per-iteration state) - [ ] Plotting (`--plot-tree`, `--plot-rtt` - parsed, return explicit error - [kb/issues/N-timetree-plot-unimplemented.md](../issues/N-timetree-plot-unimplemented.md)) diff --git a/kb/issues/H-io-auspice-v2-required-updated-missing.md b/kb/issues/H-io-auspice-v2-required-updated-missing.md deleted file mode 100644 index 5c32a61bc..000000000 --- a/kb/issues/H-io-auspice-v2-required-updated-missing.md +++ /dev/null @@ -1,18 +0,0 @@ -# Auspice v2 output omits required updated metadata - -The shared TreeIR Auspice writer declares schema version `v2` but omits required `meta.updated`. Every command using the shared writer can emit schema-invalid JSON. - -The pinned Augur export-v2 schema requires `updated` [[spec](https://github.com/nextstrain/augur/blob/d8e38736037ba9474a809f9a5a63bc2b279d2407/augur/data/schema-export-v2.json#L14-L22)], and Augur populates it during export [[src](https://github.com/nextstrain/augur/blob/d8e38736037ba9474a809f9a5a63bc2b279d2407/augur/export_v2.py#L1159-L1162)]. - -## Potential solutions - -- O1. Restore a required typed `updated` field and validate every generated document against the pinned schema. -- O2. Downgrade or change the declared schema contract. This would diverge from the selected Auspice v2 output and requires explicit approval. - -## Recommendation - -Require the output plan to generate one UTC calendar date in `YYYY-MM-DD` form and pass it into the writer as typed metadata. Validate ancestral, mugration, and timetree documents against the pinned official schema. Tests pass a fixed generation date; production supplies the current UTC date once per command run. - -## Related issues - -- [M-timetree-tree-output-inference-metadata-incomplete.md](M-timetree-tree-output-inference-metadata-incomplete.md) diff --git a/kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md b/kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md deleted file mode 100644 index d207c65bc..000000000 --- a/kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md +++ /dev/null @@ -1,25 +0,0 @@ -# UShER MAT export fabricates global reference nucleotides from branch-local alleles - -UShER `message mut` [`packages/util-usher-mat/schemas/parsimony.proto#L4-L11`](../../packages/util-usher-mat/schemas/parsimony.proto#L4-L11) defines `ref_nuc` as the reference nucleotide at a mutation's position. TreeTime's mutation-bearing command projections leave `TreeIrData.root_sequence` empty: - -- ancestral constructs default `TreeIrData` with only `has_mutations` set in [`packages/treetime/src/commands/ancestral/run.rs#L243-L263`](../../packages/treetime/src/commands/ancestral/run.rs#L243-L263); -- optimize does the same in [`packages/treetime/src/commands/optimize/run.rs#L84-L102`](../../packages/treetime/src/commands/optimize/run.rs#L84-L102); -- timetree constructs default root-sequence data in [`packages/treetime/src/commands/timetree/output/ir.rs#L22-L28`](../../packages/treetime/src/commands/timetree/output/ir.rs#L22-L28). - -`TreeIrUsherWriter` looks up the reference nucleotide and falls back to the mutation's branch-local parent allele when the reference is absent or cannot encode the position. [`packages/treetime-io/src/tree_ir/usher.rs#L75-L84`](../../packages/treetime-io/src/tree_ir/usher.rs#L75-L84) - -For a recurrent mutation, parent alleles at one coordinate can differ across branches. The exported MAT can therefore assign different `ref_nuc` values to mutations at the same coordinate even though the field denotes one global reference. - -## Required behavior - -- Populate the nucleotide root reference when projecting a mutation-bearing domain graph to TreeIR. -- Reject MAT export when the reference is absent, the coordinate is outside the reference, or the reference nucleotide is not representable as A/C/G/T. -- Verify a recurrent-mutation case where a branch parent allele differs from the global reference. - -## Related issues - -- [N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) - -## Related tickets - -- [kb/tickets/io-populate-and-validate-usher-reference-sequence.md](../tickets/io-populate-and-validate-usher-reference-sequence.md) diff --git a/kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md b/kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md index 33e5e6d6f..c53c2c73c 100644 --- a/kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md +++ b/kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md @@ -2,15 +2,17 @@ Mutation coordinates, string rendering, and tree-format projection do not share one application-wide contract. The same biological mutation can change position base, textual spelling, or representable variants depending on the code path. +> [!NOTE] +> The tree-output refactor removed the `tree_ir` layer and now writes each format directly from the graph in [`packages/treetime/src/commands/shared/tree_output.rs`](../../packages/treetime/src/commands/shared/tree_output.rs). This relocated the format-projection sites but did not, by itself, establish a single shared mutation contract. Which of the sub-points below the refactor already resolved is **not yet confirmed**; re-audit the writers in `tree_output.rs` against each point. + ## Problem - Core and sparse mutation structures mix zero-based and one-based positions. -- TreeIR substitutions use a separate string parser that accepts position zero and duplicates parsing performed elsewhere. Its UShER reader casts signed protobuf positions directly to `usize`, while its writer narrows `usize` to `i32`; zero shifts reference lookup through `saturating_sub(1)` and negative or oversized positions can wrap. - Substitution strings are assembled at multiple output sites instead of by the mutation value itself. -- Insertions and deletions reconstructed by ancestral inference do not reach every TreeIR-backed output. -- UShER conversion narrows positions through unchecked `usize`/`i32` casts and reconstructs one-based format positions ad hoc. +- Insertions and deletions reconstructed by ancestral inference may not reach every format's output. +- UShER conversion narrows positions through `usize`/`i32` casts and reconstructs one-based format positions ad hoc. -Affected boundaries include `enum Mutation` and sparse mutation types in [`packages/treetime/src/partition`](../../packages/treetime/src/partition), TreeIR mutation types in [`packages/treetime-io/src/tree_ir/types.rs`](../../packages/treetime-io/src/tree_ir/types.rs), and the Auspice, PhyloXML, and UShER adapters in [`packages/treetime-io/src/tree_ir`](../../packages/treetime-io/src/tree_ir). +Affected boundaries include `enum Mutation` and sparse mutation types in [`packages/treetime/src/partition`](../../packages/treetime/src/partition), and the Auspice, PhyloXML, and UShER writers in [`packages/treetime/src/commands/shared/tree_output.rs`](../../packages/treetime/src/commands/shared/tree_output.rs). ## Required invariant @@ -45,7 +47,7 @@ Recommendation: O1 for the canonical application string, with checked format-spe - O1. Project substitutions, insertions, and deletions only where the external format has a verified lossless representation. - O2. Define a documented lossy projection for a specific format. This changes scientific output and requires explicit approval. -Recommendation: O1 for verified mappings. Keep format-loss policy ticketless in [M-io-usher-mat-mutation-loss-is-implicit.md](M-io-usher-mat-mutation-loss-is-implicit.md) and [N-io-phyloxml-mutation-property-contract-undecided.md](N-io-phyloxml-mutation-property-contract-undecided.md) until each mapping is approved. Silent omission is invalid. +Recommendation: O1 for verified mappings. The UShER writer now rejects unrepresentable (amino-acid, indel) mutations with an explicit error rather than dropping them silently. The PhyloXML mutation-property loss policy remains open in [N-io-phyloxml-mutation-property-contract-undecided.md](N-io-phyloxml-mutation-property-contract-undecided.md) until its mapping is approved. Silent omission is invalid. ## Recommendation @@ -60,6 +62,4 @@ Use zero-based typed mutations throughout the application, put canonical renderi ## Related issues -- [M-io-usher-mat-mutation-loss-is-implicit.md](M-io-usher-mat-mutation-loss-is-implicit.md) - [N-io-phyloxml-mutation-property-contract-undecided.md](N-io-phyloxml-mutation-property-contract-undecided.md) -- [N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) diff --git a/kb/issues/M-io-auspice-entropy-perturbs-shannon-definition.md b/kb/issues/M-io-auspice-entropy-perturbs-shannon-definition.md index 450b64182..02b1236bd 100644 --- a/kb/issues/M-io-auspice-entropy-perturbs-shannon-definition.md +++ b/kb/issues/M-io-auspice-entropy-perturbs-shannon-definition.md @@ -1,8 +1,8 @@ -# Auspice entropy projection perturbs the Shannon definition +# Auspice entropy output perturbs the Shannon definition -The TreeIR projection computes entropy with an added `TINY` inside every logarithm. This changes every positive term and gives a deterministic distribution such as $(1,0)$ a nonzero entropy. +Entropy is computed with an added `TINY` inside every logarithm. This changes every positive term and gives a deterministic distribution such as $(1,0)$ a nonzero entropy. -`fn compute_entropy()` applies $p\ln(p+10^{-12})$ to every state [packages/treetime/src/commands/shared/ir_projection.rs#L207-L210](../../packages/treetime/src/commands/shared/ir_projection.rs#L207-L210). The mugration node-data path duplicates the same formula [packages/treetime/src/commands/mugration/augur_node_data.rs#L134-L138](../../packages/treetime/src/commands/mugration/augur_node_data.rs#L134-L138). +`fn compute_entropy()` applies $p\ln(p+10^{-12})$ to every state [packages/treetime/src/commands/mugration/augur_node_data.rs#L135-L137](../../packages/treetime/src/commands/mugration/augur_node_data.rs#L135-L137). Both the mugration node-data JSON [packages/treetime/src/commands/mugration/augur_node_data.rs#L113](../../packages/treetime/src/commands/mugration/augur_node_data.rs#L113) and the Auspice tree-output entropy [packages/treetime/src/commands/shared/tree_output.rs#L1385](../../packages/treetime/src/commands/shared/tree_output.rs#L1385) consume this function. For state probabilities $p_i$, Shannon entropy is diff --git a/kb/issues/M-io-auspice-trait-confidence-silently-coerced.md b/kb/issues/M-io-auspice-trait-confidence-silently-coerced.md index 02421a13f..fe9d9fe54 100644 --- a/kb/issues/M-io-auspice-trait-confidence-silently-coerced.md +++ b/kb/issues/M-io-auspice-trait-confidence-silently-coerced.md @@ -2,10 +2,12 @@ The Auspice reader treats malformed trait objects as absent and drops individual non-numeric confidence entries. A document can therefore return success with a partial probability map or without a trait that was present in the input. +> [!NOTE] +> The tree-output refactor removed the former `tree_ir` Auspice reader this issue originally cited. The current reader lives in [`packages/treetime-io/src/auspice.rs`](../../packages/treetime-io/src/auspice.rs). Whether it rejects or still silently drops malformed trait and confidence entries is **not yet confirmed**; the required contract below stands regardless. + ## Evidence -- `fn parse_trait()` [packages/treetime-io/src/tree_ir/auspice.rs#L324](../../packages/treetime-io/src/tree_ir/auspice.rs#L324) returns `Option`. A missing or non-string `value` returns `None`, and `filter_map()` silently removes confidence members whose values are not JSON numbers. -- `fn AuspiceReader::auspice_node_to_graph_components()` [packages/treetime-io/src/tree_ir/auspice.rs#L218](../../packages/treetime-io/src/tree_ir/auspice.rs#L218) inserts a trait only when `parse_trait()` returns `Some`, so malformed present input becomes indistinguishable from absence. +- The former reader returned `Option` from trait parsing: a missing or non-string `value` yielded `None`, and `filter_map()` silently removed confidence members whose values were not JSON numbers. A trait was inserted only when parsing returned `Some`, so malformed present input became indistinguishable from absence. ## Required contract @@ -32,6 +34,3 @@ Reject malformed present traits. Presence is an explicit request to parse typed - Non-object confidence and non-number entropy. - A whole-document test proving malformed present traits return an error rather than disappearing. -## Related issues - -- [N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) diff --git a/kb/issues/M-io-phyloxml-booleans-silently-coerced.md b/kb/issues/M-io-phyloxml-booleans-silently-coerced.md index 4c0c3e853..1e5fdb681 100644 --- a/kb/issues/M-io-phyloxml-booleans-silently-coerced.md +++ b/kb/issues/M-io-phyloxml-booleans-silently-coerced.md @@ -1,10 +1,13 @@ # PhyloXML boolean properties are silently coerced -The PhyloXML reader parses recognized boolean properties through equality with the string `"true"`. The valid XML Schema value `"1"` becomes false, while every invalid lexical value also becomes false without an error. +The PhyloXML reader parses recognized boolean properties (`treetime:bad_branch`, `treetime:date_inferred`) through equality with the string `"true"`. The valid XML Schema value `"1"` becomes false, while every invalid lexical value also becomes false without an error. + +> [!NOTE] +> The tree-output refactor removed the former `tree_ir` PhyloXML reader that assigned these booleans via `prop.value == "true"`. The current reader lives in [`packages/treetime-io/src/phyloxml.rs`](../../packages/treetime-io/src/phyloxml.rs). Whether recognized boolean properties are now parsed per the XSD lexical space (accepting `1`/`0`, rejecting other text) and whether `bad_branch`/`date_inferred` still round-trip is **not yet confirmed**; the required contract below stands regardless. ## Evidence -The property loop in `fn PhyloXmlReader::clade_to_graph_components()` [packages/treetime-io/src/tree_ir/phyloxml.rs#L165](../../packages/treetime-io/src/tree_ir/phyloxml.rs#L165) assigns both `REF_BAD_BRANCH` and `REF_DATE_INFERRED` using `prop.value == "true"`. +The former property loop assigned both `treetime:bad_branch` and `treetime:date_inferred` using `prop.value == "true"`, so `"1"` and every invalid value silently became false. XML Schema `boolean` permits exactly `true`, `false`, `1`, and `0`; its value space contains only true and false. [XML Schema Part 2: boolean](https://www.w3.org/TR/xmlschema-2/#boolean) @@ -30,6 +33,3 @@ Use O1: one fallible XML Schema boolean parser for both recognized boolean prope - Representative invalid values, including case variants, surrounding whitespace, an empty value, and unrelated text. - Whole-document tests proving legal values round-trip and invalid values return contextual errors. -## Related issues - -- [N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) diff --git a/kb/issues/M-io-tree-backed-output-order-inconsistent.md b/kb/issues/M-io-tree-backed-output-order-inconsistent.md deleted file mode 100644 index 908502f9f..000000000 --- a/kb/issues/M-io-tree-backed-output-order-inconsistent.md +++ /dev/null @@ -1,21 +0,0 @@ -# Tree-backed output order is inconsistent - -Commands apply topology ordering separately from TreeIR construction. Direct formats consume the ordered graph, while most TreeIR-backed formats project the original graph, so the same requested output plan can contain different node orders. - -## Evidence - -Affected command runners apply topology order under [`packages/treetime/src/commands`](../../packages/treetime/src/commands), while TreeIR projections are constructed separately under [`packages/treetime-io/src/tree_ir`](../../packages/treetime-io/src/tree_ir). `ancestral`, `clock`, `mugration`, `optimize`, and `prune` can therefore ignore the requested topology order for Auspice, PhyloXML, or UShER while direct Newick/Nexus output follows it. - -## Potential solutions - -- O1. Build one ordered graph and derive every direct output and TreeIR projection from it. -- O2. Reapply the selected order independently inside every writer. This duplicates ordering and can drift as writers are added. - -## Recommendation - -Use O1. Apply topology order once during output preparation and pass the resulting graph to every writer or projection. Method/format availability is an independent contract tracked by [N-ancestral-auspice-json-not-produced.md](N-ancestral-auspice-json-not-produced.md); this issue does not select a parsimony output policy. - -## Related issues - -- [N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) -- [N-ancestral-auspice-json-not-produced.md](N-ancestral-auspice-json-not-produced.md) diff --git a/kb/issues/M-io-usher-mat-mutation-loss-is-implicit.md b/kb/issues/M-io-usher-mat-mutation-loss-is-implicit.md deleted file mode 100644 index 010e52cbe..000000000 --- a/kb/issues/M-io-usher-mat-mutation-loss-is-implicit.md +++ /dev/null @@ -1,51 +0,0 @@ -# UShER MAT output drops unrepresentable mutations implicitly - -## Problem - -The UShER MAT adapter silently changes the requested dataset after logging at most one warning per category. It drops all indels, all amino-acid substitutions, and every substitution containing a non-ACGT state [packages/treetime-io/src/tree_ir/usher.rs#L54-L87](../../packages/treetime-io/src/tree_ir/usher.rs#L54-L87). On input, it takes only the first value from the protobuf's repeated `mut_nuc` field and casts the signed position directly to `usize` [packages/treetime-io/src/tree_ir/usher.rs#L123-L133](../../packages/treetime-io/src/tree_ir/usher.rs#L123-L133). - -UShER's protobuf mutation message has a single numeric position plus nucleotide state fields and no insertion/deletion event type [[src](https://github.com/yatisht/usher/blob/ac9c982d937c3bc43e2c16a0b73d30cf2b937118/parsimony.proto#L4-L10)]. Dropping general indels is therefore a real format limitation. Choosing the first `mut_nuc` value is a separate reader policy: UShER combines every repeated value into an ambiguity-state bitset [[src](https://github.com/yatisht/usher/blob/ac9c982d937c3bc43e2c16a0b73d30cf2b937118/src/mutation_annotated_tree.cpp#L565-L581)]. Negative positions are another distinct case: UShER treats them as masked mutations, while the TreeTime cast turns them into large ordinary coordinates [[src](https://github.com/yatisht/usher/blob/ac9c982d937c3bc43e2c16a0b73d30cf2b937118/src/matOptimize/mutation_annotated_tree_load_store.cpp#L321-L334)]. - -The current warnings do not let callers detect data loss programmatically, and a file is still produced successfully. - -## Potential solutions - -### A1. Track selection - -- **Reject mixed nucleotide/amino-acid TreeIR:** fail whenever any non-nucleotide track is present. -- **Define MAT as a nucleotide-track projection:** ignore amino-acid tracks by contract, because the target format has no gene-track model. This must be documented as selection rather than an incidental warning. - -### A2. Unrepresentable nucleotide events - -- **Warn and drop:** preserve current behavior. Successful output does not imply preservation. -- **Fail by default:** reject indels, ambiguous substitutions, and unsupported repeated-state input with an aggregate diagnostic. -- **Explicit lossy mode:** fail by default and allow a caller-selected policy that reports the number and categories of dropped events. - -### A3. Repeated `mut_nuc` input - -- **Take the first state:** deterministic but discards schema information. -- **Map the state set to an IUPAC ambiguity code:** preserves supported nucleotide ambiguity when the set has a standard code. -- **Reject multi-state values:** avoids inventing a semantic interpretation but cannot read valid ambiguous MAT data. - -### A4. Masked mutation input - -- **Reject masked events:** prevents coordinate corruption but cannot preserve valid MAT masking information. -- **Preserve a masked-event variant:** extend TreeIR with an event that carries masking without pretending it is a nucleotide substitution. -- **Explicitly drop under lossy mode:** keep topology while reporting masked-event loss through the same aggregate policy as other unsupported data. - -## Recommendation - -Treat MAT output as an explicit nucleotide-track projection, fail by default on unrepresentable nucleotide events, and aggregate every loss category before returning the error. Add a lossy mode only after its CLI/API surface and reporting contract are approved. On input, map valid nucleotide state sets to IUPAC ambiguity codes, reject invalid or empty sets, and preserve masked events as a distinct TreeIR event if TreeIR remains the accepted boundary. - -The checked numeric position conversion belongs to [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md); this issue concerns representational loss after coordinates are valid. - -The writer also reconstructs `ref_nuc` from the root sequence. That is insufficient for recurrent mutations: after `A→G` on an ancestor, a descendant `G→T` must retain reference nucleotide `A`, not the immediate parent state `G`. Non-ASCII root bytes are currently filtered out, shifting every later coordinate. These are preservation failures governed by the same explicit MAT loss policy. - -## Ticket readiness - -No implementation ticket is ready. The default loss policy, whether a lossy mode should exist, and whether TreeIR must preserve masked events are externally visible decisions. The recommendation above must be approved before an implementation ticket is created. - -## Related issues - -- [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md) -- [N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) diff --git a/kb/issues/M-io-usher-missing-branch-length-becomes-zero.md b/kb/issues/M-io-usher-missing-branch-length-becomes-zero.md new file mode 100644 index 000000000..5f8e29bfd --- /dev/null +++ b/kb/issues/M-io-usher-missing-branch-length-becomes-zero.md @@ -0,0 +1,25 @@ +# UShER input converts missing branch lengths to explicit zero + +## Problem + +The MAT contains an embedded Newick tree. Its parser correctly represents branch length as `Option`, preserving the distinction between an absent length and an explicit zero. + +`fn usher_to_graph()` then stores the parsed value in the non-optional `UsherNodeImpl.branch_length: f64` and applies `.unwrap_or(0.0)` in [`packages/treetime-io/src/usher_mat.rs`](../../packages/treetime-io/src/usher_mat.rs). The format adapter therefore converts `None` into `Some(0.0)` when a concrete graph edge is constructed. Exact round trips become impossible, and downstream code can interpret unknown evolutionary distance as observed zero change. + +This is an application bug, not a MAT limitation. MAT preserves the embedded Newick spelling, and the common graph edge API already represents branch lengths as `Option`. + +## Required behavior + +- Change `UsherNodeImpl.branch_length` to `Option`. +- Preserve `None`, `Some(0.0)`, and `Some(positive)` without an additional enum. An enum is unnecessary because the application currently has only two semantic states: absent or present; zero is an ordinary present value. +- Preserve every present finite Newick value at the format boundary. Algorithms that interpret branch length as evolutionary distance must apply their existing finite, non-negative precondition before inference; that scientific validation is separate from lossless MAT/Newick parsing. +- Preserve the same optional value when the converter constructs the graph edge. + +## Validation + +- Round-trip absent, explicit zero, and positive embedded-Newick branch lengths. +- Preserve the same negative-value behavior as the shared Newick parser, and verify that inference boundaries still reject values outside their scientific domain. + +## Related tickets + +- [kb/tickets/io-preserve-usher-missing-branch-lengths.md](../tickets/io-preserve-usher-missing-branch-lengths.md) diff --git a/kb/issues/M-io-usher-zero-branch-length-collapsed.md b/kb/issues/M-io-usher-zero-branch-length-collapsed.md deleted file mode 100644 index 8041274d2..000000000 --- a/kb/issues/M-io-usher-zero-branch-length-collapsed.md +++ /dev/null @@ -1,19 +0,0 @@ -# UShER input collapses explicit zero branch length into absence - -The UShER MAT reader converts an explicit branch length of zero to `None`. Zero substitutions and an unspecified branch length are distinct states; collapsing them prevents exact round trips and can change downstream defaulting behavior. - -The embedded Newick parser produces an optional branch length, but `fn usher_to_graph()` converts absence to `0.0` in `UsherNodeImpl` [packages/treetime-io/src/usher_mat.rs#L232-L247](../../packages/treetime-io/src/usher_mat.rs#L232-L247). `fn TreeIrUsherReader::usher_node_to_graph_components()` then maps every zero back to `None` [packages/treetime-io/src/tree_ir/usher.rs#L117-L142](../../packages/treetime-io/src/tree_ir/usher.rs#L117-L142). - -## Potential solutions - -- O1. Preserve presence independently of numeric value. -- O2. Model absent/zero/positive branch lengths as an explicit enum. This is useful only if the distinction must propagate beyond existing `Option` boundaries. - -## Recommendation - -Preserve protobuf presence independently of numeric value. `Some(0.0)` remains `Some(0.0)`, while an absent field remains `None`. Reject non-finite or negative values according to the common branch-length contract. - -## Related issues - -- [M-io-usher-mat-mutation-loss-is-implicit.md](M-io-usher-mat-mutation-loss-is-implicit.md) -- [N-core-branch-length-clamping.md](N-core-branch-length-clamping.md) diff --git a/kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md b/kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md index 5f69a0ca5..6e645568c 100644 --- a/kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md +++ b/kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md @@ -1,16 +1,17 @@ # TreeTime tree-output inference metadata is incomplete -TreeTime now produces TreeIR-backed Auspice and PhyloXML output, but its projection remains incomplete as one output contract. +TreeTime writes Auspice and PhyloXML directly from the graph in [`packages/treetime/src/commands/shared/tree_output.rs`](../../packages/treetime/src/commands/shared/tree_output.rs), but the metadata payload is not yet a complete output contract. + +> [!NOTE] +> The tree-output refactor added several of the fields this issue tracked: `treetime:gamma`, `treetime:input_branch_support`/`confidence`, `date_is_inferred`, and `genome_annotations` are now emitted. The **remaining** items below are not confirmed present; re-audit each against `tree_output.rs`. ## Problem -- Auspice substitution projection exists only where the command constructs TreeIR; insertion and deletion handling is governed by [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md). -- `branch_attrs` omits the v0 pseudo-bootstrap branch support confidence. -- `meta.colorings` omits the corresponding continuous `confidence` coloring. -- `meta.genome_annotations.nuc` omits alignment coordinates, type, and strand. -- The entropy panel metadata cannot be complete without genome annotations. -- TreeIR defines `date_is_inferred`, `date_raw_value`, and relaxed-clock `gamma`, but the TreeTime projection leaves them at defaults even though command state contains the source data. -- PhyloXML can serialize inferred-date and `gamma` properties, so command-generated relaxed-clock output remains incomplete. +- Insertion and deletion handling across formats is governed by [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md). +- `meta.colorings` may still omit the continuous `confidence` coloring that corresponds to branch support. +- `meta.genome_annotations.nuc` completeness (alignment coordinates, type, strand) is unconfirmed. +- `date_raw_value` is not emitted; confirm whether `date_is_inferred` and relaxed-clock `gamma` are populated from command state in every mode rather than left at defaults. +- PhyloXML inferred-date and `gamma` properties for relaxed-clock output are unconfirmed. The Auspice fields share one command projection and must be validated together. The already-defined non-mutation PhyloXML date and gamma properties use the same command state. PhyloXML mutation vocabulary and indel representation remain excluded pending [N-io-phyloxml-mutation-property-contract-undecided.md](N-io-phyloxml-mutation-property-contract-undecided.md). diff --git a/kb/issues/N-ancestral-auspice-json-not-produced.md b/kb/issues/N-ancestral-auspice-json-not-produced.md index 29dd4372e..7155dfbbd 100644 --- a/kb/issues/N-ancestral-auspice-json-not-produced.md +++ b/kb/issues/N-ancestral-auspice-json-not-produced.md @@ -1,6 +1,9 @@ # Ancestral Auspice output is incomplete and method-dependent -V0's `ancestral` command produces `auspice_tree.json` through `def export_sequences_and_tree()`. V1 marginal reconstruction constructs TreeIR and can produce Auspice output. Parsimony advertises TreeIR-backed formats without constructing the required projection, so a valid selection can fail after earlier outputs have been written. +V0's `ancestral` command produces `auspice_tree.json` through `def export_sequences_and_tree()`. V1 now writes Auspice output directly from the graph via `fn ancestral_to_auspice()` [`packages/treetime/src/commands/shared/tree_output.rs#L240`](../../packages/treetime/src/commands/shared/tree_output.rs#L240). + +> [!NOTE] +> The tree-output refactor removed the TreeIR projection and wired `ancestral_to_auspice` into the writer. Whether ancestral Auspice is now produced completely and for **every** method (Fitch parsimony as well as sparse/dense marginal), and whether its payload matches the v0 fields below, is **not yet confirmed** against the current writer. V0's ancestral Auspice JSON contains `node_attrs.div` (cumulative `mutation_length`), `branch_attrs.mutations.nuc` (per-branch mutations), `node_attrs.confidence` (pseudo-bootstrap), `meta.genome_annotations.nuc`, and `node_attrs.bad_branch`. It contains no dates because `timetree=false`. @@ -54,7 +57,5 @@ No implementation ticket is ready. O1 is the reference-parity recommendation, bu ## Related - [M-timetree-tree-output-inference-metadata-incomplete.md](M-timetree-tree-output-inference-metadata-incomplete.md) - shared mutation, confidence, and annotation contract -- [M-io-tree-backed-output-order-inconsistent.md](M-io-tree-backed-output-order-inconsistent.md) - shared topology ordering - [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md) - typed mutation projection -- [H-io-usher-ref-nuc-uses-parent-allele.md](H-io-usher-ref-nuc-uses-parent-allele.md) - global root-reference preservation for UShER - [kb/reports/augur-node-data-json.md](../reports/augur-node-data-json.md) - node data JSON shares the same data sources (partition mutations, annotations) diff --git a/kb/issues/N-io-auspice-number-formatting-missing-augur-golden-master.md b/kb/issues/N-io-auspice-number-formatting-missing-augur-golden-master.md index eb55f33d5..6bf2ecbd4 100644 --- a/kb/issues/N-io-auspice-number-formatting-missing-augur-golden-master.md +++ b/kb/issues/N-io-auspice-number-formatting-missing-augur-golden-master.md @@ -5,11 +5,11 @@ ## Problem -`pub fn format_number()` [packages/treetime-io/src/tree_ir/auspice.rs#L38](../../packages/treetime-io/src/tree_ir/auspice.rs#L38) reimplements Augur's significant-digit behavior for divergence and numeric dates. Its unit tests use manually written expected values attributed to `augur export_v2.format_number`; they do not capture outputs from a pinned Augur revision. +`pub(crate) fn format_number()` [packages/treetime/src/commands/shared/tree_output.rs#L1726](../../packages/treetime/src/commands/shared/tree_output.rs#L1726) reimplements Augur's significant-digit behavior for divergence and numeric dates. Its unit tests use manually written expected values attributed to `augur export_v2.format_number`; they do not capture outputs from a pinned Augur revision. -This leaves boundary behavior unverified for negative values, powers of ten, values crossing an integer-digit boundary after rounding, very small magnitudes, and ties affected by Python and Rust formatting differences. TreeIR round-trip tests cannot serve as an independent oracle because they read output produced by the same adapter. +This leaves boundary behavior unverified for negative values, powers of ten, values crossing an integer-digit boundary after rounding, very small magnitudes, and ties affected by Python and Rust formatting differences. Output-generating tests cannot serve as an independent oracle because they read output produced by the same code. -The public `i32` precision API also performs unchecked `significand + precision`. Extreme values can panic in debug, wrap in release, or request disproportionate formatting allocation. Equivalent significant-digit formatting already exists in `treetime-utils`. +The `i32` precision API also performs unchecked `significand + precision` [packages/treetime/src/commands/shared/tree_output.rs#L1736](../../packages/treetime/src/commands/shared/tree_output.rs#L1736). Extreme values can panic in debug, wrap in release, or request disproportionate formatting allocation. Equivalent significant-digit formatting already exists in `treetime-utils`. ## Potential solutions @@ -31,11 +31,10 @@ For finite nonzero $n$, let $d = \lfloor \log_{10}(\lfloor |n| \rfloor) \rfloor ## Locations -- `pub fn format_number()` [packages/treetime-io/src/tree_ir/auspice.rs#L38](../../packages/treetime-io/src/tree_ir/auspice.rs#L38) -- Unit tests [packages/treetime-io/src/tree_ir/__tests__/test_auspice.rs](../../packages/treetime-io/src/tree_ir/__tests__/test_auspice.rs) +- `pub(crate) fn format_number()` [packages/treetime/src/commands/shared/tree_output.rs#L1726](../../packages/treetime/src/commands/shared/tree_output.rs#L1726) +- Unit tests [packages/treetime/src/commands/shared/**tests**/test_tree_output.rs](../../packages/treetime/src/commands/shared/__tests__/test_tree_output.rs) ## Related KB items - [kb/decisions/multi-format-tree-io.md](../decisions/multi-format-tree-io.md) - [kb/proposals/output-format-selection.md](../proposals/output-format-selection.md) -- [kb/issues/N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) diff --git a/kb/issues/N-io-phyloxml-mutation-property-contract-undecided.md b/kb/issues/N-io-phyloxml-mutation-property-contract-undecided.md index cc29f99aa..d5ff0956e 100644 --- a/kb/issues/N-io-phyloxml-mutation-property-contract-undecided.md +++ b/kb/issues/N-io-phyloxml-mutation-property-contract-undecided.md @@ -4,7 +4,7 @@ PhyloXML does not define a native mutation element. Its generic `Property` type carries typed free text, a `prefix:value` reference key, an `applies_to` target, and an optional `id_ref` that attaches the property to a specific XML element [packages/util-phyloxml/schemas/phyloxml.xsd#L385-L418](../../packages/util-phyloxml/schemas/phyloxml.xsd#L385-L418). -TreeTime currently writes substitutions as `gene:A123T` and indels as `gene:ins|del:start:SEQ` under ad hoc `treetime:mutation` and `treetime:indel` references [packages/treetime-io/src/tree_ir/phyloxml.rs#L217-L260](../../packages/treetime-io/src/tree_ir/phyloxml.rs#L217-L260). These values round-trip through TreeTime, but the reference vocabulary is undocumented and the two event kinds use unrelated grammars. External consumers have no published contract for interpreting them. +TreeTime writes substitutions and indels under an ad hoc `treetime:*` reference vocabulary defined in [`packages/treetime/src/commands/shared/tree_output.rs#L60-L69`](../../packages/treetime/src/commands/shared/tree_output.rs#L60-L69). These values round-trip through TreeTime, but the reference vocabulary is undocumented and the event kinds use unrelated grammars. External consumers have no published contract for interpreting them. The tree-output refactor **implemented this vocabulary without a documented decision**, so the contract is now shipped-yet-unratified rather than merely undecided; it needs an explicit `kb/decisions/` entry or a redesign per the axes below. The design space was already identified in [kb/proposals/phyloxml-treetime-property-namespace.md](../proposals/phyloxml-treetime-property-namespace.md), but its actionable decisions were not represented by a known issue. @@ -38,4 +38,3 @@ No implementation ticket is ready. Reference-vocabulary identity, atomic versus - [kb/proposals/phyloxml-treetime-property-namespace.md](../proposals/phyloxml-treetime-property-namespace.md) - [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md) -- [N-io-tree-ir-architecture-unapproved.md](N-io-tree-ir-architecture-unapproved.md) diff --git a/kb/issues/N-io-tree-ir-architecture-unapproved.md b/kb/issues/N-io-tree-ir-architecture-unapproved.md deleted file mode 100644 index 34a3b3e72..000000000 --- a/kb/issues/N-io-tree-ir-architecture-unapproved.md +++ /dev/null @@ -1,57 +0,0 @@ -# Shared TreeIR architecture has not received an explicit design decision - -> [!IMPORTANT] -> **Research and discussion required.** TreeIR is implemented, but its status as an accepted v1 architecture is unresolved. No decision entry may be added without explicit approval. - -## Problem - -TreeIR provides one format-neutral graph for PhyloXML, Auspice v2, and UShER MAT serialization. Commands project domain graphs and analysis results into this representation before dispatching format-specific writers. This is a new v1 architecture rather than direct parity with a v0 abstraction. - -The implementation defines a shared loss and ownership boundary: domain-specific data must be projected into `TreeIrNode`, `TreeIrEdge`, and `TreeIrData`, while formats outside TreeIR serialize directly from the domain graph. [kb/decisions/multi-format-tree-io.md](../decisions/multi-format-tree-io.md) describes a common graph intermediate representation, but it predates the concrete TreeIR boundary, refers to the former `ConverterGraph`, and does not define TreeIR's preservation or loss invariants. The decision also requires a consent audit under the current project rule that architecture decisions must have explicit approval. - -## Decision axes - -### A1. Architectural boundary - -- O1. Keep TreeIR as an internal, format-neutral projection boundary. Commands produce one typed intermediate value and writers consume only the fields they declare. This centralizes topology and analysis-data projection without making TreeIR a public compatibility promise. -- O2. Project each domain graph directly into each output format. Each writer can model its schema precisely, but ordering, mutation conversion, and typed-field validation would be repeated across adapters. -- O3. Adopt TreeIR as a project-level architecture with a stable preservation contract. This gives every producer and consumer one durable target, but every new format requirement must remain representable in that shared contract. - -**Recommendation:** O1. Retain one internal projection boundary and document its invariants, while allowing its types to change when a format exposes a missing domain concept. This preserves centralized conversion without freezing an unshipped internal API. - -### A2. Preservation and unsupported states - -- O1. Model a lossless typed superset of every supported format and make each writer return a typed error for values outside its schema. -- O2. Model only the intersection shared by all formats and reject input that cannot enter that common subset. This makes the boundary uniform but excludes valid format-specific data. -- O3. Permit explicit lossy conversion modes whose reports enumerate every dropped value. This can complement either typed representation, but it cannot be the implicit default. - -**Recommendation:** O1 as the default contract, with O3 available only through an explicit caller-selected policy. Valid data must never disappear silently. - -### A3. Malformed typed fields - -- O1. Parse strings into typed values at the input boundary and fail the complete conversion with field and node context. -- O2. Preserve both raw and parsed representations and defer the error to writers that require the typed value. This retains malformed source text but allows invalid state deeper into the application. - -**Recommendation:** O1. Inner TreeIR values should be valid by construction; raw source retention belongs in a separate diagnostic structure when required. - -Explicit approval is required before recording any architectural decision. Until then, no implementation ticket is ready. - -## Locations - -- `type TreeIrGraph` [packages/treetime-io/src/graph.rs#L23](../../packages/treetime-io/src/graph.rs#L23) -- `pub fn write_tree_outputs()` [packages/treetime-io/src/graph.rs#L61](../../packages/treetime-io/src/graph.rs#L61) -- `mod tree_ir` adapters [packages/treetime-io/src/tree_ir/mod.rs](../../packages/treetime-io/src/tree_ir/mod.rs) -- `pub fn build_ir_with_mutations()` and sibling projections [packages/treetime/src/commands/shared/ir_projection.rs#L29](../../packages/treetime/src/commands/shared/ir_projection.rs#L29) - -## Related KB items - -- [kb/decisions/multi-format-tree-io.md](../decisions/multi-format-tree-io.md) -- [kb/proposals/output-format-selection.md](../proposals/output-format-selection.md) -- [kb/proposals/tree-format-model-embedding.md](../proposals/tree-format-model-embedding.md) -- [N-io-auspice-number-formatting-missing-augur-golden-master.md](N-io-auspice-number-formatting-missing-augur-golden-master.md) -- [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md) -- [M-io-usher-mat-mutation-loss-is-implicit.md](M-io-usher-mat-mutation-loss-is-implicit.md) -- [N-io-phyloxml-mutation-property-contract-undecided.md](N-io-phyloxml-mutation-property-contract-undecided.md) -- [M-timetree-tree-output-inference-metadata-incomplete.md](M-timetree-tree-output-inference-metadata-incomplete.md) -- [M-io-auspice-trait-confidence-silently-coerced.md](M-io-auspice-trait-confidence-silently-coerced.md) -- [M-io-phyloxml-booleans-silently-coerced.md](M-io-phyloxml-booleans-silently-coerced.md) diff --git a/kb/issues/N-mugration-confidence-rows-copied-for-output.md b/kb/issues/N-mugration-confidence-rows-copied-for-output.md new file mode 100644 index 000000000..cb1cfe9df --- /dev/null +++ b/kb/issues/N-mugration-confidence-rows-copied-for-output.md @@ -0,0 +1,23 @@ +# Mugration confidence rows are copied for output + +`PartitionMarginalDiscrete::get_confidence()` copies `node.profile.dis.row(0)` into a new `Array1`. Mugration node-data and tree-output projection call it independently, and `MugrationConfidenceOutput::new()` eagerly copies every row into a second output structure retained beside the partition. + +The posterior matrix already owns these rows for the lifetime of `MugrationGraphData`. Read-only output should borrow them. + +## Required behavior + +- Return `Option>` from `get_confidence()`. +- Accept `ArrayView1<'_, f64>` in confidence-map and entropy calculations. +- Derive both values from one borrowed row per node. +- Remove the eagerly duplicated `MugrationConfidenceOutput` profile matrix. Render confidence CSV directly from graph node names and partition row views at write time. +- Keep owned arrays only when an output value must outlive the partition. + +## Validation + +- Exact node-data, Auspice, and confidence-CSV equivalence. +- Owned, row, and non-contiguous view unit cases for consumers. +- Allocation regression coverage proving projection does not copy confidence rows. + +## Related tickets + +- [kb/tickets/mugration-borrow-confidence-rows.md](../tickets/mugration-borrow-confidence-rows.md) diff --git a/kb/issues/N-test-coverage-gaps.md b/kb/issues/N-test-coverage-gaps.md index c8272a5d5..6f886e532 100644 --- a/kb/issues/N-test-coverage-gaps.md +++ b/kb/issues/N-test-coverage-gaps.md @@ -97,7 +97,7 @@ Systematic test coverage gaps span timetree inference, clock, coalescent, ancest ## Cross-cutting scientific coverage gaps - Scientific oracle and property tests remain ignored across distribution, rerooting, coalescent, and inference paths. -- TreeIR semantic output tests cover `timetree` more thoroughly than ancestral, clock, mugration, optimize, and prune. +- Graph-backed semantic output tests cover `timetree` more thoroughly than ancestral, clock, mugration, optimize, and prune. - Parallel coverage concentrates on one sparse success case and does not establish error atomicity for marginal, optimize, or timetree passes. - Skyline tests recompute the reported formula instead of invoking the optimizer objective. - Coalescent initialization tests establish enum selection without checking topology-change sequencing or event completeness. diff --git a/kb/issues/N-timetree-dead-cli-flags.md b/kb/issues/N-timetree-dead-cli-flags.md index 8b303d377..a1d16aa9a 100644 --- a/kb/issues/N-timetree-dead-cli-flags.md +++ b/kb/issues/N-timetree-dead-cli-flags.md @@ -40,4 +40,3 @@ No aggregate ticket is ready. Create one focused ticket only after the dispositi ## Related issues - [--method-anc ignored in timetree](M-timetree-method-anc-ignored.md) `--method-anc` also dead -- [M-core-mutation-representation-and-format-projection-inconsistent.md](M-core-mutation-representation-and-format-projection-inconsistent.md) tracks the shared mutation-string API and the `--zero-based` output contract diff --git a/kb/issues/N-timetree-missing-output-files.md b/kb/issues/N-timetree-missing-output-files.md index cde3280d9..8f9682e45 100644 --- a/kb/issues/N-timetree-missing-output-files.md +++ b/kb/issues/N-timetree-missing-output-files.md @@ -19,10 +19,6 @@ v1 produces fewer output files than v0. Several output types are not implemented (see [Missing skyline output files](N-timetree-missing-skyline-output.md)) - `skyline.pdf` - skyline plot -## Produced but incomplete - -- `auspice_tree.json` - v1 produces this file with substitutions. Its remaining mutation, confidence, annotation, and panel gaps are tracked together in [M-timetree-tree-output-inference-metadata-incomplete.md](M-timetree-tree-output-inference-metadata-incomplete.md). - ## v1-only outputs - `timetree.json` - clock model in JSON format diff --git a/kb/issues/N-timetree-node-data-date-string-fp-boundary.md b/kb/issues/N-timetree-node-data-date-string-fp-boundary.md index 19f09b467..ad858bf55 100644 --- a/kb/issues/N-timetree-node-data-date-string-fp-boundary.md +++ b/kb/issues/N-timetree-node-data-date-string-fp-boundary.md @@ -3,7 +3,3 @@ The per-node `date` field in timetree node data JSON is the `"YYYY-MM-DD"` string produced from each node's `numdate` by `year_fraction_to_datestring()`. At year-fraction values that fall on a floating-point day boundary, the conversion can round to a day adjacent to the one augur reports from the same `numdate`. The discrepancy is at most one day. `augur export v2` reads `numdate` for the time tree, not the `date` string, so no downstream auspice output is affected. The mismatch only shows when comparing the `date` field directly against augur's node data JSON. - -## Related issues - -- [M-timetree-tree-output-inference-metadata-incomplete.md](M-timetree-tree-output-inference-metadata-incomplete.md) diff --git a/kb/proposals/output-format-selection.md b/kb/proposals/output-format-selection.md index 76fc6052c..b184bf7ff 100644 --- a/kb/proposals/output-format-selection.md +++ b/kb/proposals/output-format-selection.md @@ -4,11 +4,10 @@ All tree-writing commands (ancestral, clock, timetree, optimize, prune, mugratio ## Motivation -Three independent problems converge: +Two independent problems converge: -1. The shared graph writer needs one format-neutral projection with explicit preservation and loss contracts. Current gaps include [kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md](../issues/M-core-mutation-representation-and-format-projection-inconsistent.md), [kb/issues/M-io-usher-mat-mutation-loss-is-implicit.md](../issues/M-io-usher-mat-mutation-loss-is-implicit.md), and [kb/issues/N-io-phyloxml-mutation-property-contract-undecided.md](../issues/N-io-phyloxml-mutation-property-contract-undecided.md). -2. Newick annotation styles (plain, BEAST `[&...]`, NHX `[&&NHX:...]`) are a serialization choice that the user cannot control. Plain Newick and annotated Newick are both valid outputs users may need simultaneously. -3. Some commands should produce multiple tree files with different semantics -- timetree should output both time-branch-length and divergence-branch-length trees ([kb/issues/N-io-time-based-branch-lengths-not-implemented.md](../issues/N-io-time-based-branch-lengths-not-implemented.md)), and v0 does this. +1. Newick annotation styles (plain, BEAST `[&...]`, NHX `[&&NHX:...]`) are a serialization choice that the user cannot control. Plain Newick and annotated Newick are both valid outputs users may need simultaneously. +2. Some commands should produce multiple tree files with different semantics -- timetree should output both time-branch-length and divergence-branch-length trees ([kb/issues/N-io-time-based-branch-lengths-not-implemented.md](../issues/N-io-time-based-branch-lengths-not-implemented.md)), and v0 does this. A single `--nwk-style` flag cannot address all three. The root problem is that the output system lacks format selection. @@ -144,7 +143,7 @@ Both produce `BTreeMap`. The format variant dispatch applies aft ## Validation -- Each `TreeOutputFormat` variant round-trips every value in its declared vocabulary. Mutation round trips remain conditional on resolving the shared mutation vocabulary and each writer's explicit unsupported-state policy; schema-level serialization alone does not establish semantic round-trip correctness. +- Each `TreeOutputFormat` variant round-trips every value in its declared vocabulary; schema-level serialization alone does not establish semantic round-trip correctness. - Tier 1/2/3 resolution logic: unit tests for flag precedence, format selection filtering, default path generation - Per-command: verify available format set matches command's adapter traits - v0 parity: timetree produces both `timetree.nwk` and `divergence_tree.nwk` (matching v0's `export_sequences_and_tree`) diff --git a/kb/proposals/tree-format-model-embedding.md b/kb/proposals/tree-format-model-embedding.md index a44f0e6e3..844437b37 100644 --- a/kb/proposals/tree-format-model-embedding.md +++ b/kb/proposals/tree-format-model-embedding.md @@ -4,7 +4,7 @@ TreeTime infers a GTR substitution model and a molecular-clock model alongside the tree. The analysis commands write these to sidecar files (`--output-gtr`, `--output-clock-model`). Some tree formats can carry model parameters inline, which would make a single output file self-contained for downstream tools. -The TreeIR output path (PhyloXML, Auspice v2, UShER MAT) currently carries only per-node and per-branch data plus minimal graph-level metadata. It does not embed the GTR or clock model. +The graph-backed output path (PhyloXML, Auspice v2, UShER MAT) currently carries only per-node and per-branch data plus minimal graph-level metadata. It does not embed the GTR or clock model. ## Options @@ -18,4 +18,4 @@ Keep models in sidecar files (current behavior) until a concrete downstream need ## Non-goals -This proposal does not change current sidecar output. It records that inline model embedding was deliberately left out of the TreeIR output path. +This proposal does not change current sidecar output. It records that inline model embedding is absent from graph-backed tree output. diff --git a/kb/proposals/unified-input-format-support.md b/kb/proposals/unified-input-format-support.md index e36b78635..ab72b3bd5 100644 --- a/kb/proposals/unified-input-format-support.md +++ b/kb/proposals/unified-input-format-support.md @@ -10,7 +10,7 @@ The codebase already includes format adapters for unified formats where sequence - UShER MAT (protobuf): edge mutations in preorder traversal - PhyloXML: per-clade sequence elements -Schema adapters exist in `treetime-io` and are used by the `convert` command. They do not yet define a lossless shared mutation vocabulary: UShER MAT preservation policy, PhyloXML mutation properties, and shared substitution/insertion/deletion projection remain unresolved in [kb/issues/M-io-usher-mat-mutation-loss-is-implicit.md](../issues/M-io-usher-mat-mutation-loss-is-implicit.md), [kb/issues/N-io-phyloxml-mutation-property-contract-undecided.md](../issues/N-io-phyloxml-mutation-property-contract-undecided.md), and [kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md](../issues/M-core-mutation-representation-and-format-projection-inconsistent.md). Integrating a reader therefore requires an explicit preservation contract in addition to schema parsing. +Schema adapters exist in `treetime-io` and are used by the `convert` command. Integrating them into analysis commands still requires explicit preservation of embedded input data and construction of the command's partitions after parsing. ## Proposal @@ -111,8 +111,6 @@ The graph payload question affects internal organization. This proposal affects - [kb/issues/M-io-sequence-name-matching-unreliable.md](../issues/M-io-sequence-name-matching-unreliable.md) - [kb/issues/M-io-sequence-attachment-quadratic.md](../issues/M-io-sequence-attachment-quadratic.md) - [kb/issues/N-io-large-dataset-memory-constraint.md](../issues/N-io-large-dataset-memory-constraint.md) -- [kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md](../issues/M-core-mutation-representation-and-format-projection-inconsistent.md) -- [kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md](../issues/M-timetree-tree-output-inference-metadata-incomplete.md) - [kb/issues/N-io-multi-segment-genome-input.md](../issues/N-io-multi-segment-genome-input.md) ## Related documentation diff --git a/kb/tickets/io-parse-auspice-traits-fallibly.md b/kb/tickets/io-parse-auspice-traits-fallibly.md index cd7c8e948..7fe2a410e 100644 --- a/kb/tickets/io-parse-auspice-traits-fallibly.md +++ b/kb/tickets/io-parse-auspice-traits-fallibly.md @@ -2,6 +2,9 @@ Reject malformed present Auspice traits instead of dropping the trait or individual confidence entries. +> [!NOTE] +> This ticket predates the tree-output refactor, which removed the `tree_ir` layer and now writes formats directly from the graph in `packages/treetime/src/commands/shared/tree_output.rs` (readers in `packages/treetime-io/src/auspice.rs` and `packages/treetime-io/src/phyloxml.rs`). Its parent issue's current status is unconfirmed; re-derive the steps and code locations against the current code before executing. + ## Required changes - Change `parse_trait()` [packages/treetime-io/src/tree_ir/auspice.rs#L324](../../packages/treetime-io/src/tree_ir/auspice.rs#L324) to return `Result` for a present trait object. diff --git a/kb/tickets/io-parse-phyloxml-booleans-fallibly.md b/kb/tickets/io-parse-phyloxml-booleans-fallibly.md index ce585dbd0..39649f9af 100644 --- a/kb/tickets/io-parse-phyloxml-booleans-fallibly.md +++ b/kb/tickets/io-parse-phyloxml-booleans-fallibly.md @@ -2,6 +2,9 @@ Implement the XML Schema boolean lexical contract for recognized PhyloXML properties. +> [!NOTE] +> This ticket predates the tree-output refactor, which removed the `tree_ir` layer and now writes formats directly from the graph in `packages/treetime/src/commands/shared/tree_output.rs` (readers in `packages/treetime-io/src/auspice.rs` and `packages/treetime-io/src/phyloxml.rs`). Its parent issue's current status is unconfirmed; re-derive the steps and code locations against the current code before executing. + ## Required changes - Add one private fallible parser that maps `true` and `1` to true and maps `false` and `0` to false. diff --git a/kb/tickets/io-populate-and-validate-usher-reference-sequence.md b/kb/tickets/io-populate-and-validate-usher-reference-sequence.md deleted file mode 100644 index e5759ef86..000000000 --- a/kb/tickets/io-populate-and-validate-usher-reference-sequence.md +++ /dev/null @@ -1,17 +0,0 @@ -# Populate and validate the UShER MAT reference sequence - -Mutation-bearing command projections leave `TreeIrData.root_sequence` empty, and `TreeIrUsherWriter` then substitutes each mutation's branch-local parent allele for the global `ref_nuc` field. - -## Implementation - -- Obtain the nucleotide root sequence from the projected ancestral partition and store it under the `nuc` TreeIR track. -- Make MAT serialization return an actionable error when the nucleotide reference is absent, a mutation coordinate is outside it, or the reference nucleotide is not A/C/G/T. -- Remove the branch-parent fallback in [`packages/treetime-io/src/tree_ir/usher.rs#L75-L79`](../../packages/treetime-io/src/tree_ir/usher.rs#L75-L79). -- Cover ancestral, optimize, and timetree projection paths. -- Add a recurrent-mutation regression in which two mutations at the same coordinate have different parent alleles and both export the same global `ref_nuc`. -- Add parameterized writer rejection tests for a missing reference, an out-of-range coordinate, and a non-ACGT reference nucleotide, asserting actionable error context. -- Assert the exact `nuc` root sequence projected by ancestral, optimize, and timetree. - -## Related issues - -Source: [kb/issues/H-io-usher-ref-nuc-uses-parent-allele.md](../issues/H-io-usher-ref-nuc-uses-parent-allele.md) diff --git a/kb/tickets/io-preserve-usher-missing-branch-lengths.md b/kb/tickets/io-preserve-usher-missing-branch-lengths.md new file mode 100644 index 000000000..3f0bc0e9f --- /dev/null +++ b/kb/tickets/io-preserve-usher-missing-branch-lengths.md @@ -0,0 +1,20 @@ +# Preserve missing UShER branch lengths + +Retain the distinction between an absent branch length and an explicit zero in the embedded MAT Newick tree. + +## Required changes + +- Change `UsherNodeImpl.branch_length` from `f64` to `Option`. +- Remove the `.unwrap_or(0.0)` conversion in `fn usher_to_graph()`. +- Pass the optional value unchanged through every `UsherRead` implementation into `HasBranchLength::set_branch_length()` or the edge constructor. +- Preserve the shared Newick parser's present-value behavior; do not introduce MAT-only branch-length validation. + +## Validation + +- Round-trip absent, zero, and positive branch lengths as whole MAT/graph values. +- Verify that downstream algorithms retain their existing finite, non-negative validation where branch lengths are interpreted as evolutionary distances. +- Run the full lint and test suite. + +## Related issues + +Source: [kb/issues/M-io-usher-missing-branch-length-becomes-zero.md](../issues/M-io-usher-missing-branch-length-becomes-zero.md) diff --git a/kb/tickets/io-preserve-usher-zero-branch-lengths.md b/kb/tickets/io-preserve-usher-zero-branch-lengths.md deleted file mode 100644 index add616aed..000000000 --- a/kb/tickets/io-preserve-usher-zero-branch-lengths.md +++ /dev/null @@ -1,19 +0,0 @@ -# Preserve UShER zero branch lengths - -Retain the distinction between an explicit zero branch length and an absent UShER branch length. - -## Required changes - -- Decode presence without filtering numeric zero. -- Encode `Some(0.0)` as an explicit zero and `None` as absence. -- Apply the common finite/non-negative branch-length validation at the format boundary. - -## Validation - -- Round-trip absent, zero, and positive branch lengths as whole MAT/TreeIR values. -- Reject negative, NaN, and infinite inputs with field context. -- Full lint and test suite. - -## Related issues - -- Source: [kb/issues/M-io-usher-zero-branch-length-collapsed.md](../issues/M-io-usher-zero-branch-length-collapsed.md) diff --git a/kb/tickets/io-restore-and-schema-validate-auspice-updated.md b/kb/tickets/io-restore-and-schema-validate-auspice-updated.md deleted file mode 100644 index c76440fd1..000000000 --- a/kb/tickets/io-restore-and-schema-validate-auspice-updated.md +++ /dev/null @@ -1,21 +0,0 @@ -# Restore and schema-validate Auspice updated metadata - -Make every shared-writer Auspice v2 document satisfy the official pinned schema. - -## Required changes - -- Add required `meta.updated` to the typed TreeIR Auspice model. -- Generate the current UTC calendar date once per command output plan, format it as `YYYY-MM-DD`, and pass it explicitly into the writer. -- Make the generation date injectable so tests use a fixed date without normalizing output after serialization. -- Validate complete ancestral, mugration, and timetree JSON against the pinned Augur export-v2 schema. - -## Validation - -- Official-schema success for every command using the writer. -- Negative fixture proving omission of `updated` fails validation. -- Fixed injected generation date for golden tests and a production-boundary test for UTC `YYYY-MM-DD` formatting. -- Full lint and test suite. - -## Related issues - -- Source: [kb/issues/H-io-auspice-v2-required-updated-missing.md](../issues/H-io-auspice-v2-required-updated-missing.md) diff --git a/kb/tickets/io-unify-tree-output-ordering.md b/kb/tickets/io-unify-tree-output-ordering.md deleted file mode 100644 index 2c471ddea..000000000 --- a/kb/tickets/io-unify-tree-output-ordering.md +++ /dev/null @@ -1,21 +0,0 @@ -# Unify tree output ordering - -Make every tree-backed writer consume the same ordered topology. - -## Required changes - -- Apply the selected topology order once during output preparation. -- Pass the resulting graph to direct writers and every TreeIR projection. -- Remove writer-local or projection-bypassing order paths. -- Preserve the currently selected method/format availability contract; do not add or remove parsimony formats in this ticket. - -## Validation - -- Matrix tests for `ancestral`, `clock`, `mugration`, `optimize`, `prune`, and `timetree` across input/reference/list topology orders and every currently supported tree-backed format. -- Whole-output tests proving direct and TreeIR-backed formats contain the same ordered node topology. -- Full lint and test suite. - -## Related issues - -- Source: [kb/issues/M-io-tree-backed-output-order-inconsistent.md](../issues/M-io-tree-backed-output-order-inconsistent.md) -- Related: [kb/issues/N-ancestral-auspice-json-not-produced.md](../issues/N-ancestral-auspice-json-not-produced.md) diff --git a/kb/tickets/mugration-borrow-confidence-rows.md b/kb/tickets/mugration-borrow-confidence-rows.md new file mode 100644 index 000000000..b5aec6bb7 --- /dev/null +++ b/kb/tickets/mugration-borrow-confidence-rows.md @@ -0,0 +1,22 @@ +# Borrow mugration confidence rows during output + +Use ndarray row views throughout mugration output and remove the duplicated confidence matrix retained in `MugrationGraphData`. + +## Required changes + +- Return `Option>` from `PartitionMarginalDiscrete::get_confidence()`. +- Accept `ArrayView1<'_, f64>` in `build_confidence_map()` and exact Shannon entropy calculation. +- Bind one row view per node and derive confidence plus entropy from it. +- Replace `MugrationConfidenceOutput`'s owned rows with write-time CSV traversal over final graph names and partition views. +- Update every `get_confidence()` caller without adding ownership adapters. + +## Validation + +- Exact confidence CSV, Augur node-data, and Auspice output equivalence. +- Deterministic, uniform, and non-contiguous profile view cases. +- Allocation regression coverage for per-node output projection. +- Full lint and test suite. + +## Related issues + +Source: [kb/issues/N-mugration-confidence-rows-copied-for-output.md](../issues/N-mugration-confidence-rows-copied-for-output.md) diff --git a/kb/tickets/representation-standardize-mutations-and-format-projection.md b/kb/tickets/representation-standardize-mutations-and-format-projection.md index ababf33d4..abf064faf 100644 --- a/kb/tickets/representation-standardize-mutations-and-format-projection.md +++ b/kb/tickets/representation-standardize-mutations-and-format-projection.md @@ -2,6 +2,9 @@ Make zero-based positions and typed mutation variants the single application contract, then perform checked conversion only at external format boundaries. +> [!NOTE] +> This ticket predates the tree-output refactor, which removed the `tree_ir` layer and now writes formats directly from the graph in `packages/treetime/src/commands/shared/tree_output.rs` (readers in `packages/treetime-io/src/auspice.rs` and `packages/treetime-io/src/phyloxml.rs`). Its parent issue's current status is unconfirmed; re-derive the steps and code locations against the current code before executing. + ## Required changes 1. Define one mutation representation covering substitutions, insertions, and deletions with zero-based positions. @@ -24,5 +27,4 @@ For an internal position $p_0$, serialize the external position as $p_1=p_0+1$. ## Related issues - Source: [kb/issues/M-core-mutation-representation-and-format-projection-inconsistent.md](../issues/M-core-mutation-representation-and-format-projection-inconsistent.md) -- Blocked format policies: [kb/issues/M-io-usher-mat-mutation-loss-is-implicit.md](../issues/M-io-usher-mat-mutation-loss-is-implicit.md) - Blocked format policies: [kb/issues/N-io-phyloxml-mutation-property-contract-undecided.md](../issues/N-io-phyloxml-mutation-property-contract-undecided.md) diff --git a/kb/tickets/timetree-complete-tree-output-inference-metadata.md b/kb/tickets/timetree-complete-tree-output-inference-metadata.md index db5442264..bf26a9bd2 100644 --- a/kb/tickets/timetree-complete-tree-output-inference-metadata.md +++ b/kb/tickets/timetree-complete-tree-output-inference-metadata.md @@ -2,6 +2,9 @@ Produce internally consistent TreeTime Auspice metadata and populate the already-defined non-mutation PhyloXML inference properties. +> [!NOTE] +> This ticket predates the tree-output refactor, which removed the `tree_ir` layer and now writes formats directly from the graph in `packages/treetime/src/commands/shared/tree_output.rs` (readers in `packages/treetime-io/src/auspice.rs` and `packages/treetime-io/src/phyloxml.rs`). Its parent issue's current status is unconfirmed; re-derive the steps and code locations against the current code before executing. + ## Required changes 1. Consume the canonical typed nucleotide mutation projection for Auspice output. diff --git a/kb/v0-errata/auspice-nucleotide-strand-invalid.md b/kb/v0-errata/auspice-nucleotide-strand-invalid.md index 2ee5ba55c..2a9a0dd67 100644 --- a/kb/v0-errata/auspice-nucleotide-strand-invalid.md +++ b/kb/v0-errata/auspice-nucleotide-strand-invalid.md @@ -21,4 +21,4 @@ The value contains an extra colon. It is outside the Augur annotation schema and ## v1 status -V1 does not yet emit nucleotide genome annotations. [kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md](../issues/M-timetree-tree-output-inference-metadata-incomplete.md) and [kb/tickets/timetree-complete-tree-output-inference-metadata.md](../tickets/timetree-complete-tree-output-inference-metadata.md) require schema-valid `"+"` output. +V1 emits nucleotide genome annotations and validates generated documents against the pinned Augur export-v2 and annotation schemas. Completeness of the annotation fields (including strand) is tracked in [kb/issues/M-timetree-tree-output-inference-metadata-incomplete.md](../issues/M-timetree-tree-output-inference-metadata-incomplete.md).