1use faer::{Mat, MatRef};
4use std::collections::BTreeMap;
5
6use crate::prelude::*;
7
8#[derive(Clone, Debug)]
14pub struct NamedMatrix<'a, T> {
15 pub col_names: BTreeMap<String, usize>,
17 pub row_names: BTreeMap<String, usize>,
19 pub values: faer::MatRef<'a, T>,
21}
22
23impl<'a, T> NamedMatrix<'a, T>
24where
25 T: BixverseFloat,
26{
27 pub fn get_sub_mat(
41 &self,
42 rows_to_select: Option<&[&str]>,
43 cols_to_select: Option<&[&str]>,
44 ) -> Option<Mat<T>> {
45 let row_indices: Vec<usize> = match rows_to_select {
47 None => (0..self.values.nrows()).collect(), 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 let col_indices: Vec<usize> = match cols_to_select {
67 None => (0..self.values.ncols()).collect(), 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 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 pub fn get_full_mat(&self) -> Mat<T> {
102 self.get_sub_mat(None, None).unwrap()
103 }
104
105 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 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 pub fn get_col_names_refs(&self) -> Vec<&String> {
137 self.col_names.keys().collect()
138 }
139
140 pub fn get_row_names_refs(&self) -> Vec<&String> {
146 self.row_names.keys().collect()
147 }
148
149 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(), 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 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(), 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#[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 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 pub fn nrows(&self) -> usize {
268 self.row_indices.len()
269 }
270
271 pub fn ncols(&self) -> usize {
277 self.col_indices.len()
278 }
279
280 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#[cfg(test)]
300mod tests {
301 use super::*;
302 use faer::Mat;
303 use std::collections::BTreeMap;
304
305 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 #[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 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 assert!((sub[(0, 0)] - 7.0).abs() < 1e-6);
342 assert!((sub[(1, 0)] - 1.0).abs() < 1e-6);
344 }
345
346 #[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 #[test]
365 fn test_mat_slice_view() {
366 let (mat, _, _) = setup_named_matrix();
367 let r_idx = vec![0, 2]; let c_idx = vec![1, 2]; 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 assert!((owned[(0, 0)] - 1.0).abs() < 1e-6);
377 assert!((owned[(1, 1)] - 8.0).abs() < 1e-6);
379 }
380
381 #[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]; let _view = MatSliceView::new(mat.as_ref(), &r_idx, &c_idx);
389 }
390}