diff --git a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
index 8c1b094528f6d73a0932606b39625035f7e44def..29f40ffe8023d63df946d4742c020c0bfbd63025 100644
--- a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
+++ b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
@@ -105,6 +105,7 @@
     "# Plotting parameters\n",
     "skip_plots = False # exit after writing corrected files and metadata\n",
     "cell_id_preview = 1 # cell Id used for preview in single-shot plots\n",
+    "cmap = \"viridis\"  # matplolib.colormap for almost all heatmap. Other options ['plasma', 'inferno', 'magma', 'cividis', 'jet', ...]\n",
     "\n",
     "# Parallelization parameters\n",
     "chunk_size = 1000  # Size of chunk for image-wise correction\n",
@@ -114,6 +115,7 @@
     "max_nodes = 8 # Maximum number of SLURM jobs to split correction work into\n",
     "max_tasks_per_worker = 1  # the number of tasks a correction pool worker process can complete before it will exit and be replaced with a fresh worker process. Leave as -1 to keep worker alive as long as pool.\n",
     "\n",
+    "\n",
     "def balance_sequences(in_folder, run, sequences, sequences_per_node, karabo_da, max_nodes):\n",
     "    from xfel_calibrate.calibrate import balance_sequences as bs\n",
     "    return bs(in_folder, run, sequences, sequences_per_node, karabo_da, max_nodes=max_nodes)"
@@ -165,7 +167,6 @@
     "    CellRange,\n",
     "    LitFrameSelection,\n",
     ")\n",
-    "from cal_tools.ana_tools import get_range\n",
     "from cal_tools.calcat_interface import (\n",
     "    AGIPD_CalibrationData,\n",
     "    CalCatError,\n",
@@ -909,23 +910,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def do_3d_plot(data, edges, x_axis, y_axis):\n",
-    "    fig = plt.figure(figsize=(10, 10))\n",
-    "    ax = fig.gca(projection='3d')\n",
-    "\n",
-    "    # Make data.\n",
-    "    X = edges[0][:-1]\n",
-    "    Y = edges[1][:-1]\n",
-    "    X, Y = np.meshgrid(X, Y)\n",
-    "    Z = data.T\n",
-    "\n",
-    "    # Plot the surface.\n",
-    "    ax.plot_surface(X, Y, Z, cmap=colormap.coolwarm, linewidth=0, antialiased=False)\n",
-    "    ax.set_xlabel(x_axis)\n",
-    "    ax.set_ylabel(y_axis)\n",
-    "    ax.set_zlabel(\"Counts\")\n",
-    "\n",
-    "\n",
     "def do_2d_plot(data, edges, y_axis, x_axis):\n",
     "    fig = plt.figure(figsize=(10, 10))\n",
     "    ax = fig.add_subplot(111)\n",
@@ -1040,8 +1024,7 @@
     "        np.percentile(raw[:, 1, ...].astype(np.float32), [0.02, 99.8]),\n",
     "        ],\n",
     "    )\n",
