cli/commands/normalize: Add trimmed mean of M-values option

Author: zaeleus

atlas-cli/src/cli/normalize.rs

@@ -9,6 +9,8 @@ pub enum Method {
     Fpkm,
     /// Median of ratios.
     MedianOfRatios,
+    /// Trimmed mean of M-values (TMM).
+    Tmm,
     /// Transcripts per million (TPM) mapped reads
     Tpm,
 }

atlas-cli/src/commands/normalize.rs

@@ -6,7 +6,7 @@ use std::{
 };
 
 use atlas_core::{
-    counts::normalization::{fpkm, median_of_ratios, tpm},
+    counts::normalization::{fpkm, median_of_ratios, tmm, tpm},
     features::{self, Feature, ReadFeaturesError},
     StrandSpecification,
 };
@@ -73,6 +73,7 @@ pub fn normalize(args: normalize::Args) -> Result<(), NormalizeError> {
         Method::MedianOfRatios => {
             median_of_ratios::normalize_vec(sample_count, names.len(), counts)?
         }
+        Method::Tmm => tmm::normalize_vec(sample_count, names.len(), counts)?,
         Method::Tpm => {
             let lengths = calculate_feature_lengths(&features, &names)?;
 

atlas-core/src/counts/normalization/tmm.rs

@@ -8,11 +8,27 @@
 //! [edgeR]: https://bioconductor.org/packages/release/bioc/html/edgeR.html
 //! [10.1186/gb-2010-11-3-r25]: https://doi.org/10.1186/gb-2010-11-3-r25
 
-use std::collections::HashSet;
+use std::{collections::HashSet, io};
 
-use ndarray::{Array2, ArrayView1, ArrayView2, ArrayViewMut2};
+use ndarray::{Array2, ArrayView1, ArrayView2, ArrayViewMut2, Axis};
 use tracing::info;
 
+pub fn normalize_vec(
+    sample_count: usize,
+    feature_count: usize,
+    data: Vec<u32>,
+) -> io::Result<Vec<Vec<f64>>> {
+    let matrix = Array2::from_shape_vec((sample_count, feature_count), data)
+        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
+
+    let normalized_matrix = normalize(matrix);
+
+    Ok(normalized_matrix
+        .axis_iter(Axis(0))
+        .map(|row| row.to_vec())
+        .collect())
+}
+
 pub fn normalize(data: Array2<u32>) -> Array2<f64> {
     let mut data = data.mapv(f64::from);
     let mut normalized_data = data.clone();