From 635fc830d182f7b5b826f6ba8f4f6457b1ca437a Mon Sep 17 00:00:00 2001 From: Eric Coissac Date: Thu, 27 Aug 2026 06:37:48 +0200 Subject: [PATCH] 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. --- src/Cargo.lock | 4 + src/obikindex/src/layer/content_layer.rs | 79 ++++++++- src/obikindex/src/layer/tests/typed_layer.rs | 10 +- src/obikindex/src/layer/typed_layer.rs | 38 ++++- src/obikmer2/Cargo.toml | 1 + src/obikmer2/src/cmd/mod.rs | 1 + src/obikmer2/src/cmd/phylo/args.rs | 72 +++++++++ src/obikmer2/src/cmd/phylo/mod.rs | 134 +++++++++++++++ src/obikmer2/src/main.rs | 4 + src/obikphylo/Cargo.toml | 3 + src/obikphylo/src/distance.rs | 161 ++++++++++++++----- src/obikphylo/src/lib.rs | 11 +- src/obikphylo/src/matrix_store.rs | 47 ------ src/obikphylo/src/tree/mod.rs | 88 ++++++++++ src/obikphylo/src/tree/newick.rs | 34 ++++ src/obikphylo/src/tree/nj.rs | 55 +++++++ src/obikphylo/src/tree/upgma.rs | 51 ++++++ 17 files changed, 696 insertions(+), 97 deletions(-) create mode 100644 src/obikmer2/src/cmd/phylo/args.rs create mode 100644 src/obikmer2/src/cmd/phylo/mod.rs delete mode 100644 src/obikphylo/src/matrix_store.rs create mode 100644 src/obikphylo/src/tree/mod.rs create mode 100644 src/obikphylo/src/tree/newick.rs create mode 100644 src/obikphylo/src/tree/nj.rs create mode 100644 src/obikphylo/src/tree/upgma.rs diff --git a/src/Cargo.lock b/src/Cargo.lock index a2205fc6..961a2555 100644 --- a/src/Cargo.lock +++ b/src/Cargo.lock @@ -1677,6 +1677,7 @@ dependencies = [ "obikindex", "obikindexer", "obikmerge", + "obikphylo", "obikquery", "obikrebuild", "obikrope", @@ -1712,6 +1713,7 @@ dependencies = [ name = "obikphylo" version = "0.1.0" dependencies = [ + "kodama", "memmap2", "ndarray", "obicompactvec", @@ -1725,8 +1727,10 @@ dependencies = [ "obiskbuilder", "obiskio", "obisys", + "petgraph", "rand 0.10.2", "rayon", + "speedytree", "tempfile", "tracing", "tracing-subscriber", diff --git a/src/obikindex/src/layer/content_layer.rs b/src/obikindex/src/layer/content_layer.rs index 916d2a6c..e5f832ee 100644 --- a/src/obikindex/src/layer/content_layer.rs +++ b/src/obikindex/src/layer/content_layer.rs @@ -26,7 +26,8 @@ use std::path::{Path, PathBuf}; use crate::layer::utils::LAYERNAME_SUFFIX; -use obicompactvec::{PersistentBitMatrix, PersistentIntMatrix}; +use ndarray::{Array1, Array2}; +use obicompactvec::{BitPartials, ColumnWeights, CountPartials, PersistentBitMatrix, PersistentIntMatrix}; use obikseq::CanonicalKmer; use crate::index::error::OKIResult; @@ -349,3 +350,79 @@ impl KmerLayer { } } } + +// ── obicompactvec's statistics traits — read-only accessors, same spirit +// as `col_weights`/`nonzero_iter` above, just formalised as trait impls so +// generic aggregation code (`obikidxcache::IndexCache`'s callers) can rely +// on `KmerLayer: ColumnWeights`/`CountPartials`/`BitPartials` rather than +// duck-typing matching method names. `CountPartials`/`BitPartials` each +// only make sense for the content this layer actually holds — calling the +// wrong one panics, same convention as `col()`/`content()` above. + +impl ColumnWeights for KmerLayer { + fn col_weights(&self) -> Array1 { + self.col_weights() + } +} + +impl CountPartials for KmerLayer { + fn partial_bray(&self) -> Array2 { + match self { + KmerLayer::Count { layer, .. } => layer.partial_bray(), + KmerLayer::Presence { .. } => panic!("CountPartials::partial_bray() called on a Presence layer"), + KmerLayer::Empty { .. } => panic!("CountPartials::partial_bray() called on an Empty layer"), + } + } + fn partial_euclidean(&self) -> Array2 { + match self { + KmerLayer::Count { layer, .. } => layer.partial_euclidean(), + KmerLayer::Presence { .. } => panic!("CountPartials::partial_euclidean() called on a Presence layer"), + KmerLayer::Empty { .. } => panic!("CountPartials::partial_euclidean() called on an Empty layer"), + } + } + fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2, Array2) { + match self { + KmerLayer::Count { layer, .. } => layer.partial_threshold_jaccard(threshold), + KmerLayer::Presence { .. } => panic!("CountPartials::partial_threshold_jaccard() called on a Presence layer"), + KmerLayer::Empty { .. } => panic!("CountPartials::partial_threshold_jaccard() called on an Empty layer"), + } + } + fn partial_relfreq_bray(&self, global: &Array1) -> Array2 { + match self { + KmerLayer::Count { layer, .. } => layer.partial_relfreq_bray(global), + KmerLayer::Presence { .. } => panic!("CountPartials::partial_relfreq_bray() called on a Presence layer"), + KmerLayer::Empty { .. } => panic!("CountPartials::partial_relfreq_bray() called on an Empty layer"), + } + } + fn partial_relfreq_euclidean(&self, global: &Array1) -> Array2 { + match self { + KmerLayer::Count { layer, .. } => layer.partial_relfreq_euclidean(global), + KmerLayer::Presence { .. } => panic!("CountPartials::partial_relfreq_euclidean() called on a Presence layer"), + KmerLayer::Empty { .. } => panic!("CountPartials::partial_relfreq_euclidean() called on an Empty layer"), + } + } + fn partial_hellinger(&self, global: &Array1) -> Array2 { + match self { + KmerLayer::Count { layer, .. } => layer.partial_hellinger(global), + KmerLayer::Presence { .. } => panic!("CountPartials::partial_hellinger() called on a Presence layer"), + KmerLayer::Empty { .. } => panic!("CountPartials::partial_hellinger() called on an Empty layer"), + } + } +} + +impl BitPartials for KmerLayer { + fn partial_jaccard(&self) -> (Array2, Array2) { + match self { + KmerLayer::Presence { layer, .. } => layer.partial_jaccard(), + KmerLayer::Count { .. } => panic!("BitPartials::partial_jaccard() called on a Count layer"), + KmerLayer::Empty { .. } => panic!("BitPartials::partial_jaccard() called on an Empty layer"), + } + } + fn partial_hamming(&self) -> Array2 { + match self { + KmerLayer::Presence { layer, .. } => layer.partial_hamming(), + KmerLayer::Count { .. } => panic!("BitPartials::partial_hamming() called on a Count layer"), + KmerLayer::Empty { .. } => panic!("BitPartials::partial_hamming() called on an Empty layer"), + } + } +} diff --git a/src/obikindex/src/layer/tests/typed_layer.rs b/src/obikindex/src/layer/tests/typed_layer.rs index 3de5eb9b..f400dee7 100644 --- a/src/obikindex/src/layer/tests/typed_layer.rs +++ b/src/obikindex/src/layer/tests/typed_layer.rs @@ -196,11 +196,11 @@ fn count_layer_reports_count_content_and_columnar_storage() { set_k(4); let dir = tempdir().unwrap(); write_unitigs(dir.path(), &[b"AAAACGT"]); - TypedLayer::::build(dir.pat - h(), DEFAUL - T_BLOCK_BITS, &Index - Mode::Exact, |_| 1) - , + TypedLayer::::build( + dir.path(), + DEFAULT_BLOCK_BITS, + &IndexMode::Exact, + |_| 1, ).unwrap(); let layer = TypedLayer::::open(dir.path()).unwrap(); diff --git a/src/obikindex/src/layer/typed_layer.rs b/src/obikindex/src/layer/typed_layer.rs index af37d198..0b62c34e 100644 --- a/src/obikindex/src/layer/typed_layer.rs +++ b/src/obikindex/src/layer/typed_layer.rs @@ -1,8 +1,10 @@ // use crate::layer::utils::layer_dir; use obicompactvec::{ - BinaryMatrix, ColBuilder, MatrixBuilder, PersistentBitMatrix, PersistentBitMatrixBuilder, - PersistentCompactIntMatrixBuilder, PersistentIntMatrix, PersistentSparseBitMatrix, + BinaryMatrix, BitPartials, ColBuilder, CountPartials, MatrixBuilder, PersistentBitMatrix, + PersistentBitMatrixBuilder, PersistentCompactIntMatrixBuilder, PersistentIntMatrix, + PersistentSparseBitMatrix, }; +use ndarray::{Array1, Array2}; use obikseq::CanonicalKmer; use obiskio::{UnitigFileReader, UnitigFileWriter}; use std::collections::HashMap; @@ -424,6 +426,29 @@ impl TypedLayer { ) -> Box + 'a> { self.data.nonzero_iter(slots) } + + /// `CountPartials`' required primitives, delegated straight to the + /// underlying matrix — see `obicompactvec::CountPartials` for what + /// each one means and how its provided finalisation methods + /// (`bray_dist_matrix`, `jaccard_dist_matrix`, ...) consume them. + pub fn partial_bray(&self) -> Array2 { + CountPartials::partial_bray(&self.data) + } + pub fn partial_euclidean(&self) -> Array2 { + CountPartials::partial_euclidean(&self.data) + } + pub fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2, Array2) { + CountPartials::partial_threshold_jaccard(&self.data, threshold) + } + pub fn partial_relfreq_bray(&self, global: &Array1) -> Array2 { + CountPartials::partial_relfreq_bray(&self.data, global) + } + pub fn partial_relfreq_euclidean(&self, global: &Array1) -> Array2 { + CountPartials::partial_relfreq_euclidean(&self.data, global) + } + pub fn partial_hellinger(&self, global: &Array1) -> Array2 { + CountPartials::partial_hellinger(&self.data, global) + } } // ── Mode 3 — presence/absence matrix ───────────────────────────────────────── @@ -487,6 +512,15 @@ impl TypedLayer { self.data.nonzero_iter(slots) } + /// `BitPartials`' required primitives, delegated straight to the + /// underlying matrix. + pub fn partial_jaccard(&self) -> (Array2, Array2) { + BitPartials::partial_jaccard(&self.data) + } + pub fn partial_hamming(&self) -> Array2 { + BitPartials::partial_hamming(&self.data) + } + pub fn append_genome_column( layer_dir: &Path, value_of: impl Fn(usize) -> bool, diff --git a/src/obikmer2/Cargo.toml b/src/obikmer2/Cargo.toml index 97bdacd9..7409d1d7 100644 --- a/src/obikmer2/Cargo.toml +++ b/src/obikmer2/Cargo.toml @@ -23,6 +23,7 @@ obikrebuild = { path = "../obikrebuild" } obikstats = { path = "../obikstats" } obikquery = { path = "../obikquery" } obikidxcache = { path = "../obikidxcache" } +obikphylo = { path = "../obikphylo" } obikrope = { path = "../obikrope" } obifastwrite = { path = "../obifastwrite" } obiskbuilder = { path = "../obiskbuilder" } diff --git a/src/obikmer2/src/cmd/mod.rs b/src/obikmer2/src/cmd/mod.rs index f448b81c..bf381257 100644 --- a/src/obikmer2/src/cmd/mod.rs +++ b/src/obikmer2/src/cmd/mod.rs @@ -7,6 +7,7 @@ pub mod index; pub mod merge; pub mod pack; mod predicate; +pub mod phylo; pub mod query; pub mod select; pub mod superkmer; diff --git a/src/obikmer2/src/cmd/phylo/args.rs b/src/obikmer2/src/cmd/phylo/args.rs new file mode 100644 index 00000000..6b4b266a --- /dev/null +++ b/src/obikmer2/src/cmd/phylo/args.rs @@ -0,0 +1,72 @@ +use std::path::PathBuf; + +use clap::Args; +use obikphylo::DistanceMetric; + +#[derive(clap::ValueEnum, Clone, Copy, Debug)] +pub enum MetricArg { + Jaccard, + Mash, + Hamming, + BrayCurtis, + #[value(name = "relfreq-bray-curtis")] + RelfreqBrayCurtis, + Euclidean, + #[value(name = "relfreq-euclidean")] + RelfreqEuclidean, + Hellinger, + #[value(name = "hellinger-euclidean")] + HellingerEuclidean, +} + +impl From for DistanceMetric { + fn from(m: MetricArg) -> Self { + match m { + MetricArg::Jaccard => DistanceMetric::Jaccard, + MetricArg::Mash => DistanceMetric::Mash, + MetricArg::Hamming => DistanceMetric::Hamming, + MetricArg::BrayCurtis => DistanceMetric::BrayCurtis, + MetricArg::RelfreqBrayCurtis => DistanceMetric::RelfreqBrayCurtis, + MetricArg::Euclidean => DistanceMetric::Euclidean, + MetricArg::RelfreqEuclidean => DistanceMetric::RelfreqEuclidean, + MetricArg::Hellinger => DistanceMetric::Hellinger, + MetricArg::HellingerEuclidean => DistanceMetric::HellingerEuclidean, + } + } +} + +/// Partial transfer of `obikmer`'s `phylo` command: the plain distance-metric +/// path only (`--metric`/NJ/UPGMA) — everything sibling-annex-based +/// (`--sibling-annex`, `--snp`, `--sankoff`, `--tnt`/`--phyg`/`--iqtree`, ...) +/// stays in `obikmer` until `obikphylo::siblings` is reconnected (see the +/// project memory on this). +#[derive(Args)] +pub struct PhyloArgs { + /// Index directory + pub index: PathBuf, + + /// Distance metric to compute + #[arg(long, value_enum, default_value = "jaccard")] + pub metric: MetricArg, + + /// Minimum count to consider a kmer present when computing Jaccard on count indexes + #[arg(long, default_value = "1")] + pub presence_threshold: u32, + + /// Also output the shared-kmer count matrix (CSV) + #[arg(long)] + pub shared_kmers: bool, + + /// Compute and write a Neighbor-Joining tree (Newick) + #[arg(long)] + pub nj: bool, + + /// Compute and write a UPGMA tree (Newick) + #[arg(long)] + pub upgma: bool, + + /// Output prefix: _dist.csv, _shared.csv, _nj.nwk, + /// _upgma.nwk. If omitted, the distance matrix is written to stdout. + #[arg(short, long)] + pub output: Option, +} diff --git a/src/obikmer2/src/cmd/phylo/mod.rs b/src/obikmer2/src/cmd/phylo/mod.rs new file mode 100644 index 00000000..2d374a78 --- /dev/null +++ b/src/obikmer2/src/cmd/phylo/mod.rs @@ -0,0 +1,134 @@ +mod args; + +use std::io::{self, BufWriter, Write}; +use std::sync::Arc; + +use obikidxcache::index_cache::IndexCache; +use obikindex::KmerIndex; +use obikphylo::{Metrics, neighbor_joining_newick, upgma_newick}; +use obisys::{Reporter, Stage}; +use tracing::info; + +pub use args::PhyloArgs; + +pub fn run(args: PhyloArgs) { + let idx = Arc::new(KmerIndex::open(&args.index).unwrap_or_else(|e| { + eprintln!("error opening index: {e}"); + std::process::exit(1); + })); + + let labels: Vec = idx + .meta() + .genomes() + .unwrap_or_else(|e| { + eprintln!("error reading index metadata: {e}"); + std::process::exit(1); + }) + .iter() + .map(|g| g.label.clone()) + .collect(); + let n = labels.len(); + + let mut rep = Reporter::new(); + + info!("computing {:?} distances for {} genome(s)", args.metric, n); + + // Every partition/layer this needs is opened once, up front, and + // handed to `Metrics::distance` — see `obikquery`'s own use of + // `IndexCache` for the same reasoning (one open, many in-memory reads). + let cache = IndexCache::new(Arc::clone(&idx), None); + + let need_shared = args.shared_kmers || args.nj || args.upgma; + let t = Stage::start("distance"); + let result = cache + .distance(args.metric.into(), need_shared, args.presence_threshold) + .unwrap_or_else(|e| { + eprintln!("error computing distances: {e}"); + std::process::exit(1); + }); + rep.push(t.stop()); + + // ── Distance matrix → CSV ───────────────────────────────────────────────── + let write_dist_csv = |w: &mut dyn Write| { + write!(w, "genome").unwrap(); + for g in &labels { write!(w, ",{g}").unwrap(); } + writeln!(w).unwrap(); + for (i, g) in labels.iter().enumerate() { + write!(w, "{g}").unwrap(); + for j in 0..n { + write!(w, ",{:.6}", result.matrix[[i, j]]).unwrap(); + } + writeln!(w).unwrap(); + } + }; + + match &args.output { + Some(prefix) => { + let path = format!("{}_dist.csv", prefix.display()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + write_dist_csv(&mut f); + info!("distance matrix → {path}"); + } + None => { + let stdout = io::stdout(); + let mut out = BufWriter::new(stdout.lock()); + write_dist_csv(&mut out); + } + } + + // ── Shared-kmer matrix → CSV ────────────────────────────────────────────── + if args.shared_kmers { + if let Some(shared) = &result.shared_kmers { + let path = args.output.as_ref() + .map(|p| format!("{}_shared.csv", p.display())) + .unwrap_or_else(|| "shared.csv".into()); + let mut f = BufWriter::new(std::fs::File::create(&path).unwrap_or_else(|e| { + eprintln!("error creating {path}: {e}"); + std::process::exit(1); + })); + write!(f, "genome").unwrap(); + for g in &labels { write!(f, ",{g}").unwrap(); } + writeln!(f).unwrap(); + for (i, g) in labels.iter().enumerate() { + write!(f, "{g}").unwrap(); + for j in 0..n { write!(f, ",{}", shared[[i, j]]).unwrap(); } + writeln!(f).unwrap(); + } + info!("shared-kmer matrix → {path}"); + } + } + + // ── NJ tree ──────────────────────────────────────────────────────────────── + if args.nj { + let newick = neighbor_joining_newick(&result.matrix, &labels).unwrap_or_else(|e| { + eprintln!("error computing NJ tree: {e}"); + std::process::exit(1); + }); + let path = args.output.as_ref() + .map(|p| format!("{}_nj.nwk", p.display())) + .unwrap_or_else(|| "nj.nwk".into()); + std::fs::write(&path, &newick).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + info!("NJ tree → {path}"); + } + + // ── UPGMA tree ─────────────────────────────────────────────────────────────── + if args.upgma { + let newick = upgma_newick(&result.matrix, &labels); + let path = args.output.as_ref() + .map(|p| format!("{}_upgma.nwk", p.display())) + .unwrap_or_else(|| "upgma.nwk".into()); + std::fs::write(&path, &newick).unwrap_or_else(|e| { + eprintln!("error writing {path}: {e}"); + std::process::exit(1); + }); + info!("UPGMA tree → {path}"); + } + + rep.print(); +} diff --git a/src/obikmer2/src/main.rs b/src/obikmer2/src/main.rs index f7680a5e..5a587d82 100644 --- a/src/obikmer2/src/main.rs +++ b/src/obikmer2/src/main.rs @@ -39,6 +39,9 @@ enum Commands { Utils(cmd::utils::UtilsArgs), /// Convert an index's evidence representation (exact/approximate/hybrid), in place Convert(cmd::convert::ConvertArgs), + /// Genome-vs-genome distance matrix (+ optional NJ/UPGMA tree) — partial transfer, + /// sibling-annex-based operations are not yet ported (see obikmer's own `phylo`) + Phylo(cmd::phylo::PhyloArgs), } fn main() { @@ -64,5 +67,6 @@ fn main() { Commands::Annotate(args) => cmd::annotate::run(args), Commands::Utils(args) => cmd::utils::run(args), Commands::Convert(args) => cmd::convert::run(args), + Commands::Phylo(args) => cmd::phylo::run(args), } } diff --git a/src/obikphylo/Cargo.toml b/src/obikphylo/Cargo.toml index ee7a5593..24861281 100644 --- a/src/obikphylo/Cargo.toml +++ b/src/obikphylo/Cargo.toml @@ -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] diff --git a/src/obikphylo/src/distance.rs b/src/obikphylo/src/distance.rs index 00e0f3a2..511765d2 100644 --- a/src/obikphylo/src/distance.rs +++ b/src/obikphylo/src/distance.rs @@ -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 { - let n_genomes = self.meta.genomes().map_err(OKIError::Io)?.len(); +fn sum_array1(n: usize, items: impl Iterator>) -> Array1 { + items.fold(Array1::zeros(n), |acc, x| acc + x) +} +fn sum_array2_u64(n: usize, items: impl Iterator>) -> Array2 { + items.fold(Array2::zeros((n, n)), |acc, x| acc + x) +} +fn sum_array2_f64(n: usize, items: impl Iterator>) -> Array2 { + 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 { + sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights())) + } +} + +impl CountPartials for AggregatedCount<'_> { + fn partial_bray(&self) -> Array2 { + sum_array2_u64(self.n_genomes, self.layers.iter().map(|l| l.partial_bray())) + } + fn partial_euclidean(&self) -> Array2 { + sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_euclidean())) + } + fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2, Array2) { + 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) -> Array2 { + sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_bray(global))) + } + fn partial_relfreq_euclidean(&self, global: &Array1) -> Array2 { + sum_array2_f64(self.n_genomes, self.layers.iter().map(|l| l.partial_relfreq_euclidean(global))) + } + fn partial_hellinger(&self, global: &Array1) -> Array2 { + 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 { + sum_array1(self.n_genomes, self.layers.iter().map(|l| l.col_weights())) + } +} + +impl BitPartials for AggregatedBit<'_> { + fn partial_jaccard(&self) -> (Array2, Array2) { + 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 { + 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; +} + +impl Metrics for IndexCache { + fn distance(&self, metric: DistanceMetric, shared_kmers: bool, presence_threshold: u32) -> OKIResult { + 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::>()?; - 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::>()?; - 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 diff --git a/src/obikphylo/src/lib.rs b/src/obikphylo/src/lib.rs index d45e5a4b..2d2e5cce 100644 --- a/src/obikphylo/src/lib.rs +++ b/src/obikphylo/src/lib.rs @@ -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}; diff --git a/src/obikphylo/src/matrix_store.rs b/src/obikphylo/src/matrix_store.rs deleted file mode 100644 index 67054479..00000000 --- a/src/obikphylo/src/matrix_store.rs +++ /dev/null @@ -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> { - 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::>>()?; - 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> { - 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::>>()?; - Ok(LayeredStore::new(matrices)) - } -} diff --git a/src/obikphylo/src/tree/mod.rs b/src/obikphylo/src/tree/mod.rs new file mode 100644 index 00000000..4115c5fe --- /dev/null +++ b/src/obikphylo/src/tree/mod.rs @@ -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, + branch_length: Option, + bootstrap: Option, + children: Vec, +} + +/// 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, + root: Option, +} + +impl Tree { + fn new() -> Self { + Self::default() + } + + /// Create a new, detached leaf node. + fn new_leaf(&mut self, name: impl Into) -> 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] + } +} diff --git a/src/obikphylo/src/tree/newick.rs b/src/obikphylo/src/tree/newick.rs new file mode 100644 index 00000000..12e51f0f --- /dev/null +++ b/src/obikphylo/src/tree/newick.rs @@ -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 +} diff --git a/src/obikphylo/src/tree/nj.rs b/src/obikphylo/src/tree/nj.rs new file mode 100644 index 00000000..8655afef --- /dev/null +++ b/src/obikphylo/src/tree/nj.rs @@ -0,0 +1,55 @@ +//! Neighbor-Joining, via `speedytree` — converts its own tree type +//! (`petgraph::graph::UnGraph`) 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, labels: &[String]) -> OKIResult { + let n = labels.len(); + let rows: Vec> = (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::::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 +} diff --git a/src/obikphylo/src/tree/upgma.rs b/src/obikphylo/src/tree/upgma.rs new file mode 100644 index 00000000..d60bf27f --- /dev/null +++ b/src/obikphylo/src/tree/upgma.rs @@ -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, labels: &[String]) -> String { + let n = labels.len(); + let mut condensed: Vec = 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 = + 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() +}