Skip to main content

shap_rs/plot/
heatmap.rs

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}