#[cfg(feature = "rand")]
mod csr_matrix_rnd;
mod csr_matrix_row_constructor;
mod csr_matrix_row_iter_impls;
#[cfg(feature = "rayon")]
mod csr_matrix_row_par_iter_impls;
pub use self::{
csr_matrix_row_constructor::CsrMatrixRowConstructor,
csr_matrix_row_iter_impls::{CsrMatrixRowIter, CsrMatrixRowIterMut},
};
use crate::{
dim::Dim,
prelude::{Array, DynDenseStoMut, DynDenseStoRef, Matrix, StDenseStoMut, StDenseStoRef},
utils::{are_in_upper_bound, does_not_have_duplicates},
vec::{
css::{CssSlice, CssSliceMut},
VecArray,
},
};
use core::ops::Range;
#[cfg_attr(
feature = "serde1",
derive(mop_common_deps::serde::Deserialize, mop_common_deps::serde::Serialize)
)]
#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, PartialOrd)]
pub struct CsrMatrix<DS, US> {
data: DS,
dim: Dim<[usize; 2]>,
indcs: US,
ptrs: US,
}
pub type CsrMatrixArray<DA, UA> = CsrMatrix<DA, UA>;
#[cfg(feature = "std")]
pub type CsrMatrixVec<T> = CsrMatrix<Vec<T>, Vec<usize>>;
pub type CsrMatrixVecArray<DA, UA> = CsrMatrix<VecArray<DA>, VecArray<UA>>;
pub type CsrMatrixRef<'a, T> = CsrMatrix<&'a [T], &'a [usize]>;
#[cfg(feature = "std")]
impl<T> CsrMatrixVec<T> {
#[cfg(feature = "rand")]
pub fn new_rnd<F, R>(shape: [usize; 2], nnz: usize, rng: &mut R, cb: F) -> Self
where
F: FnMut(&mut R, [usize; 2]) -> T,
R: mop_common_deps::rand::Rng,
{
let dim: Dim<[usize; 2]> = shape.into();
let capacity = dim.rows() * dim.cols();
assert!(nnz <= capacity);
let mut cmr = self::csr_matrix_rnd::CsrMatrixRnd {
data: Vec::with_capacity(nnz),
dim: &dim,
indcs: Vec::with_capacity(nnz),
nnz,
ptrs: Vec::with_capacity(dim.rows() + 1),
rng,
};
cmr.fill_ptrs();
cmr.fill_indcs();
cmr.fill_data(cb);
CsrMatrixVec::new([dim.rows(), dim.cols()], cmr.data, cmr.indcs, cmr.ptrs)
}
pub fn with_vec_capacity(shape: [usize; 2], nnz: usize) -> Self {
let mut ptrs = Vec::with_capacity(shape[0] + 1);
ptrs.push(0);
CsrMatrix {
data: Vec::with_capacity(nnz),
dim: [0, shape[1]].into(),
indcs: Vec::with_capacity(nnz),
ptrs,
}
}
}
impl<DA, UA> CsrMatrixVecArray<DA, UA>
where
DA: Array,
UA: Array<Item = usize>,
{
#[cfg(feature = "rand")]
pub fn new_rnd<F, R>(shape: [usize; 2], nnz: usize, rng: &mut R, cb: F) -> Self
where
F: FnMut(&mut R, [usize; 2]) -> DA::Item,
R: mop_common_deps::rand::Rng,
{
let dim: Dim<[usize; 2]> = shape.into();
let capacity = dim.rows() * dim.cols();
assert!(nnz <= capacity);
let mut cmr = self::csr_matrix_rnd::CsrMatrixRnd {
data: VecArray::<DA>::with_capacity(),
dim: &dim,
indcs: VecArray::<UA>::with_capacity(),
nnz,
ptrs: VecArray::<UA>::with_capacity(),
rng,
};
cmr.fill_ptrs();
cmr.fill_indcs();
cmr.fill_data(cb);
CsrMatrixVecArray::new([dim.rows(), dim.cols()], cmr.data, cmr.indcs, cmr.ptrs)
}
pub fn with_vec_array(cols: usize) -> Self {
let mut ptrs_vec_array = VecArray::with_capacity();
ptrs_vec_array.push(0);
CsrMatrixVecArray {
data: VecArray::with_capacity(),
dim: [0, cols].into(),
indcs: VecArray::with_capacity(),
ptrs: ptrs_vec_array,
}
}
}
impl<DS, US> CsrMatrix<DS, US>
where
DS: StDenseStoRef,
US: StDenseStoRef<Item = usize>,
{
pub fn new(shape: [usize; 2], data: DS, indcs: US, ptrs: US) -> Self {
let dim: Dim<[usize; 2]> = shape.into();
assert!(
data.as_slice().len() == indcs.as_slice().len(),
"The length of `data` must be equal the length of `indcs`."
);
assert!(
data.as_slice().len() <= dim.rows() * dim.cols(),
"The length of `data` must be less than the number of rows times the \
number of columns."
);
assert!(
ptrs.as_slice().len() == dim.rows() + 1,
"The length of `ptrs` must be equal the number of rows plus 1."
);
assert!(
are_in_upper_bound(indcs.as_slice(), &dim.cols()),
"Column indices must be less than the number of columns."
);
assert!(
ptrs
.as_slice()
.windows(2)
.all(|x| does_not_have_duplicates(&indcs.as_slice()[x[0]..x[1]])),
"Column indices of a row must be unique."
);
assert!(
*ptrs.as_slice().last().unwrap() == data.as_slice().len(),
"Rows length must end with the nnz."
);
assert!(
ptrs.as_slice().windows(2).all(|x| x[0] <= x[1]),
"Rows length must be in ascending order."
);
unsafe { Self::new_unchecked(shape, data, indcs, ptrs) }
}
pub unsafe fn new_unchecked(shape: [usize; 2], data: DS, indcs: US, ptrs: US) -> Self {
CsrMatrix {
data,
dim: shape.into(),
indcs,
ptrs,
}
}
pub fn as_ref(&self) -> CsrMatrixRef<'_, DS::Item> {
CsrMatrix::new(
[self.dim.rows(), self.dim.cols()],
self.data.as_slice(),
self.indcs.as_slice(),
self.ptrs.as_slice(),
)
}
pub fn data(&self) -> &[DS::Item] {
self.data.as_slice()
}
pub fn indcs(&self) -> &[usize] {
&self.indcs.as_slice()
}
#[inline]
pub fn nnz(&self) -> usize {
self.data().len()
}
pub fn ptrs(&self) -> &[usize] {
&self.ptrs.as_slice()
}
pub fn row(&self, row_idx: usize) -> CssSlice<'_, DS::Item> {
let range = self.ptr_range_of_row(row_idx);
CssSlice::new(
self.dim.cols(),
&self.data.as_slice()[range.clone()],
&self.indcs.as_slice()[range],
)
}
pub fn row_iter(&self) -> CsrMatrixRowIter<'_, DS::Item> {
CsrMatrixRowIter::new(
[self.rows(), self.cols()],
self.data().as_ptr(),
self.indcs(),
self.ptrs(),
)
}
#[cfg(feature = "rayon")]
pub fn row_par_iter(
&self,
) -> crate::utils::ParallelIteratorWrapper<CsrMatrixRowIter<'_, DS::Item>> {
crate::utils::ParallelIteratorWrapper(self.row_iter())
}
pub fn value(&self, row_idx: usize, col_idx: usize) -> Option<&DS::Item> {
self
.data_idx_of_coords(row_idx, col_idx)
.map(|x| &self.data()[x])
}
#[inline]
fn data_idx_of_coords(&self, row_idx: usize, col_idx: usize) -> Option<usize> {
assert!(row_idx < self.rows() && col_idx < self.cols());
let row_start_ptr = self.ptrs()[row_idx];
let row_end_ptr = self.ptrs()[row_idx + 1];
if let Ok(x) = self.indcs()[row_start_ptr..row_end_ptr].binary_search(&col_idx) {
Some(row_start_ptr + x)
} else {
None
}
}
#[inline]
fn ptr_range_of_row(&self, row_idx: usize) -> Range<usize> {
assert!(
row_idx < self.dim.rows(),
"`row_idx` must be less than the number of rows."
);
let lower_bound = self.ptrs()[row_idx] - self.ptrs()[0];
let upper_bound = self.ptrs()[row_idx + 1] - self.ptrs()[0];
lower_bound..upper_bound
}
}
#[cfg(feature = "rand")]
impl<DS, US> CsrMatrix<DS, US>
where
DS: StDenseStoRef,
US: StDenseStoRef<Item = usize>,
{
}
impl<DS, US> CsrMatrix<DS, US>
where
DS: StDenseStoMut,
US: StDenseStoRef<Item = usize>,
{
pub fn data_mut(&mut self) -> &mut [DS::Item] {
self.data.as_mut_slice()
}
pub fn parts_with_data_mut(&mut self) -> (&mut [DS::Item], &[usize], &[usize]) {
(
self.data.as_mut_slice(),
self.indcs.as_slice(),
self.ptrs.as_slice(),
)
}
pub fn row_mut(&mut self, row_idx: usize) -> CssSliceMut<'_, DS::Item> {
let range = self.ptr_range_of_row(row_idx);
CssSliceMut::new(
self.dim.cols(),
&mut self.data.as_mut_slice()[range.clone()],
&self.indcs.as_slice()[range],
)
}
pub fn row_iter_mut(&mut self) -> CsrMatrixRowIterMut<'_, DS::Item> {
CsrMatrixRowIterMut::new(
[self.rows(), self.cols()],
self.data_mut().as_mut_ptr(),
self.indcs(),
self.ptrs(),
)
}
#[cfg(feature = "rayon")]
pub fn row_par_iter_mut(
&mut self,
) -> crate::utils::ParallelIteratorWrapper<CsrMatrixRowIterMut<'_, DS::Item>> {
crate::utils::ParallelIteratorWrapper(self.row_iter_mut())
}
pub fn split_at_mut(
&mut self,
row_idx: usize,
) -> (CssSliceMut<'_, DS::Item>, CssSliceMut<'_, DS::Item>) {
assert!(row_idx <= self.rows());
let split_at = self.ptrs()[row_idx];
let (first_data, second_data) = self.data.as_mut_slice().split_at_mut(split_at);
let first = CssSliceMut::new(
self.dim.cols(),
first_data,
&self.indcs.as_slice()[0..split_at],
);
let second = CssSliceMut::new(
self.dim.cols(),
second_data,
&self.indcs.as_slice()[split_at..],
);
(first, second)
}
pub fn value_mut(&mut self, row_idx: usize, col_idx: usize) -> Option<&mut DS::Item> {
if let Some(x) = self.data_idx_of_coords(row_idx, col_idx) {
Some(&mut self.data_mut()[x])
} else {
None
}
}
pub fn swap(&mut self, a: [usize; 2], b: [usize; 2]) {
let a_data_idx = self.data_idx_of_coords(a[0], a[1]).unwrap();
let b_data_idx = self.data_idx_of_coords(b[0], b[1]).unwrap();
self.data_mut().swap(a_data_idx, b_data_idx);
}
}
impl<DS, US> CsrMatrix<DS, US>
where
DS: DynDenseStoRef,
US: DynDenseStoRef<Item = usize>,
{
#[inline]
pub fn rows_capacity(&self) -> usize {
self.ptrs.capacity() - 1
}
}
impl<DS, US> CsrMatrix<DS, US>
where
DS: DynDenseStoMut,
US: DynDenseStoMut<Item = usize>,
{
pub fn clear(&mut self) {
self.data.clear();
self.indcs.clear();
self.ptrs.clear();
self.ptrs.push(0);
*self.dim.rows_mut() = 0;
}
pub fn extend(&mut self, other: &Self)
where
DS::Item: Copy,
US::Item: Copy,
{
assert!(
self.cols() == other.cols(),
"The number of columns of `Self` must be equal the number of columns of `other`."
);
let mut current_ptr = *self.ptrs.as_slice().last().unwrap();
let mut last_other_ptr = *other.ptrs().first().unwrap();
other.ptrs().iter().skip(1).for_each(|&x| {
current_ptr += x - last_other_ptr;
last_other_ptr = x;
self.ptrs.push(current_ptr);
});
self.data.extend(other.data());
self.indcs.extend(other.indcs());
*self.dim.rows_mut() += other.rows();
}
pub fn row_constructor(&mut self) -> CsrMatrixRowConstructor<'_, DS, US> {
CsrMatrixRowConstructor::new(
&mut self.dim,
&mut self.data,
&mut self.indcs,
&mut self.ptrs,
)
}
pub fn truncate(&mut self, until_row_idx: usize) {
self.data.truncate(until_row_idx);
self.indcs.truncate(until_row_idx);
self.ptrs.truncate(until_row_idx + 1);
*self.dim.rows_mut() = until_row_idx;
}
}
impl<DS, US> Matrix for CsrMatrix<DS, US> {
#[inline]
fn cols(&self) -> usize {
self.dim.cols()
}
#[inline]
fn rows(&self) -> usize {
self.dim.rows()
}
}
#[cfg(all(feature = "quickcheck", feature = "rand"))]
use mop_common_deps::{
quickcheck::{Arbitrary, Gen},
rand::{
distributions::{Distribution, Standard},
Rng,
},
};
#[cfg(all(feature = "quickcheck", feature = "rand"))]
impl<T> Arbitrary for CsrMatrixVec<T>
where
Standard: Distribution<T>,
T: Arbitrary,
{
#[inline]
fn arbitrary<G>(g: &mut G) -> Self
where
G: Gen,
{
let rows = g.gen_range(0, g.size());
let cols = g.gen_range(0, g.size());
let nnz = g.gen_range(0, rows * cols + 1);
CsrMatrixVec::new_rnd([rows, cols], nnz, g, |g, _| g.gen())
}
}