rudrani-rane/ATIS
0
1"""
2ML Model Explainability Module
3Provides SHAP values, attention weights, and feature importance analysis for GNN predictions
4"""
5
6import torch
7import numpy as np
8from typing import Dict, List, Tuple, Optional
9import pandas as pd
10from torch_geometric.data import Data
11from pathlib import Path
12import sys
13
14# Add src to path for imports
15sys.path.append(str(Path(__file__).parent.parent.parent))
16
17from src.models.gnn_model import ATISGNN
18from src.risk.threat_engine import compute_threat_scores
19
20
21class ModelExplainer:
22 """Explains GNN model predictions using various techniques"""
23
24 def __init__(self, model_path: str = "outputs/best_model.pth", device: str = "cpu"):
25 self.device = device
26 self.model = None
27
28 # Load trained model
29 if not Path(model_path).exists():
30 print(f"ℹ️ No trained model at {model_path} — explainer running in fallback mode")
31 self.model = ATISGNN(in_channels=15).to(device)
32 return
33 try:
34 checkpoint = torch.load(model_path, map_location=device)
35 in_ch = checkpoint.get('in_channels', 15)
36 self.model = ATISGNN(
37 in_channels=in_ch,
38 hidden_channels=64,
39 latent_dim=32,
40 heads=4
41 ).to(device)
42 self.model.load_state_dict(checkpoint['model_state_dict'], strict=False)
43 self.model.eval()
44 print(f"✓ Loaded trained GNN model from {model_path}")
45 except Exception as e:
46 print(f"⚠️ Could not load model weights from {model_path}: {e}")
47 print(" Explainer running with untrained model weights")
48 if self.model is None:
49 self.model = ATISGNN(in_channels=15).to(device)
50
51 def extract_attention_weights(self, graph_data: Data) -> Dict[str, List]:
52 """
53 Extract simplified attention statistics
54 Returns basic layer information
55 """
56 # Simplified version to avoid slow computation
57 return {
58 'layer_1': [0.5], # Placeholder attention weights
59 'layer_2': [0.3]
60 }
61
62 def _approximate_attention(self, x_in: torch.Tensor, x_out: torch.Tensor,
63 edge_index: torch.Tensor) -> np.ndarray:
64 """Approximate attention weights from input/output changes"""
65 # Calculate magnitude of change for each node
66 change = torch.norm(x_out - x_in, dim=1)
67
68 # Edge attention is the average of source and target node changes
69 src, dst = edge_index
70 edge_attention = (change[src] + change[dst]) / 2
71
72 # Normalize to [0, 1]
73 edge_attention = edge_attention / (edge_attention.max() + 1e-8)
74
75 return edge_attention.cpu().numpy()
76
77 def compute_feature_importance(self, graph_data: Data,
78 target_node: int) -> Dict[str, float]:
79 """
80 Compute feature importance using gradient-based attribution
81 Higher values indicate features that strongly influence the prediction
82 """
83 self.model.eval()
84
85 from src.models.gnn_model import ORBITAL_FEATURE_START
86 # Enable gradient computation for orbital-only input (strips neo/pha to match training)
87 x = graph_data.x[:, ORBITAL_FEATURE_START:].clone().detach().requires_grad_(True).to(self.device)
88 edge_index = graph_data.edge_index.to(self.device)
89
90 # Forward pass
91 mu, sigma, _pha_logit = self.model(x, edge_index)
92
93 # Compute threat score (need gradients)
94 latent_risk = torch.norm(mu, dim=1)
95 target_output = latent_risk[target_node]
96
97 # Backward pass
98 target_output.backward()
99
100 # Feature importance is the gradient magnitude
101 importance = torch.abs(x.grad[target_node]).cpu().numpy()
102
103 # Map to feature names
104 feature_names = [
105 'eccentricity', 'semi_major_axis', 'inclination',
106 'longitude_ascending', 'argument_perihelion', 'mean_anomaly',
107 'perihelion_distance', 'aphelion_distance', 'orbital_period',
108 'mean_motion', 'absolute_magnitude', 'diameter'
109 ]
110
111 importance_dict = {
112 name: float(imp) for name, imp in zip(feature_names, importance)
113 }
114
115 # Normalize to percentages
116 total = sum(importance_dict.values())
117 if total > 0:
118 importance_dict = {k: v/total * 100 for k, v in importance_dict.items()}
119
120 return importance_dict
121
122 def compute_shap_values(self, graph_data: Data, target_node: int,
123 num_samples: int = 10) -> Dict[str, float]:
124 """
125 Compute SHAP-like values using permutation importance
126 Reduced sample count for faster computation
127 """
128 self.model.eval()
129
130 from src.models.gnn_model import ORBITAL_FEATURE_START
131 with torch.no_grad():
132 # Get baseline prediction
133 x = graph_data.x[:, ORBITAL_FEATURE_START:].to(self.device)
134 edge_index = graph_data.edge_index.to(self.device)
135
136 mu, sigma, _pha_logit = self.model(x, edge_index)
137 threat_scores = compute_threat_scores(mu, sigma, graph_data)
138 baseline_pred = float(threat_scores[target_node].item())
139
140 feature_names = [
141 'eccentricity', 'semi_major_axis', 'inclination',
142 'longitude_ascending', 'argument_perihelion', 'mean_anomaly',
143 'perihelion_distance', 'aphelion_distance', 'orbital_period',
144 'mean_motion', 'absolute_magnitude', 'diameter'
145 ]
146
147 shap_values = {}
148
149 # For each feature, compute impact of masking it
150 for i, feature_name in enumerate(feature_names):
151 impacts = []
152
153 for _ in range(num_samples):
154 # Create perturbed input by replacing feature with random value
155 x_perturbed = x.clone()
156
157 # Random perturbation from same distribution
158 random_idx = np.random.randint(0, x.shape[0])
159 x_perturbed[target_node, i] = x[random_idx, i]
160
161 # Prediction with perturbed feature
162 mu_p, sigma_p, _pha_p = self.model(x_perturbed, edge_index)
163 threat_scores_p = compute_threat_scores(mu_p, sigma_p, graph_data)
164 perturbed_pred = float(threat_scores_p[target_node].item())
165
166 # Impact is the difference
167 impacts.append(baseline_pred - perturbed_pred)
168
169 # SHAP value is the average impact
170 shap_values[feature_name] = np.mean(impacts)
171
172 return shap_values
173
174 def explain_prediction(self, graph_data: Data, target_node: int,
175 asteroid_id: str = None) -> Dict:
176 """
177 Generate comprehensive explanation for a prediction
178 Combines multiple explainability techniques
179 """
180 # Get prediction
181 from src.models.gnn_model import ORBITAL_FEATURE_START
182 self.model.eval()
183 with torch.no_grad():
184 x = graph_data.x[:, ORBITAL_FEATURE_START:].to(self.device)
185 edge_index = graph_data.edge_index.to(self.device)
186
187 mu, sigma, _pha_logit = self.model(x, edge_index)
188
189 # Compute threat score for this asteroid
190 threat_scores = compute_threat_scores(mu, sigma, graph_data)
191 prediction = float(threat_scores[target_node].item())
192
193 # Get various explanations
194 feature_importance = self.compute_feature_importance(graph_data, target_node)
195 shap_values = self.compute_shap_values(graph_data, target_node, num_samples=5) # Reduced samples
196 attention_weights = self.extract_attention_weights(graph_data)
197
198 # Identify most influential features
199 sorted_importance = sorted(feature_importance.items(), key=lambda x: x[1], reverse=True)
200 top_features = sorted_importance[:5]
201
202 # Generate human-readable explanation
203 explanation_text = self._generate_explanation_text(
204 prediction, top_features, shap_values
205 )
206
207 return {
208 'asteroid_id': asteroid_id,
209 'prediction': float(prediction),
210 'prediction_label': 'High Threat' if prediction > 0.7 else 'Medium Threat' if prediction > 0.4 else 'Low Threat',
211 'confidence': self._calculate_confidence(prediction),
212 'feature_importance': feature_importance,
213 'shap_values': shap_values,
214 'attention_weights': attention_weights, # Already in list format
215 'top_influential_features': [
216 {'feature': name, 'importance': value} for name, value in top_features
217 ],
218 'explanation_text': explanation_text
219 }
220
221 def _calculate_confidence(self, prediction: float) -> float:
222 """Calculate prediction confidence based on distance from decision boundary"""
223 # Confidence is higher when prediction is far from 0.5
224 distance_from_boundary = abs(prediction - 0.5)
225 confidence = min(1.0, distance_from_boundary * 2)
226 return float(confidence)
227
228 def _generate_explanation_text(self, prediction: float,
229 top_features: List[Tuple[str, float]],
230 shap_values: Dict[str, float]) -> str:
231 """Generate human-readable explanation"""
232 threat_level = 'high' if prediction > 0.7 else 'medium' if prediction > 0.4 else 'low'
233
234 explanation = f"This asteroid has a {threat_level} threat score of {prediction:.2%}. "
235
236 # Identify key factors
237 positive_factors = [f for f, v in shap_values.items() if v > 0]
238 negative_factors = [f for f, v in shap_values.items() if v < 0]
239
240 if positive_factors:
241 top_positive = sorted([(f, shap_values[f]) for f in positive_factors],
242 key=lambda x: x[1], reverse=True)[:3]
243 explanation += "The threat score is primarily driven by: "
244 explanation += ", ".join([f.replace('_', ' ') for f, _ in top_positive])
245 explanation += ". "
246
247 if negative_factors:
248 top_negative = sorted([(f, shap_values[f]) for f in negative_factors],
249 key=lambda x: x[1])[:2]
250 explanation += "Factors reducing the threat include: "
251 explanation += ", ".join([f.replace('_', ' ') for f, _ in top_negative])
252 explanation += ". "
253
254 # Add context from top features
255 if top_features:
256 explanation += f"The most influential factor is {top_features[0][0].replace('_', ' ')} "
257 explanation += f"({top_features[0][1]:.1f}% of prediction influence)."
258
259 return explanation
260
261
262# Singleton instance
263explainer = ModelExplainer()
264
265
266def get_explanation(graph_data: Data, target_node: int,
267 asteroid_id: str = None) -> Dict:
268 """Convenience function to get prediction explanation"""
269 return explainer.explain_prediction(graph_data, target_node, asteroid_id)
270