added new top n genes function

Author: Ian

src/sparse/csc.rs

@@ -7,6 +7,7 @@ use std::iter::Sum;
 use std::ops::Add;
 use std::ops::AddAssign;
 
+use crate::sparse::MatrixNTop;
 use crate::utils::Normalize;
 
 use super::{
@@ -1027,6 +1028,41 @@ impl<M: NumericOps + NumCast> BatchMatrixMean for CscMatrix<M> {
     }
 }
 
+impl<M: NumericOps + NumCast> MatrixNTop for CscMatrix<M> {
+    type Item = M;
+
+    fn sum_row_n_top<T>(&self, n: usize) -> anyhow::Result<Vec<T>>
+    where
+        T: Float + NumCast + AddAssign + Sum {
+        let mut result = vec![T::zero(); self.nrows()];
+        
+        let mut row_values: Vec<Vec<T>> = vec![Vec::new(); self.nrows()];
+        
+        for col_idx in 0..self.ncols() {
+            let col_start = self.col_offsets()[col_idx];
+            let col_end = self.col_offsets()[col_idx + 1];
+            
+            for idx in col_start..col_end {
+                let row_idx = self.row_indices()[idx];
+                if let Some(val) = T::from(self.values()[idx]) {
+                    row_values[row_idx].push(val);
+                }
+            }
+        }
+        
+        for (row_idx, mut values) in row_values.into_iter().enumerate() {
+            if values.len() <= n {
+                result[row_idx] = values.into_iter().sum();
+            } else {
+                values.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
+                result[row_idx] = values.into_iter().take(n).sum();
+            }
+        }
+        
+        Ok(result)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use Direction;

src/sparse/csr.rs

@@ -6,6 +6,7 @@ use std::ops::{Add, AddAssign};
 use super::{
     BatchMatrixMean, BatchMatrixVariance, MatrixMinMax, MatrixNonZero, MatrixSum, MatrixVariance,
 };
+use crate::sparse::MatrixNTop;
 use crate::utils::Normalize;
 use crate::utils::{BatchIdentifier, Log1P};
 use anyhow::{anyhow, Ok};
@@ -1033,6 +1034,36 @@ impl<M: NumericOps + NumCast> BatchMatrixMean for CsrMatrix<M> {
     }
 }
 
+impl<M: NumericOps + NumCast> MatrixNTop for CsrMatrix<M> {
+    type Item = M;
+
+    fn sum_row_n_top<T>(&self, n: usize) -> anyhow::Result<Vec<T>>
+    where
+        T: Float + NumCast + AddAssign + Sum {
+        let mut result = vec![T::zero(); self.nrows()];
+
+        for row_idx in 0..self.nrows() {
+            let row_start = self.row_offsets()[row_idx];
+            let row_end = self.row_offsets()[row_idx + 1];
+
+            let mut row_values: Vec<T> = Vec::new();
+            for idx in row_start..row_end {
+                if let Some(val) = T::from(self.values()[idx]) {
+                    row_values.push(val);
+                }
+            }
+
+            if row_values.len() <= n {
+                result[row_idx] = row_values.into_iter().sum();
+            } else {
+                row_values.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal));
+                result[row_idx] = row_values.into_iter().take(n).sum();
+            }
+        }
+        Ok(result)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use Direction;

src/sparse/mod.rs

@@ -164,4 +164,12 @@ pub trait BatchMatrixMean {
         B: BatchIdentifier;
 }
 
+pub trait MatrixNTop {
+    type Item: NumCast;
+
+    fn sum_row_n_top<T>(&self, n: usize) -> anyhow::Result<Vec<T>>
+    where
+        T: Float + NumCast + AddAssign + std::iter::Sum;
+}
+