use std::collections::HashMap;
use ndarray::{Array1, Array2, ArrayBase, AsArray, Axis, Dimension, ViewRepr};
use rayon::prelude::*;
use crate::prelude::*;
#[inline]
pub fn roi_cloud_map<'a, A, D>(labels: A, threads: Option<usize>) -> HashMap<u64, Array2<usize>>
where
A: AsArray<'a, u64, D>,
D: Dimension,
{
let data: ArrayBase<ViewRepr<&'a u64>, D> = labels.into();
let vec_to_arr = |k: u64, v: Vec<Vec<usize>>| {
let arr = Array2::from_shape_vec((v.len(), v[0].len()), v.into_iter().flatten().collect())
.expect("Failed to reshape ROI point cloud into an Array2<usize>.");
(k, arr)
};
let labels_to_map_seq = || {
let mut cloud_map: HashMap<u64, Vec<Vec<usize>>> = HashMap::new();
data.view()
.into_dyn()
.indexed_iter()
.filter(|&(_, &v)| v != 0)
.for_each(|(p, &v)| {
cloud_map
.entry(v)
.or_default()
.push(p.as_array_view().to_vec());
});
cloud_map
};
let labels_to_map_par = || {
data.view()
.into_dyn()
.indexed_iter()
.par_bridge()
.filter(|&(_, &v)| v != 0)
.fold(
HashMap::new,
|mut map: HashMap<u64, Vec<Vec<usize>>>, (p, &v)| {
map.entry(v).or_default().push(p.as_array_view().to_vec());
map
},
)
.reduce(HashMap::new, |mut map_a, map_b| {
map_b.into_iter().for_each(|(k, mut v)| {
map_a.entry(k).or_insert_with(Vec::new).append(&mut v);
});
map_a
})
};
let cloud_map = par!(threads,
seq_exp: labels_to_map_seq(),
par_exp: labels_to_map_par());
cloud_map
.into_par_iter()
.map(|(k, v)| vec_to_arr(k, v))
.collect()
}
#[inline]
pub fn roi_data_map<'a, T, A, B, D>(
data: A,
labels: B,
threads: Option<usize>,
) -> Result<HashMap<u64, Array1<T>>, ImgalError>
where
A: AsArray<'a, T, D>,
B: AsArray<'a, u64, D>,
D: Dimension,
T: 'a + AsNumeric,
{
let data: ArrayBase<ViewRepr<&'a T>, D> = data.into();
let labels: ArrayBase<ViewRepr<&'a u64>, D> = labels.into();
if data.shape() != labels.shape() {
return Err(ImgalError::MismatchedArrayShapes {
a_arr_name: "data",
a_shape: data.shape().to_vec(),
b_arr_name: "labels",
b_shape: labels.shape().to_vec(),
});
}
let data = data.into_dyn();
let rcm = roi_cloud_map(labels, threads);
let mut rdm: HashMap<u64, Array1<T>> = HashMap::new();
rcm.iter().for_each(|(&k, c)| {
let cloud_lns = c.lanes(Axis(1));
let cloud_data = cloud_lns
.into_iter()
.map(|l| match l.as_slice() {
Some(s) => *data.get(s).unwrap(),
None => {
let coords = l.to_vec();
*data.get(coords.as_slice()).unwrap()
}
})
.collect::<Vec<T>>();
rdm.insert(k, Array1::from_vec(cloud_data));
});
Ok(rdm)
}