diff --git a/nalgebra-sparse/src/coo.rs b/nalgebra-sparse/src/coo.rs index 6f9698143..c0627eb71 100644 --- a/nalgebra-sparse/src/coo.rs +++ b/nalgebra-sparse/src/coo.rs @@ -311,4 +311,146 @@ impl CooMatrix { pub fn disassemble(self) -> (Vec, Vec, Vec) { (self.row_indices, self.col_indices, self.values) } + + /// Returns a new COO matrix containing only entries specified by the given predicate. The cost + /// of this operation is O(nnz) and allocates space for the new matrix. + #[inline] + pub fn filter(&self, predicate: F) -> Self + where + F: Fn(usize, usize, T) -> bool, + T: Copy, + { + let ((row_indices, col_indices), values) = self + .row_indices + .iter() + .zip(self.col_indices.iter()) + .zip(self.values.iter()) + .filter(|((i, j), v)| predicate(**i, **j, **v)) + .unzip(); + Self { + nrows: self.nrows, + ncols: self.ncols, + row_indices, + col_indices, + values, + } + } + + /// Removes the `i`th row from the matrix. Beware the cost of the operation is `O(nnz)` and + /// causes a reallocation. + /// + /// Panics + /// ------- + /// + /// Panics if `i` is out of bounds. + /// + /// Examples + /// -------- + /// + /// ``` + /// # use nalgebra_sparse::coo::CooMatrix; + /// let row_indices = vec![0, 1]; + /// let col_indices = vec![1, 2]; + /// let values = vec![1.0, 2.0]; + /// let mut coo = CooMatrix::try_from_triplets(2, 3, row_indices, col_indices, values) + /// .unwrap(); + /// + /// coo.remove_row(0); + /// ``` + pub fn remove_row(&self, i: usize) -> Self + where + T: Copy, + { + assert!(i < self.nrows); + let mut new_coo = self.filter(|row, _, _| row != i); + new_coo.row_indices.iter_mut().for_each(|idx| { + if *idx > i { + *idx -= 1; + } + }); + new_coo.nrows -= 1; + new_coo + } + + /// Removes the `i`th column from the matrix. Beware the cost of the operation is `O(nnz)` and + /// causes a reallocation. + /// + /// Panics + /// ------- + /// + /// Panics if `i` is out of bounds. + /// + /// Examples + /// -------- + /// + /// ``` + /// # use nalgebra_sparse::coo::CooMatrix; + /// let row_indices = vec![0, 1]; + /// let col_indices = vec![1, 2]; + /// let values = vec![1.0, 2.0]; + /// let mut coo = CooMatrix::try_from_triplets(2, 3, row_indices, col_indices, values) + /// .unwrap(); + /// + /// coo.remove_column(0); + /// ``` + pub fn remove_column(&self, i: usize) -> Self + where + T: Copy, + { + assert!(i < self.ncols); + let mut new_coo = self.filter(|_, col, _| col != i); + new_coo.col_indices.iter_mut().for_each(|idx| { + if *idx > i { + *idx -= 1; + } + }); + new_coo.ncols -= 1; + new_coo + } + + /// Removes the `i`th row and the `j`th column from the matrix. Beware the cost of the operation + /// is `O(nnz)` and causes a reallocation. Note that a reallocation can be saved calling this + /// function rather than successive calls to `remove_row` or `remove_column`. + /// + /// Panics + /// ------- + /// + /// Panics if `i` or `j` are out of bounds. + /// + /// Examples + /// -------- + /// + /// ``` + /// # use nalgebra_sparse::coo::CooMatrix; + /// let row_indices = vec![0, 1]; + /// let col_indices = vec![1, 2]; + /// let values = vec![1.0, 2.0]; + /// let mut coo = CooMatrix::try_from_triplets(2, 3, row_indices, col_indices, values) + /// .unwrap(); + /// + /// coo.remove_row_column(0, 1); + /// ``` + pub fn remove_row_column(&self, i: usize, j: usize) -> Self + where + T: Copy, + { + assert!(i < self.nrows); + assert!(j < self.ncols); + let mut new_coo = self.filter(|row, col, _| row != i && col != j); + new_coo + .row_indices + .iter_mut() + .zip(new_coo.col_indices.iter_mut()) + .for_each(|(row_idx, col_idx)| { + if *row_idx > i { + *row_idx -= 1; + } + if *col_idx > j { + *col_idx -= 1; + } + }); + new_coo.nrows -= 1; + new_coo.ncols -= 1; + new_coo + } } diff --git a/nalgebra-sparse/tests/unit_tests/coo.rs b/nalgebra-sparse/tests/unit_tests/coo.rs index 89b2d26ed..60819259f 100644 --- a/nalgebra-sparse/tests/unit_tests/coo.rs +++ b/nalgebra-sparse/tests/unit_tests/coo.rs @@ -428,3 +428,185 @@ fn coo_push_matrix_out_of_bounds_entries() { assert_panics!(CooMatrix::new(3, 3).push_matrix(2, 2, &inserted)); } } + +#[test] +fn coo_remove_row_valid() { + let mut coo = CooMatrix::new(3, 3); + + coo.push(0, 0, 1); + coo.push(0, 1, 2); + coo.push(0, 2, 3); + + let mut removed_coo = coo.remove_row(0); + assert_eq!(removed_coo.nrows(), 2); + assert_eq!(removed_coo.ncols(), 3); + assert_eq!(removed_coo.nnz(), 0); + + assert_eq!(removed_coo.triplet_iter().collect::>(), vec![]); + + #[rustfmt::skip] + let expected_dense = DMatrix::from_row_slice(2, 3, &[ + 0, 0, 0, + 0, 0, 0, + ]); + assert_eq!(DMatrix::from(&removed_coo), expected_dense); + + // makes sure resulting COO matrix still works. This will push to the new + // matrices row 0. + removed_coo.push(0, 0, 1); + removed_coo.push(0, 1, 2); + removed_coo.push(0, 2, 3); + assert_eq!(removed_coo.nnz(), 3); + + assert_panics!(removed_coo.clone().push(2, 0, 4)); + + // makes sure original matrix is untouched. + assert_eq!( + coo.triplet_iter().collect::>(), + vec![(0, 0, &1), (0, 1, &2), (0, 2, &3)] + ); + + #[rustfmt::skip] + let expected_dense = DMatrix::from_row_slice(2, 3, &[ + 1, 2, 3, + 0, 0, 0, + ]); + assert_eq!(DMatrix::from(&removed_coo), expected_dense); +} + +#[test] +fn coo_remove_row_out_of_bounds() { + let mut coo = CooMatrix::new(3, 3); + + coo.push(0, 0, 1); + coo.push(0, 1, 2); + coo.push(0, 2, 3); + + // Push past col-dim + assert_panics!(coo.clone().remove_row(3)); +} + +#[test] +fn coo_remove_column_valid() { + let mut coo = CooMatrix::new(3, 3); + + coo.push(0, 0, 1); + coo.push(0, 1, 2); + coo.push(0, 2, 3); + + let mut removed_coo = coo.remove_column(1); + assert_eq!(removed_coo.ncols(), 2); + assert_eq!(removed_coo.nrows(), 3); + assert_eq!(removed_coo.nnz(), 2); + + assert_eq!( + removed_coo.triplet_iter().collect::>(), + vec![(0, 0, &1), (0, 1, &3)] + ); + + #[rustfmt::skip] + let expected_dense = DMatrix::from_row_slice(3, 2, &[ + 1, 3, + 0, 0, + 0, 0, + ]); + assert_eq!(DMatrix::from(&removed_coo), expected_dense); + + // makes sure resulting COO matrix still works. + removed_coo.push(0, 1, 2); + removed_coo.push(2, 1, 4); + assert_eq!(removed_coo.nnz(), 4); + + assert_panics!(removed_coo.clone().push(0, 2, 4)); + + // makes sure original matrix is untouched. + assert_eq!( + coo.triplet_iter().collect::>(), + vec![(0, 0, &1), (0, 1, &2), (0, 2, &3)] + ); + + #[rustfmt::skip] + let expected_dense = DMatrix::from_row_slice(3, 2, &[ + 1, 5, + 0, 0, + 0, 4, + ]); + assert_eq!(DMatrix::from(&removed_coo), expected_dense); +} + +#[test] +fn coo_remove_column_out_of_bounds() { + let mut coo = CooMatrix::new(3, 3); + + coo.push(0, 0, 1); + coo.push(0, 1, 2); + coo.push(0, 2, 3); + + // Push past col-dim + assert_panics!(coo.clone().remove_column(3)); +} + +#[test] +fn coo_remove_row_column_valid() { + let mut coo = CooMatrix::new(3, 3); + + coo.push(0, 0, 1); + coo.push(0, 1, 2); + coo.push(0, 2, 3); + coo.push(1, 0, 4); + coo.push(1, 1, 5); + coo.push(1, 2, 6); + coo.push(2, 0, 7); + coo.push(2, 1, 8); + coo.push(2, 2, 9); + + let mut removed_coo = coo.remove_row_column(1, 1); + assert_eq!(removed_coo.ncols(), 2); + assert_eq!(removed_coo.nrows(), 2); + assert_eq!(removed_coo.nnz(), 4); + + assert_eq!( + removed_coo.triplet_iter().collect::>(), + vec![(0, 0, &1), (0, 1, &3), (1, 0, &7), (1, 1, &9)] + ); + + #[rustfmt::skip] + let expected_dense = DMatrix::from_row_slice(2, 2, &[ + 1, 3, + 7, 9 + ]); + assert_eq!(DMatrix::from(&removed_coo), expected_dense); + + // makes sure resulting COO matrix still works. + removed_coo.push(0, 0, 1); + removed_coo.push(0, 1, 0); + removed_coo.push(1, 0, 0); + removed_coo.push(1, 1, 4); + assert_eq!(removed_coo.nnz(), 8); + + assert_panics!(removed_coo.clone().push(0, 2, 4)); + assert_panics!(removed_coo.clone().push(2, 1, 4)); + + // makes sure original matrix is untouched. + assert_eq!( + coo.triplet_iter().collect::>(), + vec![ + (0, 0, &1), + (0, 1, &2), + (0, 2, &3), + (1, 0, &4), + (1, 1, &5), + (1, 2, &6), + (2, 0, &7), + (2, 1, &8), + (2, 2, &9) + ] + ); + + #[rustfmt::skip] + let expected_dense = DMatrix::from_row_slice(2, 2, &[ + 2, 3, + 7, 13 + ]); + assert_eq!(DMatrix::from(&removed_coo), expected_dense); +}