1use crate::{Explanation, Result, ShapError};
2use ndarray::Array2;
3#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
4pub struct HeatmapData {
5 pub values: Array2<f64>,
6 pub feature_order: Vec<usize>,
7}
8pub fn data(e: &Explanation, output: usize) -> Result<HeatmapData> {
9 if output >= e.n_outputs() {
10 return Err(ShapError::InvalidOutputIndex {
11 index: output,
12 n_outputs: e.n_outputs(),
13 });
14 }
15 let order = crate::plot::bar::data(e)
16 .into_iter()
17 .map(|x| x.0)
18 .collect::<Vec<_>>();
19 let values = Array2::from_shape_fn((e.n_samples(), e.n_features()), |(i, j)| {
20 e.values()[[i, order[j], output]]
21 });
22 Ok(HeatmapData {
23 values,
24 feature_order: order,
25 })
26}