shap_rs/
evaluation.rs

1use crate::{coalition, Masker, Predict, Result, ShapError};
2use ndarray::{Array2, ArrayView1, Axis, Slice};
3use std::collections::{HashMap, HashSet, VecDeque};
4
5/// Shared limits for model-agnostic explainer evaluation.
6#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
7#[serde(try_from = "EvaluationConfigPayload")]
8pub struct EvaluationConfig {
9    /// Maximum coalitions combined into one model call.
10    pub coalition_batch_size: usize,
11    /// Maximum cached coalition results for one explained sample.
12    pub cache_capacity: usize,
13    /// Optional hard limit on total model rows evaluated per sample.
14    pub max_model_rows: Option<usize>,
15}
16#[derive(serde::Deserialize)]
17struct EvaluationConfigPayload {
18    coalition_batch_size: usize,
19    cache_capacity: usize,
20    max_model_rows: Option<usize>,
21}
22impl TryFrom<EvaluationConfigPayload> for EvaluationConfig {
23    type Error = ShapError;
24    fn try_from(payload: EvaluationConfigPayload) -> Result<Self> {
25        Self {
26            coalition_batch_size: payload.coalition_batch_size,
27            cache_capacity: payload.cache_capacity,
28            max_model_rows: payload.max_model_rows,
29        }
30        .validate()
31    }
32}
33impl Default for EvaluationConfig {
34    fn default() -> Self {
35        Self {
36            coalition_batch_size: 64,
37            cache_capacity: 4096,
38            max_model_rows: None,
39        }
40    }
41}
42impl EvaluationConfig {
43    pub fn validate(self) -> Result<Self> {
44        if self.coalition_batch_size == 0 || self.cache_capacity == 0 {
45            return Err(ShapError::InvalidConfiguration(
46                "coalition batch size and cache capacity must be positive".into(),
47            ));
48        }
49        if self.coalition_batch_size > self.cache_capacity {
50            return Err(ShapError::InvalidConfiguration(
51                "cache capacity must be at least the coalition batch size".into(),
52            ));
53        }
54        if self.max_model_rows == Some(0) {
55            return Err(ShapError::InvalidConfiguration(
56                "model row evaluation limit must be positive when configured".into(),
57            ));
58        }
59        Ok(self)
60    }
61}
62
63pub(crate) struct CoalitionEvaluator<'a, M, K> {
64    model: &'a M,
65    masker: &'a K,
66    config: EvaluationConfig,
67    cache: HashMap<u64, Vec<f64>>,
68    cache_order: VecDeque<u64>,
69    rows_evaluated: usize,
70    outputs: Option<usize>,
71}
72impl<'a, M: Predict, K: Masker> CoalitionEvaluator<'a, M, K> {
73    pub(crate) fn new(model: &'a M, masker: &'a K, config: EvaluationConfig) -> Result<Self> {
74        Ok(Self {
75            model,
76            masker,
77            config: config.validate()?,
78            cache: HashMap::new(),
79            cache_order: VecDeque::new(),
80            rows_evaluated: 0,
81            outputs: None,
82        })
83    }
84    pub(crate) fn evaluate(
85        &mut self,
86        sample: ArrayView1<'_, f64>,
87        masks: &[u64],
88    ) -> Result<Vec<Vec<f64>>> {
89        let mut results = HashMap::new();
90        let mut missing = Vec::new();
91        let mut missing_set = HashSet::new();
92        for &mask in masks {
93            if let Some(value) = self.cache.get(&mask).cloned() {
94                self.touch(mask);
95                results.insert(mask, value);
96            } else if missing_set.insert(mask) {
97                missing.push(mask);
98            }
99        }
100        for chunk in missing.chunks(self.config.coalition_batch_size) {
101            self.evaluate_chunk(sample, chunk)?;
102            for &mask in chunk {
103                let value =
104                    self.cache.get(&mask).cloned().ok_or_else(|| {
105                        ShapError::Other("coalition cache invariant failed".into())
106                    })?;
107                results.insert(mask, value);
108            }
109        }
110        masks
111            .iter()
112            .map(|m| {
113                results
114                    .get(m)
115                    .cloned()
116                    .ok_or_else(|| ShapError::Other("coalition cache invariant failed".into()))
117            })
118            .collect()
119    }
120    fn evaluate_chunk(&mut self, sample: ArrayView1<'_, f64>, masks: &[u64]) -> Result<()> {
121        if masks.is_empty() {
122            return Ok(());
123        }
124        if self.masker.streams_masked_batches() {
125            for &mask in masks {
126                self.evaluate_streaming_mask(sample, mask)?;
127            }
128            return Ok(());
129        }
130        let mut masked = Vec::with_capacity(masks.len());
131        let mut rows = 0usize;
132        for &mask in masks {
133            let part = self
134                .masker
135                .mask(sample, &coalition::members(mask, self.masker.n_features()))?;
136            if part.nrows() == 0 {
137                return Err(ShapError::MaskerError("masker returned no rows".into()));
138            }
139            rows = rows.checked_add(part.nrows()).ok_or_else(|| {
140                ShapError::InvalidConfiguration("coalition batch is too large".into())
141            })?;
142            masked.push(part)
143        }
144        if self
145            .config
146            .max_model_rows
147            .is_some_and(|limit| self.rows_evaluated.saturating_add(rows) > limit)
148        {
149            return Err(ShapError::InvalidConfiguration(
150                "model row evaluation limit exceeded".into(),
151            ));
152        }
153        crate::error::checked_f64_shape(
154            &[rows, self.masker.n_input_features()],
155            "masked coalition batch",
156        )?;
157        let mut batch = Array2::zeros((rows, self.masker.n_input_features()));
158        let mut offset = 0;
159        for part in &masked {
160            let end = offset + part.nrows();
161            batch
162                .slice_axis_mut(Axis(0), Slice::from(offset..end))
163                .assign(part);
164            offset = end
165        }
166        let predictions = self.model.predict_owned(batch)?;
167        if predictions.nrows() != rows || predictions.ncols() == 0 {
168            return Err(ShapError::DimensionMismatch {
169                expected: format!("({rows}, outputs>0)"),
170                found: format!("{:?}", predictions.dim()),
171            });
172        }
173        if self.outputs.is_some_and(|o| o != predictions.ncols()) {
174            return Err(ShapError::OutputDimensionMismatch {
175                expected: self.outputs.unwrap(),
176                found: predictions.ncols(),
177            });
178        }
179        if predictions.iter().any(|v| !v.is_finite()) {
180            return Err(ShapError::ModelError(
181                "prediction contains a non-finite value".into(),
182            ));
183        }
184        self.outputs = Some(predictions.ncols());
185        self.rows_evaluated += rows;
186        while masks.len() > self.config.cache_capacity.saturating_sub(self.cache.len()) {
187            let Some(key) = self.cache_order.pop_front() else {
188                break;
189            };
190            self.cache.remove(&key);
191        }
192        let mut offset = 0;
193        for (i, &mask) in masks.iter().enumerate() {
194            let end = offset + masked[i].nrows();
195            let value = predictions
196                .slice_axis(Axis(0), Slice::from(offset..end))
197                .mean_axis(Axis(0))
198                .unwrap()
199                .to_vec();
200            offset = end;
201            self.cache.insert(mask, value);
202            self.cache_order.push_back(mask);
203        }
204        Ok(())
205    }
206
207    fn evaluate_streaming_mask(&mut self, sample: ArrayView1<'_, f64>, mask: u64) -> Result<()> {
208        let members = coalition::members(mask, self.masker.n_features());
209        let model = self.model;
210        let starting_rows = self.rows_evaluated;
211        let row_limit = self.config.max_model_rows;
212        let mut rows = 0usize;
213        let mut outputs = self.outputs;
214        let mut sums = Vec::<f64>::new();
215        self.masker
216            .for_each_masked_batch(sample, &members, &mut |batch| {
217                let batch_rows = batch.nrows();
218                let next_rows = rows.checked_add(batch_rows).ok_or_else(|| {
219                    ShapError::InvalidConfiguration("streaming model row count overflow".into())
220                })?;
221                if row_limit.is_some_and(|limit| starting_rows.saturating_add(next_rows) > limit) {
222                    return Err(ShapError::InvalidConfiguration(
223                        "model row evaluation limit exceeded".into(),
224                    ));
225                }
226                let predictions = model.predict_owned(batch)?;
227                if predictions.nrows() != batch_rows || predictions.ncols() == 0 {
228                    return Err(ShapError::DimensionMismatch {
229                        expected: format!("({batch_rows}, outputs>0)"),
230                        found: format!("{:?}", predictions.dim()),
231                    });
232                }
233                if let Some(expected) = outputs {
234                    if expected != predictions.ncols() {
235                        return Err(ShapError::OutputDimensionMismatch {
236                            expected,
237                            found: predictions.ncols(),
238                        });
239                    }
240                } else {
241                    outputs = Some(predictions.ncols());
242                    sums.resize(predictions.ncols(), 0.0);
243                }
244                if predictions.iter().any(|value| !value.is_finite()) {
245                    return Err(ShapError::ModelError(
246                        "prediction contains a non-finite value".into(),
247                    ));
248                }
249                if sums.is_empty() {
250                    sums.resize(predictions.ncols(), 0.0);
251                }
252                for prediction in predictions.rows() {
253                    for (sum, value) in sums.iter_mut().zip(prediction) {
254                        *sum += *value;
255                    }
256                }
257                rows = next_rows;
258                Ok(())
259            })?;
260        if rows == 0 {
261            return Err(ShapError::MaskerError(
262                "streaming masker returned no rows".into(),
263            ));
264        }
265        let value = sums
266            .into_iter()
267            .map(|sum| sum / rows as f64)
268            .collect::<Vec<_>>();
269        if value.iter().any(|value| !value.is_finite()) {
270            return Err(ShapError::ModelError(
271                "streaming prediction mean is non-finite".into(),
272            ));
273        }
274        self.rows_evaluated = starting_rows.checked_add(rows).ok_or_else(|| {
275            ShapError::InvalidConfiguration("model row evaluation count overflow".into())
276        })?;
277        self.outputs = outputs;
278        while self.cache.len() >= self.config.cache_capacity {
279            let Some(key) = self.cache_order.pop_front() else {
280                break;
281            };
282            self.cache.remove(&key);
283        }
284        self.cache.insert(mask, value);
285        self.cache_order.push_back(mask);
286        Ok(())
287    }
288    fn touch(&mut self, mask: u64) {
289        if let Some(position) = self.cache_order.iter().position(|&key| key == mask) {
290            self.cache_order.remove(position);
291        }
292        self.cache_order.push_back(mask);
293    }
294    #[cfg(test)]
295    pub(crate) fn rows_evaluated(&self) -> usize {
296        self.rows_evaluated
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303    use crate::{Background, FnModel, FnStreamingMasker, IndependentMasker};
304    use ndarray::{array, Array2, ArrayView1, ArrayView2};
305    use std::cell::Cell;
306    #[test]
307    fn batches_and_deduplicates_coalitions() {
308        let calls = Cell::new(0);
309        let model = FnModel::new(|x: ArrayView2<'_, f64>| {
310            calls.set(calls.get() + 1);
311            Ok(Array2::from_shape_fn((x.nrows(), 1), |(i, _)| {
312                x.row(i).sum()
313            }))
314        });
315        let masker = IndependentMasker::new(Background::new(array![[0., 0.], [1., 1.]]).unwrap());
316        let config = EvaluationConfig {
317            coalition_batch_size: 8,
318            cache_capacity: 8,
319            max_model_rows: None,
320        };
321        let mut evaluator = CoalitionEvaluator::new(&model, &masker, config).unwrap();
322        let values = evaluator
323            .evaluate(array![2., 3.].view(), &[0, 1, 2, 3, 1])
324            .unwrap();
325        assert_eq!(calls.get(), 1);
326        assert_eq!(values[1], values[4]);
327        assert_eq!(evaluator.rows_evaluated(), 8);
328    }
329    #[test]
330    fn coalition_batch_uses_owned_prediction_fast_path_once() {
331        struct OwnedTrackingModel {
332            borrowed: Cell<usize>,
333            owned: Cell<usize>,
334        }
335        impl Predict for OwnedTrackingModel {
336            fn predict(&self, x: ArrayView2<'_, f64>) -> Result<Array2<f64>> {
337                self.borrowed.set(self.borrowed.get() + 1);
338                Ok(x.sum_axis(Axis(1)).insert_axis(Axis(1)))
339            }
340            fn predict_owned(&self, x: Array2<f64>) -> Result<Array2<f64>> {
341                self.owned.set(self.owned.get() + 1);
342                Ok(x.sum_axis(Axis(1)).insert_axis(Axis(1)))
343            }
344        }
345        let model = OwnedTrackingModel {
346            borrowed: Cell::new(0),
347            owned: Cell::new(0),
348        };
349        let masker = IndependentMasker::new(Background::new(array![[0., 0.], [1., 1.]]).unwrap());
350        let mut evaluator = CoalitionEvaluator::new(
351            &model,
352            &masker,
353            EvaluationConfig {
354                coalition_batch_size: 4,
355                cache_capacity: 4,
356                max_model_rows: None,
357            },
358        )
359        .unwrap();
360        evaluator
361            .evaluate(array![2., 3.].view(), &[0, 1, 2, 3])
362            .unwrap();
363        assert_eq!(model.owned.get(), 1);
364        assert_eq!(model.borrowed.get(), 0);
365    }
366    #[test]
367    fn bounded_cache_does_not_limit_request_size() {
368        let model = FnModel::new(|x: ArrayView2<'_, f64>| {
369            Ok(Array2::from_shape_fn((x.nrows(), 1), |(i, _)| {
370                x.row(i).sum()
371            }))
372        });
373        let masker = IndependentMasker::new(Background::new(array![[0., 0.]]).unwrap());
374        let mut evaluator = CoalitionEvaluator::new(
375            &model,
376            &masker,
377            EvaluationConfig {
378                coalition_batch_size: 2,
379                cache_capacity: 2,
380                max_model_rows: None,
381            },
382        )
383        .unwrap();
384        let values = evaluator
385            .evaluate(array![2., 3.].view(), &[0, 1, 2, 3])
386            .unwrap();
387        assert_eq!(values.len(), 4);
388        assert_eq!(evaluator.cache.len(), 2);
389    }
390    #[test]
391    fn cache_uses_deterministic_lru_eviction() {
392        let model =
393            FnModel::new(|x: ArrayView2<'_, f64>| Ok(x.sum_axis(Axis(1)).insert_axis(Axis(1))));
394        let masker = IndependentMasker::new(Background::new(array![[0., 0.]]).unwrap());
395        let mut evaluator = CoalitionEvaluator::new(
396            &model,
397            &masker,
398            EvaluationConfig {
399                coalition_batch_size: 1,
400                cache_capacity: 2,
401                max_model_rows: None,
402            },
403        )
404        .unwrap();
405        evaluator.evaluate(array![2., 3.].view(), &[0, 1]).unwrap();
406        evaluator.evaluate(array![2., 3.].view(), &[0]).unwrap();
407        evaluator.evaluate(array![2., 3.].view(), &[2]).unwrap();
408        assert!(evaluator.cache.contains_key(&0));
409        assert!(evaluator.cache.contains_key(&2));
410        assert!(!evaluator.cache.contains_key(&1));
411    }
412
413    #[test]
414    fn rejects_zero_model_row_budget() {
415        assert!(EvaluationConfig {
416            coalition_batch_size: 1,
417            cache_capacity: 1,
418            max_model_rows: Some(0),
419        }
420        .validate()
421        .is_err());
422    }
423
424    #[test]
425    fn consumes_streaming_masker_batches_without_collecting_background() {
426        let calls = Cell::new(0usize);
427        let model = FnModel::new(|x: ArrayView2<'_, f64>| {
428            calls.set(calls.get() + 1);
429            Ok(x.sum_axis(Axis(1)).insert_axis(Axis(1)))
430        });
431        let masker = FnStreamingMasker::new(
432            2,
433            |sample: ArrayView1<'_, f64>,
434             present: &[bool],
435             visitor: &mut dyn FnMut(Array2<f64>) -> Result<()>| {
436                for mut batch in [array![[0., 0.], [2., 2.]], array![[4., 4.]]] {
437                    for (feature, enabled) in present.iter().copied().enumerate() {
438                        if enabled {
439                            batch.column_mut(feature).fill(sample[feature]);
440                        }
441                    }
442                    visitor(batch)?;
443                }
444                Ok(())
445            },
446        )
447        .unwrap();
448        let mut evaluator = CoalitionEvaluator::new(
449            &model,
450            &masker,
451            EvaluationConfig {
452                coalition_batch_size: 4,
453                cache_capacity: 4,
454                max_model_rows: None,
455            },
456        )
457        .unwrap();
458        let values = evaluator
459            .evaluate(array![10., 20.].view(), &[0, 3])
460            .unwrap();
461        assert_eq!(values, vec![vec![4.], vec![30.]]);
462        assert_eq!(calls.get(), 4);
463        assert_eq!(evaluator.rows_evaluated(), 6);
464    }
465
466    #[test]
467    fn streaming_batches_respect_the_total_model_row_budget() {
468        let model = FnModel::new(|x: ArrayView2<'_, f64>| Ok(x.to_owned()));
469        let masker = FnStreamingMasker::new(
470            1,
471            |_: ArrayView1<'_, f64>,
472             _: &[bool],
473             visitor: &mut dyn FnMut(Array2<f64>) -> Result<()>| {
474                visitor(array![[0.]])?;
475                visitor(array![[1.]])
476            },
477        )
478        .unwrap();
479        let mut evaluator = CoalitionEvaluator::new(
480            &model,
481            &masker,
482            EvaluationConfig {
483                coalition_batch_size: 1,
484                cache_capacity: 1,
485                max_model_rows: Some(1),
486            },
487        )
488        .unwrap();
489        assert!(evaluator.evaluate(array![2.].view(), &[0]).is_err());
490    }
491}
shap_rs/evaluation.rs

shap_rs/
evaluation.rs
