pub mod matrix_io;
pub mod run_log;
use anyhow::{Context, Result, anyhow};
use fixedbitset::FixedBitSet;
use log::{debug, info};
use ndarray::Array2;
use std::{
fs,
path::{Path, PathBuf},
time::Instant,
};
use crate::algorithms::equal_angle::{EqualAngleOpts, equal_angle_apply};
use crate::cli::NeighbourNetArgs;
use crate::data::splits_blocks::SplitsBlock;
use crate::neighbour_net::matrix_io::ParseMeta;
use crate::neighbour_net::run_log::{MatrixMeta, NNLSMeta, RunLog, RunTimings, system_stats};
use crate::nexus::Nexus;
use crate::nexus::nexus_writer::{NexusProperties, write_nexus_all_to_path};
use crate::ordering::OrderingMethod;
use crate::ordering::ordering_huson2023::compute_order_huson_2023;
use crate::ordering::ordering_multiway::compute_order_multiway;
use crate::phylo::phylo_splits_graph::PhyloSplitsGraph;
use crate::splits::asplit::ASplit;
use crate::utils::compute_least_squares_fit;
use crate::weights::InferenceMethod;
use crate::weights::active_set_weights::{NNLSParams, compute_asplits};
use crate::weights::cg_weights::{
CGSolveStats, DEFAULT_CUTOFF, compute_splits as compute_cg_splits,
};
struct PipelineResult {
params: NNLSParams,
splits_block: SplitsBlock,
graph: PhyloSplitsGraph,
fit: f32,
cg_stats: Option<CGSolveStats>,
cycle_sec: f64,
nnls_sec: f64,
fit_sec: f64,
splits_sec: f64,
graph_sec: f64,
}
pub struct NeighbourNet {
out_dir: String,
args: NeighbourNetArgs,
distance_matrix: Array2<f64>,
labels: Vec<String>,
parse_meta: ParseMeta,
}
impl NeighbourNet {
pub fn new(out_dir: String, args: NeighbourNetArgs) -> Result<Self> {
let (distance_matrix, labels, parse_meta) =
matrix_io::load_distance_matrix(&args.input).context("loading distance matrix")?;
Ok(NeighbourNet {
out_dir,
args,
distance_matrix,
labels,
parse_meta,
})
}
pub fn from_distance_matrix(
mut dist: Array2<f64>,
labels: Vec<String>,
args: NeighbourNetArgs,
) -> Result<Self> {
let n = dist.nrows();
if n != dist.ncols() {
anyhow::bail!("Distance matrix must be square ({}x{})", n, dist.ncols());
}
if labels.len() != n {
anyhow::bail!(
"Label count ({}) does not match matrix size ({})",
labels.len(),
n
);
}
let symmetry_pairs_fixed = matrix_io::symmetrize(&mut dist, n);
Ok(NeighbourNet {
out_dir: "output".to_string(),
args,
distance_matrix: dist,
labels,
parse_meta: ParseMeta {
delimiter: ',',
has_header: false,
has_index: false,
symmetry_pairs_fixed,
load_sec: 0.0,
},
})
}
pub fn load_distance_matrix(path: &str) -> Result<(Array2<f64>, Vec<String>, ParseMeta)> {
matrix_io::load_distance_matrix(path)
}
pub fn run(&self) -> Result<()> {
let t0 = Instant::now();
let pipeline = self.compute_pipeline()?;
let t_out = Instant::now();
self.output_results(
pipeline
.splits_block
.cycle()
.ok_or_else(|| anyhow!("cycle should be set after pipeline"))?,
&self.labels,
&pipeline.splits_block,
&self.distance_matrix,
&pipeline.graph,
pipeline.fit,
)
.context("writing outputs")?;
let out_sec = t_out.elapsed().as_secs_f64();
info!("Wrote outputs in {:.3}s", out_sec);
let run_log_path = Path::new(&self.out_dir).join("run_log.json");
let stats = self.build_run_stats(
&self.parse_meta,
&self.labels,
&pipeline.splits_block,
&pipeline.params,
pipeline.fit,
pipeline.cg_stats,
RunTimings {
load_sec: self.parse_meta.load_sec,
cycle_sec: pipeline.cycle_sec,
nnls_sec: pipeline.nnls_sec,
fit_sec: pipeline.fit_sec,
splits_sec: pipeline.splits_sec,
graph_sec: pipeline.graph_sec,
output_sec: out_sec,
total_sec: t0.elapsed().as_secs_f64(),
},
);
fs::write(&run_log_path, serde_json::to_string_pretty(&stats)?)?;
info!("Run log written: {}", run_log_path.display());
info!("Done in {:.3}s total.", t0.elapsed().as_secs_f64());
Ok(())
}
pub fn into_nexus(self) -> Result<Nexus> {
let pipeline = self.compute_pipeline()?;
Ok(Nexus::new(
self.labels,
self.distance_matrix,
pipeline.splits_block,
pipeline.graph,
))
}
pub fn generate_nexus(&self) -> Result<Nexus> {
let pipeline = self.compute_pipeline()?;
Ok(Nexus::new(
self.labels.clone(),
self.distance_matrix.clone(),
pipeline.splits_block,
pipeline.graph,
))
}
fn compute_pipeline(&self) -> Result<PipelineResult> {
let n = self.distance_matrix.nrows();
info!("Loaded distance matrix: {}x{}", n, n);
let t_cycle = Instant::now();
let cycle = self.get_ordering().context("computing cycle")?;
let cycle_sec = t_cycle.elapsed().as_secs_f64();
info!("Computed cycle in {:.3}s", cycle_sec);
debug!("Cycle (1-based): {:?}", &cycle[1..]);
let t_nnls = Instant::now();
let (params, splits, cg_stats) = self.compute_asplits(&cycle).context("ASplits solved")?;
let nnls_sec = t_nnls.elapsed().as_secs_f64();
info!(
"Estimated {} splits (cutoff = {}) in {:.3}s",
splits.len(),
self.effective_cutoff(),
nnls_sec
);
let t_fit = Instant::now();
let fit = compute_least_squares_fit(&self.distance_matrix, &splits);
let fit_sec = t_fit.elapsed().as_secs_f64();
info!(
"Least-squares fit: {:.4} % (computed in {:.3}s)",
fit, fit_sec
);
let t_spl = Instant::now();
let splits_block = self.create_splits_block(splits, fit, cycle)?;
let splits_sec = t_spl.elapsed().as_secs_f64();
info!(
"Created splits block with {} splits in {:.3}s",
splits_block.nsplits(),
splits_sec
);
let t_graph = Instant::now();
let graph = self.create_graph(&splits_block)?;
let graph_sec = t_graph.elapsed().as_secs_f64();
info!("Created phylogenetic splits graph in {:.3}s", graph_sec);
Ok(PipelineResult {
params,
splits_block,
graph,
fit,
cg_stats,
cycle_sec,
nnls_sec,
fit_sec,
splits_sec,
graph_sec,
})
}
pub fn get_ordering(&self) -> Result<Vec<usize>> {
let mut cycle = match self.args.ordering {
OrderingMethod::ClosestPair => {
compute_order_huson_2023(&self.distance_matrix).context("computing cycle")?
}
OrderingMethod::Multiway => {
compute_order_multiway(&self.distance_matrix).context("computing cycle")?
}
};
if cycle.first().copied() != Some(0) {
cycle = std::iter::once(0usize)
.chain(cycle.into_iter().map(|i| i + 1))
.collect();
}
Ok(cycle)
}
pub fn compute_asplits(
&self,
cycle: &[usize],
) -> Result<(NNLSParams, Vec<ASplit>, Option<CGSolveStats>)> {
match self.args.inference {
InferenceMethod::ActiveSet => {
let mut params = self.args.nnls_params.clone();
let splits = compute_asplits(cycle, &self.distance_matrix, &mut params, None)
.context("ASplits solved")?;
Ok((params, splits, None))
}
InferenceMethod::CG => {
let (splits, solve_stats) =
compute_cg_splits(cycle, &self.distance_matrix).context("CG weights solved")?;
let mut params = self.args.nnls_params.clone();
params.cutoff = DEFAULT_CUTOFF;
Ok((params, splits, Some(solve_stats)))
}
}
}
fn effective_cutoff(&self) -> f64 {
match self.args.inference {
InferenceMethod::ActiveSet => self.args.nnls_params.cutoff,
InferenceMethod::CG => DEFAULT_CUTOFF,
}
}
pub fn create_splits_block(
&self,
splits: Vec<ASplit>,
fit: f32,
cycle: Vec<usize>,
) -> Result<SplitsBlock> {
let mut splits_blocks = SplitsBlock::new();
splits_blocks.set_splits(splits);
splits_blocks.set_fit(fit);
splits_blocks.set_cycle(cycle, true)?;
Ok(splits_blocks)
}
pub fn create_graph(&self, splits_blocks: &SplitsBlock) -> Result<PhyloSplitsGraph> {
let mut graph = PhyloSplitsGraph::new();
let mut used_splits = FixedBitSet::with_capacity(splits_blocks.nsplits() + 1);
equal_angle_apply(
EqualAngleOpts::default(),
&self.labels,
splits_blocks,
&mut graph,
None,
&mut used_splits,
)?;
graph.create_node_ids();
Ok(graph)
}
fn output_results(
&self,
cycle: &[usize],
labels: &[String],
splits: &SplitsBlock,
distances: &Array2<f64>,
graph: &PhyloSplitsGraph,
fit: f32,
) -> Result<()> {
let out_dir = self.out_dir.clone();
fs::create_dir_all(&out_dir).with_context(|| format!("creating {}", out_dir))?;
let meta_path = Path::new(&out_dir).join("meta.json");
let meta = serde_json::json!({
"fit_percent": fit,
"ntax": labels.len(),
"cutoff": self.effective_cutoff(),
"input": self.args.input,
"inference": self.args.inference.as_str(),
});
fs::write(&meta_path, serde_json::to_string_pretty(&meta)?)?;
let props = NexusProperties {
splits_properties: Some("leastsquares cyclic".to_string()),
total_angle_deg: Some(360.0),
network_root_split: None,
distances_triangle_both: true,
};
let nexus_path = self.resolve_nexus_path();
if let Some(parent) = nexus_path.parent()
&& !parent.as_os_str().is_empty()
{
fs::create_dir_all(parent).with_context(|| format!("creating {}", parent.display()))?;
}
write_nexus_all_to_path(
&nexus_path,
labels,
Some(distances),
Some(splits),
Some(cycle),
Some(graph),
props,
)?;
info!("Outputs:");
info!(" {}", meta_path.display());
info!(" {}", nexus_path.display());
Ok(())
}
fn resolve_nexus_path(&self) -> PathBuf {
let prefix_path = Path::new(&self.args.output_prefix);
let has_explicit_path = prefix_path.is_absolute()
|| prefix_path
.parent()
.is_some_and(|p| !p.as_os_str().is_empty());
let mut nexus_path = if has_explicit_path {
prefix_path.to_path_buf()
} else {
Path::new(&self.out_dir).join(prefix_path)
};
if nexus_path.extension().is_none() {
nexus_path.set_extension("nex");
}
nexus_path
}
#[allow(clippy::too_many_arguments)] fn build_run_stats(
&self,
parse: &ParseMeta,
labels: &[String],
splits: &SplitsBlock,
nnls: &NNLSParams,
fit_percent: f32,
cg_solve_stats: Option<CGSolveStats>,
timings: RunTimings,
) -> RunLog {
let n = labels.len();
let npairs = n * (n - 1) / 2;
let (num_trivial, num_nontrivial, sum_weights) =
splits
.splits()
.fold((0usize, 0usize, 0.0f64), |mut acc, s| {
if s.is_trivial() {
acc.0 += 1;
} else {
acc.1 += 1;
}
acc.2 += s.weight;
acc
});
RunLog {
input: self.args.input.clone(),
out_dir: self.out_dir.clone(),
matrix: MatrixMeta {
n,
npairs,
delimiter: parse.delimiter.to_string(),
has_header: parse.has_header,
has_index: parse.has_index,
symmetry_pairs_fixed: parse.symmetry_pairs_fixed,
},
nnls: NNLSMeta {
cutoff: self.effective_cutoff(),
proj_grad_bound: nnls.proj_grad_bound,
cgnr_iterations: nnls.cgnr_iterations,
cgnr_tolerance: nnls.cgnr_tolerance,
active_set_rho: nnls.active_set_rho,
num_splits: splits.nsplits(),
num_trivial_splits: num_trivial,
num_nontrivial_splits: num_nontrivial,
sum_weights,
},
fit_percent,
cg_solve_stats,
timings,
system: system_stats(),
inference: self.args.inference.as_str().to_string(),
}
}
}