fix: SHAP explain - torch import, reduce samples, fix visualizer scale

explain/shap_wrapper.py

@@ -12,21 +12,17 @@ class GNNFeatureWrapper:
 
     def __call__(self, binary_coalitions: np.ndarray) -> np.ndarray:
         predictions = []
+        x = self.data_obj.x.to(self.device)
+        edge_index = self.data_obj.edge_index.to(self.device)
+        edge_attr = self.data_obj.edge_attr.to(self.device)
+        batch = torch.zeros(x.shape[0], dtype=torch.long, device=self.device)
+
         with torch.no_grad():
             for coalition in binary_coalitions:
-                x_perturbed = self.data_obj.x.clone().to(self.device)
-                mask_tensor = torch.tensor(
-                    coalition, dtype=torch.float32, device=self.device
-                ).unsqueeze(1)
-                x_perturbed = x_perturbed * mask_tensor
-
-                logits = self.model(
-                    x_perturbed,
-                    self.data_obj.edge_index.to(self.device),
-                    self.data_obj.edge_attr.to(self.device),
-                    torch.zeros(x_perturbed.shape[0], dtype=torch.long, device=self.device),
-                )
-                probs = torch.sigmoid(logits)
-                predictions.append(probs[0, self.target_task_idx].cpu().item())
+                mask = torch.tensor(coalition, dtype=torch.float32, device=self.device).unsqueeze(1)
+                x_perturbed = x * mask
+                logits = self.model(x_perturbed, edge_index, edge_attr, batch)
+                prob = torch.sigmoid(logits[0, self.target_task_idx]).cpu().item()
+                predictions.append(prob)
 
         return np.array(predictions)

explain/visualizer.py

@@ -13,12 +13,18 @@ def generate_similarity_map(smiles: str, attributions: list, target_assay: str)
     if mol is None:
         raise ValueError(f"Invalid SMILES: {smiles}")
 
+    if mol.GetNumAtoms() != len(attributions):
+        raise ValueError(
+            f"Atom count mismatch: mol has {mol.GetNumAtoms()} atoms, "
+            f"but {len(attributions)} attributions provided"
+        )
+
     fig, ax = plt.subplots(figsize=(8, 6))
     SimilarityMaps.GetSimilarityMapFromWeights(
         mol,
         attributions,
         colorMap='RdYlBu_r',
-        scale=-1,
+        scale=1,
         alpha=0.45,
     )
     ax.set_title(f"SHAP Attributions — {target_assay}", fontsize=14, fontweight="bold")

routes/explain.py

@@ -1,4 +1,5 @@
 import numpy as np
+import torch
 from fastapi import APIRouter, Header, HTTPException
 from pydantic import BaseModel
 from firebase_auth import verify_token
@@ -43,20 +44,23 @@ async def explain(request: ExplainRequest, authorization: str = Header(...)):
         raise HTTPException(status_code=400, detail=str(e))
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    import torch
     data = data.to(device)
 
     num_nodes = data.x.shape[0]
     predict_fn = GNNFeatureWrapper(model, data, request.target_task, device)
-    background_reference = np.zeros((100, num_nodes))
+
+    random_states = np.random.randint(0, 2, size=(20, num_nodes))
+    background_reference = np.vstack([np.zeros((1, num_nodes)), random_states])
     active_state = np.ones((1, num_nodes))
 
     import shap
     explainer = shap.KernelExplainer(predict_fn, background_reference)
-    shap_values = explainer.shap_values(active_state, nsamples=min(1000, 2 ** num_nodes))
+    nsamples = min(500, 2 ** num_nodes)
+    shap_values = explainer.shap_values(active_state, nsamples=nsamples)
 
-    attributions = shap_values[0].tolist()
-    high_impact = np.where(np.abs(shap_values[0]) > 0.05)[0].tolist()
+    attributions = np.array(shap_values[0]) if isinstance(shap_values, list) else np.array(shap_values)
+    attributions = attributions.flatten().tolist()
+    high_impact = np.where(np.abs(np.array(attributions)) > 0.05)[0].tolist()
 
     try:
         vis_base64 = generate_similarity_map(