first untested sparse implementation

This commit is contained in:
Magnus Ulimoen 2020-06-12 21:04:18 +02:00
parent 9261f1b239
commit 4f772b8dc5
6 changed files with 802 additions and 40 deletions

View File

@ -13,6 +13,7 @@ members = [
"euler",
"maxwell",
"shallow_water",
"sparse",
]
[profile.bench]
@ -24,3 +25,4 @@ debug = true
[patch]
[patch.crates-io]
hdf5 = { git = "https://github.com/mulimoen/hdf5-rust.git", branch = "master" }
sprs = { git = "https://github.com/vbarrielle/sprs.git" }

View File

@ -1,40 +0,0 @@
#! /usr/bin/env python3
import numpy as np
A = np.array([[0, 0, 0], [0, 0, -1], [0, -1, 0]])
B = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 0]])
def similarity_transform(matrix):
L, S = np.linalg.eig(matrix)
L = np.diag(L)
S = S.transpose()
assert np.allclose(np.matmul(S.transpose(), np.matmul(L, S)), matrix)
return L, S
def plusminus(matrix):
L, S = similarity_transform(matrix)
def signed(op):
return 0.5 * np.matmul(S.transpose(), np.matmul(op(L, np.abs(L)), S))
plus = signed(np.add)
minus = signed(np.subtract)
assert np.allclose(matrix, plus + minus)
return plus, minus
Aplus, Aminus = plusminus(A)
Bplus, Bminus = plusminus(B)
print("A+")
print(Aplus)
print("A-")
print(Aminus)
print()
print("B+")
print(Bplus)
print("B-")
print(Bminus)

View File

@ -9,10 +9,12 @@ ndarray = { version = "0.13.1", features = ["approx"] }
approx = "0.3.2"
packed_simd = "0.3.3"
rayon = { version = "1.3.0", optional = true }
sprs = { version = "0.7.1", optional = true }
[features]
# Use f32 as precision, default is f64
f32 = []
sparse = ["sprs"]
[dev-dependencies]
criterion = "0.3.1"

View File

@ -11,6 +11,14 @@ pub trait SbpOperator1d: Send + Sync {
fn is_h2(&self) -> bool {
false
}
#[cfg(feature = "sparse")]
fn diff_matrix(&self, n: usize) -> sprs::CsMat<Float> {
unimplemented!()
}
#[cfg(feature = "sparse")]
fn h_matrix(&self, n: usize) -> sprs::CsMat<Float> {
unimplemented!()
}
}
pub trait SbpOperator2d: Send + Sync {
@ -517,6 +525,95 @@ fn diff_op_row(
}
}
#[cfg(feature = "sparse")]
fn sparse_from_block(
block: &[&[Float]],
diag: &[Float],
symmetry: Symmetry,
optype: OperatorType,
n: usize,
) -> sprs::CsMat<Float> {
assert!(n >= 2 * block.len());
let nnz = {
let block_elems = block.iter().fold(0, |acc, x| {
acc + x
.iter()
.fold(0, |acc, &x| if x != 0.0 { acc + 1 } else { acc })
});
let diag_elems = diag
.iter()
.fold(0, |acc, &x| if x != 0.0 { acc + 1 } else { acc });
2 * block_elems + (n - 2 * block.len()) * diag_elems
};
let mut mat = sprs::TriMat::with_capacity((n, n), nnz);
let dx = if optype == OperatorType::H2 {
1.0 / (n - 2) as Float
} else {
1.0 / (n - 1) as Float
};
let idx = 1.0 / dx;
for (j, bl) in block.iter().enumerate() {
for (i, &b) in bl.iter().enumerate() {
if b == 0.0 {
continue;
}
mat.add_triplet(j, i, b * idx);
}
}
for j in block.len()..n - block.len() {
let half_diag_len = diag.len() / 2;
for (&d, i) in diag.iter().zip(j - half_diag_len..) {
if d == 0.0 {
continue;
}
mat.add_triplet(j, i, d * idx);
}
}
for (bl, j) in block.iter().zip((0..n).rev()).rev() {
for (&b, i) in bl.iter().zip((0..n).rev()).rev() {
if b == 0.0 {
continue;
}
if symmetry == Symmetry::AntiSymmetric {
mat.add_triplet(j, i, -b * idx);
} else {
mat.add_triplet(j, i, b * idx);
}
}
}
mat.to_csr()
}
#[cfg(feature = "sparse")]
fn h_matrix(diag: &[Float], n: usize, is_h2: bool) -> sprs::CsMat<Float> {
let h = if is_h2 {
1.0 / (n - 2) as Float
} else {
1.0 / (n - 1) as Float
};
let nmiddle = n - 2 * diag.len();
let iter = diag
.iter()
.chain(std::iter::repeat(&1.0).take(nmiddle))
.chain(diag.iter().rev())
.map(|&x| h * x);
let mut mat = sprs::TriMat::with_capacity((n, n), n);
for (i, d) in iter.enumerate() {
mat.add_triplet(i, i, d);
}
mat.to_csr()
}
mod upwind4;
pub use upwind4::Upwind4;
mod upwind9;

