Skip to main content

bixverse_core/core/
mat_struct.rs

1//! Helper structure to store named matrices
2
3use faer::{Mat, MatRef};
4use std::collections::BTreeMap;
5
6use crate::prelude::*;
7
8//////////////
9// Matrices //
10//////////////
11
12/// Structure to store named matrices
13#[derive(Clone, Debug)]
14pub struct NamedMatrix<'a, T> {
15    /// A BTreeMap representing the column names and the col indices
16    pub col_names: BTreeMap<String, usize>,
17    /// A BTreeMap representing the row names and the row indices
18    pub row_names: BTreeMap<String, usize>,
19    /// A faer matrix reference representing the matrix values.
20    pub values: faer::MatRef<'a, T>,
21}
22
23impl<'a, T> NamedMatrix<'a, T>
24where
25    T: BixverseFloat,
26{
27    /// Return a submatrix based on the row names and columns to select.
28    ///
29    /// If no rows or columns are specified, returns the full matrix. If empty slices are
30    /// provided, returns None.
31    ///
32    /// ### Params
33    ///
34    /// * `rows_to_select` - Names of the row features
35    /// * `cols_to_select` - Names of the column features
36    ///
37    /// ### Returns
38    ///
39    /// The faer Matrix with the values from the rows and columns.
40    pub fn get_sub_mat(
41        &self,
42        rows_to_select: Option<&[&str]>,
43        cols_to_select: Option<&[&str]>,
44    ) -> Option<Mat<T>> {
45        // Determine which rows to select
46        let row_indices: Vec<usize> = match rows_to_select {
47            None => (0..self.values.nrows()).collect(), // All rows
48            Some(rows) => {
49                if rows.is_empty() {
50                    return None;
51                }
52                let mut indices = Vec::new();
53                for &row_name in rows {
54                    if let Some(&index) = self.row_names.get(row_name) {
55                        indices.push(index);
56                    }
57                }
58                if indices.is_empty() {
59                    return None;
60                }
61                indices
62            }
63        };
64
65        // Determine which columns to select
66        let col_indices: Vec<usize> = match cols_to_select {
67            None => (0..self.values.ncols()).collect(), // All columns
68            Some(cols) => {
69                if cols.is_empty() {
70                    return None;
71                }
72                let mut indices = Vec::new();
73                for &col_name in cols {
74                    if let Some(&index) = self.col_names.get(col_name) {
75                        indices.push(index);
76                    }
77                }
78                if indices.is_empty() {
79                    return None;
80                }
81                indices
82            }
83        };
84
85        // Create new matrix by copying values
86        let mut result = Mat::<T>::zeros(row_indices.len(), col_indices.len());
87        for (new_col, &old_col) in col_indices.iter().enumerate() {
88            for (new_row, &old_row) in row_indices.iter().enumerate() {
89                result[(new_row, new_col)] = self.values[(old_row, old_col)];
90            }
91        }
92
93        Some(result)
94    }
95
96    /// Convenience method to get the full matrix
97    ///
98    /// ### Returns
99    ///
100    /// Faer matrix representing the full data.
101    pub fn get_full_mat(&self) -> Mat<T> {
102        self.get_sub_mat(None, None).unwrap()
103    }
104
105    /// Convenience method to get submatrix with only row selection
106    ///
107    /// ### Params
108    ///
109    /// * `rows_to_select` - Names of the row features
110    ///
111    /// ### Returns
112    ///
113    /// Faer matrix with the selected rows.
114    pub fn get_rows(&self, rows_to_select: &[&str]) -> Option<Mat<T>> {
115        self.get_sub_mat(Some(rows_to_select), None)
116    }
117
118    /// Convenience method to get submatrix with only column selection
119    ///
120    /// ### Params
121    ///
122    /// * `cols_to_select` - Names of the column features
123    ///
124    /// ### Returns
125    ///
126    /// Faer matrix with the selected columns.
127    pub fn get_cols(&self, cols_to_select: &[&str]) -> Option<Mat<T>> {
128        self.get_sub_mat(None, Some(cols_to_select))
129    }
130
131    /// Get column names as references (for temporary use within same scope)
132    ///
133    /// ### Returns
134    ///
135    /// The column names as a reference
136    pub fn get_col_names_refs(&self) -> Vec<&String> {
137        self.col_names.keys().collect()
138    }
139
140    /// Get row names as references (for temporary use within same scope)
141    ///
142    /// ### Returns
143    ///
144    /// The row names as a reference
145    pub fn get_row_names_refs(&self) -> Vec<&String> {
146        self.row_names.keys().collect()
147    }
148
149    /// Get the row indices based on names
150    ///
151    /// ### Params
152    ///
153    /// * `rows_to_select` - Names of the row features
154    pub fn get_row_indices(&self, rows_to_select: Option<&[&str]>) -> Vec<usize> {
155        let row_indices: Vec<usize> = match rows_to_select {
156            None => (0..self.values.nrows()).collect(), // All rows
157            Some(rows) => {
158                if rows.is_empty() {
159                    return vec![];
160                }
161                let mut indices = Vec::new();
162                for &row_name in rows {
163                    if let Some(&index) = self.row_names.get(row_name) {
164                        indices.push(index);
165                    }
166                }
167                if indices.is_empty() {
168                    return vec![];
169                }
170                indices
171            }
172        };
173
174        row_indices
175    }
176
177    /// Get the col indices based on names
178    ///
179    /// ### Params
180    ///
181    /// * `rows_to_select` - Names of the row features
182    pub fn get_col_indices(&self, cols_to_select: Option<&[&str]>) -> Vec<usize> {
183        let col_indices: Vec<usize> = match cols_to_select {
184            None => (0..self.values.ncols()).collect(), // All columns
185            Some(cols) => {
186                if cols.is_empty() {
187                    return vec![];
188                }
189                let mut indices = Vec::new();
190                for &col_name in cols {
191                    if let Some(&index) = self.col_names.get(col_name) {
192                        indices.push(index);
193                    }
194                }
195                if indices.is_empty() {
196                    return vec![];
197                }
198                indices
199            }
200        };
201
202        col_indices
203    }
204}
205
206/// Matrix slice view
207///
208/// Structure to help creating sub slices of a given matrix when needed.
209/// Due to memory structure, slicing creates deep copies, but this function
210/// avoids generating them until the last possible moment.
211///
212/// ### Fields
213///
214/// * `data` - The faer MatRef (original matrix)
215/// * `row_indices` - The row indices you want to slice out.
216/// * `col_indices` - The col indices you want to slice out.
217#[derive(Clone, Debug)]
218pub struct MatSliceView<'a, 'r, 'c, T> {
219    data: MatRef<'a, T>,
220    row_indices: &'r [usize],
221    col_indices: &'c [usize],
222}
223
224impl<'a, 'r, 'c, T> MatSliceView<'a, 'r, 'c, T>
225where
226    T: BixverseFloat,
227{
228    /// Generate a new MatSliceView
229    ///
230    /// This function will panic if you try to select indices larger than the
231    /// underlying matrix.
232    ///
233    /// ### Params
234    ///
235    /// * `data` - The original MatRef from which you want to slice out data
236    /// * `row_indices` - The row indices you want to slice out.
237    /// * `col_indices` - The col indices you want to slice out.
238    pub fn new(
239        data: MatRef<'a, T>,
240        row_indices: &'r [usize],
241        col_indices: &'c [usize],
242    ) -> MatSliceView<'a, 'r, 'c, T> {
243        let max_col_index = col_indices.iter().copied().fold(0, usize::max);
244        let max_row_index = row_indices.iter().copied().fold(0, usize::max);
245
246        assert!(
247            max_col_index < data.ncols(),
248            "You selected indices larger than ncol."
249        );
250        assert!(
251            max_row_index < data.nrows(),
252            "You selected indices larger than nrow."
253        );
254
255        Self {
256            data,
257            row_indices,
258            col_indices,
259        }
260    }
261
262    /// Return the number of rows
263    ///
264    /// ### Returns
265    ///
266    /// Number of rows
267    pub fn nrows(&self) -> usize {
268        self.row_indices.len()
269    }
270
271    /// Return the number of columns
272    ///
273    /// ### Returns
274    ///
275    /// Number of columns
276    pub fn ncols(&self) -> usize {
277        self.col_indices.len()
278    }
279
280    /// Return an owned matrix.
281    ///
282    /// Deep copying cannot be circumvented due to memory accessing at this point.
283    /// Subsequent matrix algebra needs a continouos view into memory.
284    ///
285    /// ### Returns
286    ///
287    /// Owned sliced matrix for subsequent usage.
288    pub fn to_owned(&self) -> Mat<T> {
289        Mat::from_fn(self.nrows(), self.ncols(), |i, j| {
290            *self.data.get(self.row_indices[i], self.col_indices[j])
291        })
292    }
293}
294
295///////////
296// Tests //
297///////////
298
299#[cfg(test)]
300mod tests {
301    use super::*;
302    use faer::Mat;
303    use std::collections::BTreeMap;
304
305    /// A 3 by 3 matrix holding 0 to 8 in row-major order, named r1 to r3 and c1 to c3.
306    fn setup_named_matrix() -> (Mat<f64>, BTreeMap<String, usize>, BTreeMap<String, usize>) {
307        let mat: Mat<f64> = Mat::from_fn(3, 3, |i, j| (i * 3 + j) as f64);
308
309        let mut row_names = BTreeMap::new();
310        row_names.insert("r1".to_string(), 0);
311        row_names.insert("r2".to_string(), 1);
312        row_names.insert("r3".to_string(), 2);
313
314        let mut col_names = BTreeMap::new();
315        col_names.insert("c1".to_string(), 0);
316        col_names.insert("c2".to_string(), 1);
317        col_names.insert("c3".to_string(), 2);
318
319        (mat, row_names, col_names)
320    }
321
322    /// Selection follows the order of the requested names, not the stored index order.
323    #[test]
324    fn test_named_matrix_sub_mat() {
325        let (mat, row_names, col_names) = setup_named_matrix();
326        let named_mat = NamedMatrix {
327            col_names,
328            row_names,
329            values: mat.as_ref(),
330        };
331
332        // Select row 3 and row 1 (in that order), and column 2
333        let sub = named_mat
334            .get_sub_mat(Some(&["r3", "r1"]), Some(&["c2"]))
335            .unwrap();
336
337        assert_eq!(sub.nrows(), 2);
338        assert_eq!(sub.ncols(), 1);
339
340        // "r3", "c2" -> original (2, 1) -> value 7.0
341        assert!((sub[(0, 0)] - 7.0).abs() < 1e-6);
342        // "r1", "c2" -> original (0, 1) -> value 1.0
343        assert!((sub[(1, 0)] - 1.0).abs() < 1e-6);
344    }
345
346    /// An empty selection and an all-unknown selection both give None, never a zero-sized matrix.
347    #[test]
348    fn test_named_matrix_empty_selection() {
349        let (mat, row_names, col_names) = setup_named_matrix();
350        let named_mat = NamedMatrix {
351            col_names,
352            row_names,
353            values: mat.as_ref(),
354        };
355
356        let empty_rows = named_mat.get_sub_mat(Some(&[]), None);
357        assert!(empty_rows.is_none());
358
359        let invalid_rows = named_mat.get_sub_mat(Some(&["r_invalid"]), None);
360        assert!(invalid_rows.is_none());
361    }
362
363    /// The deferred view reports the sliced shape and copies the right cells when realised.
364    #[test]
365    fn test_mat_slice_view() {
366        let (mat, _, _) = setup_named_matrix();
367        let r_idx = vec![0, 2]; // row 0, row 2
368        let c_idx = vec![1, 2]; // col 1, col 2
369
370        let view = MatSliceView::new(mat.as_ref(), &r_idx, &c_idx);
371        assert_eq!(view.nrows(), 2);
372        assert_eq!(view.ncols(), 2);
373
374        let owned = view.to_owned();
375        // original (0, 1) -> 1.0
376        assert!((owned[(0, 0)] - 1.0).abs() < 1e-6);
377        // original (2, 2) -> 8.0
378        assert!((owned[(1, 1)] - 8.0).abs() < 1e-6);
379    }
380
381    /// Out-of-range indices panic at construction, not later during `to_owned`.
382    #[test]
383    #[should_panic(expected = "You selected indices larger than ncol.")]
384    fn test_mat_slice_view_out_of_bounds_col() {
385        let (mat, _, _) = setup_named_matrix();
386        let r_idx = vec![0];
387        let c_idx = vec![5]; // Out of bounds
388        let _view = MatSliceView::new(mat.as_ref(), &r_idx, &c_idx);
389    }
390}