Merge branch 'feat/timetree-scaffolding' into rust

Author: ivan-aksamentov

Cargo.toml

@@ -30,7 +30,6 @@ authors = [
 ]
 
 
-
 [patch.crates-io]
 lapack-sys = { git = "https://github.com/numrs/lapack-sys", rev = "d3b1ca9" }
 lax = { git = "https://github.com/numrs/ndarray-linalg", rev = "ac7052e" }
@@ -319,6 +318,10 @@ iter_over_hash_type = "allow"
 needless_range_loop = "allow"
 non_std_lazy_statics = "allow" # bugs clap annotations
 
+# Remove once amount of unimplemented code is reduced
+todo = "allow"
+let_underscore_untyped = "allow"
+
 [workspace.lints.rustdoc]
 broken_intra_doc_links = "warn"
 bare_urls = "warn"

packages/treetime-cli/src/bin/treetime.rs

@@ -7,7 +7,7 @@ use treetime::commands::homoplasy::run_homoplasy::run_homoplasy;
 use treetime::commands::mugration::run_mugration::run_mugration;
 use treetime::commands::optimize::run::run_optimize;
 use treetime::commands::prune::run::run_prune;
-use treetime::commands::timetree::run_timetree_estimation::run_timetree_estimation;
+use treetime::commands::timetree::run::run_timetree_estimation;
 use treetime::utils::global_init::global_init;
 use treetime::utils::openblas::get_openblas_info_str;
 use treetime_cli::cli::treetime_cli::{TreetimeCommands, generate_shell_completions, treetime_parse_cli_args};

packages/treetime-cli/src/cli/treetime_cli.rs

@@ -15,7 +15,7 @@ use treetime::commands::homoplasy::homoplasy_args::TreetimeHomoplasyArgs;
 use treetime::commands::mugration::mugration_args::TreetimeMugrationArgs;
 use treetime::commands::optimize::args::TreetimeOptimizeArgs;
 use treetime::commands::prune::args::TreetimePruneArgs;
-use treetime::commands::timetree::timetree_args::TreetimeTimetreeArgs;
+use treetime::commands::timetree::args::TreetimeTimetreeArgs;
 use treetime::utils::clap_styles::styles;
 use treetime::utils::global_init::setup_logger;
 

packages/treetime/examples/convolution.rs

@@ -141,7 +141,7 @@ fn main() -> eyre::Result<()> {
     _ => {
       eprintln!("Unknown algorithm: {}", args.algorithm);
       return Ok(());
-    }
+    },
   };
 
   // Compute analytical expected result
@@ -220,14 +220,7 @@ fn main() -> eyre::Result<()> {
       format!("{rel_err:.2}")
     };
 
-    println!(
-      "{:>8.2} {:>12.6} {:>12.6} {:>12} {:>9}%",
-      x_val,
-      actual_val,
-      expected_val,
-      diff_str,
-      rel_err_str
-    );
+    println!("{x_val:>8.2} {actual_val:>12.6} {expected_val:>12.6} {diff_str:>12} {rel_err_str:>9}%");
   }
 
   // Create and save results structure
@@ -259,7 +252,11 @@ fn compute_domain_agreement_metrics(
   Ok(metrics)
 }
 
-fn plot_input_functions(f: &treetime::distribution::reference::grid_fn::GridFn, g: &treetime::distribution::reference::grid_fn::GridFn, args: &Args) -> eyre::Result<()> {
+fn plot_input_functions(
+  f: &treetime::distribution::reference::grid_fn::GridFn,
+  g: &treetime::distribution::reference::grid_fn::GridFn,
+  args: &Args,
+) -> eyre::Result<()> {
   let output_path = format!("{}/input_functions.svg", args.output_dir);
   let root = SVGBackend::new(&output_path, (800, 600)).into_drawing_area();
   root.fill(&WHITE)?;
@@ -307,7 +304,10 @@ fn plot_convolution_results(
   let max_val = actual.iter().chain(expected.iter()).fold(0.0_f64, |a, &b| a.max(b));
 
   let mut chart = ChartBuilder::on(&root)
-    .caption(&format!("Convolution Results: (f * g)(x) [{}]", args.algorithm), ("Arial", 24))
+    .caption(
+      format!("Convolution Results: (f * g)(x) [{}]", args.algorithm),
+      ("Arial", 24),
+    )
     .margin(20)
     .x_label_area_size(40)
     .y_label_area_size(50)
@@ -318,7 +318,7 @@ fn plot_convolution_results(
   let actual_data: Vec<(f64, f64)> = x_grid.iter().zip(actual.iter()).map(|(&x, &y)| (x, y)).collect();
   chart
     .draw_series(LineSeries::new(actual_data, BLUE.stroke_width(2)))?
-    .label(&format!("Actual ({})", args.algorithm))
+    .label(format!("Actual ({})", args.algorithm))
     .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 10, y)], BLUE));
 
   let expected_data: Vec<(f64, f64)> = x_grid.iter().zip(expected.iter()).map(|(&x, &y)| (x, y)).collect();
