use std::{collections::HashSet, ops::Mul};
use anyhow::Result;
use nalgebra::DVector;
use nalgebra_sparse::{pattern::SparsityPattern, CooMatrix, CscMatrix};
use crate::{scalar::Scale, IndexType, Scalar};
use super::{Matrix, MatrixCommon, MatrixSparsity};
impl<T: Scalar> MatrixCommon for CscMatrix<T> {
type V = DVector<T>;
type T = T;
fn ncols(&self) -> IndexType {
self.ncols()
}
fn nrows(&self) -> IndexType {
self.nrows()
}
}
impl<T: Scalar> Mul<Scale<T>> for CscMatrix<T> {
type Output = CscMatrix<T>;
fn mul(self, rhs: Scale<T>) -> Self::Output {
self * rhs.value()
}
}
impl MatrixSparsity for SparsityPattern {
type Index = Vec<IndexType>;
fn get_index(&self, rows: &[IndexType], cols: &[IndexType]) -> Self::Index {
let mut index = Vec::with_capacity(rows.len());
for (&i, &j) in rows.iter().zip(cols) {
let offset = self.major_offsets()[j];
let lane = self.lane(j);
let lane_i = lane.iter().position(|&x| x == i).unwrap();
index.push(offset + lane_i);
}
index
}
fn nrows(&self) -> IndexType {
self.minor_dim()
}
fn ncols(&self) -> IndexType {
self.major_dim()
}
fn is_sparse(&self) -> bool {
true
}
fn indices(&self) -> Vec<(IndexType, IndexType)> {
let mut indices = Vec::with_capacity(self.nnz());
for (j, &offset) in self.major_offsets().iter().enumerate() {
let next_offset = self
.major_offsets()
.get(j + 1)
.copied()
.unwrap_or(self.minor_indices().len());
for i in offset..next_offset {
indices.push((self.minor_indices()[i], j));
}
}
indices
}
fn union(&self, other: &Self) -> Result<Self> {
let max_nnz = self.nnz().max(other.nnz());
let min_nnz = self.nnz().min(other.nnz());
let mut minor_indices = Vec::with_capacity(self.nnz() + max_nnz - min_nnz);
let mut major_offsets = Vec::with_capacity(self.major_dim());
let mut offset = 0;
for j in 0..self.major_dim() {
let lane = self.lane(j);
let other_lane = other.lane(j);
let set: HashSet<usize> =
HashSet::from_iter(lane.iter().chain(other_lane.iter()).cloned());
let mut set = set.into_iter().collect::<Vec<_>>();
major_offsets.push(offset);
offset += set.len();
minor_indices.append(&mut set);
}
SparsityPattern::try_from_offsets_and_indices(
self.major_dim(),
self.minor_dim(),
major_offsets,
minor_indices,
)
.map_err(anyhow::Error::new)
}
fn try_from_indices(
nrows: IndexType,
ncols: IndexType,
indices: Vec<(IndexType, IndexType)>,
) -> Result<Self> {
let major_dim = ncols;
let minor_dim = nrows;
let mut indices = indices;
indices.sort_unstable_by_key(|&(_, j)| j);
let mut curr_col = 0;
let mut major_offsets = Vec::with_capacity(major_dim + 1);
let mut minor_indices = Vec::with_capacity(indices.len());
for (i, j) in indices {
while curr_col < j {
major_offsets.push(minor_indices.len());
curr_col += 1;
}
minor_indices.push(i);
}
major_offsets.push(minor_indices.len());
SparsityPattern::try_from_offsets_and_indices(
major_dim,
minor_dim,
major_offsets,
minor_indices,
)
.map_err(anyhow::Error::new)
}
fn new_diagonal(n: IndexType) -> Self {
let mut major_offsets = Vec::with_capacity(n + 1);
let mut minor_indices = Vec::with_capacity(n);
for i in 0..n {
major_offsets.push(i);
minor_indices.push(i);
}
major_offsets.push(n);
SparsityPattern::try_from_offsets_and_indices(n, n, major_offsets, minor_indices).unwrap()
}
}
impl<T: Scalar> Matrix for CscMatrix<T> {
type Sparsity = SparsityPattern;
fn sparsity(&self) -> Option<&Self::Sparsity> {
Some(self.pattern())
}
fn set_data_with_indices(
&mut self,
dst_indices: &<Self::Sparsity as MatrixSparsity>::Index,
src_indices: &<Self::V as crate::vector::Vector>::Index,
data: &Self::V,
) {
let values = self.values_mut();
for (&dst_i, &src_i) in dst_indices.iter().zip(src_indices.iter()) {
values[dst_i] = data[src_i];
}
}
fn triplet_iter(&self) -> impl Iterator<Item = (IndexType, IndexType, &Self::T)> {
self.triplet_iter()
}
fn add_column_to_vector(&self, j: IndexType, v: &mut Self::V) {
let col = self.col(j);
for (&i, &val) in col.row_indices().iter().zip(col.values().iter()) {
v[i] += val;
}
}
fn try_from_triplets(
nrows: IndexType,
ncols: IndexType,
triplets: Vec<(IndexType, IndexType, T)>,
) -> Result<Self> {
let mut coo = CooMatrix::new(nrows, ncols);
for (i, j, v) in triplets {
coo.push(i, j, v);
}
Ok(CscMatrix::from(&coo))
}
fn zeros(nrows: IndexType, ncols: IndexType) -> Self {
Self::zeros(nrows, ncols)
}
fn copy_from(&mut self, other: &Self) {
self.clone_from(other);
}
fn gemv(&self, alpha: Self::T, x: &Self::V, beta: Self::T, y: &mut Self::V) {
let mut tmp = self * x;
tmp *= alpha;
y.axpy(alpha, &tmp, beta);
}
fn from_diagonal(v: &DVector<T>) -> Self {
let nrows = v.len();
let ncols = v.len();
let mut coo = CooMatrix::<T>::new(nrows, ncols);
for (i, &v) in v.into_iter().enumerate() {
coo.push(i, i, v);
}
CscMatrix::from(&coo)
}
fn diagonal(&self) -> Self::V {
let mut ret = DVector::zeros(self.nrows());
for (i, _j, &v) in self.diagonal_as_csc().triplet_iter() {
ret[i] = v;
}
ret
}
fn set_column(&mut self, j: IndexType, v: &Self::V) {
assert_eq!(v.len(), self.nrows());
let mut col = self.col_mut(j);
let (dst_row_indices, dst_values) = col.rows_and_values_mut();
for (&dst_i, dst_v) in dst_row_indices.iter().zip(dst_values.iter_mut()) {
*dst_v = v[dst_i];
}
}
fn scale_add_and_assign(&mut self, x: &Self, beta: Self::T, y: &Self) {
*self = x + y * beta;
}
fn new_from_sparsity(
nrows: IndexType,
ncols: IndexType,
sparsity: Option<&Self::Sparsity>,
) -> Self {
if let Some(sparsity) = sparsity {
let values = vec![T::zero(); sparsity.nnz()];
CscMatrix::try_from_pattern_and_values(sparsity.clone(), values).unwrap()
} else {
CscMatrix::zeros(nrows, ncols)
}
}
}