Push zunrplorkwkt #70

Merged
coissac merged 93 commits from push-zunrplorkwkt into main 2026-08-28 23:15:38 +00:00
17 changed files with 696 additions and 97 deletions
Showing only changes of commit 635fc830d1 - Show all commits
+4
View File
@@ -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",
+78 -1
View File
@@ -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<u64> {
self.col_weights()
}
}
impl CountPartials for KmerLayer {
fn partial_bray(&self) -> Array2<u64> {
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<f64> {
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<u64>, Array2<u64>) {
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<u64>) -> Array2<f64> {
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<u64>) -> Array2<f64> {
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<u64>) -> Array2<f64> {
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<u64>, Array2<u64>) {
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<u64> {
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"),
}
}
}
+5 -5
View File
@@ -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::<PersistentIntMatrix>::build(dir.pat
h(), DEFAUL
T_BLOCK_BITS, &Index
Mode::Exact, |_| 1)
,
TypedLayer::<PersistentIntMatrix>::build(
dir.path(),
DEFAULT_BLOCK_BITS,
&IndexMode::Exact,
|_| 1,
).unwrap();
let layer = TypedLayer::<PersistentIntMatrix>::open(dir.path()).unwrap();
+36 -2
View File
@@ -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<PersistentIntMatrix> {
) -> Box<dyn Iterator<Item = (usize, usize, u32)> + '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<u64> {
CountPartials::partial_bray(&self.data)
}
pub fn partial_euclidean(&self) -> Array2<f64> {
CountPartials::partial_euclidean(&self.data)
}
pub fn partial_threshold_jaccard(&self, threshold: u32) -> (Array2<u64>, Array2<u64>) {
CountPartials::partial_threshold_jaccard(&self.data, threshold)
}
pub fn partial_relfreq_bray(&self, global: &Array1<u64>) -> Array2<f64> {
CountPartials::partial_relfreq_bray(&self.data, global)
}
pub fn partial_relfreq_euclidean(&self, global: &Array1<u64>) -> Array2<f64> {
CountPartials::partial_relfreq_euclidean(&self.data, global)
}
pub fn partial_hellinger(&self, global: &Array1<u64>) -> Array2<f64> {
CountPartials::partial_hellinger(&self.data, global)
}
}
// ── Mode 3 — presence/absence matrix ─────────────────────────────────────────
@@ -487,6 +512,15 @@ impl TypedLayer<PersistentBitMatrix> {
self.data.nonzero_iter(slots)
}
/// `BitPartials`' required primitives, delegated straight to the
/// underlying matrix.
pub fn partial_jaccard(&self) -> (Array2<u64>, Array2<u64>) {
BitPartials::partial_jaccard(&self.data)
}
pub fn partial_hamming(&self) -> Array2<u64> {
BitPartials::partial_hamming(&self.data)
}
pub fn append_genome_column(
layer_dir: &Path,
value_of: impl Fn(usize) -> bool,
+1
View File
@@ -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" }
+1
View File
@@ -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;
+72
View File
@@ -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<MetricArg> 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: <prefix>_dist.csv, <prefix>_shared.csv, <prefix>_nj.nwk,
/// <prefix>_upgma.nwk. If omitted, the distance matrix is written to stdout.
#[arg(short, long)]
pub output: Option<PathBuf>,
}
+134
View File
@@ -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<String> = 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();
}
+4
View File
@@ -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),
}
}
+3
View File
@@ -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
View File
@@ -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
+6 -5
View File
@@ -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};
-47
View File
@@ -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))
}
}
+88
View File
@@ -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]
}
}
+34
View File
@@ -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
}
+55
View File
@@ -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
}
+51
View File
@@ -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()
}