diff --git a/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb b/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
index 3832165cc90bd938ebeeac61ddf12032c9de7f0b..fd89a19f3ac12f8391f72534a13ddee296c9e184 100644
--- a/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
+++ b/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
@@ -70,6 +70,7 @@
    "outputs": [],
    "source": [
     "import datetime\n",
+    "import warnings\n",
     "from functools import partial\n",
     "\n",
     "import h5py\n",
@@ -89,6 +90,9 @@
     "    get_dir_creation_date,\n",
     ")\n",
     "from iCalibrationDB import Conditions, Constants\n",
+    "from XFELDetAna.plotting.heatmap import heatmapPlot\n",
+    "\n",
+    "warnings.filterwarnings('ignore')\n",
     "\n",
     "%matplotlib inline"
    ]
@@ -353,6 +357,9 @@
     "        dc = dc.select(instr_mod_src, require_all=True)\n",
     "        data = dc[instr_mod_src, \"data.adc\"].ndarray()\n",
     "        gain = dc[instr_mod_src, \"data.gain\"].ndarray()\n",
+    "        step_timer.done_step(\"preparing raw data\")\n",
+    "        dshape = data.shape\n",
+    "\n",
     "        step_timer.start()\n",
     "\n",
     "        # Allocate shared arrays.\n",
@@ -444,43 +451,59 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "display(Markdown(\"### Mean RAW across pulses for one train:\"))\n",
+    "display(Markdown(\"### Mean RAW and CORRECTED across pulses for one train:\"))\n",
     "display(Markdown(f\"Train: {tid}\"))\n",
     "\n",
     "step_timer.start()\n",
     "for mod in karabo_da:\n",
+    "\n",
+    "    fig, (ax1, ax2) = plt.subplots(nrows=2, figsize=(20, 10))\n",
     "    raw_data = mod_dcs[mod][\"train_raw_data\"]\n",
