diff --git a/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb b/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
index a994b180c6cc2d039212df3ed26ce1e147f1c13e..a91f1fe78d6ccef2da245dc561c9553702e800a9 100644
--- a/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
+++ b/notebooks/Gotthard2/Correction_Gotthard2_NBC.ipynb
@@ -39,6 +39,9 @@
     "cal_db_interface = \"tcp://max-exfl016:8017#8025\" # the database interface to use\n",
     "cal_db_timeout = 180000  # timeout on caldb requests\n",
     "\n",
+    "# Parameters affecting corrected data.\n",
+    "relative_gain = True  # do relative gain correction\n",
+    "\n",
     "# Parameters for plotting\n",
     "skip_plots = False  # exit after writing corrected files\n",
     "pulse_preview = 3 # pulseId to preview. The following even/odd pulseId is used for preview.\n",
@@ -201,8 +204,13 @@
     "def correct_train(wid, index, d):\n",
     "    g = gain[index]\n",
     "    d_corr = np.zeros_like(d, dtype=np.float32)\n",
+    "    gotthard2algs.convert_to_10bit(d, lut, d_corr)\n",
     "    gotthard2algs.correct_train(\n",
-    "        d, g, lut, offset_map.astype(np.float32), gain_map.astype(np.float32), d_corr)\n",
+    "        d_corr, g,\n",
+    "        offset_map.astype(np.float32),\n",
+    "        gain_map.astype(np.float32),\n",
+    "        apply_relgain=relative_gain,\n",
+    "    )\n",
     "    data_corr[index, ...] = d_corr"
    ]
   },
@@ -235,12 +243,13 @@
     "        # Allocate shared arrays.\n",
     "        data_corr = context.alloc(shape=dshape, dtype=np.float32)\n",
     "\n",
-    "        context.map(correct_train, data) \n",
+    "        context.map(correct_train, data)\n",
     "        step_timer.done_step(\"Correcting one sequence file\")\n",
     "        step_timer.start()\n",
     "\n",
     "        # Provided PSI gain map has 0 values. Set inf values to nan.\n",
     "        data_corr[np.isinf(data_corr)] = np.nan\n",
+    "\n",
     "        # Create CORR files and add corrected data sources.\n",
     "        # Exclude raw data images (data/adc)\n",
     "        with h5py.File(out_file, 'w') as ofile:\n",
@@ -255,7 +264,7 @@
     "                chunks=((1,) + dshape[1:]),  # 1 chunk == 1 image\n",
     "                dtype=np.float32,\n",
     "            )\n",
-    "            step_timer.done_step(\"Storing data\")"
+    "            step_timer.done_step(\"Storing data\")\n"
    ]
   },
   {
diff --git a/src/cal_tools/gotthard2algs.pyx b/src/cal_tools/gotthard2algs.pyx
index 6003ed93773b2bef9a4cdde866e97c3e79d97ada..0d4e8de126550ba0d2b8834d3aeaecebae92adad 100644
--- a/src/cal_tools/gotthard2algs.pyx
+++ b/src/cal_tools/gotthard2algs.pyx
@@ -1,32 +1,48 @@
 import numpy as np
 
 from cython cimport boundscheck, wraparound, cdivision
-from cython.parallel import prange
 cimport numpy as cnp
 
 
+@boundscheck(False)
+@wraparound(False)
+def convert_to_10bit(
+    unsigned short[:, :] data,
+    unsigned short[:, :, :] lut,
+    float[:, :] data_10bit,
+):
+    """Convert 12bit RAW data to 10bit data."""
+
+    cdef:
+        unsigned short d_10bit, pulse, train, x, raw_val
+
+    for pulse in range(data.shape[0]):
+        cell = pulse % 2
+        for x in range(data.shape[1]):
+            raw_val = data[pulse, x]
+            d_10bit = lut[cell, raw_val, x]
+            data_10bit[pulse, x] = <float>d_10bit
+
+
 @boundscheck(False)
 @wraparound(False)
 @cdivision(True)
 def correct_train(
-    unsigned short[:, :] data,
+    float[:, :] data,
     unsigned char[:, :] gain,
-    unsigned short[:, :, :] lut,
-    float[:, :, :] offset,
-    float[:, :, :] relgain,
-    float[:, :] corr_data,
+    float[:, :, :] offset_map,
+    float[:, :, :] relgain_map,
+    unsigned short short apply_relgain = 1,
 ):
-    """
-    correct Gotthard2 raw data.
-    1. Convert 12 bit data to 10 bit data.
-    2. Offset correction.
-    3. Gain correction.
+    """Correct Gotthard2 raw data.
+    1. Offset correction.
+    2. Gain correction.
     """
 
-    cdef :
-        cdef float corr_val
-        cdef unsigned short d_10bit, pulse, train, x, raw_val, g
-    
+    cdef:
+        float raw_val
+        unsigned short d_10bit, pulse, train, x, g
+
     for pulse in range(data.shape[0]):
         cell = pulse % 2
         for x in range(data.shape[1]):
@@ -34,8 +50,7 @@ def correct_train(
             if g == 3:
                 g = 2
             raw_val = data[pulse, x]
-            d_10bit = lut[cell, raw_val, x]
-            corr_val = <float>d_10bit
-            corr_val -= offset[g, cell, x]
-            corr_val /= relgain[g, cell, x]
-            corr_data[pulse, x] = corr_val
+            raw_val -= offset_map[g, cell, x]
+            if apply_relgain == 1:
+                raw_val /= relgain_map[g, cell, x]
+            data[pulse, x] = raw_val