Merge pull request #1288 from JulianKnodt/dev

Add `try_from_triplets_iter`
This commit is contained in:
Sébastien Crozet 2023-09-30 17:55:41 +02:00 committed by GitHub
commit 25749e3b89
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 44 additions and 0 deletions

View File

@ -160,6 +160,25 @@ impl<T> CooMatrix<T> {
}
}
/// Try to construct a COO matrix from the given dimensions and a finite iterator of
/// (i, j, v) triplets.
///
/// Returns an error if either row or column indices contain indices out of bounds.
/// Note that the COO format inherently supports duplicate entries, but they are not
/// eagerly summed.
///
/// Implementation note:
/// Calls try_from_triplets so each value is scanned twice.
pub fn try_from_triplets_iter(
nrows: usize,
ncols: usize,
triplets: impl IntoIterator<Item = (usize, usize, T)>,
) -> Result<Self, SparseFormatError> {
let (row_indices, (col_indices, values)) =
triplets.into_iter().map(|(r, c, v)| (r, (c, v))).unzip();
Self::try_from_triplets(nrows, ncols, row_indices, col_indices, values)
}
/// An iterator over triplets (i, j, v).
// TODO: Consider giving the iterator a concrete type instead of impl trait...?
pub fn triplet_iter(&self) -> impl Iterator<Item = (usize, usize, &T)> {

View File

@ -184,6 +184,31 @@ fn coo_try_from_triplets_reports_out_of_bounds_indices() {
}
}
#[test]
fn coo_try_from_triplets_iter() {
// Check that try_from_triplets_iter panics when the triplet vectors have different lengths
macro_rules! assert_errs {
($result:expr) => {
assert!(matches!(
$result.unwrap_err().kind(),
SparseFormatErrorKind::IndexOutOfBounds
))
};
}
assert_errs!(CooMatrix::<f32>::try_from_triplets_iter(
3,
5,
vec![(0, 6, 3.0)].into_iter(),
));
assert!(CooMatrix::<f32>::try_from_triplets_iter(
3,
5,
vec![(0, 3, 3.0), (1, 2, 2.0), (0, 3, 1.0),].into_iter(),
)
.is_ok());
}
#[test]
fn coo_try_from_triplets_panics_on_mismatched_vectors() {
// Check that try_from_triplets panics when the triplet vectors have different lengths