use crate::compressed::{check_compressed, transpose};
use crate::index::{Col, Row};
use crate::real::Real;
use crate::{CscMatrix, SparseError};
#[derive(Clone, Debug, PartialEq)]
pub struct CsrMatrix<T: Real = f64> {
pub n_cols: usize,
pub row_ptr: Vec<usize>,
pub col_idx: Vec<u32>,
pub val: Vec<T>,
}
impl<T: Real> CsrMatrix<T> {
pub fn from_raw_parts(
n_cols: usize,
row_ptr: Vec<usize>,
col_idx: Vec<u32>,
val: Vec<T>,
) -> Result<Self, SparseError> {
let r = Self {
n_cols,
row_ptr,
col_idx,
val,
};
r.check_invariants()?;
Ok(r)
}
#[inline]
pub fn n_rows(&self) -> usize {
self.row_ptr.len() - 1
}
#[inline]
pub fn n_cols(&self) -> usize {
self.n_cols
}
#[inline]
pub fn nnz(&self) -> usize {
self.val.len()
}
#[inline]
pub fn row(&self, i: Row) -> (&[u32], &[T]) {
let r = self.row_ptr[i.index()]..self.row_ptr[i.index() + 1];
(&self.col_idx[r.clone()], &self.val[r])
}
pub fn get(&self, i: Row, j: Col) -> T {
let (cols, vals) = self.row(i);
cols.binary_search(&j.0).map_or(T::ZERO, |p| vals[p])
}
pub fn to_csc(&self) -> CscMatrix<T> {
let (col_ptr, row_idx, val) =
transpose(self.n_cols, &self.row_ptr, &self.col_idx, &self.val);
let a = CscMatrix {
n_rows: self.n_rows(),
col_ptr,
row_idx,
val,
};
debug_assert!(a.check_invariants().is_ok(), "{:?}", a.check_invariants());
a
}
pub fn check_invariants(&self) -> Result<(), SparseError> {
check_compressed(
"row",
"column",
self.n_cols,
&self.row_ptr,
&self.col_idx,
&self.val,
)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::real::{approx_eq, is_exact_zero};
#[test]
fn row_access() {
let r = CsrMatrix::from_raw_parts(3, vec![0, 2, 2, 3], vec![0, 2, 1], vec![1.0, 2.0, 3.0])
.unwrap();
assert_eq!(r.n_rows(), 3);
assert_eq!(r.row(Row(0)), (&[0_u32, 2][..], &[1.0, 2.0][..]));
assert!(approx_eq(r.get(Row(2), Col(1)), 3.0, 0.0));
assert!(is_exact_zero(r.get(Row(1), Col(1))));
assert!(CsrMatrix::from_raw_parts(3, vec![0, 1], vec![3], vec![1.0_f64]).is_err());
}
}