Fix SHAP compatibility — handle both old list format and new 3D array format

explainable_ids_full_pipeline.ipynb

@@ -59,7 +59,16 @@
     "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
     "print(f'Device: {DEVICE}')\n",
     "if DEVICE.type == 'cuda':\n",
-    "    print(f'GPU: {torch.cuda.get_device_name(0)}')"
+    "    print(f'GPU: {torch.cuda.get_device_name(0)}')\n",
+    "\n",
+    "def get_shap_for_class(shap_values, class_idx=0):\n",
+    "    \"\"\"Handle both old SHAP format (list of arrays) and new format (3D array).\"\"\"\n",
+    "    if isinstance(shap_values, list):\n",
+    "        return shap_values[class_idx]\n",
+    "    elif shap_values.ndim == 3:\n",
+    "        return shap_values[:, :, class_idx]\n",
+    "    else:\n",
+    "        return shap_values"
    ]
   },
   {
@@ -252,7 +261,6 @@
     "            optimizer.step()\n",
     "            total_loss += loss.item() * len(yb)\n",
     "        \n",
-    "        # Evaluate\n",
     "        model.eval()\n",
     "        preds, probs, labels = [], [], []\n",
     "        with torch.no_grad():\n",
@@ -284,7 +292,6 @@
     "    \n",
     "    dt = time.time() - t0\n",
     "    \n",
-    "    # Load best and final eval\n",
     "    model.load_state_dict(best_state)\n",
     "    model.eval()\n",
     "    preds, probs, labels = [], [], []\n",
@@ -307,9 +314,8 @@
     "    print('Confusion Matrix:')\n",
     "    print(confusion_matrix(labels, preds))\n",
     "    \n",
-    "    return model, {'f1': best_f1, 'roc_auc': roc, 'pr_auc': pr, 'time': dt, 'history': history, 'preds': preds, 'probs': probs, 'labels': labels}\n",
+    "    return model, {'f1': best_f1, 'roc_auc': roc, 'pr_auc': pr, 'time': dt, 'history': history}\n",
     "\n",
-    "# Train all 3\n",
     "models = {}\n",
     "results = {}\n",
     "for name, cls in [('mlp', MLP_IDS), ('lstm', LSTM_IDS), ('cnn1d', CNN1D_IDS)]:\n",
@@ -373,8 +379,11 @@
     "\n",
     "explainer = shap.KernelExplainer(predict_fn, X_train[bg_idx])\n",
     "print('Computing SHAP values for 150 test samples (this takes a few minutes)...')\n",
-    "shap_values = explainer.shap_values(X_test[exp_idx], nsamples=200, silent=True)\n",
-    "print('Done!')"
+    "shap_values_raw = explainer.shap_values(X_test[exp_idx], nsamples=200, silent=True)\n",
+    "\n",
+    "# Get SHAP values for anomaly class (class 0) — works with any SHAP version\n",
+    "shap_vals_anomaly = get_shap_for_class(shap_values_raw, class_idx=0)\n",
+    "print(f'Done! SHAP values shape: {shap_vals_anomaly.shape}')"
    ]
   },
   {
@@ -384,7 +393,7 @@
    "outputs": [],
    "source": [
     "# Global feature importance (anomaly class)\n",
-    "mean_abs_shap = np.abs(shap_values[0]).mean(axis=0)\n",
+    "mean_abs_shap = np.abs(shap_vals_anomaly).mean(axis=0)\n",
     "feature_importance = sorted(zip(FEATURE_NAMES, mean_abs_shap), key=lambda x: x[1], reverse=True)\n",
     "\n",
     "print('Top 15 features by mean |SHAP| (anomaly class):')\n",
@@ -399,7 +408,7 @@
    "outputs": [],
    "source": [
     "# SHAP summary plot\n",
-    "shap.summary_plot(shap_values[0], X_test[exp_idx], feature_names=FEATURE_NAMES, max_display=15)"
+    "shap.summary_plot(shap_vals_anomaly, X_test[exp_idx], feature_names=FEATURE_NAMES, max_display=15)"
    ]
   },
   {
@@ -429,7 +438,10 @@
     "pred = predict_fn(X_test[exp_idx[idx:idx+1]])\n",
     "print(f'Sample prediction: anomaly={pred[0][0]:.3f}, normal={pred[0][1]:.3f}')\n",
     "print(f'True label: {class_names[y_test[exp_idx[idx]]]}')\n",
-    "shap.force_plot(explainer.expected_value[0], shap_values[0][idx], X_test[exp_idx[idx]], feature_names=FEATURE_NAMES, matplotlib=True)"
+    "\n",
+    "ev = explainer.expected_value\n",
+    "ev0 = ev[0] if isinstance(ev, (list, np.ndarray)) else ev\n",
+    "shap.force_plot(ev0, shap_vals_anomaly[idx], X_test[exp_idx[idx]], feature_names=FEATURE_NAMES, matplotlib=True)"
    ]
   },
   {
@@ -479,13 +491,11 @@
     "# LIME vs SHAP comparison\n",
     "fig, axes = plt.subplots(1, 2, figsize=(16, 6))\n",
     "\n",
-    "# SHAP\n",
     "top10_shap = feature_importance[:10]\n",
     "axes[0].barh(range(10), [v for _, v in top10_shap][::-1], color='steelblue')\n",
     "axes[0].set_yticks(range(10)); axes[0].set_yticklabels([f for f, _ in top10_shap][::-1])\n",
     "axes[0].set_xlabel('Mean |SHAP value|'); axes[0].set_title('SHAP Top 10')\n",
     "\n",
-    "# LIME\n",
     "top10_lime = lime_sorted[:10]\n",
     "axes[1].barh(range(10), [v for _, v in top10_lime][::-1], color='coral')\n",
     "axes[1].set_yticks(range(10)); axes[1].set_yticklabels([f for f, _ in top10_lime][::-1])\n",
@@ -519,15 +529,15 @@
     "def compute_shap_stability(explainer, sample, epsilon, n_perturbs=10):\n",
     "    \"\"\"Compute SENS_MAX and PCC for one sample.\"\"\"\n",
     "    rng = np.random.RandomState(SEED)\n",
-    "    base = np.array(explainer.shap_values(sample.reshape(1,-1), nsamples=100, silent=True))\n",
-    "    base = base[0].flatten() if isinstance(base, list) else base.flatten()\n",
+    "    base_raw = explainer.shap_values(sample.reshape(1,-1), nsamples=100, silent=True)\n",
+    "    base = get_shap_for_class(base_raw, 0).flatten()\n",
     "    \n",
     "    max_delta, pccs = 0, []\n",
     "    for _ in range(n_perturbs):\n",
     "        noise = rng.uniform(-epsilon, epsilon, sample.shape)\n",
     "        perturbed = np.clip(sample + noise, 0, 1)\n",
-    "        p_shap = np.array(explainer.shap_values(perturbed.reshape(1,-1), nsamples=100, silent=True))\n",
-    "        p_shap = p_shap[0].flatten() if isinstance(p_shap, list) else p_shap.flatten()\n",
+    "        p_raw = explainer.shap_values(perturbed.reshape(1,-1), nsamples=100, silent=True)\n",
+    "        p_shap = get_shap_for_class(p_raw, 0).flatten()\n",
     "        max_delta = max(max_delta, np.linalg.norm(p_shap - base))\n",
     "        if np.std(base) > 1e-8 and np.std(p_shap) > 1e-8:\n",
     "            pccs.append(pearsonr(base, p_shap)[0])\n",
@@ -597,8 +607,8 @@
     "\n",
     "for idx in stability_idx[:10]:\n",
     "    sample = X_test[idx]\n",
-    "    sv = np.array(explainer.shap_values(sample.reshape(1,-1), nsamples=100, silent=True))\n",
-    "    sv = sv[0].flatten() if isinstance(sv, list) else sv.flatten()\n",
+    "    sv_raw = explainer.shap_values(sample.reshape(1,-1), nsamples=100, silent=True)\n",
+    "    sv = get_shap_for_class(sv_raw, 0).flatten()\n",
     "    \n",
     "    base_conf = predict_fn(sample.reshape(1,-1))[0]\n",
     "    pred_cls = np.argmax(base_conf)\n",
@@ -622,13 +632,11 @@
     "# Stability summary plot\n",
     "fig, axes = plt.subplots(1, 3, figsize=(16, 5))\n",
     "\n",
-    "# SENS_MAX\n",
     "eps_list = list(stability_results.keys())\n",
     "axes[0].plot(eps_list, [stability_results[e]['sens_max'] for e in eps_list], 'o-', color='steelblue', markersize=8)\n",
     "axes[0].set_xlabel('Perturbation epsilon'); axes[0].set_ylabel('SENS_MAX')\n",
     "axes[0].set_title('SHAP Sensitivity (lower = more stable)'); axes[0].grid(alpha=0.3)\n",
     "\n",
-    "# PCC\n",
     "pcc_vals = [stability_results[e]['pcc'] for e in eps_list]\n",
     "colors = ['green' if p > 0.6 else 'red' for p in pcc_vals]\n",
     "axes[1].bar(range(len(eps_list)), pcc_vals, color=colors)\n",
@@ -636,7 +644,6 @@
     "axes[1].axhline(y=0.6, color='gray', linestyle='--', label='Threshold (0.6)')\n",
     "axes[1].set_ylabel('Mean PCC'); axes[1].set_title('SHAP Stability'); axes[1].legend()\n",
     "\n",
-    "# Faithfulness\n",
     "ks = list(faith_results.keys())\n",
     "axes[2].bar(range(len(ks)), [np.mean(faith_results[k]) for k in ks],\n",
     "            yerr=[np.std(faith_results[k]) for k in ks], color='coral', capsize=5)\n",
@@ -660,15 +667,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Analyze which top SHAP features are attacker-manipulable\n",
     "manipulable = {'src_bytes', 'dst_bytes', 'hot', 'num_failed_logins', 'duration', 'num_compromised',\n",
     "               'root_shell', 'su_attempted', 'num_root', 'num_file_creations', 'num_shells', 'num_access_files'}\n",
     "partial = {'count', 'srv_count', 'serror_rate', 'srv_serror_rate', 'rerror_rate', 'srv_rerror_rate',\n",
     "           'protocol_type', 'flag', 'service'}\n",
-    "non_manip = {'dst_host_count', 'dst_host_srv_count', 'dst_host_same_srv_rate', 'dst_host_diff_srv_rate',\n",
-    "             'dst_host_same_src_port_rate', 'dst_host_srv_diff_host_rate', 'dst_host_serror_rate',\n",
-    "             'dst_host_srv_serror_rate', 'dst_host_rerror_rate', 'dst_host_srv_rerror_rate',\n",
-    "             'same_srv_rate', 'diff_srv_rate', 'srv_diff_host_rate'}\n",
     "\n",
     "print('SECURITY ANALYSIS: Top 15 Features by Manipulability')\n",
     "print('='*70)\n",
@@ -698,7 +700,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Final summary\n",
     "print('\\n' + '='*60)\n",
     "print('FINAL RESULTS SUMMARY')\n",
     "print('='*60)\n",