diff --git a/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb b/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
index bca6e72574b225e3f2995c4a843eb19550277735..1f3d6ac8bffbb567af01976fd2c1759abc95f773 100644
--- a/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
+++ b/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
@@ -340,6 +340,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "parallel_num_procs = 6\n",
+    "parallel_num_threads = multiprocessing.cpu_count() // parallel_num_procs\n",
+    "\n",
     "def characterize_module(\n",
     "    fast_data_filename: str, channel: int, gain_index: int\n",
     ") -> Tuple[np.array, np.array, np.array, np.array, int, np.array, int, float]:\n",
@@ -390,22 +393,20 @@
     "            ga = im[:, 1, ...]\n",
     "        im = im[:, 0, ...].astype(np.float32)\n",
     "\n",
-    "    im = np.rollaxis(im, 2)\n",
-    "    im = np.rollaxis(im, 2, 1)\n",
-    "\n",
+    "    im = np.transpose(im)\n",
     "    if not fixed_gain_mode:\n",
-    "        ga = np.rollaxis(ga, 2)\n",
-    "        ga = np.rollaxis(ga, 2, 1)\n",
+    "        ga = np.transpose(ga)\n",
     "    \n",
-    "    offset = psh.alloc(shape=(im.shape[0], im.shape[1], num_cells), dtype=np.float64)\n",
-    "    noise = psh.alloc(like=offset)\n",
+    "    context = psh.context.ThreadContext(num_workers=parallel_num_threads)\n",
+    "    offset = context.alloc(shape=(im.shape[0], im.shape[1], num_cells), dtype=np.float64)\n",
+    "    noise = context.alloc(like=offset)\n",
     "\n",
     "    if fixed_gain_mode:\n",
     "        gains = None\n",
     "        gains_std = None\n",
     "    else:\n",
-    "        gains = psh.alloc(like=offset)\n",
-    "        gains_std = psh.alloc(like=offset)\n",
+    "        gains = context.alloc(like=offset)\n",
+    "        gains_std = context.alloc(like=offset)\n",
     "\n",
     "    def process_cell(worker_id, array_index, cell_number):\n",
     "        cell_slice_index = (cellIds == cell_number)\n",
@@ -416,10 +417,10 @@
     "            ga_slice = ga[..., cell_slice_index]\n",
     "            gains[..., cell_number] = np.median(ga_slice, axis=2)\n",
     "            gains_std[..., cell_number] = np.std(ga_slice, axis=2)\n",
-    "    psh.map(process_cell, np.arange(num_cells))\n",
+    "    context.map(process_cell, np.arange(num_cells))\n",
     "\n",
     "    # bad pixels\n",
-    "    bp = np.zeros(offset.shape, np.uint32)\n",
+    "    bp = np.zeros_like(offset, dtype=np.uint32)\n",
     "    # offset related bad pixels\n",
     "    offset_mn = np.nanmedian(offset, axis=(0,1))\n",
     "    offset_std = np.nanstd(offset, axis=(0,1))\n",
@@ -441,17 +442,6 @@
     "    return offset, noise, gains, gains_std, bp, num_cells, local_acq_rate"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "parallel_num_procs = 6\n",
-    "parallel_num_threads = multiprocessing.cpu_count() // parallel_num_procs\n",
-    "psh.set_default_context(\"threads\", num_workers=parallel_num_threads)"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,