Feature: add filter and remove functions to CooMatrix

nalgebra-sparse/src/coo.rs

@@ -311,4 +311,146 @@ impl<T> CooMatrix<T> {
     pub fn disassemble(self) -> (Vec<usize>, Vec<usize>, Vec<T>) {
         (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<F>(&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
+    }
 }

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<_>>(), 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<_>>(),
+        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<_>>(),
+        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<_>>(),
+        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<_>>(),
+        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<_>>(),
+        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);
+}