-    "do_2d_plot(hist, (bins_x, bins_y), \"Signal (ADU)\", \"Analogue gain (ADU)\")\n",
-    "do_3d_plot(hist, (bins_x, bins_y), \"Signal (ADU)\", \"Analogue gain (ADU)\")"
+    "do_2d_plot(hist, (bins_x, bins_y), \"Signal (ADU)\", \"Analogue gain (ADU)\")"
    ]
   },
   {
@@ -1113,8 +1096,7 @@
     "    ],\n",
     ")\n",
     "\n",
-    "do_2d_plot(hist, (bins_x, bins_y), \"Signal (ADU)\", \"Pulse id\")\n",
-    "do_3d_plot(hist, (bins_x, bins_y), \"Signal (ADU)\", \"Pulse id\")"
+    "do_2d_plot(hist, (bins_x, bins_y), \"Signal (ADU)\", \"Pulse id\")"
    ]
   },
   {
@@ -1183,7 +1165,8 @@
     "    ax = fig.add_subplot(111)\n",
     "    data = np.mean(raw[slice(*cell_sel.crange), 0, ...], axis=0)\n",
     "    vmin, vmax = np.percentile(data, [5, 95])\n",
-    "    ax = geom.plot_data_fast(data, ax=ax, vmin=vmin, vmax=vmax)\n",
+    "    ax = geom.plot_data_fast(data, ax=ax, vmin=vmin, vmax=vmax, cmap=cmap)\n",
+    "    pass\n",
     "else:\n",
     "    print(\"Skipping mean RAW preview for single memory cell, \"\n",
     "          f\"see single shot image for selected cell ID {cell_id_preview}.\")"
@@ -1199,7 +1182,9 @@
     "fig = plt.figure(figsize=(20, 10))\n",
     "ax = fig.add_subplot(111)\n",
     "vmin, vmax = np.percentile(raw[cell_idx_preview, 0, ...], [5, 95])\n",
-    "ax = geom.plot_data_fast(raw[cell_idx_preview, 0, ...], ax=ax, vmin=vmin, vmax=vmax)"
+    "ax = geom.plot_data_fast(\n",
+    "    raw[cell_idx_preview, 0, ...], ax=ax, vmin=vmin, vmax=vmax, cmap=cmap)\n",
+    "pass"
    ]
   },
   {
@@ -1215,8 +1200,9 @@
     "    fig = plt.figure(figsize=(20, 10))\n",
     "    ax = fig.add_subplot(111)\n",
     "    data = np.mean(corrected, axis=0)\n",
-    "    vmin, vmax = np.percentile(data, [5, 95])\n",
-    "    ax = geom.plot_data_fast(data, ax=ax, vmin=vmin, vmax=vmax)\n",
+    "    vmin, vmax = np.percentile(data, [5, 99.9])\n",
+    "    ax = geom.plot_data_fast(data, ax=ax, vmin=vmin, vmax=vmax, cmap=cmap)\n",
+    "    pass\n",
     "else:\n",
     "    print(\"Skipping mean CORRECTED preview for single memory cell, \"\n",
     "          f\"see single shot image for selected cell ID {cell_id_preview}.\")"
@@ -1231,8 +1217,10 @@
     "display(Markdown(f'A single shot of the CORRECTED image from cell {cell_id_preview} \\n'))\n",
     "fig = plt.figure(figsize=(20, 10))\n",
     "ax = fig.add_subplot(111)\n",
-    "vmin, vmax = np.percentile(corrected[cell_idx_preview], [5, 95])\n",
-    "ax = geom.plot_data_fast(corrected[cell_idx_preview], ax=ax, vmin=vmin, vmax=vmax)"
+    "vmin, vmax = np.percentile(corrected[cell_idx_preview], [5, 99.9])\n",
+    "ax = geom.plot_data_fast(\n",
+    "    corrected[cell_idx_preview], ax=ax, vmin=vmin, vmax=vmax, cmap=cmap)\n",
+    "pass"
    ]
   },
   {
@@ -1243,13 +1231,16 @@
    "source": [
     "fig = plt.figure(figsize=(20, 10))\n",
     "ax = fig.add_subplot(111)\n",
-    "nbins = np.int((vmax + 50) / 2)\n",
+    "vmin, vmax = np.percentile(corrected[cell_idx_preview], [5, 100])\n",
+    "nbins = int((vmax + 50) / 2)\n",
     "h = ax.hist(corrected[cell_idx_preview].flatten(),\n",
     "            bins=nbins, range=(vmin, vmax),\n",
     "            histtype='stepfilled', log=True)\n",
     "plt.xlabel('[ADU]')\n",
     "plt.ylabel('Counts')\n",
-    "ax.grid()"
+    "ax.grid()\n",
+    "plt.title(f'Log-scaled histogram for corrected data for cell {cell_idx_preview}')\n",
+    "pass"
    ]
   },
   {
@@ -1260,11 +1251,10 @@
    "source": [
     "fig = plt.figure(figsize=(20, 10))\n",
     "ax = fig.add_subplot(111)\n",
-    "vmin, vmax = np.percentile(corrected[cell_idx_preview], [5, 95])\n",
-    "vmax = np.nanmax(corrected)\n",
+    "vmin, vmax = np.nanpercentile(corrected, [5, 100])\n",
     "if vmax > 50000:\n",
-    "    vmax=50000\n",
-    "nbins = np.int((vmax + 100) / 5)\n",
+    "    vmax = 50000\n",
+    "nbins = int((vmax + 100) / 5)\n",
     "hist_range=(vmin, vmax)\n",
     "h = ax.hist(corrected.flatten(), bins=nbins,\n",
     "            range=hist_range, histtype='step', log=True, label = 'All')\n",
@@ -1277,7 +1267,9 @@
     "ax.legend()\n",
     "ax.grid()\n",
     "plt.xlabel('[ADU]')\n",
-    "plt.ylabel('Counts')"
+    "plt.ylabel('Counts')\n",
+    "plt.title(f'Overlaid Histograms for corrected data for multiple gains')\n",
+    "pass"
    ]
   },
   {
@@ -1299,7 +1291,9 @@
     "fig = plt.figure(figsize=(20, 10))\n",
     "ax = fig.add_subplot(111)\n",
     "vmin, vmax = np.percentile(np.mean(gains, axis=0), [0, 100])\n",
-    "ax = geom.plot_data_fast(np.mean(gains, axis=0), ax=ax, vmin=vmin, vmax=vmax)"
+    "ax = geom.plot_data_fast(\n",
+    "    np.mean(gains, axis=0), ax=ax, vmin=vmin, vmax=vmax, cmap=cmap)\n",
+    "pass"
    ]
   },
   {
@@ -1341,7 +1335,9 @@
    "source": [
     "fig = plt.figure(figsize=(20, 10))\n",
     "ax = fig.add_subplot(111)\n",
-    "geom.plot_data_fast(np.log2(mask[cell_idx_preview]), ax=ax, vmin=0, vmax=32)"
+    "geom.plot_data_fast(\n",
+    "    np.log2(mask[cell_idx_preview]), ax=ax, vmin=0, vmax=32, cmap=cmap)\n",
+    "pass"
    ]
   },
   {
@@ -1389,7 +1385,9 @@
    "source": [
     "fig = plt.figure(figsize=(20, 10))\n",
     "ax = fig.add_subplot(111)\n",
-    "geom.plot_data_fast(np.mean(mask>0, axis=0), vmin=0, ax=ax, vmax=1)"
+    "geom.plot_data_fast(\n",
+    "    np.mean(mask>0, axis=0), vmin=0, ax=ax, vmax=1, cmap=cmap)\n",
+    "pass"
    ]
   },
   {
@@ -1409,7 +1407,9 @@
     "ax = fig.add_subplot(111)\n",
     "cm = np.copy(mask)\n",
     "cm[cm > BadPixels.NO_DARK_DATA.value] = 0\n",
-    "ax = geom.plot_data_fast(np.mean(cm>0, axis=0), vmin=0, ax=ax, vmax=1)"
+    "ax = geom.plot_data_fast(\n",
+    "    np.mean(cm>0, axis=0), vmin=0, ax=ax, vmax=1, cmap=cmap)\n",
+    "pass"
    ]
   }
  ],