feat: add phylogenetic analysis CLI command and lazy distance traits
Introduce the `obikphylo` crate to support genome-vs-genome distance matrix computation and phylogenetic tree inference via Neighbor-Joining and UPGMA algorithms. Extend the `obikmer2` CLI with a new `phylo` command that exposes configurable metrics, presence thresholds, and CSV/Newick output options. Refactor `obikindex` to expose trait-based partial aggregation for efficient distance metric finalization without full matrix materialization. Update dependency graphs and remove obsolete storage modules.
This commit is contained in:
Eric Coissac
@@ -12,10 +12,13 @@ obicompactvec = { path = "../obicompactvec" }
|
||||
obikidxcache = { path = "../obikidxcache" }
|
||||
obiskbuilder = { path = "../obiskbuilder" }
|
||||
obipipeline = { path = "../obipipeline" }
|
||||
kodama = "0.3.0"
|
||||
memmap2 = "0.9"
|
||||
ndarray = "0.17"
|
||||
petgraph = "0.6.4"
|
||||
rand = "0.10"
|
||||
rayon = "1"
|
||||
speedytree = "0.1"
|
||||
tracing = "0.1.44"
|
||||
|
||||
[dev-dependencies]
|
||||
|
||||
+124
-37
@@ -1,10 +1,23 @@
|
||||
use ndarray::Array2;
|
||||
use obicompactvec::traits::{BitPartials, CountPartials};
|
||||
use obikidxcache::LayeredStore;
|
||||
use rayon::prelude::*;
|
||||
//! Genome-vs-genome distance matrices over a whole index — the aggregation
|
||||
//! (summing per-layer partials into one global partial) and the final
|
||||
//! distance calculation are both phylogenetic domain logic, so both live
|
||||
//! here, not in `obikidxcache` (which only ever hands out already-opened
|
||||
//! `KmerLayer`s via `IndexCache::iter()` — plain exposition, no
|
||||
//! computation of its own).
|
||||
//!
|
||||
//! [`AggregatedCount`]/[`AggregatedBit`] exist purely to "hang"
|
||||
//! `obicompactvec`'s `CountPartials`/`BitPartials` traits on a borrowed
|
||||
//! collection of `KmerLayer`s, so their already-written finalisation
|
||||
//! methods (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) can be reused
|
||||
//! as-is instead of re-derived here — see `KmerLayer`'s own
|
||||
//! `ColumnWeights`/`CountPartials`/`BitPartials` impls (`obikindex`) for
|
||||
//! the per-layer primitives these sum.
|
||||
|
||||
use ndarray::{Array1, Array2};
|
||||
use obicompactvec::{BitPartials, ColumnWeights, CountPartials};
|
||||
use obikidxcache::index_cache::IndexCache;
|
||||
use obikindex::layer::{KmerLayer, LayerContent};
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use obikindex::KmerIndex;
|
||||
|
||||
// ── Public API ────────────────────────────────────────────────────────────────
|
||||
|
||||
@@ -51,43 +64,122 @@ impl DistanceMetric {
|
||||
}
|
||||
}
|
||||
|
||||
// ── KmerIndex::distance ───────────────────────────────────────────────────────
|
||||
// ── Aggregators — sum every cached layer's own partials into one global
|
||||
// partial, then hand the result off to `CountPartials`/`BitPartials`'s own
|
||||
// provided finalisation methods.
|
||||
|
||||
impl KmerIndex {
|
||||
pub fn distance(&self, metric: DistanceMetric, shared_kmers: bool, presence_threshold: u32) -> OKIResult<DistanceOutput> {
|
||||
let n_genomes = self.meta.genomes().map_err(OKIError::Io)?.len();
|
||||
fn sum_array1(n: usize, items: impl Iterator<Item = Array1<u64>>) -> Array1<u64> {
|
||||
items.fold(Array1::zeros(n), |acc, x| acc + x)
|
||||
}
|
||||
fn sum_array2_u64(n: usize, items: impl Iterator<Item = Array2<u64>>) -> Array2<u64> {
|
||||
items.fold(Array2::zeros((n, n)), |acc, x| acc + x)
|
||||
}
|
||||
fn sum_array2_f64(n: usize, items: impl Iterator<Item = Array2<f64>>) -> Array2<f64> {
|
||||
items.fold(Array2::zeros((n, n)), |acc, x| acc + x)
|
||||
}
|
||||
|
||||
struct AggregatedCount<'a> {
|
||||
n_genomes: usize,
|
||||
layers: Vec<&'a KmerLayer>,
|
||||
}
|
||||
|
||||
impl ColumnWeights for AggregatedCount<'_> {
|
||||
fn col_weights(&self) -> Array1<u64> {
|
||||
sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights()))
|
||||
}
|
||||
}
|
||||
|
||||
impl CountPartials for AggregatedCount<'_> {
|
||||
fn partial_bray(&self) -> Array2<u64> {
|
||||
sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_bray()))
|
||||
}
|
||||
fn partial_euclidean(&self) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_euclidean()))
|
||||
}
|
||||
fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
|
||||
let n = self.n_genomes;
|
||||
self.layers.iter().map(|l| l.partial_threshold_jaccard(threshold)).fold(
|
||||
(Array2::zeros((n, n)), Array2::zeros((n, n))),
|
||||
|(inter, union), (i, u)| (inter + i, union + u),
|
||||
)
|
||||
}
|
||||
fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_bray(global)))
|
||||
}
|
||||
fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_euclidean(global)))
|
||||
}
|
||||
fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
|
||||
sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_hellinger(global)))
|
||||
}
|
||||
}
|
||||
|
||||
struct AggregatedBit<'a> {
|
||||
n_genomes: usize,
|
||||
layers: Vec<&'a KmerLayer>,
|
||||
}
|
||||
|
||||
impl ColumnWeights for AggregatedBit<'_> {
|
||||
fn col_weights(&self) -> Array1<u64> {
|
||||
sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights()))
|
||||
}
|
||||
}
|
||||
|
||||
impl BitPartials for AggregatedBit<'_> {
|
||||
fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
|
||||
let n = self.n_genomes;
|
||||
self.layers.iter().map(|l| l.partial_jaccard()).fold(
|
||||
(Array2::zeros((n, n)), Array2::zeros((n, n))),
|
||||
|(inter, union), (i, u)| (inter + i, union + u),
|
||||
)
|
||||
}
|
||||
fn partial_hamming(&self) -> Array2<u64> {
|
||||
sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_hamming()))
|
||||
}
|
||||
}
|
||||
|
||||
// ── Metrics — the domain-facing entry point ────────────────────────────────────
|
||||
|
||||
/// Genome-vs-genome distance computation over a whole (already-cached)
|
||||
/// index. `IndexCache` is foreign to this crate, so this is an extension
|
||||
/// trait — same reasoning as every other `obikindex`/`obikidxcache`
|
||||
/// extension trait in this codebase.
|
||||
pub trait Metrics {
|
||||
fn distance(&self, metric: DistanceMetric, shared_kmers: bool, presence_threshold: u32) -> OKIResult<DistanceOutput>;
|
||||
}
|
||||
|
||||
impl Metrics for IndexCache {
|
||||
fn distance(&self, metric: DistanceMetric, shared_kmers: bool, presence_threshold: u32) -> OKIResult<DistanceOutput> {
|
||||
let n_genomes = self.meta().genomes().len();
|
||||
if n_genomes < 2 {
|
||||
return Err(OKIError::InvalidInput(
|
||||
"distance requires at least 2 genomes in the index".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let use_counts = self.meta.config.with_counts;
|
||||
let use_counts = self.meta().config().with_counts;
|
||||
if metric.requires_counts() && !use_counts {
|
||||
return Err(OKIError::InvalidInput(format!(
|
||||
"{metric:?} requires a count index (with_counts = true)"
|
||||
)));
|
||||
}
|
||||
|
||||
let n_parts = self.n_partitions();
|
||||
let kmer_size = self.meta().config().kmer_size;
|
||||
|
||||
if use_counts {
|
||||
let stores: Vec<_> = (0..n_parts)
|
||||
.into_par_iter()
|
||||
.map(|i| self.count_store(i).map_err(OKIError::Partition))
|
||||
.collect::<OKIResult<_>>()?;
|
||||
let global = LayeredStore::new(stores);
|
||||
let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Count).collect();
|
||||
let agg = AggregatedCount { n_genomes, layers };
|
||||
|
||||
let matrix = match metric {
|
||||
DistanceMetric::BrayCurtis => CountPartials::bray_dist_matrix(&global),
|
||||
DistanceMetric::RelfreqBrayCurtis => CountPartials::relfreq_bray_dist_matrix(&global),
|
||||
DistanceMetric::Euclidean => CountPartials::euclidean_dist_matrix(&global),
|
||||
DistanceMetric::RelfreqEuclidean => CountPartials::relfreq_euclidean_dist_matrix(&global),
|
||||
DistanceMetric::Hellinger => CountPartials::hellinger_dist_matrix(&global),
|
||||
DistanceMetric::HellingerEuclidean => CountPartials::hellinger_euclidean_dist_matrix(&global),
|
||||
DistanceMetric::Jaccard => CountPartials::threshold_jaccard_dist_matrix(&global, presence_threshold),
|
||||
DistanceMetric::Mash => CountPartials::threshold_mash_dist_matrix(&global, self.kmer_size(), presence_threshold),
|
||||
DistanceMetric::Hamming => {
|
||||
DistanceMetric::BrayCurtis => agg.bray_dist_matrix(),
|
||||
DistanceMetric::RelfreqBrayCurtis => agg.relfreq_bray_dist_matrix(),
|
||||
DistanceMetric::Euclidean => agg.euclidean_dist_matrix(),
|
||||
DistanceMetric::RelfreqEuclidean => agg.relfreq_euclidean_dist_matrix(),
|
||||
DistanceMetric::Hellinger => agg.hellinger_dist_matrix(),
|
||||
DistanceMetric::HellingerEuclidean => agg.hellinger_euclidean_dist_matrix(),
|
||||
DistanceMetric::Jaccard => agg.threshold_jaccard_dist_matrix(presence_threshold),
|
||||
DistanceMetric::Mash => agg.threshold_mash_dist_matrix(kmer_size, presence_threshold),
|
||||
DistanceMetric::Hamming => {
|
||||
return Err(OKIError::InvalidInput(
|
||||
"Hamming is only available for presence/absence indexes".into(),
|
||||
));
|
||||
@@ -95,7 +187,7 @@ impl KmerIndex {
|
||||
};
|
||||
|
||||
let shared = if shared_kmers {
|
||||
let (inter, _) = CountPartials::partial_threshold_jaccard(&global, presence_threshold);
|
||||
let (inter, _) = agg.partial_threshold_jaccard(presence_threshold);
|
||||
Some(inter)
|
||||
} else {
|
||||
None
|
||||
@@ -103,18 +195,13 @@ impl KmerIndex {
|
||||
|
||||
Ok(DistanceOutput { matrix, shared_kmers: shared })
|
||||
} else {
|
||||
let stores: Vec<_> = (0..n_parts)
|
||||
.into_par_iter()
|
||||
.map(|i| self.presence_store(i).map_err(OKIError::Partition))
|
||||
.collect::<OKIResult<_>>()?;
|
||||
let global = LayeredStore::new(stores);
|
||||
let layers: Vec<&KmerLayer> = self.iter().filter(|l| l.content() == LayerContent::Presence).collect();
|
||||
let agg = AggregatedBit { n_genomes, layers };
|
||||
|
||||
let matrix = match metric {
|
||||
DistanceMetric::Jaccard => BitPartials::jaccard_dist_matrix(&global),
|
||||
DistanceMetric::Mash => BitPartials::mash_dist_matrix(&global, self.kmer_size()),
|
||||
DistanceMetric::Hamming => {
|
||||
BitPartials::hamming_dist_matrix(&global).mapv(|v| v as f64)
|
||||
}
|
||||
DistanceMetric::Jaccard => agg.jaccard_dist_matrix(),
|
||||
DistanceMetric::Mash => agg.mash_dist_matrix(kmer_size),
|
||||
DistanceMetric::Hamming => agg.hamming_dist_matrix().mapv(|v| v as f64),
|
||||
other => {
|
||||
return Err(OKIError::InvalidInput(format!(
|
||||
"{other:?} requires a count index; use --metric jaccard or --metric hamming"
|
||||
@@ -123,7 +210,7 @@ impl KmerIndex {
|
||||
};
|
||||
|
||||
let shared = if shared_kmers {
|
||||
let (inter, _) = BitPartials::partial_jaccard(&global);
|
||||
let (inter, _) = agg.partial_jaccard();
|
||||
Some(inter)
|
||||
} else {
|
||||
None
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
//! Library-level phylogenetic functionality for `obikmer`, built as
|
||||
//! extension traits over `obikindex::KmerIndex` and `obikindex::layer`'s
|
||||
//! generic layer types — the `phylo` CLI command is a consumer of this
|
||||
//! crate, not the owner of this logic (see `DevDocMD/architecture/siblings.md`).
|
||||
//! extension traits over `obikindex`/`obikidxcache`'s foreign types — the
|
||||
//! `phylo` CLI command is a consumer of this crate, not the owner of this
|
||||
//! logic (see `DevDocMD/architecture/siblings.md`).
|
||||
//!
|
||||
//! Starts with [`siblings`] (family presence-mask annex, SNP distance,
|
||||
//! cardinality, pseudo-alignment); further phylo-domain functionality
|
||||
@@ -15,7 +15,8 @@
|
||||
//! this crate until `siblings` comes back.
|
||||
|
||||
mod distance;
|
||||
mod matrix_store;
|
||||
mod tree;
|
||||
// pub mod siblings; // temporarily disconnected — see module doc above.
|
||||
|
||||
pub use distance::{DistanceMetric, DistanceOutput};
|
||||
pub use distance::{DistanceMetric, DistanceOutput, Metrics};
|
||||
pub use tree::{neighbor_joining_newick, upgma_newick};
|
||||
|
||||
@@ -1,47 +0,0 @@
|
||||
use obicompactvec::{PersistentBitMatrix, PersistentIntMatrix};
|
||||
use obikidxcache::LayeredStore;
|
||||
use obikindex::OKIResult;
|
||||
use obikindex::layer::open_data;
|
||||
|
||||
use obikindex::KmerIndex;
|
||||
use obikindex::load_meta;
|
||||
|
||||
impl KmerIndex {
|
||||
/// Open all count matrices for partition `part`, one per layer.
|
||||
/// Layers without a `counts/` directory are skipped.
|
||||
pub fn count_store(&self, part: usize) -> OKIResult<LayeredStore<PersistentIntMatrix>> {
|
||||
let index_dir = self.index_dir(part);
|
||||
if !index_dir.exists() {
|
||||
return Ok(LayeredStore::new(vec![]));
|
||||
}
|
||||
let n_layers = self.n_layers();
|
||||
let matrices = (0..n_layers)
|
||||
.filter_map(|l| {
|
||||
self.layer_dir(part, l)
|
||||
.join("counts")
|
||||
.exists()
|
||||
.then(|| open_data(&index_dir, l))
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
Ok(LayeredStore::new(matrices))
|
||||
}
|
||||
|
||||
/// Open all presence matrices for partition `part`, one per layer.
|
||||
/// Layers without a `presence/` directory are skipped.
|
||||
pub fn presence_store(&self, part: usize) -> OKIResult<LayeredStore<PersistentBitMatrix>> {
|
||||
let index_dir = self.index_dir(part);
|
||||
if !index_dir.exists() {
|
||||
return Ok(LayeredStore::new(vec![]));
|
||||
}
|
||||
let n_layers = load_meta(&index_dir)?.n_layers;
|
||||
let matrices = (0..n_layers)
|
||||
.filter_map(|l| {
|
||||
self.layer_dir(part, l)
|
||||
.join("presence")
|
||||
.exists()
|
||||
.then(|| open_data(&index_dir, l))
|
||||
})
|
||||
.collect::<OKIResult<Vec<_>>>()?;
|
||||
Ok(LayeredStore::new(matrices))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
//! A minimal, write-only phylogenetic tree: just enough structure to unify
|
||||
//! whatever a reconstruction algorithm produces (NJ, UPGMA, ...) into one
|
||||
//! shared Newick serialisation, instead of each algorithm hand-rolling its
|
||||
//! own Newick writer. No reader — nothing in this crate ever needs to parse
|
||||
//! a Newick string back in, so there is no `from_newick`.
|
||||
|
||||
mod newick;
|
||||
mod nj;
|
||||
mod upgma;
|
||||
|
||||
pub use nj::neighbor_joining_newick;
|
||||
pub use upgma::upgma_newick;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
struct NodeId(usize);
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct Node {
|
||||
name: Option<String>,
|
||||
branch_length: Option<f64>,
|
||||
bootstrap: Option<f64>,
|
||||
children: Vec<NodeId>,
|
||||
}
|
||||
|
||||
/// Arena-based tree: nodes are created detached ([`new_leaf`](Tree::new_leaf)/
|
||||
/// [`new_node`](Tree::new_node)) and wired together with
|
||||
/// [`attach`](Tree::attach) — the shape every bottom-up reconstruction
|
||||
/// algorithm (UPGMA's successive merges, NJ's already-built graph) needs:
|
||||
/// a node's parent is only known once *it* already exists.
|
||||
#[derive(Debug, Default)]
|
||||
struct Tree {
|
||||
nodes: Vec<Node>,
|
||||
root: Option<NodeId>,
|
||||
}
|
||||
|
||||
impl Tree {
|
||||
fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Create a new, detached leaf node.
|
||||
fn new_leaf(&mut self, name: impl Into<String>) -> NodeId {
|
||||
let id = self.new_node();
|
||||
self.nodes[id.0].name = Some(name.into());
|
||||
id
|
||||
}
|
||||
|
||||
/// Create a new, detached internal node (no name).
|
||||
fn new_node(&mut self) -> NodeId {
|
||||
let id = NodeId(self.nodes.len());
|
||||
self.nodes.push(Node::default());
|
||||
id
|
||||
}
|
||||
|
||||
/// Attach `child` under `parent`, recording the branch length between them.
|
||||
fn attach(&mut self, parent: NodeId, child: NodeId, branch_length: f64) {
|
||||
self.nodes[child.0].branch_length = Some(branch_length);
|
||||
self.nodes[parent.0].children.push(child);
|
||||
}
|
||||
|
||||
/// Record a bootstrap/support value for `node` — meaningful on an
|
||||
/// internal node only; printed right after its closing parenthesis, per
|
||||
/// Newick convention.
|
||||
#[allow(dead_code)] // no current caller sets one (NJ/UPGMA don't produce support values); kept for the next algorithm that does.
|
||||
fn set_bootstrap(&mut self, node: NodeId, value: f64) {
|
||||
self.nodes[node.0].bootstrap = Some(value);
|
||||
}
|
||||
|
||||
/// Mark `node` as the tree's root. Required exactly once before
|
||||
/// [`to_newick`](Tree::to_newick).
|
||||
fn set_root(&mut self, node: NodeId) {
|
||||
self.root = Some(node);
|
||||
}
|
||||
|
||||
/// Serialise to Newick (terminated with `;`). Panics if
|
||||
/// [`set_root`](Tree::set_root) was never called — a tree with no
|
||||
/// designated root is a builder bug, not a recoverable input.
|
||||
fn to_newick(&self) -> String {
|
||||
let root = self.root.expect("Tree::to_newick: root not set");
|
||||
let mut out = newick::write_node(self, root);
|
||||
out.push(';');
|
||||
out
|
||||
}
|
||||
|
||||
fn node(&self, id: NodeId) -> &Node {
|
||||
&self.nodes[id.0]
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,34 @@
|
||||
//! The one Newick writer every tree-reconstruction algorithm in this
|
||||
//! module shares.
|
||||
|
||||
use super::{Node, NodeId, Tree};
|
||||
|
||||
pub(super) fn write_node(tree: &Tree, id: NodeId) -> String {
|
||||
let node: &Node = tree.node(id);
|
||||
let mut s = String::new();
|
||||
|
||||
if node.children.is_empty() {
|
||||
if let Some(name) = &node.name {
|
||||
s.push_str(name);
|
||||
}
|
||||
} else {
|
||||
s.push('(');
|
||||
for (i, &child) in node.children.iter().enumerate() {
|
||||
if i > 0 {
|
||||
s.push(',');
|
||||
}
|
||||
s.push_str(&write_node(tree, child));
|
||||
}
|
||||
s.push(')');
|
||||
if let Some(bootstrap) = node.bootstrap {
|
||||
s.push_str(&format!("{bootstrap}"));
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(branch_length) = node.branch_length {
|
||||
s.push(':');
|
||||
s.push_str(&format!("{branch_length:.6}"));
|
||||
}
|
||||
|
||||
s
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
//! Neighbor-Joining, via `speedytree` — converts its own tree type
|
||||
//! (`petgraph::graph::UnGraph<String, f64>`) into this crate's [`super::Tree`]
|
||||
//! so its Newick output goes through the same writer as every other
|
||||
//! algorithm here.
|
||||
|
||||
use ndarray::Array2;
|
||||
use obikindex::{OKIError, OKIResult};
|
||||
use petgraph::graph::NodeIndex;
|
||||
use speedytree::{DistanceMatrix, Hybrid, NeighborJoiningSolver};
|
||||
|
||||
use super::{NodeId, Tree};
|
||||
|
||||
/// Compute a Neighbor-Joining tree from a symmetric `n×n` distance `matrix`
|
||||
/// (genomes in `labels` order), returning it as a Newick string.
|
||||
pub fn neighbor_joining_newick(matrix: &Array2<f64>, labels: &[String]) -> OKIResult<String> {
|
||||
let n = labels.len();
|
||||
let rows: Vec<Vec<f64>> = (0..n).map(|i| (0..n).map(|j| matrix[[i, j]]).collect()).collect();
|
||||
let dm = DistanceMatrix::build(rows, labels.to_vec())
|
||||
.map_err(|e| OKIError::InvalidInput(format!("building NJ distance matrix: {e}")))?;
|
||||
let graph = NeighborJoiningSolver::<Hybrid>::default(dm)
|
||||
.solve()
|
||||
.map_err(|e| OKIError::InvalidInput(format!("computing NJ tree: {e}")))?;
|
||||
|
||||
// speedytree's own Newick writer expects the same convention: the
|
||||
// unrooted result's "display root" is whichever node still has all
|
||||
// three of its original neighbours (every other node lost one to
|
||||
// becoming somebody's child during the walk below).
|
||||
let root = graph
|
||||
.node_indices()
|
||||
.find(|&node| graph.neighbors(node).count() == 3)
|
||||
.ok_or_else(|| OKIError::InvalidInput("NJ tree has no trifurcating root".into()))?;
|
||||
|
||||
let mut tree = Tree::new();
|
||||
let root_id = convert(&graph, root, root, &mut tree);
|
||||
tree.set_root(root_id);
|
||||
Ok(tree.to_newick())
|
||||
}
|
||||
|
||||
/// Recursively convert `node` (and everything below it, i.e. every
|
||||
/// neighbour except `parent`) into this crate's `Tree`. Excluding just the
|
||||
/// immediate parent is enough to avoid revisiting — `graph` is a tree
|
||||
/// (no cycles), so that's the only edge that could lead back.
|
||||
fn convert(graph: &speedytree::Tree, node: NodeIndex, parent: NodeIndex, tree: &mut Tree) -> NodeId {
|
||||
let label = &graph[node];
|
||||
let id = if label.is_empty() { tree.new_node() } else { tree.new_leaf(label.clone()) };
|
||||
|
||||
for child in graph.neighbors(node).filter(|&c| c != parent) {
|
||||
let edge = graph.find_edge(node, child).expect("adjacent nodes must share an edge");
|
||||
let branch_length = *graph.edge_weight(edge).expect("edge must carry a weight");
|
||||
let child_id = convert(graph, child, node, tree);
|
||||
tree.attach(id, child_id, branch_length);
|
||||
}
|
||||
|
||||
id
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
//! UPGMA (average-linkage hierarchical clustering), via `kodama` — converts
|
||||
//! its dendrogram (a flat list of merge steps) into this crate's
|
||||
//! [`super::Tree`], bottom-up, so its Newick output goes through the same
|
||||
//! writer as every other algorithm here.
|
||||
|
||||
use kodama::{Method, linkage};
|
||||
use ndarray::Array2;
|
||||
|
||||
use super::Tree;
|
||||
|
||||
/// Compute a UPGMA tree from a symmetric `n×n` distance `matrix` (genomes in
|
||||
/// `labels` order), returning it as a Newick string.
|
||||
pub fn upgma_newick(matrix: &Array2<f64>, labels: &[String]) -> String {
|
||||
let n = labels.len();
|
||||
let mut condensed: Vec<f64> = Vec::with_capacity(n * (n - 1) / 2);
|
||||
for i in 0..n {
|
||||
for j in (i + 1)..n {
|
||||
condensed.push(matrix[[i, j]]);
|
||||
}
|
||||
}
|
||||
let dendro = linkage(&mut condensed, n, Method::Average);
|
||||
|
||||
let mut tree = Tree::new();
|
||||
// `kodama` numbers clusters `0..n` for the original leaves and `n..`
|
||||
// for each successive merge, in the order `dendro.steps()` yields them
|
||||
// — so appending to `node_of` as we go keeps `node_of[cluster_id]`
|
||||
// valid the moment a later step references it.
|
||||
let mut node_of: Vec<super::NodeId> =
|
||||
labels.iter().map(|label| tree.new_leaf(label.clone())).collect();
|
||||
// height of each node: leaves = 0, internal = dissimilarity/2 (UPGMA is
|
||||
// ultrametric: both children of a merge are equidistant from it).
|
||||
let mut heights = vec![0.0f64; 2 * n - 1];
|
||||
|
||||
let mut root = None;
|
||||
for (k, step) in dendro.steps().iter().enumerate() {
|
||||
let new_node = n + k;
|
||||
let h = step.dissimilarity / 2.0;
|
||||
heights[new_node] = h;
|
||||
|
||||
let parent = tree.new_node();
|
||||
for &cluster in &[step.cluster1, step.cluster2] {
|
||||
let branch_length = (h - heights[cluster]).max(0.0);
|
||||
tree.attach(parent, node_of[cluster], branch_length);
|
||||
}
|
||||
node_of.push(parent);
|
||||
root = Some(parent);
|
||||
}
|
||||
|
||||
tree.set_root(root.expect("at least one merge step for n >= 2 genomes"));
|
||||
tree.to_newick()
|
||||
}
|
||||
Reference in New Issue
Block a user