@@ -353,7 +353,10 @@ fn plot_error_analysis(
   let max_error = absolute_errors.iter().fold(0.0_f64, |a, &b| a.max(b));
 
   let mut chart = ChartBuilder::on(&root)
-    .caption(&format!("Absolute Error: |Actual - Expected| [{}]", args.algorithm), ("Arial", 24))
+    .caption(
+      format!("Absolute Error: |Actual - Expected| [{}]", args.algorithm),
+      ("Arial", 24),
+    )
     .margin(20)
     .x_label_area_size(40)
     .y_label_area_size(50)

packages/treetime/src/commands/ancestral/marginal_unified.rs

@@ -12,7 +12,7 @@ use parking_lot::RwLock;
 use std::sync::Arc;
 
 /// Main entry point for marginal reconstruction
-pub fn run_marginal<P: PartitionMarginalOps + HasLogLh>(
+pub fn run_marginal<P: PartitionMarginalOps + HasLogLh + ?Sized>(
   graph: &GraphAncestral,
   partitions: &[Arc<RwLock<P>>],
   aln: Option<&[FastaRecord]>,
@@ -32,7 +32,7 @@ pub fn run_marginal<P: PartitionMarginalOps + HasLogLh>(
 }
 
 /// Ancestral sequence reconstruction
-pub fn ancestral_reconstruction_marginal<P: PartitionMarginalOps + HasLogLh>(
+pub fn ancestral_reconstruction_marginal<P: PartitionMarginalOps + HasLogLh + ?Sized>(
   graph: &GraphAncestral,
   include_leaves: bool,
   partitions: &[Arc<RwLock<P>>],
@@ -59,7 +59,7 @@ pub fn ancestral_reconstruction_marginal<P: PartitionMarginalOps + HasLogLh>(
 }
 
 /// Backward pass: calculates ingroup profiles
-fn marginal_backward<P: PartitionMarginalOps + HasLogLh>(
+fn marginal_backward<P: PartitionMarginalOps + HasLogLh + ?Sized>(
   graph: &GraphAncestral,
   partitions: &[Arc<RwLock<P>>],
 ) -> Result<(), Report> {
@@ -70,7 +70,7 @@ fn marginal_backward<P: PartitionMarginalOps + HasLogLh>(
   Ok(())
 }
 
-fn run_marginal_backward<P: PartitionMarginalOps + HasLogLh>(
+fn run_marginal_backward<P: PartitionMarginalOps + HasLogLh + ?Sized>(
   partitions: &[Arc<RwLock<P>>],
   node: &GraphNodeBackward<NodeAncestral, EdgeAncestral, ()>,
 ) -> Result<(), Report> {
@@ -82,7 +82,7 @@ fn run_marginal_backward<P: PartitionMarginalOps + HasLogLh>(
 }
 
 /// Forward pass: calculates outgroup profiles
-fn marginal_forward<P: PartitionMarginalOps + HasLogLh>(
+fn marginal_forward<P: PartitionMarginalOps + HasLogLh + ?Sized>(
   graph: &GraphAncestral,
   partitions: &[Arc<RwLock<P>>],
 ) -> Result<(), Report> {
@@ -93,7 +93,7 @@ fn marginal_forward<P: PartitionMarginalOps + HasLogLh>(
   Ok(())
 }
 
-fn run_marginal_forward<P: PartitionMarginalOps + HasLogLh>(
+fn run_marginal_forward<P: PartitionMarginalOps + HasLogLh + ?Sized>(
   graph: &GraphAncestral,
   partitions: &[Arc<RwLock<P>>],
   node: &GraphNodeForward<NodeAncestral, EdgeAncestral, ()>,

packages/treetime/src/commands/clock/clock_args.rs

@@ -3,7 +3,7 @@ use crate::commands::clock::clock_regression::ClockOptions;
 use crate::commands::clock::find_best_root::params::{
   BrentParams, GoldenSectionParams, GridSearchParams, OptimizationMethod,
 };
-use crate::commands::timetree::timetree_args::{BranchLengthMode, RerootMode};
+use crate::commands::timetree::args::{BranchLengthMode, RerootMode};
 use crate::gtr::get_gtr::GtrModelName;
 use clap::{Args, Parser, ValueHint};
 use smart_default::SmartDefault;

packages/treetime/src/commands/optimize/optimize_dense.rs

@@ -24,6 +24,7 @@ use crate::commands::optimize::optimize_unified::OptimizationMetrics;
 use crate::gtr::gtr::GTR;
 use crate::representation::graph_ancestral::GraphAncestral;
 use crate::representation::graph_dense::DenseSeqDis;
+use crate::representation::partition_marginal::PartitionMarginalOps;
 use crate::representation::partition_marginal_dense::PartitionMarginalDense;
 use eyre::Report;
 use ndarray::{Array2, Axis};
@@ -76,7 +77,10 @@ pub fn run_optimize_dense(
   graph: &GraphAncestral,
   partitions: &[Arc<RwLock<PartitionMarginalDense>>],
 ) -> Result<(), Report> {
-  let total_length: usize = partitions.iter().map(|part| part.read_arc().length).sum();
+  let total_length: usize = partitions
+    .iter()
+    .map(|part| part.read_arc().get_sequence_length().unwrap_or(0))
+    .sum();
   let one_mutation = 1.0 / total_length as f64;
   let n_partitions = partitions.len();
   graph.get_edges().iter().for_each(|edge_ref| {

packages/treetime/src/commands/optimize/optimize_sparse.rs

@@ -22,7 +22,9 @@
 //!
 use crate::commands::optimize::optimize_unified::OptimizationMetrics;
 use crate::graph::edge::GraphEdgeKey;
-use crate::representation::{graph_ancestral::GraphAncestral, partition_marginal_sparse::PartitionMarginalSparse};
+use crate::representation::graph_ancestral::GraphAncestral;
+use crate::representation::partition_marginal::PartitionMarginalOps;
+use crate::representation::partition_marginal_sparse::PartitionMarginalSparse;
 use crate::seq::mutation::Sub;
 use eyre::{OptionExt, Report};
 use itertools::Itertools;
@@ -141,7 +143,10 @@ pub fn run_optimize_sparse(
   graph: &GraphAncestral,
   partitions: &[Arc<RwLock<PartitionMarginalSparse>>],
 ) -> Result<(), Report> {
-  let total_length: usize = partitions.iter().map(|part| part.read_arc().length).sum();
+  let total_length: usize = partitions
+    .iter()
+    .map(|part| part.read_arc().get_sequence_length().unwrap_or(0))
+    .sum();
   let one_mutation = 1.0 / total_length as f64;
   let n_partitions = partitions.len();
   graph.get_edges().iter().try_for_each(|edge_ref| {

packages/treetime/src/commands/optimize/optimize_unified.rs

@@ -2,6 +2,7 @@ use crate::commands::optimize::optimize_dense;
 use crate::commands::optimize::optimize_sparse;
 use crate::graph::edge::GraphEdgeKey;
 use crate::representation::graph_ancestral::GraphAncestral;
+use crate::representation::partition_marginal::PartitionMarginalOps;
 use crate::representation::partition_marginal_dense::PartitionMarginalDense;
 use crate::representation::partition_marginal_sparse::PartitionMarginalSparse;
 use eyre::Report;
@@ -220,8 +221,12 @@ pub fn run_optimize_mixed(
 ) -> Result<(), Report> {
   let total_length: usize = dense_partitions
     .iter()
-    .map(|part| part.read_arc().length)
-    .chain(sparse_partitions.iter().map(|part| part.read_arc().length))
+    .map(|part| part.read_arc().get_sequence_length().unwrap_or(0))
+    .chain(
+      sparse_partitions
+        .iter()
+        .map(|part| part.read_arc().get_sequence_length().unwrap_or(0)),
+    )
     .sum();
 
   let one_mutation = 1.0 / total_length as f64;
@@ -342,8 +347,12 @@ pub fn initial_guess_mixed(
 
     let total_length: usize = dense_partitions
       .iter()
-      .map(|part| part.read_arc().length)
-      .chain(sparse_partitions.iter().map(|part| part.read_arc().length))
+      .map(|part| part.read_arc().get_sequence_length().unwrap_or(0))
+      .chain(
+        sparse_partitions
+          .iter()
+          .map(|part| part.read_arc().get_sequence_length().unwrap_or(0)),
+      )
       .sum();
 
     let branch_length = (differences as f64) / (total_length as f64);

packages/treetime/src/commands/timetree/timetree_args.rs renamed to packages/treetime/src/commands/timetree/args.rs

@@ -1,6 +1,8 @@
+use crate::alphabet::alphabet::AlphabetName;
 use crate::commands::ancestral::anc_args::MethodAncestral;
 use crate::gtr::get_gtr::GtrModelName;
 use clap::{Parser, ValueEnum, ValueHint};
+use smart_default::SmartDefault;
 use std::fmt::Debug;
 use std::path::PathBuf;
 
@@ -49,7 +51,7 @@ impl Default for RerootMode {
   }
 }
 
-#[derive(Parser, Debug)]
+#[derive(Parser, Debug, SmartDefault)]
 pub struct TreetimeTimetreeArgs {
   /// Path to one or multiple FASTA files with aligned input sequences
   ///
@@ -130,6 +132,10 @@ pub struct TreetimeTimetreeArgs {
   #[clap(long)]
   pub keep_polytomies: bool,
 
+  /// Resolve polytomies using temporal information
+  #[clap(long)]
+  pub resolve_polytomies: bool,
+
   /// use an autocorrelated molecular clock. Strength of the gaussian priors on branch specific rate
   /// deviation and the coupling of parent and offspring rates can be specified e.g. as --relax 1.0
   /// 0.5. Values around 1.0 correspond to weak priors, larger values constrain rate deviations more
@@ -139,8 +145,9 @@ pub struct TreetimeTimetreeArgs {
 
   /// maximal number of iterations the inference cycle is run. Note that for polytomy resolution and
   /// coalescence models max_iter should be at least 2
-  #[clap(long)]
-  pub max_iter: Option<usize>,
+  #[default = 2]
+  #[clap(long, default_value_t = TreetimeTimetreeArgs::default().max_iter)]
+  pub max_iter: usize,
 
   /// coalescent time scale -- sensible values are on the order of the average hamming distance of
   /// contemporaneous sequences. In addition, 'opt' 'skyline' are valid options and estimate a
@@ -183,6 +190,10 @@ pub struct TreetimeTimetreeArgs {
   #[clap(long, default_value = "3.0")]
   pub clock_filter: f64,
 
+  /// Number of IQD (interquartile distance) for clock filter outlier detection
+  #[clap(long)]
+  pub n_iqd: Option<f64>,
+
   /// Reroot the tree using root-to-tip regression. Valid choices are 'min_dev', 'least-squares',
   /// and 'oldest'. 'least-squares' adjusts the root to minimize residuals of the root-to-tip vs
   /// sampling time regression, 'min_dev' minimizes variance of root-to-tip distances. 'least-
@@ -207,6 +218,10 @@ pub struct TreetimeTimetreeArgs {
   #[clap(long)]
   pub covariation: bool,
 
+  /// Estimate timetree with rate variation to assess sensitivity to clock rate uncertainty
+  #[clap(long)]
+  pub vary_rate: bool,
+
   /// GTR model to use
   ///
   /// '--gtr infer' will infer a model from the data. Alternatively, specify the model type. If the specified model requires additional options, use '--gtr-params' to specify those.
@@ -225,6 +240,14 @@ pub struct TreetimeTimetreeArgs {
   #[clap(long, value_enum, default_value_t = MethodAncestral::default())]
   pub method_anc: MethodAncestral,
 
+  /// Alphabet to use for sequences
+  #[clap(long, value_enum, default_value_t = AlphabetName::default())]
+  pub alphabet: AlphabetName,
+
+  /// Use dense representation for sequences (store full probability distributions)
+  #[clap(long)]
+  pub dense: Option<bool>,
+
   /// Use aminoacid alphabet
   #[clap(long)]
   pub aa: bool,
@@ -249,7 +272,18 @@ pub struct TreetimeTimetreeArgs {
   #[clap(long, short = 'O')]
   pub outdir: PathBuf,
 
+  /// Write iteration statistics to tracelog CSV file for monitoring convergence
+  #[clap(long)]
+  #[clap(value_hint = ValueHint::FilePath)]
+  pub tracelog: Option<PathBuf>,
+
   /// Random seed
   #[clap(long)]
   pub seed: Option<u64>,
 }
+
+impl TreetimeTimetreeArgs {
+  pub fn clock_filter_enabled(&self) -> bool {
+    self.clock_filter > 0.0
+  }
+}