bixverse_methods/methods/nmf_hals/
nmf_preprocessing.rs

1//! Pre-processing methods for NMF on sparse and dense data.
2
3use faer::{Mat, MatRef};
4use rayon::prelude::*;
5
6use crate::prelude::*;
7use bixverse_core::core::math::matrix_helpers::col_sds;
8use bixverse_core::core::math::sparse::sparse_col_sds_csc;
9
10//////////
11// Enum //
12//////////
13
14use crate::prelude::BixverseFloat;
15
16/// Options to do NmfPreprocessing
17#[derive(Clone, Copy, Debug, Default)]
18pub enum NmfPreprocessing {
19    /// No pre-processing.
20    #[default]
21    None,
22    /// Scale each column (features) by the standard deviation.
23    SdScaling,
24    /// Scale each column (features) by the squared standard deviation.
25    SqrtSdScaling,
26}
27
28/// Transpose sparse counts to CSC if they arrived as CSR.
29///
30/// [`nmf_process_sparse`] indexes `indptr` per column, so CSR input scales the
31/// wrong axis and returns quiet nonsense rather than an error. The R side always
32/// hands over CSR, so this is the common path, not the exception.
33///
34/// ### Params
35///
36/// * `data` - The counts, in either orientation.
37/// * `verbosity` - Resolved verbosity, used only for the transpose notice.
38///
39/// ### Returns
40///
41/// The same data in CSC orientation.
42pub fn ensure_csc(
43    data: CompressedSparseData2<f32>,
44    verbosity: Verbosity,
45) -> CompressedSparseData2<f32> {
46    if data.cs_type.is_csr() {
47        if verbosity.detailed_verbosity() {
48            println!("NMF: data was provided as CSR. Transposing to CSC.")
49        }
50        data.transform()
51    } else {
52        data
53    }
54}
55
56/// Parse the NMF pre-processing
57///
58/// ### Params
59///
60/// * `s` - String to parse. `"none"` -> leave data as is; `"sd" | "scaling"`
61///   -> divide the respective features by their standard deviation; `"sqrt_sd"`
62///   -> divide the respective features by the square root of their standard
63///   deviation.
64///
65/// ### Returns
66///
67/// The option of [NmfPreprocessing]
68pub fn parse_nmf_processing(s: &str) -> Option<NmfPreprocessing> {
69    match s.to_lowercase().as_str() {
70        "none" => Some(NmfPreprocessing::None),
71        "sd" | "scaling" => Some(NmfPreprocessing::SdScaling),
72        "sqrt_sd" => Some(NmfPreprocessing::SqrtSdScaling),
73        _ => None,
74    }
75}
76
77///////////
78// Dense //
79///////////
80
81/// Process dense matrices of samples x features
82///
83/// ### Params
84///
85/// * `x` - Matrix to preprocess
86/// * `process` - Which approach of NMF pre-processing to apply, see
87///   [NmfPreprocessing]
88///
89/// ### Returns
90///
91/// The pre-processed matrix
92pub fn nmf_process_dense<T: BixverseFloat>(x: MatRef<T>, process: &NmfPreprocessing) -> Mat<T> {
93    let epsilon = T::from_f32(1e-6).unwrap();
94
95    match process {
96        NmfPreprocessing::None => x.to_owned(),
97        NmfPreprocessing::SdScaling => {
98            let mut sds = col_sds(x);
99            // safe guard against tiny values...
100            sds.par_iter_mut()
101                .for_each(|x| *x = if *x <= epsilon { T::one() } else { *x });
102
103            Mat::from_fn(x.nrows(), x.ncols(), |i, j| x[(i, j)] / sds[j])
104        }
105        NmfPreprocessing::SqrtSdScaling => {
106            let mut sds = col_sds(x);
107            sds.par_iter_mut()
108                .for_each(|x| *x = if *x <= epsilon { T::one() } else { x.sqrt() });
109
110            Mat::from_fn(x.nrows(), x.ncols(), |i, j| x[(i, j)] / sds[j])
111        }
112    }
113}
114
115////////////
116// Sparse //
117////////////
118
119/// Process sparse matrices of samples x features (need to be in CSC format!)
120///
121/// ### Params
122///
123/// * `x` - Sparse matrix to preprocess, see [CompressedSparseData2].
124/// * `process` - Which approach of NMF pre-processing to apply, see
125///   [NmfPreprocessing].
126/// * `use_second_layer` - Use the second data layer.
127///
128/// ### Returns
129///
130/// The pre-processed matrix
131pub fn nmf_process_sparse<T>(
132    x: &CompressedSparseData2<T>,
133    process: &NmfPreprocessing,
134    use_second_layer: bool,
135) -> Result<CompressedSparseData2<T>, MethodsErrors>
136where
137    T: BixverseFloat + BixverseSimd,
138{
139    let epsilon = T::from_f32(1e-6).unwrap();
140
141    match process {
142        NmfPreprocessing::None => Ok(x.clone()),
143        NmfPreprocessing::SdScaling | NmfPreprocessing::SqrtSdScaling => {
144            let use_sqrt = matches!(process, NmfPreprocessing::SqrtSdScaling);
145            let mut sds = sparse_col_sds_csc(x, use_second_layer)?;
146
147            sds.par_iter_mut().for_each(|s| {
148                *s = if *s <= epsilon {
149                    T::one()
150                } else if use_sqrt {
151                    s.sqrt()
152                } else {
153                    *s
154                };
155            });
156
157            let mut out = x.clone();
158            let active_data: &mut Vec<T> = if use_second_layer {
159                out.data_2.as_mut().ok_or(CoreErrors::Data2NotAvailable)?
160            } else {
161                &mut out.data
162            };
163
164            let ncols = x.shape().1;
165            for j in 0..ncols {
166                let start = x.indptr[j] as usize;
167                let end = x.indptr[j + 1] as usize;
168                let divisor = sds[j];
169                for v in &mut active_data[start..end] {
170                    *v /= divisor;
171                }
172            }
173
174            Ok(out)
175        }
176    }
177}
bixverse_methods/methods/nmf_hals/nmf_preprocessing.rs

bixverse_methods/methods/nmf_hals/
nmf_preprocessing.rs
