fn main() {
use mpi::traits::Root;
use mpi::{self, traits::Communicator};
use rlst::dense::array::DynArray;
use rlst::{
self, Shape,
io::matrix_market::{read_array_mm, read_coordinate_mm},
};
use rlst::{AijIteratorByValue, AsMatrixApply, FromAijDistributed, assert_array_relative_eq};
use std::path::PathBuf;
use std::rc::Rc;
fn relative_file(filename: &str) -> String {
let file = PathBuf::from(file!());
let dir = file.parent().unwrap();
format!("{}/{filename}", dir.display())
}
let universe = mpi::initialize().unwrap();
let world = universe.world();
let rank = world.rank();
let (dist_mat, dist_x) = {
if rank == 0 {
println!("Loading CSR matrix example on process {rank}");
let sparse_mat = read_coordinate_mm::<f64>(&relative_file("mat_507_313.mm")).unwrap();
let mut shape = sparse_mat.shape();
world.this_process().broadcast_into(shape.as_mut_slice());
let domain_layout = Rc::new(
rlst::distributed_tools::IndexLayout::from_equidistributed_chunks(
shape[1], 1, &world,
),
);
let range_layout = Rc::new(
rlst::distributed_tools::IndexLayout::from_equidistributed_chunks(
shape[0], 1, &world,
),
);
let dist_mat = rlst::sparse::distributed_csr_mat::DistributedCsrMatrix::from_aij_iter(
domain_layout.clone(),
range_layout.clone(),
sparse_mat.iter_aij_value(),
);
let x = read_array_mm::<f64>(&relative_file("x_313.mm")).unwrap();
let dist_x = x.slice(1, 0).scatter_from_one_root(domain_layout.clone());
(dist_mat, dist_x)
} else {
let mut shape = [0, 0];
world
.process_at_rank(0)
.broadcast_into(shape.as_mut_slice());
let domain_layout = Rc::new(
rlst::distributed_tools::IndexLayout::from_equidistributed_chunks(
shape[1], 1, &world,
),
);
let range_layout = Rc::new(
rlst::distributed_tools::IndexLayout::from_equidistributed_chunks(
shape[0], 1, &world,
),
);
let dist_mat =
rlst::sparse::distributed_csr_mat::DistributedCsrMatrix::<'_, f64, _>::from_aij(
domain_layout.clone(),
range_layout.clone(),
Default::default(),
Default::default(),
Default::default(),
);
let dist_x = DynArray::<f64, 1>::scatter_from_one(0, domain_layout.clone());
(dist_mat, dist_x)
}
};
let mut dist_y = rlst::dist_vec!(f64, dist_mat.range_layout().clone());
dist_mat.apply(1.0, &dist_x, 0.0, &mut dist_y);
if rank == 0 {
let result = dist_y.gather_to_one_root();
let y_expected = read_array_mm::<f64>(&relative_file("/y_507.mm"))
.unwrap()
.slice(1, 0);
assert_array_relative_eq!(result, y_expected, 1E-10);
println!("Test passed.")
} else {
dist_y.gather_to_one(0);
}
}