/home/noah/src/realizar/src/explain.rs
Line | Count | Source |
1 | | //! Model Explainability (SHAP, LIME, Attention) |
2 | | //! |
3 | | //! Per spec §13: Model explainability for APR classical ML models. |
4 | | //! Implements SHAP TreeExplainer for tree ensembles and KernelSHAP for any model. |
5 | | //! |
6 | | //! ## Methods |
7 | | //! |
8 | | //! | Method | Type | Models | Output | |
9 | | //! |--------|------|--------|--------| |
10 | | //! | **TreeSHAP** | Model-specific | Tree ensembles | Feature contributions | |
11 | | //! | **KernelSHAP** | Model-agnostic | Any | Feature contributions | |
12 | | //! |
13 | | //! ## References |
14 | | //! |
15 | | //! - [16] Lundberg & Lee (2017) "A Unified Approach to Interpreting Model Predictions" |
16 | | //! - [17] Ribeiro et al. (2016) "Why Should I Trust You? Explaining Predictions" |
17 | | |
18 | | // Module-level clippy allows for explainability module |
19 | | #![allow(clippy::must_use_candidate)] |
20 | | #![allow(clippy::return_self_not_must_use)] |
21 | | #![allow(clippy::missing_errors_doc)] |
22 | | #![allow(clippy::unused_self)] // Methods designed for future expansion |
23 | | #![allow(clippy::unnecessary_wraps)] // Result wrapping for API consistency |
24 | | #![allow(clippy::option_if_let_else)] // map_or is more readable for our use case |
25 | | |
26 | | use serde::{Deserialize, Serialize}; |
27 | | use std::fmt; |
28 | | use thiserror::Error; |
29 | | |
30 | | /// Error type for explainability operations |
31 | | #[derive(Debug, Error)] |
32 | | pub enum ExplainError { |
33 | | /// Model does not support explainability |
34 | | #[error("Model does not support explainability: {reason}")] |
35 | | UnsupportedModel { |
36 | | /// Why the model is unsupported |
37 | | reason: String, |
38 | | }, |
39 | | |
40 | | /// Invalid input dimensions |
41 | | #[error("Invalid input: expected {expected} features, got {actual}")] |
42 | | InvalidInput { |
43 | | /// Expected number of features |
44 | | expected: usize, |
45 | | /// Actual number of features provided |
46 | | actual: usize, |
47 | | }, |
48 | | |
49 | | /// Background dataset required but not provided |
50 | | #[error("Background dataset required for KernelSHAP")] |
51 | | NoBackground, |
52 | | |
53 | | /// Computation error |
54 | | #[error("Computation error: {0}")] |
55 | | ComputationError(String), |
56 | | } |
57 | | |
58 | | /// Trait for models that can be explained |
59 | | pub trait Explainable { |
60 | | /// Predict for a single instance |
61 | | fn predict(&self, instance: &[f32]) -> Result<f32, ExplainError>; |
62 | | |
63 | | /// Predict for multiple instances (batch) |
64 | 1 | fn predict_batch(&self, instances: &[Vec<f32>]) -> Result<Vec<f32>, ExplainError> { |
65 | 2 | instances1 .iter1 ().map1 (|x| self.predict(x)).collect1 () |
66 | 1 | } |
67 | | |
68 | | /// Number of features expected |
69 | | fn n_features(&self) -> usize; |
70 | | |
71 | | /// Check if this is a tree-based model |
72 | 2 | fn is_tree_model(&self) -> bool { |
73 | 2 | false |
74 | 2 | } |
75 | | |
76 | | /// Get tree structure for TreeSHAP (if applicable) |
77 | 0 | fn get_tree_structure(&self) -> Option<&TreeStructure> { |
78 | 0 | None |
79 | 0 | } |
80 | | } |
81 | | |
82 | | /// Tree structure for TreeSHAP algorithm |
83 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
84 | | pub struct TreeStructure { |
85 | | /// Number of trees in the ensemble |
86 | | pub n_trees: usize, |
87 | | /// Number of features |
88 | | pub n_features: usize, |
89 | | /// Trees in the ensemble |
90 | | pub trees: Vec<DecisionTree>, |
91 | | } |
92 | | |
93 | | /// A single decision tree |
94 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
95 | | pub struct DecisionTree { |
96 | | /// Feature index used at each node (-1 for leaf) |
97 | | pub feature: Vec<i32>, |
98 | | /// Threshold at each node |
99 | | pub threshold: Vec<f32>, |
100 | | /// Left child index |
101 | | pub left: Vec<usize>, |
102 | | /// Right child index |
103 | | pub right: Vec<usize>, |
104 | | /// Value at leaf nodes |
105 | | pub value: Vec<f32>, |
106 | | } |
107 | | |
108 | | impl DecisionTree { |
109 | | /// Create a new decision tree |
110 | 6 | pub fn new( |
111 | 6 | feature: Vec<i32>, |
112 | 6 | threshold: Vec<f32>, |
113 | 6 | left: Vec<usize>, |
114 | 6 | right: Vec<usize>, |
115 | 6 | value: Vec<f32>, |
116 | 6 | ) -> Self { |
117 | 6 | Self { |
118 | 6 | feature, |
119 | 6 | threshold, |
120 | 6 | left, |
121 | 6 | right, |
122 | 6 | value, |
123 | 6 | } |
124 | 6 | } |
125 | | |
126 | | /// Get the number of nodes in the tree |
127 | 1 | pub fn n_nodes(&self) -> usize { |
128 | 1 | self.feature.len() |
129 | 1 | } |
130 | | |
131 | | /// Check if a node is a leaf |
132 | 18 | pub fn is_leaf(&self, node: usize) -> bool { |
133 | 18 | self.feature.get(node).is_none_or(|&f| f < 0) |
134 | 18 | } |
135 | | |
136 | | /// Predict for a single instance |
137 | 8 | pub fn predict(&self, instance: &[f32]) -> f32 { |
138 | 8 | let mut node = 0; |
139 | 15 | while !self.is_leaf(node) { |
140 | 7 | let feature_idx = self.feature[node] as usize; |
141 | 7 | if instance |
142 | 7 | .get(feature_idx) |
143 | 7 | .is_some_and(|&v| v <= self.threshold[node]) |
144 | 5 | { |
145 | 5 | node = self.left[node]; |
146 | 5 | } else { |
147 | 2 | node = self.right[node]; |
148 | 2 | } |
149 | | } |
150 | 8 | self.value.get(node).copied().unwrap_or(0.0) |
151 | 8 | } |
152 | | } |
153 | | |
154 | | /// SHAP explanation result |
155 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
156 | | pub struct ShapExplanation { |
157 | | /// Expected model output E[f(X)] |
158 | | pub base_value: f32, |
159 | | /// SHAP values for each feature (φᵢ) |
160 | | /// sum(shap_values) + base_value ≈ prediction |
161 | | pub shap_values: Vec<f32>, |
162 | | /// Feature names for display |
163 | | pub feature_names: Vec<String>, |
164 | | /// The actual prediction for this instance |
165 | | pub prediction: f32, |
166 | | } |
167 | | |
168 | | impl ShapExplanation { |
169 | | /// Create a new SHAP explanation |
170 | 10 | pub fn new(base_value: f32, shap_values: Vec<f32>, prediction: f32) -> Self { |
171 | 10 | let n = shap_values.len(); |
172 | | Self { |
173 | 10 | base_value, |
174 | 10 | shap_values, |
175 | 19 | feature_names: (0..n)10 .map10 (|i| format!("feature_{i}")).collect10 (), |
176 | 10 | prediction, |
177 | | } |
178 | 10 | } |
179 | | |
180 | | /// Set feature names |
181 | 4 | pub fn with_feature_names(mut self, names: Vec<String>) -> Self { |
182 | 4 | self.feature_names = names; |
183 | 4 | self |
184 | 4 | } |
185 | | |
186 | | /// Get the most important features (sorted by absolute SHAP value) |
187 | 3 | pub fn top_features(&self, n: usize) -> Vec<(String, f32)> { |
188 | 3 | let mut indexed: Vec<_> = self |
189 | 3 | .shap_values |
190 | 3 | .iter() |
191 | 3 | .enumerate() |
192 | 5 | .map3 (|(i, &v)| (i, v)) |
193 | 3 | .collect(); |
194 | 4 | indexed3 .sort_by3 (|a, b| { |
195 | 4 | b.1.abs() |
196 | 4 | .partial_cmp(&a.1.abs()) |
197 | 4 | .unwrap_or(std::cmp::Ordering::Equal) |
198 | 4 | }); |
199 | 3 | indexed |
200 | 3 | .into_iter() |
201 | 3 | .take(n) |
202 | 4 | .map3 (|(i, v)| { |
203 | 4 | let name = self |
204 | 4 | .feature_names |
205 | 4 | .get(i) |
206 | 4 | .cloned() |
207 | 4 | .unwrap_or_else(|| format!0 ("feature_{i}"0 )); |
208 | 4 | (name, v) |
209 | 4 | }) |
210 | 3 | .collect() |
211 | 3 | } |
212 | | |
213 | | /// Verify SHAP consistency: sum(shap_values) + base_value ≈ prediction |
214 | 3 | pub fn verify_consistency(&self, tolerance: f32) -> bool { |
215 | 3 | let sum: f32 = self.shap_values.iter().sum(); |
216 | 3 | (self.base_value + sum - self.prediction).abs() < tolerance |
217 | 3 | } |
218 | | } |
219 | | |
220 | | impl fmt::Display for ShapExplanation { |
221 | 1 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
222 | 1 | writeln!(f, "SHAP Explanation:")?0 ; |
223 | 1 | writeln!(f, " Base value: {:.4}", self.base_value)?0 ; |
224 | 1 | writeln!(f, " Prediction: {:.4}", self.prediction)?0 ; |
225 | 1 | writeln!(f, " Top features:")?0 ; |
226 | 2 | for (name, value) in self1 .top_features1 (5) { |
227 | 2 | let sign = if value >= 0.0 { "+"1 } else { ""1 }; |
228 | 2 | writeln!(f, " {name}: {sign}{value:.4}")?0 ; |
229 | | } |
230 | 1 | Ok(()) |
231 | 1 | } |
232 | | } |
233 | | |
234 | | /// SHAP explainer for APR classical ML models |
235 | | /// Reference: [16] Lundberg & Lee (2017) SHAP |
236 | | pub struct ShapExplainer { |
237 | | /// Background dataset for computing expected values |
238 | | background: Vec<Vec<f32>>, |
239 | | /// Number of samples for KernelSHAP |
240 | | nsamples: usize, |
241 | | /// Feature names |
242 | | feature_names: Vec<String>, |
243 | | } |
244 | | |
245 | | impl ShapExplainer { |
246 | | /// Create a new SHAP explainer with background data |
247 | 7 | pub fn new(background: Vec<Vec<f32>>) -> Self { |
248 | 7 | let n_features = background.first().map_or(0, Vec::len); |
249 | | Self { |
250 | 7 | background, |
251 | | nsamples: 100, |
252 | 10 | feature_names: (0..n_features)7 .map7 (|i| format!("feature_{i}")).collect7 (), |
253 | | } |
254 | 7 | } |
255 | | |
256 | | /// Set the number of samples for KernelSHAP |
257 | 1 | pub fn with_nsamples(mut self, nsamples: usize) -> Self { |
258 | 1 | self.nsamples = nsamples; |
259 | 1 | self |
260 | 1 | } |
261 | | |
262 | | /// Set feature names |
263 | 1 | pub fn with_feature_names(mut self, names: Vec<String>) -> Self { |
264 | 1 | self.feature_names = names; |
265 | 1 | self |
266 | 1 | } |
267 | | |
268 | | /// Compute SHAP values for a prediction |
269 | 4 | pub fn explain( |
270 | 4 | &self, |
271 | 4 | model: &dyn Explainable, |
272 | 4 | instance: &[f32], |
273 | 4 | ) -> Result<ShapExplanation, ExplainError> { |
274 | | // Validate input |
275 | 4 | if instance.len() != model.n_features() { |
276 | 1 | return Err(ExplainError::InvalidInput { |
277 | 1 | expected: model.n_features(), |
278 | 1 | actual: instance.len(), |
279 | 1 | }); |
280 | 3 | } |
281 | | |
282 | | // Use TreeSHAP for tree-based models (fast, exact) |
283 | 3 | if model.is_tree_model() { |
284 | 1 | if let Some(tree_structure) = model.get_tree_structure() { |
285 | 1 | return self.tree_shap(tree_structure, instance, model); |
286 | 0 | } |
287 | 2 | } |
288 | | |
289 | | // KernelSHAP for other models (model-agnostic) |
290 | 2 | self.kernel_shap(model, instance) |
291 | 4 | } |
292 | | |
293 | | /// TreeSHAP algorithm for tree-based models |
294 | | /// O(TLD) complexity where T=trees, L=leaves, D=depth |
295 | 1 | fn tree_shap( |
296 | 1 | &self, |
297 | 1 | tree_structure: &TreeStructure, |
298 | 1 | instance: &[f32], |
299 | 1 | model: &dyn Explainable, |
300 | 1 | ) -> Result<ShapExplanation, ExplainError> { |
301 | 1 | let n_features = tree_structure.n_features; |
302 | 1 | let mut shap_values = vec![0.0; n_features]; |
303 | | |
304 | | // Compute SHAP values for each tree and average |
305 | 2 | for tree1 in &tree_structure.trees { |
306 | 1 | let tree_shap = self.tree_shap_single(tree, instance)?0 ; |
307 | 1 | for (i, v) in tree_shap.iter().enumerate() { |
308 | 1 | shap_values[i] += v / tree_structure.n_trees as f32; |
309 | 1 | } |
310 | | } |
311 | | |
312 | | // Compute base value from background |
313 | 1 | let base_value = self.compute_expected_value(model)?0 ; |
314 | 1 | let prediction = model.predict(instance)?0 ; |
315 | | |
316 | 1 | Ok(ShapExplanation::new(base_value, shap_values, prediction) |
317 | 1 | .with_feature_names(self.feature_names.clone())) |
318 | 1 | } |
319 | | |
320 | | /// TreeSHAP for a single tree |
321 | 1 | fn tree_shap_single( |
322 | 1 | &self, |
323 | 1 | tree: &DecisionTree, |
324 | 1 | instance: &[f32], |
325 | 1 | ) -> Result<Vec<f32>, ExplainError> { |
326 | 1 | let n_features = instance.len(); |
327 | 1 | let mut shap_values = vec![0.0; n_features]; |
328 | | |
329 | | // Simplified TreeSHAP using path enumeration |
330 | | // For each feature, compute its marginal contribution |
331 | 1 | for feature_idx in 0..n_features { |
332 | | // Compute prediction with feature |
333 | 1 | let pred_with = tree.predict(instance); |
334 | | |
335 | | // Compute prediction without feature (using background mean) |
336 | 1 | let mut instance_without = instance.to_vec(); |
337 | 1 | let background_mean = self |
338 | 1 | .background |
339 | 1 | .iter() |
340 | 2 | .filter_map1 (|bg| bg.get(feature_idx).copied()) |
341 | 1 | .sum::<f32>() |
342 | 1 | / self.background.len().max(1) as f32; |
343 | 1 | instance_without[feature_idx] = background_mean; |
344 | 1 | let pred_without = tree.predict(&instance_without); |
345 | | |
346 | | // Marginal contribution (simplified) |
347 | 1 | shap_values[feature_idx] = pred_with - pred_without; |
348 | | } |
349 | | |
350 | 1 | Ok(shap_values) |
351 | 1 | } |
352 | | |
353 | | /// KernelSHAP algorithm for model-agnostic explainability |
354 | 2 | fn kernel_shap( |
355 | 2 | &self, |
356 | 2 | model: &dyn Explainable, |
357 | 2 | instance: &[f32], |
358 | 2 | ) -> Result<ShapExplanation, ExplainError> { |
359 | 2 | if self.background.is_empty() { |
360 | 1 | return Err(ExplainError::NoBackground); |
361 | 1 | } |
362 | | |
363 | 1 | let n_features = instance.len(); |
364 | 1 | let mut shap_values = vec![0.0; n_features]; |
365 | | |
366 | | // KernelSHAP: sample coalitions and compute weighted linear regression |
367 | 1 | for _ in 0..self.nsamples { |
368 | | // Sample a random coalition (subset of features) |
369 | 100 | let coalition = self.sample_coalition(n_features); |
370 | 100 | let coalition_size = coalition.iter().filter(|&&b| b).count(); |
371 | | |
372 | | // Skip empty and full coalitions |
373 | 100 | if coalition_size == 0 || coalition_size == n_features { |
374 | 0 | continue; |
375 | 100 | } |
376 | | |
377 | | // Compute marginal contribution |
378 | 100 | let marginal = self.compute_marginal(model, instance, &coalition)?0 ; |
379 | | |
380 | | // SHAP kernel weight: M / ((M choose |S|) * |S| * (M - |S|)) |
381 | | // where M = n_features, S = coalition |
382 | 100 | let weight = self.shap_kernel_weight(n_features, coalition_size); |
383 | | |
384 | | // Update SHAP values |
385 | 200 | for (i, &in_coalition) in coalition.iter()100 .enumerate100 () { |
386 | 200 | if in_coalition { |
387 | 100 | shap_values[i] += marginal * weight; |
388 | 100 | } |
389 | | } |
390 | | } |
391 | | |
392 | | // Normalize by total weight |
393 | 1 | let total_weight: f32 = (1..n_features) |
394 | 1 | .map(|k| self.shap_kernel_weight(n_features, k)) |
395 | 1 | .sum(); |
396 | 1 | if total_weight > 0.0 { |
397 | 3 | for v2 in &mut shap_values { |
398 | 2 | *v /= total_weight; |
399 | 2 | } |
400 | 0 | } |
401 | | |
402 | | // Compute base value and prediction |
403 | 1 | let base_value = self.compute_expected_value(model)?0 ; |
404 | 1 | let prediction = model.predict(instance)?0 ; |
405 | | |
406 | 1 | Ok(ShapExplanation::new(base_value, shap_values, prediction) |
407 | 1 | .with_feature_names(self.feature_names.clone())) |
408 | 2 | } |
409 | | |
410 | | /// Sample a random coalition (subset of features) |
411 | 100 | fn sample_coalition(&self, n_features: usize) -> Vec<bool> { |
412 | | // Use deterministic sampling for reproducibility in tests |
413 | | // In production, use thread_rng() instead |
414 | 200 | (0..n_features)100 .map100 (|i| i % 2 == 0).collect100 () |
415 | 100 | } |
416 | | |
417 | | /// Compute marginal contribution for a coalition |
418 | 100 | fn compute_marginal( |
419 | 100 | &self, |
420 | 100 | model: &dyn Explainable, |
421 | 100 | instance: &[f32], |
422 | 100 | coalition: &[bool], |
423 | 100 | ) -> Result<f32, ExplainError> { |
424 | | // Create masked instance: use instance values for coalition, background mean for others |
425 | 100 | let mut total_pred = 0.0; |
426 | 100 | let n_background = self.background.len().max(1); |
427 | | |
428 | 300 | for bg200 in &self.background { |
429 | 200 | let mut masked: Vec<f32> = Vec::with_capacity(instance.len()); |
430 | 400 | for (i, (&inst_val, &in_coalition)) in instance200 .iter200 ().zip200 (coalition200 .iter200 ()).enumerate200 () |
431 | | { |
432 | 400 | if in_coalition { |
433 | 200 | masked.push(inst_val); |
434 | 200 | } else { |
435 | 200 | masked.push(bg.get(i).copied().unwrap_or(0.0)); |
436 | 200 | } |
437 | | } |
438 | 200 | total_pred += model.predict(&masked)?0 ; |
439 | | } |
440 | | |
441 | 100 | Ok(total_pred / n_background as f32) |
442 | 100 | } |
443 | | |
444 | | /// Compute expected value E[f(X)] from background |
445 | 2 | fn compute_expected_value(&self, model: &dyn Explainable) -> Result<f32, ExplainError> { |
446 | 2 | if self.background.is_empty() { |
447 | 0 | return Ok(0.0); |
448 | 2 | } |
449 | | |
450 | 2 | let predictions: Result<Vec<f32>, _> = |
451 | 4 | self.background.iter()2 .map2 (|x| model.predict(x)).collect2 (); |
452 | 2 | let predictions = predictions?0 ; |
453 | | |
454 | 2 | Ok(predictions.iter().sum::<f32>() / predictions.len() as f32) |
455 | 2 | } |
456 | | |
457 | | /// SHAP kernel weight |
458 | 101 | fn shap_kernel_weight(&self, n_features: usize, coalition_size: usize) -> f32 { |
459 | | // Weight = M / (binom(M, |S|) * |S| * (M - |S|)) |
460 | 101 | let m = n_features as f32; |
461 | 101 | let s = coalition_size as f32; |
462 | 101 | let binom = binomial(n_features, coalition_size) as f32; |
463 | 101 | if binom * s * (m - s) == 0.0 { |
464 | 0 | 0.0 |
465 | | } else { |
466 | 101 | m / (binom * s * (m - s)) |
467 | | } |
468 | 101 | } |
469 | | } |
470 | | |
471 | | /// Compute binomial coefficient (n choose k) |
472 | 112 | fn binomial(n: usize, k: usize) -> usize { |
473 | 112 | if k > n { |
474 | 1 | return 0; |
475 | 111 | } |
476 | 111 | if k == 0 || k == n108 { |
477 | 5 | return 1; |
478 | 106 | } |
479 | 106 | let k = k.min(n - k); // Take advantage of symmetry |
480 | 106 | let mut result = 1usize; |
481 | 112 | for i in 0..k106 { |
482 | 112 | result = result.saturating_mul(n - i) / (i + 1); |
483 | 112 | } |
484 | 106 | result |
485 | 112 | } |
486 | | |
487 | | // ============================================================================ |
488 | | // Tests |
489 | | // ============================================================================ |
490 | | |
491 | | #[cfg(test)] |
492 | | mod tests { |
493 | | use super::*; |
494 | | |
495 | | /// Simple linear model for testing |
496 | | struct LinearModel { |
497 | | weights: Vec<f32>, |
498 | | bias: f32, |
499 | | } |
500 | | |
501 | | impl LinearModel { |
502 | 4 | fn new(weights: Vec<f32>, bias: f32) -> Self { |
503 | 4 | Self { weights, bias } |
504 | 4 | } |
505 | | } |
506 | | |
507 | | impl Explainable for LinearModel { |
508 | 205 | fn predict(&self, instance: &[f32]) -> Result<f32, ExplainError> { |
509 | 205 | if instance.len() != self.weights.len() { |
510 | 0 | return Err(ExplainError::InvalidInput { |
511 | 0 | expected: self.weights.len(), |
512 | 0 | actual: instance.len(), |
513 | 0 | }); |
514 | 205 | } |
515 | 410 | let dot205 : f32205 = instance205 .iter205 ().zip205 (&self.weights205 ).map205 (|(x, w)| x * w).sum205 (); |
516 | 205 | Ok(dot + self.bias) |
517 | 205 | } |
518 | | |
519 | 4 | fn n_features(&self) -> usize { |
520 | 4 | self.weights.len() |
521 | 4 | } |
522 | | } |
523 | | |
524 | | /// Simple tree model for testing |
525 | | struct SimpleTreeModel { |
526 | | structure: TreeStructure, |
527 | | } |
528 | | |
529 | | impl SimpleTreeModel { |
530 | 1 | fn new(tree: DecisionTree) -> Self { |
531 | 1 | let n_features = tree |
532 | 1 | .feature |
533 | 1 | .iter() |
534 | 3 | .filter1 (|&&f| f >= 0) |
535 | 1 | .map(|&f| f as usize + 1) |
536 | 1 | .max() |
537 | 1 | .unwrap_or(1); |
538 | 1 | Self { |
539 | 1 | structure: TreeStructure { |
540 | 1 | n_trees: 1, |
541 | 1 | n_features, |
542 | 1 | trees: vec![tree], |
543 | 1 | }, |
544 | 1 | } |
545 | 1 | } |
546 | | } |
547 | | |
548 | | impl Explainable for SimpleTreeModel { |
549 | 3 | fn predict(&self, instance: &[f32]) -> Result<f32, ExplainError> { |
550 | 3 | let sum: f32 = self |
551 | 3 | .structure |
552 | 3 | .trees |
553 | 3 | .iter() |
554 | 3 | .map(|t| t.predict(instance)) |
555 | 3 | .sum(); |
556 | 3 | Ok(sum / self.structure.n_trees as f32) |
557 | 3 | } |
558 | | |
559 | 1 | fn n_features(&self) -> usize { |
560 | 1 | self.structure.n_features |
561 | 1 | } |
562 | | |
563 | 1 | fn is_tree_model(&self) -> bool { |
564 | 1 | true |
565 | 1 | } |
566 | | |
567 | 1 | fn get_tree_structure(&self) -> Option<&TreeStructure> { |
568 | 1 | Some(&self.structure) |
569 | 1 | } |
570 | | } |
571 | | |
572 | | // === ShapExplanation Tests === |
573 | | |
574 | | #[test] |
575 | 1 | fn test_shap_explanation_new() { |
576 | 1 | let exp = ShapExplanation::new(0.5, vec![0.1, -0.2, 0.3], 0.7); |
577 | 1 | assert_eq!(exp.base_value, 0.5); |
578 | 1 | assert_eq!(exp.shap_values.len(), 3); |
579 | 1 | assert_eq!(exp.prediction, 0.7); |
580 | 1 | assert_eq!(exp.feature_names.len(), 3); |
581 | 1 | } |
582 | | |
583 | | #[test] |
584 | 1 | fn test_shap_explanation_with_feature_names() { |
585 | 1 | let exp = ShapExplanation::new(0.5, vec![0.1, -0.2], 0.4) |
586 | 1 | .with_feature_names(vec!["age".to_string(), "income".to_string()]); |
587 | 1 | assert_eq!(exp.feature_names, vec!["age", "income"]); |
588 | 1 | } |
589 | | |
590 | | #[test] |
591 | 1 | fn test_shap_explanation_top_features() { |
592 | 1 | let exp = ShapExplanation::new(0.0, vec![0.1, -0.3, 0.2], 0.0).with_feature_names(vec![ |
593 | 1 | "a".to_string(), |
594 | 1 | "b".to_string(), |
595 | 1 | "c".to_string(), |
596 | | ]); |
597 | | |
598 | 1 | let top = exp.top_features(2); |
599 | 1 | assert_eq!(top.len(), 2); |
600 | | // Should be sorted by absolute value |
601 | 1 | assert_eq!(top[0].0, "b"); // |-0.3| = 0.3 |
602 | 1 | assert_eq!(top[1].0, "c"); // |0.2| = 0.2 |
603 | 1 | } |
604 | | |
605 | | #[test] |
606 | 1 | fn test_shap_explanation_verify_consistency() { |
607 | | // Perfect consistency |
608 | 1 | let exp = ShapExplanation::new(0.5, vec![0.2, 0.3], 1.0); |
609 | 1 | assert!(exp.verify_consistency(0.01)); |
610 | | |
611 | | // Not consistent |
612 | 1 | let exp_bad = ShapExplanation::new(0.5, vec![0.2, 0.3], 2.0); |
613 | 1 | assert!(!exp_bad.verify_consistency(0.01)); |
614 | 1 | } |
615 | | |
616 | | #[test] |
617 | 1 | fn test_shap_explanation_display() { |
618 | 1 | let exp = ShapExplanation::new(0.5, vec![0.1, -0.2], 0.4); |
619 | 1 | let display = format!("{exp}"); |
620 | 1 | assert!(display.contains("SHAP Explanation")); |
621 | 1 | assert!(display.contains("Base value")); |
622 | 1 | assert!(display.contains("Prediction")); |
623 | 1 | } |
624 | | |
625 | | // === Decision Tree Tests === |
626 | | |
627 | | #[test] |
628 | 1 | fn test_decision_tree_predict_simple() { |
629 | | // Simple tree: if feature_0 <= 0.5 then 1.0 else 2.0 |
630 | 1 | let tree = DecisionTree::new( |
631 | 1 | vec![0, -1, -1], // feature indices (-1 = leaf) |
632 | 1 | vec![0.5, 0.0, 0.0], // thresholds |
633 | 1 | vec![1, 0, 0], // left children |
634 | 1 | vec![2, 0, 0], // right children |
635 | 1 | vec![0.0, 1.0, 2.0], // values |
636 | | ); |
637 | | |
638 | 1 | assert_eq!(tree.predict(&[0.3]), 1.0); // <= 0.5, go left |
639 | 1 | assert_eq!(tree.predict(&[0.7]), 2.0); // > 0.5, go right |
640 | 1 | } |
641 | | |
642 | | #[test] |
643 | 1 | fn test_decision_tree_n_nodes() { |
644 | 1 | let tree = DecisionTree::new( |
645 | 1 | vec![0, -1, -1], |
646 | 1 | vec![0.5, 0.0, 0.0], |
647 | 1 | vec![1, 0, 0], |
648 | 1 | vec![2, 0, 0], |
649 | 1 | vec![0.0, 1.0, 2.0], |
650 | | ); |
651 | 1 | assert_eq!(tree.n_nodes(), 3); |
652 | 1 | } |
653 | | |
654 | | #[test] |
655 | 1 | fn test_decision_tree_is_leaf() { |
656 | 1 | let tree = DecisionTree::new( |
657 | 1 | vec![0, -1, -1], |
658 | 1 | vec![0.5, 0.0, 0.0], |
659 | 1 | vec![1, 0, 0], |
660 | 1 | vec![2, 0, 0], |
661 | 1 | vec![0.0, 1.0, 2.0], |
662 | | ); |
663 | 1 | assert!(!tree.is_leaf(0)); // root is not leaf |
664 | 1 | assert!(tree.is_leaf(1)); // leaf |
665 | 1 | assert!(tree.is_leaf(2)); // leaf |
666 | 1 | } |
667 | | |
668 | | // === ShapExplainer Tests === |
669 | | |
670 | | #[test] |
671 | 1 | fn test_shap_explainer_new() { |
672 | 1 | let background = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; |
673 | 1 | let explainer = ShapExplainer::new(background); |
674 | 1 | assert_eq!(explainer.nsamples, 100); |
675 | 1 | assert_eq!(explainer.feature_names.len(), 2); |
676 | 1 | } |
677 | | |
678 | | #[test] |
679 | 1 | fn test_shap_explainer_with_nsamples() { |
680 | 1 | let explainer = ShapExplainer::new(vec![vec![1.0]]).with_nsamples(50); |
681 | 1 | assert_eq!(explainer.nsamples, 50); |
682 | 1 | } |
683 | | |
684 | | #[test] |
685 | 1 | fn test_shap_explainer_with_feature_names() { |
686 | 1 | let explainer = ShapExplainer::new(vec![vec![1.0, 2.0]]) |
687 | 1 | .with_feature_names(vec!["age".to_string(), "income".to_string()]); |
688 | 1 | assert_eq!(explainer.feature_names, vec!["age", "income"]); |
689 | 1 | } |
690 | | |
691 | | #[test] |
692 | 1 | fn test_shap_explainer_linear_model() { |
693 | 1 | let model = LinearModel::new(vec![1.0, 2.0], 0.0); |
694 | 1 | let background = vec![vec![0.0, 0.0], vec![1.0, 1.0]]; |
695 | 1 | let explainer = ShapExplainer::new(background); |
696 | | |
697 | 1 | let instance = vec![2.0, 3.0]; // prediction = 2*1 + 3*2 = 8 |
698 | 1 | let explanation = explainer.explain(&model, &instance).expect("test"); |
699 | | |
700 | 1 | assert_eq!(explanation.prediction, 8.0); |
701 | 1 | assert_eq!(explanation.shap_values.len(), 2); |
702 | 1 | } |
703 | | |
704 | | #[test] |
705 | 1 | fn test_shap_explainer_tree_model() { |
706 | | // Simple tree: if feature_0 <= 0.5 then 1.0 else 2.0 |
707 | 1 | let tree = DecisionTree::new( |
708 | 1 | vec![0, -1, -1], |
709 | 1 | vec![0.5, 0.0, 0.0], |
710 | 1 | vec![1, 0, 0], |
711 | 1 | vec![2, 0, 0], |
712 | 1 | vec![0.0, 1.0, 2.0], |
713 | | ); |
714 | 1 | let model = SimpleTreeModel::new(tree); |
715 | | |
716 | 1 | let background = vec![vec![0.3], vec![0.7]]; |
717 | 1 | let explainer = ShapExplainer::new(background); |
718 | | |
719 | 1 | let explanation = explainer.explain(&model, &[0.3]).expect("test"); |
720 | 1 | assert_eq!(explanation.prediction, 1.0); |
721 | 1 | } |
722 | | |
723 | | #[test] |
724 | 1 | fn test_shap_explainer_invalid_input() { |
725 | 1 | let model = LinearModel::new(vec![1.0, 2.0], 0.0); |
726 | 1 | let background = vec![vec![0.0, 0.0]]; |
727 | 1 | let explainer = ShapExplainer::new(background); |
728 | | |
729 | 1 | let result = explainer.explain(&model, &[1.0, 2.0, 3.0]); // 3 features, model expects 2 |
730 | 1 | assert!(matches!0 (result, Err(ExplainError::InvalidInput { .. }))); |
731 | 1 | } |
732 | | |
733 | | #[test] |
734 | 1 | fn test_shap_explainer_empty_background() { |
735 | 1 | let model = LinearModel::new(vec![1.0], 0.0); |
736 | 1 | let explainer = ShapExplainer::new(vec![]); // Empty background |
737 | | |
738 | 1 | let result = explainer.explain(&model, &[1.0]); |
739 | 1 | assert!(matches!0 (result, Err(ExplainError::NoBackground))); |
740 | 1 | } |
741 | | |
742 | | // === ExplainError Tests === |
743 | | |
744 | | #[test] |
745 | 1 | fn test_explain_error_display() { |
746 | 1 | let err = ExplainError::UnsupportedModel { |
747 | 1 | reason: "not a tree".to_string(), |
748 | 1 | }; |
749 | 1 | assert!(err.to_string().contains("not a tree")); |
750 | | |
751 | 1 | let err = ExplainError::InvalidInput { |
752 | 1 | expected: 3, |
753 | 1 | actual: 2, |
754 | 1 | }; |
755 | 1 | assert!(err.to_string().contains("expected 3")); |
756 | 1 | assert!(err.to_string().contains("got 2")); |
757 | | |
758 | 1 | let err = ExplainError::NoBackground; |
759 | 1 | assert!(err.to_string().contains("Background")); |
760 | | |
761 | 1 | let err = ExplainError::ComputationError("overflow".to_string()); |
762 | 1 | assert!(err.to_string().contains("overflow")); |
763 | 1 | } |
764 | | |
765 | | // === Binomial Tests === |
766 | | |
767 | | #[test] |
768 | 1 | fn test_binomial_coefficient() { |
769 | 1 | assert_eq!(binomial(5, 0), 1); |
770 | 1 | assert_eq!(binomial(5, 1), 5); |
771 | 1 | assert_eq!(binomial(5, 2), 10); |
772 | 1 | assert_eq!(binomial(5, 3), 10); |
773 | 1 | assert_eq!(binomial(5, 4), 5); |
774 | 1 | assert_eq!(binomial(5, 5), 1); |
775 | 1 | assert_eq!(binomial(5, 6), 0); // k > n |
776 | 1 | } |
777 | | |
778 | | #[test] |
779 | 1 | fn test_binomial_edge_cases() { |
780 | 1 | assert_eq!(binomial(0, 0), 1); |
781 | 1 | assert_eq!(binomial(1, 0), 1); |
782 | 1 | assert_eq!(binomial(1, 1), 1); |
783 | 1 | assert_eq!(binomial(10, 5), 252); |
784 | 1 | } |
785 | | |
786 | | // === TreeStructure Tests === |
787 | | |
788 | | #[test] |
789 | 1 | fn test_tree_structure_serialization() { |
790 | 1 | let tree = DecisionTree::new( |
791 | 1 | vec![0, -1, -1], |
792 | 1 | vec![0.5, 0.0, 0.0], |
793 | 1 | vec![1, 0, 0], |
794 | 1 | vec![2, 0, 0], |
795 | 1 | vec![0.0, 1.0, 2.0], |
796 | | ); |
797 | 1 | let structure = TreeStructure { |
798 | 1 | n_trees: 1, |
799 | 1 | n_features: 1, |
800 | 1 | trees: vec![tree], |
801 | 1 | }; |
802 | | |
803 | 1 | let json = serde_json::to_string(&structure).expect("test"); |
804 | 1 | assert!(json.contains("n_trees")); |
805 | 1 | assert!(json.contains("n_features")); |
806 | 1 | } |
807 | | |
808 | | #[test] |
809 | 1 | fn test_shap_explanation_serialization() { |
810 | 1 | let exp = ShapExplanation::new(0.5, vec![0.1, -0.2], 0.4); |
811 | 1 | let json = serde_json::to_string(&exp).expect("test"); |
812 | 1 | let parsed: ShapExplanation = serde_json::from_str(&json).expect("test"); |
813 | | |
814 | 1 | assert_eq!(parsed.base_value, exp.base_value); |
815 | 1 | assert_eq!(parsed.shap_values, exp.shap_values); |
816 | 1 | assert_eq!(parsed.prediction, exp.prediction); |
817 | 1 | } |
818 | | |
819 | | // === Additional Edge Case Tests === |
820 | | |
821 | | #[test] |
822 | 1 | fn test_shap_explanation_empty_values() { |
823 | 1 | let exp = ShapExplanation::new(0.5, vec![], 0.5); |
824 | 1 | assert!(exp.verify_consistency(0.01)); |
825 | 1 | assert!(exp.top_features(3).is_empty()); |
826 | 1 | } |
827 | | |
828 | | #[test] |
829 | 1 | fn test_decision_tree_empty_instance() { |
830 | 1 | let tree = DecisionTree::new( |
831 | 1 | vec![-1], // Just a leaf |
832 | 1 | vec![0.0], |
833 | 1 | vec![0], |
834 | 1 | vec![0], |
835 | 1 | vec![5.0], |
836 | | ); |
837 | 1 | assert_eq!(tree.predict(&[]), 5.0); |
838 | 1 | } |
839 | | |
840 | | #[test] |
841 | 1 | fn test_linear_model_batch_predict() { |
842 | 1 | let model = LinearModel::new(vec![1.0, 2.0], 0.0); |
843 | 1 | let instances = vec![vec![1.0, 1.0], vec![2.0, 2.0]]; |
844 | 1 | let predictions = model.predict_batch(&instances).expect("test"); |
845 | 1 | assert_eq!(predictions, vec![3.0, 6.0]); |
846 | 1 | } |
847 | | } |