-    "    fig, ax = plt.subplots(figsize=(20, 10))\n",
-    "    im = ax.plot(np.mean(raw_data, axis=0))\n",
-    "    ax.set_title(f\"Module {mod}\")\n",
-    "    ax.set_xlabel(\"Strip #\", size=15)\n",
-    "    ax.set_ylabel(\"12-bit ADC output\", size=15)\n",
-    "    plt.show()\n",
-    "step_timer.done_step(\"Plotting raw data\")"
+    "    \n",
+    "    im = ax1.plot(np.mean(raw_data, axis=0))\n",
+    "    ax1.set_title(f\"Module {mod}\")\n",
+    "    ax1.set_xlabel(\"Strip #\", size=15)\n",
+    "    ax1.set_ylabel(\"12-bit ADC output\", size=15)\n",
+    "\n",
+    "    corr_data = mod_dcs[mod][\"train_corr_data\"]\n",
+    "    im = ax2.plot(np.mean(corr_data, axis=0))\n",
+    "    ax2.set_title(f\"Module {mod}\")\n",
+    "    ax2.set_xlabel(\"Strip #\", size=20)\n",
+    "    ax2.set_ylabel(\"10-bit KeV. output\", size=20)\n",
+    "    plt.xticks(fontsize=15)\n",
+    "    plt.yticks(fontsize=15)\n",
+    "    pass\n",
+    "step_timer.done_step(\"Plotting mean data\")\n"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "11337d95",
+   "id": "58a6a276",
    "metadata": {},
    "outputs": [],
    "source": [
-    "display(Markdown(\"### Mean CORRECTED across pulses for one train:\"))\n",
-    "display(Markdown(f\"Train: {tid}\"))\n",
+    "display(Markdown(\"### RAW and CORRECTED strips across pulses for one train:\"))\n",
     "\n",
     "step_timer.start()\n",
     "for mod in karabo_da:\n",
-    "    corr_data = mod_dcs[mod][\"train_corr_data\"]\n",
-    "    fig, ax = plt.subplots(figsize=(20, 10))\n",
-    "    im = ax.plot(np.mean(corr_data, axis=0))\n",
-    "    ax.set_title(f\"Module {mod}\")\n",
-    "    ax.set_xlabel(\"Strip #\", size=20)\n",
-    "    ax.set_ylabel(\"10-bit KeV. output\", size=20)\n",
-    "    plt.xticks(fontsize=15)\n",
-    "    plt.yticks(fontsize=15)\n",
-    "    plt.show()\n",
-    "step_timer.done_step(\"Plotting corrected data\")"
+    "    fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(20, 10))\n",
+    "    heatmapPlot(\n",
+    "        mod_dcs[mod][\"train_raw_data\"],\n",
+    "        y_label=\"Pulses\",\n",
+    "        x_label=\"Strips\",\n",
+    "        title=f\"Train {tid} RAW\",\n",
+    "        use_axis=ax1\n",
+    "    )\n",
+    "    heatmapPlot(\n",
+    "        mod_dcs[mod][\"train_corr_data\"],\n",
+    "        y_label=\"Pulses\",\n",
+    "        x_label=\"Strips\",\n",
+    "        title=f\"Train {tid} CORRECTED\",\n",
+    "        use_axis=ax2\n",
+    "    )\n",
+    "    pass\n",
+    "step_timer.done_step(\"Plotting RAW and CORRECTED data for one train\")\n"
    ]
   },
   {
@@ -521,45 +544,32 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "display(Markdown(\"### RAW even/odd pulses for one train:\"))\n",
-    "display(Markdown(f\"RAW train: {tid}\"))\n",
+    "display(Markdown(\"### RAW and CORRECTED even/odd pulses for one train:\"))\n",
+    "display(Markdown(f\"Train: {tid}\"))\n",
     "for mod in karabo_da:\n",
+    "    fig, (ax1, ax2) = plt.subplots(nrows=2, figsize=(20, 10))\n",
     "    raw_data = mod_dcs[mod][\"train_raw_data\"]\n",
-    "    fig, ax = plt.subplots(figsize=(20, 10))\n",
-    "    ax.plot(raw_data[odd_pulse], label=f\"Odd Pulse {odd_pulse}\")\n",
+    "    corr_data = mod_dcs[mod][\"train_corr_data\"]\n",
+    "\n",
+    "    ax1.plot(raw_data[odd_pulse], label=f\"Odd Pulse {odd_pulse}\")\n",
+    "    ax2.plot(corr_data[odd_pulse], label=f\"Odd Pulse {odd_pulse}\")\n",
     "    if even_pulse:\n",
-    "        ax.plot(raw_data[even_pulse], label=f\"Even Pulse {even_pulse}\")\n",
-    "    ax.set_title(f\"Module {mod}\")\n",
-    "    ax.set_xlabel(\"Strip #\", size=20)\n",
-    "    ax.set_ylabel(\"12-bit ADC RAW\", size=20)\n",
+    "        ax1.plot(raw_data[even_pulse], label=f\"Even Pulse {even_pulse}\")\n",
+    "        ax2.plot(corr_data[even_pulse], label=f\"Even Pulse {even_pulse}\")\n",
+    "\n",
+    "    ax1.set_title(f\"Module {mod}\")\n",
+    "    ax1.set_xlabel(\"Strip #\", size=20)\n",
+    "    ax1.set_ylabel(\"12-bit ADC RAW\", size=20)\n",
     "    plt.xticks(fontsize=15)\n",
     "    plt.yticks(fontsize=15)\n",
-    "    ax.legend()\n",
-    "    plt.show()\n",
-    "step_timer.done_step(\"Plotting RAW odd/even pulses.\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "96900ba8-96a3-44b0-b489-d948f19298dc",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "display(Markdown(\"### CORRECTED even/odd pulses for one train:\"))\n",
-    "display(Markdown(f\"CORRECTED train: {tid}\"))\n",
-    "for mod in karabo_da:\n",
-    "    corr_data = mod_dcs[mod][\"train_corr_data\"]\n",
-    "    fig, ax = plt.subplots(figsize=(20, 10))\n",
-    "    ax.plot(corr_data[odd_pulse], label=f\"Odd Pulse {odd_pulse}\")\n",
-    "    if even_pulse:\n",
-    "        ax.plot(corr_data[even_pulse], label=f\"Even Pulse {even_pulse}\")\n",
-    "    ax.set_title(f\"Module {mod}\")\n",
-    "    ax.set_xlabel(\"Strip #\", size=15)\n",
-    "    ax.set_ylabel(\"10-bit KeV CORRECTED\", size=15)\n",
-    "    ax.legend()\n",
-    "    plt.show()\n",
-    "step_timer.done_step(\"Plotting CORRECTED odd/even pulses.\")"
+    "    ax1.legend()\n",
+    "\n",
+    "    ax2.set_title(f\"Module {mod}\")\n",
+    "    ax2.set_xlabel(\"Strip #\", size=15)\n",
+    "    ax2.set_ylabel(\"10-bit KeV CORRECTED\", size=15)\n",
+    "    ax2.legend()\n",
+    "    pass\n",
+    "step_timer.done_step(\"Plotting RAW and CORRECTED odd/even pulses.\")"
    ]
   }
  ],