View File

@ -207,6 +207,20 @@ impl SbpOperator1d for Upwind4 {
fn h(&self) -> &'static [Float] {
Self::HBLOCK
}
#[cfg(feature = "sparse")]
fn diff_matrix(&self, n: usize) -> sprs::CsMat<Float> {
super::sparse_from_block(
Self::BLOCK,
Self::DIAG,
super::Symmetry::AntiSymmetric,
super::OperatorType::Normal,
n,
)
}
#[cfg(feature = "sparse")]
fn h_matrix(&self, n: usize) -> sprs::CsMat<Float> {
super::h_matrix(Self::DIAG, n, self.is_h2())
}
}
impl<SBP: SbpOperator1d> SbpOperator2d for (&SBP, &Upwind4) {

687
sparse/src/main.rs Normal file
View File

@ -0,0 +1,687 @@
use maxwell::Field;
use ndarray::Array2;
use sbp::{operators::SbpOperator1d, Float};
struct SparseMaxwellSystem {
x: Array2<Float>,
y: Array2<Float>,
rhs: sprs::CsMat<Float>,
lhs_implicit: Option<sprs::CsMat<Float>>,
now: Field,
next: Field,
k: [Field; 4],
}
impl SparseMaxwellSystem {
fn new(ny: usize, nx: usize) -> Self {
let x = ndarray::Array::from_shape_fn((ny, nx), |(_j, i)| {
i as Float * (1.0 / (nx - 1) as Float)
});
let y = ndarray::Array::from_shape_fn((ny, nx), |(j, _i)| {
j as Float * (1.0 / (nx - 1) as Float)
});
let rhs = Self::make_matrix(ny, nx);
let mut now = Field::new(ny, nx);
let mut next = Field::new(ny, nx);
let mut k = [now.clone(), now.clone(), now.clone(), now.clone()];
Self {
x,
y,
rhs,
now,
next,
k,
lhs_implicit: None,
}
}
fn nx(&self) -> usize {
self.x.shape()[1]
}
fn ny(&self) -> usize {
self.x.shape()[0]
}
fn max_dt(&self) -> Float {
1.0 / std::cmp::max(self.nx(), self.ny()) as Float
}
fn make_matrix(ny: usize, nx: usize) -> sprs::CsMat<Float> {
let d1_x = sbp::operators::Upwind4.diff_matrix(nx);
let d1_y = sbp::operators::Upwind4.diff_matrix(ny);
let ix = sprs::CsMat::<Float>::eye(nx);
let iy = sprs::CsMat::eye(ny);
let dx = sparse_sparse_outer_product(iy.view(), d1_x.view());
let dy = sparse_sparse_outer_product(d1_y.view(), ix.view());
let mut a_flux = sprs::TriMat::new((3, 3));
a_flux.add_triplet(1, 2, -1.0);
a_flux.add_triplet(2, 1, -1.0);
let a_flux = a_flux.to_csr();
let mut b_flux = sprs::TriMat::new((3, 3));
b_flux.add_triplet(0, 1, 1.0);
b_flux.add_triplet(1, 0, 1.0);
let b_flux = b_flux.to_csr();
let f = &sparse_sparse_outer_product(a_flux.view(), dx.view())
+ &sparse_sparse_outer_product(b_flux.view(), dy.view());
let mut hx =
sparse_sparse_outer_product(iy.view(), sbp::operators::Upwind4.h_matrix(nx).view());
hx.map_inplace(|h| 1.0 / h);
let ihx = hx;
let mut hy =
sparse_sparse_outer_product(sbp::operators::Upwind4.h_matrix(ny).view(), ix.view());
hy.map_inplace(|h| 1.0 / h);
let ihy = hy;
let f = {
// West boundary
let mut aminus = sprs::TriMat::new((3, 3));
aminus.add_triplet(1, 1, -0.5);
aminus.add_triplet(1, 2, -0.5);
aminus.add_triplet(2, 1, -0.5);
aminus.add_triplet(2, 2, -0.5);
let aminus = aminus.to_csr();
let mut e0x = sprs::TriMat::new((nx, 1));
e0x.add_triplet(0, 0, 1.0);
let e0x = e0x.to_csr();
let e0x_nt = &e0x * &e0x.transpose_view();
let e0x_nt = sparse_sparse_outer_product(iy.view(), e0x_nt.view());
let sat0 = &ihx * &e0x_nt;
let mut sat0 = sparse_sparse_outer_product(aminus.view(), sat0.view());
let tau = 1.0;
sat0.map_inplace(|x| tau * x);
&f + &sat0
};
let f = {
// East boundary
let mut aplus = sprs::TriMat::new((3, 3));
aplus.add_triplet(1, 1, 0.5);
aplus.add_triplet(1, 2, -0.5);
aplus.add_triplet(2, 1, -0.5);
aplus.add_triplet(2, 2, 0.5);
let aplus = aplus.to_csr();
let mut enx = sprs::TriMat::new((nx, 1));
enx.add_triplet(nx - 1, 0, 1.0);
let enx = enx.to_csr();
let enx_nt = &enx * &enx.transpose_view();
let enx_nt = sparse_sparse_outer_product(iy.view(), enx_nt.view());
let satn = &ihx * &enx_nt;
let mut satn = sparse_sparse_outer_product(aplus.view(), satn.view());
let tau = -1.0;
satn.map_inplace(|x| tau * x);
&f + &satn
};
let f = {
// South boundary
let mut bminus = sprs::TriMat::new((3, 3));
bminus.add_triplet(0, 0, -0.5);
bminus.add_triplet(0, 1, 0.5);
bminus.add_triplet(1, 0, 0.5);
bminus.add_triplet(1, 1, -0.5);
let bminus = bminus.to_csr();
let mut e0y = sprs::TriMat::new((ny, 1));
e0y.add_triplet(0, 0, 1.0);
let e0y = e0y.to_csr();
let e0y_nt = &e0y * &e0y.transpose_view();
let e0y_nt = sparse_sparse_outer_product(e0y_nt.view(), ix.view());
let sat0 = &ihx * &e0y_nt;
let mut sat0 = sparse_sparse_outer_product(bminus.view(), sat0.view());
let tau = 1.0;
sat0.map_inplace(|x| tau * x);
&f + &sat0
};
let f = {
// North boundary
let mut bplus = sprs::TriMat::new((3, 3));
bplus.add_triplet(0, 0, 0.5);
bplus.add_triplet(0, 1, 0.5);
bplus.add_triplet(1, 0, 0.5);
bplus.add_triplet(1, 1, 0.5);
let bplus = bplus.to_csr();
let mut eny = sprs::TriMat::new((ny, 1));
eny.add_triplet(ny - 1, 0, 1.0);
let eny = eny.to_csr();
let eny_nt = &eny * &eny.transpose_view();
let eny_nt = sparse_sparse_outer_product(eny_nt.view(), ix.view());
let satn = &ihy * &eny_nt;
let mut satn = sparse_sparse_outer_product(bplus.view(), satn.view());
let tau = -1.0;
satn.map_inplace(|x| tau * x);
&f + &satn
};
// Setting up the periodic boundaries
let f = {
// West
let mut aminus = sprs::TriMat::new((3, 3));
aminus.add_triplet(1, 1, -0.5);
aminus.add_triplet(1, 2, -0.5);
aminus.add_triplet(2, 1, -0.5);
aminus.add_triplet(2, 2, -0.5);
let aminus = aminus.to_csr();
let mut e0x = sprs::TriMat::new((nx, 1));
e0x.add_triplet(0, 0, 1.0);
let e0x = e0x.to_csr();
let mut enx = sprs::TriMat::new((nx, 1));
enx.add_triplet(nx - 1, 0, 1.0);
let enx = enx.to_csr();
let e0nx_nt = &e0x * &enx.transpose_view();
let e0nx_nt = sparse_sparse_outer_product(iy.view(), e0nx_nt.view());
let sat0 = &ihx * &e0nx_nt;
let mut sat0 = sparse_sparse_outer_product(aminus.view(), sat0.view());
let tau = 1.0;
// Negative => subtracting this boundary
sat0.map_inplace(|x| -tau * x);
&f + &sat0
};
let f = {
// East boundary
let mut aplus = sprs::TriMat::new((3, 3));
aplus.add_triplet(1, 1, 0.5);
aplus.add_triplet(1, 2, -0.5);
aplus.add_triplet(2, 1, -0.5);
aplus.add_triplet(2, 2, 0.5);
let aplus = aplus.to_csr();
let mut enx = sprs::TriMat::new((nx, 1));
enx.add_triplet(nx - 1, 0, 1.0);
let enx = enx.to_csr();
let mut e0x = sprs::TriMat::new((nx, 1));
e0x.add_triplet(0, 0, 1.0);
let e0x = e0x.to_csr();
let en0x_nt = &enx * &e0x.transpose_view();
let en0x_nt = sparse_sparse_outer_product(iy.view(), en0x_nt.view());
let satn = &ihx * &en0x_nt;
let mut satn = sparse_sparse_outer_product(aplus.view(), satn.view());
let tau = -1.0;
satn.map_inplace(|x| -tau * x);
&f + &satn
};
let f = {
// South boundary
let mut bminus = sprs::TriMat::new((3, 3));
bminus.add_triplet(0, 0, -0.5);
bminus.add_triplet(0, 1, 0.5);
bminus.add_triplet(1, 0, 0.5);
bminus.add_triplet(1, 1, -0.5);
let bminus = bminus.to_csr();
let mut e0y = sprs::TriMat::new((ny, 1));
e0y.add_triplet(0, 0, 1.0);
let e0y = e0y.to_csr();
let mut eny = sprs::TriMat::new((ny, 1));
eny.add_triplet(ny - 1, 0, 1.0);
let eny = eny.to_csr();
let e0ny_nt = &e0y * &eny.transpose_view();
let e0ny_nt = sparse_sparse_outer_product(e0ny_nt.view(), ix.view());
let sat0 = &ihx * &e0ny_nt;
let mut sat0 = sparse_sparse_outer_product(bminus.view(), sat0.view());
let tau = 1.0;
sat0.map_inplace(|x| -tau * x);
&f + &sat0
};
let f = {
// North boundary
let mut bplus = sprs::TriMat::new((3, 3));
bplus.add_triplet(0, 0, 0.5);
bplus.add_triplet(0, 1, 0.5);
bplus.add_triplet(1, 0, 0.5);
bplus.add_triplet(1, 1, 0.5);
let bplus = bplus.to_csr();
let mut eny = sprs::TriMat::new((ny, 1));
eny.add_triplet(ny - 1, 0, 1.0);
let eny = eny.to_csr();
let mut e0y = sprs::TriMat::new((ny, 1));
e0y.add_triplet(0, 0, 1.0);
let e0y = e0y.to_csr();
let en0y_nt = &eny * &e0y.transpose_view();
let en0y_nt = sparse_sparse_outer_product(en0y_nt.view(), ix.view());
let satn = &ihy * &en0y_nt;
let mut satn = sparse_sparse_outer_product(bplus.view(), satn.view());
let tau = -1.0;
satn.map_inplace(|x| -tau * x);
&f + &satn
};
f
}
fn advance(&mut self) {
let max_dt = self.max_dt();
let rhs = self.rhs.view();
let rhs_f = |next: &mut Field, now: &Field, _t: Float| {
next.fill(0.0);
sprs::prod::mul_acc_mat_vec_csr(
rhs,
now.as_slice().unwrap(),
next.as_slice_mut().unwrap(),
);
};
sbp::integrate::integrate::<sbp::integrate::Rk4, _, _>(
rhs_f,
&self.now,
&mut self.next,
&mut 0.0,
max_dt,
&mut self.k[..],
);
std::mem::swap(&mut self.now, &mut self.next);
}
fn advance_implicit(&mut self) {
if self.lhs_implicit.is_none() {
self.lhs_implicit = Some({
let i_kyx = sprs::CsMat::eye(3 * self.ny() * self.nx());
let dt = self.max_dt();
let f = self.rhs.map(|x| x * dt);
&i_kyx - &f
});
}
let b = self.now.clone();
let tnow = std::time::Instant::now();
jacobi_method(
self.lhs_implicit.as_ref().unwrap().view(),
b.as_slice().unwrap(),
self.now.as_slice_mut().unwrap(),
self.next.as_slice_mut().unwrap(),
10,
);
let elapsed = tnow.elapsed();
println!("{:?}", elapsed);
}
}
fn main() {
let nx = 64;
let ny = 64;
let mut sys = SparseMaxwellSystem::new(ny, nx);
let to_image = |mat: sprs::CsMatView<Float>, path: &str| {
let sparsity = sprs::visu::nnz_image(mat);
let im: image::ImageBuffer<image::Luma<u8>, _> = sparsity
.as_slice()
.map(|slice| {
image::ImageBuffer::from_raw((3 * nx * ny) as u32, (3 * nx * ny) as u32, slice)
.expect("failed to create image from slice")
})
.unwrap();
im.save(path).unwrap();
};
let tnow = std::time::Instant::now();
for _ in 0..100 {
sys.advance();
}
let elapsed = tnow.elapsed();
println!("{:?}", elapsed.div_f64(100.0));
}
/// A x = b
/// with A and b known
/// x should contain a first guess of
fn jacobi_method(
a: sprs::CsMatView<Float>,
b: &[Float],
x: &mut [Float],
tmp: &mut [Float],
iter_count: usize,
) {
for _ in 0..iter_count {
jacobi_step(a, b, x, tmp);
x.copy_from_slice(tmp);
}
}
fn jacobi_step(a: sprs::CsMatView<Float>, b: &[Float], x0: &[Float], x: &mut [Float]) {
let n = a.shape().0;
assert_eq!(n, a.shape().1);
let b = &b[..n];
let x0 = &x0[..n];
let x = &mut x[..n];
for (((i, ai), xi), &bi) in a
.outer_iterator()
.enumerate()
.zip(x.iter_mut())
.zip(b.iter())
{
let mut summa = 0.0;
let mut aii = None;
for (j, aij) in ai.iter() {
if i == j {
aii = Some(aij);
continue;
}
summa += aij * x0[j];
}
*xi = 1.0 / aii.unwrap() * (bi - summa);
}
}
#[test]
fn test_jacobi_2x2() {
let mut a = sprs::CsMat::zero((2, 2));
a.insert(0, 0, 2.0);
a.insert(0, 1, 1.0);
a.insert(1, 0, 5.0);
a.insert(1, 1, 7.0);
let b = ndarray::arr1(&[11.0, 13.0]);
let mut x0 = ndarray::arr1(&[1.0; 2]);
let mut tmp = x0.clone();
jacobi_method(
a.view(),
b.as_slice().unwrap(),
x0.as_slice_mut().unwrap(),
tmp.as_slice_mut().unwrap(),
25,
);
approx::assert_abs_diff_eq!(x0, ndarray::arr1(&[7.111, -3.222]), epsilon = 1e-2);
}
#[test]
fn test_jacobi_4x4() {
let mut a = sprs::CsMat::zero((4, 4));
a.insert(0, 0, 10.0);
a.insert(0, 1, -1.0);
a.insert(0, 2, 2.0);
a.insert(1, 0, -1.0);
a.insert(1, 1, 11.0);
a.insert(1, 2, -1.0);
a.insert(1, 3, 3.0);
a.insert(2, 0, 2.0);
a.insert(2, 1, -1.0);
a.insert(2, 2, 10.0);
a.insert(2, 3, -1.0);
a.insert(3, 1, 3.0);
a.insert(3, 2, -1.0);
a.insert(3, 3, 8.0);
let b = ndarray::arr1(&[6.0, 25.0, -11.0, 15.0]);
let mut x0 = ndarray::Array::zeros(b.len());
let mut tmp = x0.clone();
for iter in 0.. {
jacobi_step(
a.view(),
b.as_slice().unwrap(),
x0.as_slice().unwrap(),
tmp.as_slice_mut().unwrap(),
);
x0.as_slice_mut()
.unwrap()
.copy_from_slice(tmp.as_slice().unwrap());
match iter {
0 => approx::assert_abs_diff_eq!(
x0,
ndarray::arr1(&[0.6, 2.27272, -1.1, 1.875]),
epsilon = 1e-4
),
1 => approx::assert_abs_diff_eq!(
x0,
ndarray::arr1(&[1.04727, 1.7159, -0.80522, 0.88522]),
epsilon = 1e-4
),
2 => approx::assert_abs_diff_eq!(
x0,
ndarray::arr1(&[0.93263, 2.05330, -1.0493, 1.13088]),
epsilon = 1e-4
),
3 => approx::assert_abs_diff_eq!(
x0,
ndarray::arr1(&[1.01519, 1.95369, -0.9681, 0.97384]),
epsilon = 1e-4
),
4 => approx::assert_abs_diff_eq!(
x0,
ndarray::arr1(&[0.98899, 2.0114, -1.0102, 1.02135]),
epsilon = 1e-4
),
_ => break,
}
}
}
/// Computes the sparse kronecker product
/// M = A \kron B
#[allow(non_snake_case)]
#[must_use]
fn sparse_sparse_outer_product<
N: num_traits::Num + Copy + Default,
I: sprs::SpIndex,
Iptr: sprs::SpIndex,
>(
A: sprs::CsMatViewI<N, I, Iptr>,
B: sprs::CsMatViewI<N, I, Iptr>,
) -> sprs::CsMatI<N, I, Iptr> {
match (A.storage(), B.storage()) {
(sprs::CompressedStorage::CSR, sprs::CompressedStorage::CSR) => {
let nnz = A.nnz() * B.nnz();
let a_shape = A.shape();
let b_shape = B.shape();
let shape = (a_shape.0 * b_shape.0, a_shape.1 * b_shape.1);
let mut mat = sprs::CsMatI::zero(shape);
mat.reserve_nnz_exact(nnz);
for (aj, a) in A.outer_iterator().enumerate() {
for (bj, b) in B.outer_iterator().enumerate() {
for (ai, &a) in a.iter() {
for (bi, &b) in b.iter() {
let i = ai * b_shape.1 + bi;
let j = aj * b_shape.0 + bj;
mat.insert(j, i, a * b)
}
}
}
}
debug_assert_eq!(mat.nnz(), nnz);
mat
}
(sprs::CompressedStorage::CSC, sprs::CompressedStorage::CSC) => {
let nnz = A.nnz() * B.nnz();
let a_shape = A.shape();
let b_shape = B.shape();
let shape = (a_shape.0 * b_shape.0, a_shape.1 * b_shape.1);
let mat = sprs::CsMatI::zero(shape);
let mut mat = mat.to_csc();
for (ai, a) in A.outer_iterator().enumerate() {
for (bi, b) in B.outer_iterator().enumerate() {
for (aj, &a) in a.iter() {
for (bj, &b) in b.iter() {
let i = ai * b_shape.1 + bi;
let j = aj * b_shape.0 + bj;
mat.insert(j, i, a * b)
}
}
}
}
debug_assert_eq!(mat.nnz(), nnz);
mat
}
(sprs::CompressedStorage::CSR, sprs::CompressedStorage::CSC) => {
let nnz = A.nnz() * B.nnz();
let a_shape = A.shape();
let b_shape = B.shape();
let shape = (a_shape.0 * b_shape.0, a_shape.1 * b_shape.1);
let mut mat = sprs::CsMatI::zero(shape);
mat.reserve_nnz_exact(nnz);
for (aj, a) in A.outer_iterator().enumerate() {
for (bi, b) in B.outer_iterator().enumerate() {
for (ai, &a) in a.iter() {
for (bj, &b) in b.iter() {
let i = ai * b_shape.1 + bi;
let j = aj * b_shape.0 + bj;
mat.insert(j, i, a * b)
}
}
}
}
debug_assert_eq!(mat.nnz(), nnz);
mat
}
(sprs::CompressedStorage::CSC, sprs::CompressedStorage::CSR) => {
let nnz = A.nnz() * B.nnz();
let a_shape = A.shape();
let b_shape = B.shape();
let shape = (a_shape.0 * b_shape.0, a_shape.1 * b_shape.1);
let mat = sprs::CsMatI::zero(shape);
let mut mat = mat.to_csc();
for (aj, a) in A.outer_iterator().enumerate() {
for (bi, b) in B.outer_iterator().enumerate() {
for (ai, &a) in a.iter() {
for (bj, &b) in b.iter() {
let i = ai * b_shape.1 + bi;
let j = aj * b_shape.0 + bj;
mat.insert(j, i, a * b)
}
}
}
}
debug_assert_eq!(mat.nnz(), nnz);
mat
}
}
}
#[test]
fn test_outer_product() {
let mut a = sprs::TriMat::new((2, 3));
a.add_triplet(0, 1, 2);
a.add_triplet(0, 2, 3);
a.add_triplet(1, 0, 6);
a.add_triplet(1, 2, 8);
let a = a.to_csr();
let mut b = sprs::TriMat::new((3, 2));
b.add_triplet(0, 0, 1);
b.add_triplet(1, 0, 2);
b.add_triplet(2, 0, 3);
b.add_triplet(2, 1, -3);
let b = b.to_csr();
let c = sparse_sparse_outer_product(a.view(), b.view());
for (&n, (j, i)) in c.iter() {
match (j, i) {
(0, 2) => assert_eq!(n, 2),
(0, 4) => assert_eq!(n, 3),
(1, 2) => assert_eq!(n, 4),
(1, 4) => assert_eq!(n, 6),
(2, 2) => assert_eq!(n, 6),
(2, 3) => assert_eq!(n, -6),
(2, 4) => assert_eq!(n, 9),
(2, 5) => assert_eq!(n, -9),
(3, 0) => assert_eq!(n, 6),
(3, 4) => assert_eq!(n, 8),
(4, 0) => assert_eq!(n, 12),
(4, 4) => assert_eq!(n, 16),
(5, 0) => assert_eq!(n, 18),
(5, 1) => assert_eq!(n, -18),
(5, 4) => assert_eq!(n, 24),
(5, 5) => assert_eq!(n, -24),
_ => panic!("index ({},{}) should be 0, found {}", j, i, n),
}
}
}
#[test]
fn test_outer_product_csc() {
let mut a = sprs::TriMat::new((2, 3));
a.add_triplet(0, 1, 2);
a.add_triplet(0, 2, 3);
a.add_triplet(1, 0, 6);
a.add_triplet(1, 2, 8);
let a = a.to_csc();
let mut b = sprs::TriMat::new((3, 2));
b.add_triplet(0, 0, 1);
b.add_triplet(1, 0, 2);
b.add_triplet(2, 0, 3);
b.add_triplet(2, 1, -3);
let b = b.to_csc();
let c = sparse_sparse_outer_product(a.view(), b.view());
for (&n, (j, i)) in c.iter() {
match (j, i) {
(0, 2) => assert_eq!(n, 2),
(0, 4) => assert_eq!(n, 3),
(1, 2) => assert_eq!(n, 4),
(1, 4) => assert_eq!(n, 6),
(2, 2) => assert_eq!(n, 6),
(2, 3) => assert_eq!(n, -6),
(2, 4) => assert_eq!(n, 9),
(2, 5) => assert_eq!(n, -9),
(3, 0) => assert_eq!(n, 6),
(3, 4) => assert_eq!(n, 8),
(4, 0) => assert_eq!(n, 12),
(4, 4) => assert_eq!(n, 16),
(5, 0) => assert_eq!(n, 18),
(5, 1) => assert_eq!(n, -18),
(5, 4) => assert_eq!(n, 24),
(5, 5) => assert_eq!(n, -24),
_ => panic!("index ({},{}) should be 0, found {}", j, i, n),
}
}
}
#[test]
fn test_outer_product_2() {
let mut e0 = sprs::CsMat::zero((10, 1));
e0.insert(0, 0, 1);
let mut en = sprs::CsMat::zero((11, 1));
en.insert(10, 0, 1);
let v = sparse_sparse_outer_product(e0.view(), en.transpose_view());
for (&val, (j, i)) in v.iter() {
match (j, i) {
(0, 10) => assert_eq!(val, 1),
_ => panic!("Unexpected element: ({},{}): {}", j, i, val),
}
}
}