diff --git a/docs/requirements.txt b/docs/requirements.txt
index d7504541fbdcba87ddbb589179a85ce40658a118..5676a86030e171f069e80e5dc6a8ef71593eca63 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,2 +1 @@
-iCalibrationDB @ git+https://xcalgitlab:${GITHUB_TOKEN}@git.xfel.eu/gitlab/detectors/cal_db_interactive.git@2.0.9
-XFELDetectorAnalysis @ git+https://xcalgitlab:${GITHUB_TOKEN}@git.xfel.eu/gitlab/karaboDevices/pyDetLib.git@2.7.0
+iCalibrationDB @ git+https://xcalgitlab:${GITHUB_TOKEN}@git.xfel.eu/gitlab/detectors/cal_db_interactive.git@2.1.0
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 4a732c00cda5ce1383e8b032179750b8af17778a..fcb9b13da661bbef10a49bf52f74c8f4d05119b3 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -13,6 +13,8 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
+import glob
+
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
@@ -20,6 +22,23 @@
 # import os
 # import sys
 # sys.path.insert(0, os.path.abspath('.'))
+import os
+import shutil
+import sys
+import textwrap
+from datetime import datetime
+from subprocess import Popen, check_output
+from textwrap import dedent, indent
+from uuid import uuid4
+
+import nbformat
+import tabulate
+from dateutil.parser import parse
+from lxml import etree
+from nbconvert import RSTExporter
+
+# generate the list of available notebooks
+from xfel_calibrate import notebooks
 
 # -- General configuration ------------------------------------------------
 
@@ -39,11 +58,6 @@ extensions = [
     'sphinx.ext.viewcode',
 ]
 
-import glob
-import os
-import sys
-from subprocess import Popen
-
 sys.path.append(os.path.abspath("../pycalibration/"))
 p = Popen(["./makeAllDocs.sh"])
 p.communicate()
@@ -374,16 +388,6 @@ except:
     check_call(["wget", pandoc_url])
     check_call(["dpkg", "-i", pandoc_pack])
 
-import os
-from subprocess import check_output
-from textwrap import dedent, indent
-
-import nbformat
-from nbconvert import RSTExporter
-
-# generate the list of available notebooks
-from xfel_calibrate import notebooks
-
 rst_exporter = RSTExporter()
 with open("available_notebooks.rst", "w") as f:
     f.write(dedent("""
@@ -409,6 +413,8 @@ with open("available_notebooks.rst", "w") as f:
 
         for caltype in sorted(values.keys()):
             data = values[caltype]
+            if data.get("notebook", None) is None:
+                continue
             nbpath = os.path.abspath("{}/../../../{}".format(__file__, data["notebook"]))
             with open(nbpath, "r") as nf:
                 nb = nbformat.read(nf, as_version=4)
@@ -443,15 +449,6 @@ with open("available_notebooks.rst", "w") as f:
 # add test results
 test_artefact_dir = os.path.realpath("../../tests/legacy/artefacts")
 
-import shutil
-import textwrap
-from datetime import datetime
-from uuid import uuid4
-
-import tabulate
-from dateutil.parser import parse
-from lxml import etree
-
 
 def xml_to_rst_report(xml, git_tag, reports=[]):
     e = etree.fromstring(xml.encode())
@@ -609,5 +606,7 @@ with open("test_results.rst", "w") as f:
                     rst = xml_to_rst_report(xs, commit, reports=reports.get(rloc, []))
                     fr.write(rst)
             f.write("   test_rsts/{}\n".format(commit))
+
+
 def setup(app):
     app.add_stylesheet('css/test_decorators.css')
diff --git a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
index 668f0c3cd8f7a134ed6d6f3f679de3f261b78eb3..962d0d6cb514b0edd8a34f8867ad557595b2e879 100644
--- a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
+++ b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
@@ -17,37 +17,37 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "in_folder = \"/gpfs/exfel/exp/HED/202031/p900174/raw\" # the folder to read data from, required\n",
-    "out_folder = \"/gpfs/exfel/data/scratch/ahmedk/test/hibef_agipd2\"  # the folder to output to, required\n",
+    "in_folder = \"/gpfs/exfel/exp/SPB/202131/p900230/raw\" # the folder to read data from, required\n",
+    "out_folder = \"/gpfs/exfel/data/scratch/ahmedk/test/remove/agipd_resolve_conf\"  # the folder to output to, required\n",
     "sequences = [-1] # sequences to correct, set to -1 for all, range allowed\n",
     "modules = [-1] # modules to correct, set to -1 for all, range allowed\n",
     "train_ids = [-1] # train IDs to correct, set to -1 for all, range allowed\n",
-    "run = 155 # runs to process, required\n",
+    "run = 275 # runs to process, required\n",
     "\n",
-    "karabo_id = \"HED_DET_AGIPD500K2G\" # karabo karabo_id\n",
+    "karabo_id = \"SPB_DET_AGIPD1M-1\" # karabo karabo_id\n",
     "karabo_da = ['-1']  # a list of data aggregators names, Default [-1] for selecting all data aggregators\n",
-    "receiver_id = \"{}CH0\" # inset for receiver devices\n",
+    "receiver_template = \"{}CH0\" # inset for receiver devices\n",
     "path_template = 'RAW-R{:04d}-{}-S{:05d}.h5' # the template to use to access data\n",
-    "h5path = 'INSTRUMENT/{}/DET/{}:xtdf/' # path in the HDF5 file to images\n",
-    "h5path_idx = 'INDEX/{}/DET/{}:xtdf/' # path in the HDF5 file to images\n",
-    "h5path_ctrl = '/CONTROL/{}/MDL/FPGA_COMP' # path to control information\n",
-    "karabo_id_control = \"HED_EXP_AGIPD500K2G\" # karabo-id for control device\n",
-    "karabo_da_control = 'AGIPD500K2G00' # karabo DA for control infromation\n",
+    "instrument_source_template = '{}/DET/{}:xtdf'  # path in the HDF5 file to images\n",
+    "index_source_template = 'INDEX/{}/DET/{}:xtdf/'  # path in the HDF5 file to images\n",
+    "ctrl_source_template = '{}/MDL/FPGA_COMP'  # path to control information\n",
+    "karabo_id_control = \"SPB_IRU_AGIPD1M1\" # karabo-id for control device\n",
     "\n",
     "slopes_ff_from_files = \"\" # Path to locally stored SlopesFF and BadPixelsFF constants\n",
     "\n",
     "use_dir_creation_date = True # use the creation data of the input dir for database queries\n",
     "cal_db_interface = \"tcp://max-exfl016:8015#8045\" # the database interface to use\n",
-    "cal_db_timeout = 30000 # in milli seconds\n",
+    "cal_db_timeout = 30000 # in milliseconds\n",
     "creation_date_offset = \"00:00:00\" # add an offset to creation date, e.g. to get different constants\n",
     "\n",
     "use_ppu_device = ''  # Device ID for a pulse picker device to only process picked trains, empty string to disable\n",
     "ppu_train_offset = 0  # When using the pulse picker, offset between the PPU's sequence start and actually picked train\n",
     "\n",
-    "max_cells = 0 # number of memory cells used, set to 0 to automatically infer\n",
-    "bias_voltage = 300 # Bias voltage\n",
+    "mem_cells = 0  # Number of memory cells used, set to 0 to automatically infer\n",
+    "bias_voltage = 0  # bias voltage, set to 0 to use stored value in slow data.\n",
     "acq_rate = 0. # the detector acquisition rate, use 0 to try to auto-determine\n",
-    "gain_setting = 0.1 # the gain setting, use 0.1 to try to auto-determine\n",
+    "gain_setting = -1  # the gain setting, use -1 to use value stored in slow data.\n",
+    "gain_mode = -1  # gain mode (0: adaptive, 1-3 fixed high/med/low, -1: read from CONTROL data)\n",
     "photon_energy = 9.2 # photon energy in keV\n",
     "overwrite = True # set to True if existing data should be overwritten\n",
     "max_pulses = [0, 352, 1] # range list [st, end, step] of memory cell indices to be processed within a train. 3 allowed maximum list input elements.\n",
@@ -76,7 +76,7 @@
     "zero_nans = False # set NaN values in corrected data to 0\n",
     "zero_orange = False # set to 0 very negative and very large values in corrected data\n",
     "blc_set_min = False # Shift to 0 negative medium gain pixels after offset corr\n",
-    "corr_asic_diag = False # if set, diagonal drop offs on ASICs are correted\n",
+    "corr_asic_diag = False # if set, diagonal drop offs on ASICs are corrected\n",
     "force_hg_if_below = False # set high gain if mg offset subtracted value is below hg_hard_threshold\n",
     "force_mg_if_below = False # set medium gain if mg offset subtracted value is below mg_hard_threshold\n",
     "mask_noisy_adc = False # Mask entire ADC if they are noise above a relative threshold\n",
@@ -85,17 +85,20 @@
     "mask_zero_std = False # Mask pixels with zero standard deviation across train\n",
     "low_medium_gap = False # 5 sigma separation in thresholding between low and medium gain\n",
     "\n",
+    "use_litframe_device = '' # Device ID for a lit frame finder device to only process illuminated frames, empty string to disable\n",
+    "energy_threshold = -1000 # The low limit for the energy (uJ) exposed by frames subject to processing. If -1000, selection by pulse energy is disabled\n",
+    "\n",
     "# 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",
     "\n",
     "# Paralellization parameters\n",
-    "chunk_size = 1000 # Size of chunk for image-weise correction\n",
-    "chunk_size_idim = 1  # chunking size of imaging dimension, adjust if user software is sensitive to this.\n",
+    "chunk_size = 1000  # Size of chunk for image-wise correction\n",
     "n_cores_correct = 16 # Number of chunks to be processed in parallel\n",
     "n_cores_files = 4 # Number of files to be processed in parallel\n",
-    "sequences_per_node = 2 # number of sequence files per cluster node if run as slurm job, set to 0 to not run SLURM parallel\n",
-    "max_nodes = 8 # Maximum number of Slurm jobs to split correction work into\n",
+    "sequences_per_node = 2 # number of sequence files per cluster node if run as SLURM job, set to 0 to not run SLURM parallel\n",
+    "max_nodes = 8 # Maximum number of SLURM jobs to split correction work into\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",
@@ -127,7 +130,7 @@
     "import matplotlib\n",
     "import matplotlib.pyplot as plt\n",
     "import yaml\n",
-    "from extra_data import RunDirectory, stack_detector_data, by_id\n",
+    "from extra_data import H5File, RunDirectory, stack_detector_data, by_id\n",
     "from extra_geom import AGIPD_1MGeometry, AGIPD_500K2GGeometry\n",
     "from matplotlib import cm as colormap\n",
     "from matplotlib.colors import LogNorm\n",
@@ -146,14 +149,12 @@
     "from cal_tools import agipdalgs as calgs\n",
     "from cal_tools.agipdlib import (\n",
     "    AgipdCorrections,\n",
-    "    get_acq_rate,\n",
-    "    get_gain_mode,\n",
-    "    get_integration_time,\n",
-    "    get_gain_setting,\n",
-    "    get_num_cells,\n",
+    "    AgipdCtrl,\n",
+    "    CellRange,\n",
+    "    LitFrameSelection,\n",
     ")\n",
     "from cal_tools.ana_tools import get_range\n",
-    "from cal_tools.enums import BadPixels\n",
+    "from cal_tools.enums import AgipdGainMode, BadPixels\n",
     "from cal_tools.step_timing import StepTimer\n",
     "\n",
     "sns.set()\n",
@@ -168,7 +169,8 @@
    "outputs": [],
    "source": [
     "in_folder = Path(in_folder)\n",
-    "out_folder = Path(out_folder)"
+    "out_folder = Path(out_folder)\n",
+    "run_folder = in_folder / f'r{run:04d}'"
    ]
   },
   {
@@ -234,12 +236,11 @@
     "if sequences == [-1]:\n",
     "    sequences = None\n",
     "\n",
-    "control_fn = in_folder / f'r{run:04d}' / f'RAW-R{run:04d}-{karabo_da_control}-S00000.h5'\n",
-    "h5path_ctrl = h5path_ctrl.format(karabo_id_control)\n",
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
-    "h5path_idx = h5path_idx.format(karabo_id, receiver_id)\n",
+    "dc = RunDirectory(run_folder)\n",
     "\n",
-    "print(f'Path to control file {control_fn}')"
+    "ctrl_src = ctrl_source_template.format(karabo_id_control)\n",
+    "instrument_src = instrument_source_template.format(karabo_id, receiver_template)\n",
+    "index_src = index_source_template.format(karabo_id, receiver_template)"
    ]
   },
   {
@@ -285,7 +286,6 @@
    "source": [
     "if use_ppu_device:\n",
     "    # Obtain trains to process if using a pulse picker device.\n",
-    "    dc = RunDirectory(in_folder / f'r{run:04d}')\n",
     "\n",
     "    # Will throw an uncaught exception if the device is wrong.\n",
     "    seq_start = dc[use_ppu_device, 'trainTrigger.sequenceStart.value'].ndarray()\n",
@@ -295,21 +295,6 @@
     "    train_ids = np.unique(seq_start)[1:] + ppu_train_offset\n",
     "\n",
     "    print(f'PPU device {use_ppu_device} triggered for {len(train_ids)} train(s)')\n",
-    "    \n",
-    "    # Since we got the DataCollection already, narrow down the files we open.\n",
-    "    # This hardcodes the receiver_id and path_template parameters currently, but this\n",
-    "    # will disappear with moving the entire notebook to EXtra-data.\n",
-    "    subdc = dc.select_trains(by_id[train_ids]).select(f'{karabo_id}/DET/*CH0:xtdf')\n",
-    "    subseq = {int(f.filename[-8:-3]) for f in subdc.files}\n",
-    "    \n",
-    "    if sequences is None:\n",
-    "        # All sequences were meant to be processed by this job, so take the entire\n",
-    "        # subset of sequences.\n",
-    "        sequences = sorted(subseq)\n",
-    "    else:\n",
-    "        # If explicit sequences were specified (e.g. due to job balancing by xfel-calibrate)\n",
-    "        # only work on the intersection between that and what the PPU device offers.\n",
-    "        sequences = sorted(set(sequences) & subseq)\n",
     "\n",
     "elif train_ids != [-1]:\n",
     "    # Specific trains passed by parameter, convert to ndarray.\n",
@@ -326,9 +311,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# set everything up filewise\n",
@@ -362,25 +345,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "filename = file_list[0]\n",
-    "channel = int(re.findall(r\".*-AGIPD([0-9]+)-.*\", filename)[0])\n",
-    "\n",
-    "# Evaluate number of memory cells\n",
-    "mem_cells = get_num_cells(filename, karabo_id, channel)\n",
-    "if mem_cells is None:\n",
-    "    raise ValueError(f\"No raw images found in {filename}\")\n",
-    "\n",
-    "mem_cells_db = mem_cells if mem_cells_db == 0 else mem_cells_db\n",
-    "max_cells = mem_cells if max_cells == 0 else max_cells\n",
-    "\n",
-    "fast_paths = (filename, karabo_id, channel)\n",
-    "slow_paths = (control_fn, karabo_id_control)\n",
-    "\n",
-    "# Evaluate aquisition rate\n",
-    "if acq_rate == 0:\n",
-    "    acq_rate = get_acq_rate(fast_paths, slow_paths)\n",
-    "\n",
-    "print(f\"Maximum memory cells to calibrate: {max_cells}\")"
+    "first_mod_channel = sorted(modules)[0]\n",
+    "\n",
+    "instrument_src_mod = [\n",
+    "    s for s in list(dc.all_sources) if f\"{first_mod_channel}CH\" in s][0]\n",
+    "mod_channel = int(re.findall(rf\".*{first_mod_channel}CH([0-9]+):.*\", instrument_src_mod)[0])\n",
+    "\n",
+    "agipd_cond = AgipdCtrl(\n",
+    "    run_dc=dc,\n",
+    "    image_src=instrument_src_mod,\n",
+    "    ctrl_src=ctrl_src,\n",
+    "    raise_error=False,  # to be able to process very old data without gain_setting value\n",
+    ")"
    ]
   },
   {
@@ -397,26 +373,35 @@
     "    delta = timedelta(hours=offset.hour, minutes=offset.minute, seconds=offset.second)\n",
     "    creation_time += delta\n",
     "\n",
-    "# Evaluate gain setting\n",
-    "if gain_setting == 0.1:\n",
-    "    if creation_time.replace(tzinfo=None) < parser.parse('2020-01-31'):\n",
-    "        print(\"Set gain-setting to None for runs taken before 2020-01-31\")\n",
-    "        gain_setting = None\n",
-    "    else:\n",
-    "        try:\n",
-    "            gain_setting = get_gain_setting(str(control_fn), h5path_ctrl)\n",
-    "        except Exception as e:\n",
-    "            print(f'ERROR: while reading gain setting from: \\n{control_fn}')\n",
-    "            print(e)\n",
-    "            print(\"Set gain setting to 0\")\n",
-    "            gain_setting = 0\n",
-    "\n",
-    "# Evaluate gain mode (operation mode)\n",
-    "gain_mode = get_gain_mode(control_fn, h5path_ctrl)\n",
-    "\n",
-    "# Evaluate integration time\n",
-    "if integration_time < 0:\n",
-    "    integration_time = get_integration_time(control_fn, h5path_ctrl)"
+    "if acq_rate == 0.:\n",
+    "    acq_rate = agipd_cond.get_acq_rate()\n",
+    "if mem_cells == 0.:\n",
+    "    mem_cells = agipd_cond.get_num_cells()\n",
+    "# TODO: look for alternative for passing creation_time\n",
+    "if gain_setting == -1:\n",
+    "    gain_setting = agipd_cond.get_gain_setting(creation_time)\n",
+    "if bias_voltage == 0.:\n",
+    "    bias_voltage = agipd_cond.get_bias_voltage(karabo_id_control)\n",
+    "if integration_time == -1:\n",
+    "    integration_time = agipd_cond.get_integration_time()\n",
+    "if gain_mode == -1:\n",
+    "    gain_mode = agipd_cond.get_gain_mode()\n",
+    "else:\n",
+    "    gain_mode = AgipdGainMode(gain_mode)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if mem_cells is None:\n",
+    "    raise ValueError(f\"No raw images found in {filename}\")\n",
+    "\n",
+    "mem_cells_db = mem_cells if mem_cells_db == 0 else mem_cells_db\n",
+    "\n",
+    "print(f\"Maximum memory cells to calibrate: {mem_cells}\")"
    ]
   },
   {
@@ -449,6 +434,30 @@
     "            corr_bools[to_disable] = False"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if use_litframe_device:\n",
+    "    # check run for the AgipdLitFrameFinder device\n",
+    "\n",
+    "    if use_litframe_device + ':output' in dc.instrument_sources:\n",
+    "        # Use selection provided by the AgipdLitFrameFinder (if the device is recorded)\n",
+    "        cell_sel = LitFrameSelection(use_litframe_device, dc, train_ids, max_pulses, energy_threshold)\n",
+    "        train_ids = cell_sel.train_ids\n",
+    "    else:\n",
+    "        # Use range selection (if the device is not recorded)\n",
+    "        print(f\"WARNING: LitFrameFinder {use_litframe_device} device is not found.\")\n",
+    "        cell_sel = CellRange(max_pulses, max_cells=mem_cells)\n",
+    "else:\n",
+    "    # Use range selection\n",
+    "    cell_sel = CellRange(max_pulses, max_cells=mem_cells)\n",
+    "\n",
+    "print(cell_sel.msg())"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -463,10 +472,10 @@
    "outputs": [],
    "source": [
     "agipd_corr = AgipdCorrections(\n",
-    "    max_cells,\n",
-    "    max_pulses,\n",
-    "    h5_data_path=h5path,\n",
-    "    h5_index_path=h5path_idx,\n",
+    "    mem_cells,\n",
+    "    cell_sel,\n",
+    "    h5_data_path=instrument_src,\n",
+    "    h5_index_path=index_src,\n",
     "    corr_bools=corr_bools,\n",
     "    gain_mode=gain_mode,\n",
     "    comp_threads=os.cpu_count() // n_cores_files,\n",
@@ -556,7 +565,7 @@
    "outputs": [],
    "source": [
     "# allocate memory for images and hists\n",
-    "n_images_max = max_cells * 256\n",
+    "n_images_max = mem_cells * 256\n",
     "data_shape = (n_images_max, 512, 128)\n",
     "agipd_corr.allocate_images(data_shape, n_cores_files)"
    ]
@@ -586,6 +595,8 @@
     "\n",
     "    Yields (file data slot, start index, stop index)\n",
     "    \"\"\"\n",
+    "    \n",
+    "    \n",
     "    for i_proc, n_img in enumerate(img_counts):\n",
     "        n_chunks = math.ceil(n_img / chunk_size)\n",
     "        for i in range(n_chunks):\n",
@@ -615,8 +626,11 @@
     "    for file_batch in batches(file_list, n_cores_files):\n",
     "        # TODO: Move some printed output to logging or similar\n",
     "        print(f\"Processing next {len(file_batch)} files\")\n",
-    "        img_counts = pool.starmap(agipd_corr.read_file, zip(range(len(file_batch)), file_batch,\n",
-    "                                                                  [not common_mode]*len(file_batch)))\n",
+    "        step_timer.start()\n",
+    "        img_counts = pool.starmap(\n",
+    "            agipd_corr.read_file,\n",
+    "            zip(range(len(file_batch)), file_batch, [not common_mode]*len(file_batch))\n",
+    "        )\n",
     "        step_timer.done_step(f'Loading data from files')\n",
     "\n",
     "        if img_counts == 0:\n",
@@ -626,9 +640,12 @@
     "\n",
     "        if mask_zero_std:\n",
     "            # Evaluate zero-data-std mask\n",
-    "            pool.starmap(agipd_corr.mask_zero_std, itertools.product(\n",
-    "                range(len(file_batch)), np.array_split(np.arange(agipd_corr.max_cells), n_cores_correct)\n",
-    "            ))\n",
+    "            pool.starmap(\n",
+    "                agipd_corr.mask_zero_std, itertools.product(\n",
+    "                    range(len(file_batch)),\n",
+    "                    np.array_split(np.arange(agipd_corr.max_cells), n_cores_correct)\n",
+    "                )\n",
+    "            )\n",
     "            step_timer.done_step('Mask 0 std')\n",
     "\n",
     "        # Perform offset image-wise correction\n",
@@ -641,6 +658,8 @@
     "            step_timer.done_step(\"Base-line shift correction\")\n",
     "\n",
     "        if common_mode:\n",
+    "            # In common mode corrected is enabled.\n",
+    "            # Cell selection is only activated after common mode correction.\n",
     "            # Perform cross-file correction parallel over asics\n",
     "            pool.starmap(agipd_corr.cm_correction, itertools.product(\n",
     "                range(len(file_batch)), range(16)  # 16 ASICs per module\n",
@@ -648,11 +667,12 @@
     "            step_timer.done_step(\"Common-mode correction\")\n",
     "\n",
     "            img_counts = pool.map(agipd_corr.apply_selected_pulses, range(len(file_batch)))\n",
-    "            step_timer.done_step(\"Applying selected pulses after common mode correction\")\n",
+    "            step_timer.done_step(\"Applying selected cells after common mode correction\")\n",
     "\n",
     "        # Perform image-wise correction\n",
-    "        pool.starmap(agipd_corr.gain_correction, imagewise_chunks(img_counts))\n",
-    "        step_timer.done_step(\"Gain corrections\")\n",
+    "        if any(agipd_corr.pc_bools):\n",
+    "            pool.starmap(agipd_corr.gain_correction, imagewise_chunks(img_counts))\n",
+    "            step_timer.done_step(\"Gain corrections\")\n",
     "\n",
     "        # Save corrected data\n",
     "        pool.starmap(agipd_corr.write_file, [\n",
@@ -780,17 +800,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def get_trains_data(run_folder, source, include, detector_id, tid=None, modules=16, fillvalue=None):\n",
+    "def get_trains_data(data_folder, source, include, detector_id, tid=None, modules=16, fillvalue=None):\n",
     "    \"\"\"Load single train for all module\n",
     "\n",
-    "    :param run_folder: Path to folder with data\n",
+    "    :param data_folder: Path to folder with data\n",
     "    :param source: Data source to be loaded\n",
     "    :param include: Inset of file name to be considered\n",
     "    :param detector_id: The karabo id of the detector to get data for\n",
     "    :param tid: Train Id to be loaded. First train is considered if None is given\n",
     "    :param path: Path to find image data inside h5 file\n",
     "    \"\"\"\n",
-    "    run_data = RunDirectory(run_folder, include)\n",
+    "    run_data = RunDirectory(data_folder, include)\n",
     "    if tid is not None:\n",
     "        tid, data = run_data.select(f'{detector_id}/DET/*', source).train_from_id(tid)\n",
     "    else:\n",
@@ -830,7 +850,7 @@
     "_, blshift = get_trains_data(out_folder, 'image.blShift', include, karabo_id, tid, modules=nmods)\n",
     "_, cellId = get_trains_data(out_folder, 'image.cellId', include, karabo_id, tid, modules=nmods)\n",
     "_, pulseId = get_trains_data(out_folder, 'image.pulseId', include, karabo_id, tid, modules=nmods, fillvalue=0)\n",
-    "_, raw = get_trains_data(f'{in_folder}/r{run:04d}/', 'image.data', include, karabo_id, tid, modules=nmods)"
+    "_, raw = get_trains_data(run_folder, 'image.data', include, karabo_id, tid, modules=nmods)"
    ]
   },
   {
@@ -1216,9 +1236,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.8.12"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
-}
\ No newline at end of file
+ "nbformat_minor": 4
+}
diff --git a/notebooks/AGIPD/AGIPD_Retrieve_Constants_Precorrection.ipynb b/notebooks/AGIPD/AGIPD_Retrieve_Constants_Precorrection.ipynb
index 53ecb323c0e52497a23b0b0ad003f4d5a833c081..2cd11eb0b59cbb97f5c72c5e2ac9670e5573b95e 100644
--- a/notebooks/AGIPD/AGIPD_Retrieve_Constants_Precorrection.ipynb
+++ b/notebooks/AGIPD/AGIPD_Retrieve_Constants_Precorrection.ipynb
@@ -26,9 +26,10 @@
     "karabo_id = \"SPB_DET_AGIPD1M-1\" # karabo karabo_id\n",
     "karabo_da = ['-1']  # a list of data aggregators names, Default [-1] for selecting all data aggregators\n",
     "path_template = 'RAW-R{:04d}-{}-S{:05d}.h5' # the template to use to access data\n",
-    "h5path_ctrl = '/CONTROL/{}/MDL/FPGA_COMP_TEST' # path to control information\n",
+    "ctrl_source_template = '{}/MDL/FPGA_COMP_TEST'  # path to control information\n",
+    "instrument_source_template = '{}/DET/{}:xtdf'  # path in the HDF5 file to images\n",
+    "receiver_template = \"{}CH0\" # inset for receiver devices\n",
     "karabo_id_control = \"SPB_IRU_AGIPD1M1\" # karabo-id for control device\n",
-    "karabo_da_control = 'AGIPD1MCTRL00' # karabo DA for control infromation\n",
     "\n",
     "use_dir_creation_date = True # use the creation data of the input dir for database queries\n",
     "cal_db_interface = \"tcp://max-exfl016:8015#8045\" # the database interface to use\n",
@@ -38,12 +39,11 @@
     "calfile =  \"\" # path to calibration file. Leave empty if all data should come from DB\n",
     "nodb = False # if set only file-based constants will be used\n",
     "mem_cells = 0 # number of memory cells used, set to 0 to automatically infer\n",
-    "bias_voltage = 300\n",
+    "bias_voltage = 0  # bias voltage, set to 0 to use stored value in slow data.\n",
     "acq_rate = 0. # the detector acquisition rate, use 0 to try to auto-determine\n",
-    "gain_setting = 0.1 # the gain setting, use 0.1 to try to auto-determine\n",
+    "gain_setting = -1  # the gain setting, use -1 to use value stored in slow data.\n",
+    "gain_mode = -1  # gain mode (0: adaptive, 1-3 fixed high/med/low, -1: read from CONTROL data)\n",
     "photon_energy = 9.2 # photon energy in keV\n",
-    "max_cells_db_dark = 0  # set to a value different than 0 to use this value for dark data DB queries\n",
-    "max_cells_db = 0 # set to a value different than 0 to use this value for DB queries\n",
     "integration_time = -1 # integration time, negative values for auto-detection.\n",
     "\n",
     "# Correction Booleans\n",
@@ -85,19 +85,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from pathlib import Path\n",
     "from typing import List, Tuple\n",
     "\n",
     "import matplotlib\n",
+    "import matplotlib.pyplot as plt\n",
+    "import multiprocessing\n",
     "import numpy as np\n",
+    "from datetime import timedelta\n",
+    "from dateutil import parser\n",
+    "from extra_data import RunDirectory\n",
     "\n",
     "matplotlib.use(\"agg\")\n",
-    "import multiprocessing\n",
-    "from datetime import timedelta\n",
-    "from pathlib import Path\n",
     "\n",
-    "import matplotlib.pyplot as plt\n",
     "from cal_tools import agipdlib, tools\n",
-    "from dateutil import parser\n",
+    "from cal_tools.enums import AgipdGainMode\n",
     "from iCalibrationDB import Conditions, Constants, Detectors"
    ]
   },
@@ -119,8 +121,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "max_cells = mem_cells\n",
-    "\n",
     "creation_time = None\n",
     "if use_dir_creation_date:\n",
     "    creation_time = tools.get_dir_creation_date(str(in_folder), run)\n",
@@ -144,50 +144,21 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "control_fn = in_folder / f'r{run:04d}' / f'RAW-R{run:04d}-{karabo_da_control}-S00000.h5'\n",
-    "h5path_ctrl = h5path_ctrl.format(karabo_id_control)\n",
-    "slow_paths = (control_fn, karabo_id_control)\n",
-    "if gain_setting == 0.1:\n",
-    "    if creation_time.replace(tzinfo=None) < parser.parse('2020-01-31'):\n",
-    "        print(\"Set gain-setting to None for runs taken before 2020-01-31\")\n",
-    "        gain_setting = None\n",
-    "    else:\n",
-    "        try:\n",
-    "            gain_setting = agipdlib.get_gain_setting(str(control_fn), h5path_ctrl)\n",
-    "        except Exception as e:\n",
-    "            print(f'ERROR: while reading gain setting from: \\n{control_fn}')\n",
-    "            print(e)\n",
-    "            print(\"Set gain setting to 0\")\n",
-    "            gain_setting = 0\n",
-    "\n",
-    "# Evaluate gain mode (operation mode)\n",
-    "gain_mode = agipdlib.get_gain_mode(control_fn, h5path_ctrl)\n",
-    "\n",
-    "# Evaluate integration time\n",
-    "if integration_time < 0:\n",
-    "    integration_time = agipdlib.get_integration_time(control_fn, h5path_ctrl)\n",
-    "            \n",
-    "print(f\"Gain setting: {gain_setting}\")\n",
-    "print(f\"Gain mode: {gain_mode.name}\")\n",
-    "print(f\"Detector in use is {karabo_id}\")\n",
+    "ctrl_src = ctrl_source_template.format(karabo_id_control)\n",
     "\n",
+    "print(f\"Detector in use is {karabo_id}\")\n",
     "\n",
     "# Extracting Instrument string\n",
     "instrument = karabo_id.split(\"_\")[0]\n",
     "# Evaluate detector instance for mapping\n",
     "if instrument == \"SPB\":\n",
-    "    dinstance = \"AGIPD1M1\"\n",
     "    nmods = 16\n",
     "elif instrument == \"MID\":\n",
-    "    dinstance = \"AGIPD1M2\"\n",
     "    nmods = 16\n",
     "elif instrument == \"HED\":\n",
-    "    dinstance = \"AGIPD500K\"\n",
     "    nmods = 8\n",
     "\n",
     "print(f\"Instrument {instrument}\")\n",
-    "print(f\"Detector instance {dinstance}\")\n",
-    "\n",
     "\n",
     "if karabo_da[0] == '-1':\n",
     "    if modules[0] == -1:\n",
@@ -197,6 +168,48 @@
     "    modules = [int(x[-2:]) for x in karabo_da]"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "run_dc = RunDirectory(in_folder / f\"r{run:04d}\")\n",
+    "\n",
+    "# set everything up filewise\n",
+    "mapped_files, _, _, _, _ = tools.map_modules_from_folder(\n",
+    "    str(in_folder), run, path_template, karabo_da, sequences=[0]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Read AGIPD conditions from the 1st sequence of 1st module and slow data.\n",
+    "instrument_src = instrument_source_template.format(karabo_id, receiver_template)\n",
+    "instrument_src_mod = instrument_src.format(0)\n",
+    "\n",
+    "agipd_cond = agipdlib.AgipdCtrl(\n",
+    "    run_dc=run_dc,\n",
+    "    image_src=None,  # Not need, as we wont read mem_cells or acq_rate.\n",
+    "    ctrl_src=ctrl_src,\n",
+    ")\n",
+    "\n",
+    "if gain_setting == -1:\n",
+    "    gain_setting = agipd_cond.get_gain_setting(creation_time)\n",
+    "if bias_voltage == 0.:\n",
+    "    bias_voltage = agipd_cond.get_bias_voltage(karabo_id_control)\n",
+    "if integration_time == -1:\n",
+    "    integration_time = agipd_cond.get_integration_time()\n",
+    "if gain_mode == -1:\n",
+    "    gain_mode = agipd_cond.get_gain_mode()\n",
+    "else:\n",
+    "    gain_mode = AgipdGainMode(gain_mode)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -211,34 +224,33 @@
    "outputs": [],
    "source": [
     "def retrieve_constants(\n",
-    "    qm_files: List[Path], qm: str, karabo_da: str, idx: int\n",
+    "    karabo_da: str, idx: int\n",
     ") -> Tuple[str, str, float, float, str, dict]:\n",
     "    \"\"\"\n",
     "    Retrieve constants for a module.\n",
     "\n",
     "    :return:\n",
-    "            qm: module virtual name i.e. Q1M1.\n",
     "            karabo_da: karabo data aggregator.\n",
     "            acq_rate: acquisition rate parameter.\n",
-    "            max_cells: number of memory cells.\n",
+    "            mem_cells: number of memory cells.\n",
     "            mdata_dict: (DICT) dictionary with the metadata for the retrieved constants.\n",
     "    \"\"\"\n",
-    "    if max_cells != 0:\n",
-    "        # either use overriding notebook parameter\n",
-    "        local_max_cells = max_cells\n",
+    "    if mem_cells == 0:\n",
+    "        # either or look around in sequence files\n",
+    "        agipd_cond.image_src = instrument_src.format(idx)\n",
+    "        local_mem_cells = agipd_cond.get_num_cells()\n",
     "    else:\n",
-    "        # or look around in sequence files\n",
-    "        for f in qm_files:\n",
-    "            local_max_cells = agipdlib.get_num_cells(f, karabo_id, idx)\n",
-    "            if local_max_cells is not None:\n",
-    "                break\n",
+    "        # or use overriding notebook parameter\n",
+    "        local_mem_cells = mem_cells\n",
+    "\n",
     "    # maybe we never found this in a sequence file...\n",
-    "    if local_max_cells is None:\n",
-    "        raise ValueError(f\"No raw images found for {qm} for all sequences\")\n",
+    "    if local_mem_cells is None:\n",
+    "        raise ValueError(\n",
+    "            \"No raw images found for \"\n",
+    "            f\"{tools.module_index_to_qm(module_index)}({karabo_da}) for all sequences\")\n",
     "\n",
-    "    if acq_rate == 0:\n",
-    "        local_acq_rate = agipdlib.get_acq_rate(\n",
-    "            fast_paths=(f, karabo_id, idx), slow_paths=slow_paths)\n",
+    "    if acq_rate == 0.:\n",
+    "        local_acq_rate = agipd_cond.get_acq_rate()\n",
     "    else:\n",
     "        local_acq_rate = acq_rate\n",
     "\n",
@@ -249,7 +261,7 @@
     "    const_dict = agipdlib.assemble_constant_dict(\n",
     "        corr_bools,\n",
     "        pc_bools,\n",
-    "        local_max_cells,\n",
+    "        local_mem_cells,\n",
     "        bias_voltage,\n",
     "        gain_setting,\n",
     "        local_acq_rate,\n",
@@ -266,7 +278,7 @@
     "    mdata_dict[\"constants\"] = dict()\n",
     "    mdata_dict[\"physical-detector-unit\"] = None  # initialization\n",
     "\n",
-    "    for const_name, (const_init_fun, const_shape, (cond_type, cond_param)) in const_dict.items():\n",
+    "    for const_name, (const_init_fun, const_shape, (cond_type, cond_param)) in const_dict.items():  # noqa\n",
     "        if gain_mode and const_name in (\"ThresholdsDark\",):\n",
     "            continue\n",
     "        \n",
@@ -275,11 +287,14 @@
     "        mdata_dict[\"constants\"][const_name] = const_mdata\n",
     "\n",
     "        if slopes_ff_from_files and const_name in [\"SlopesFF\", \"BadPixelsFF\"]:\n",
-    "            const_mdata[\"file-path\"] = f\"{slopes_ff_from_files}/slopesff_bpmask_module_{qm}.h5\"\n",
+    "            const_mdata[\"file-path\"] = (\n",
+    "                f\"{slopes_ff_from_files}/slopesff_bpmask_module_{tools.module_index_to_qm(module_index)}.h5\")  # noqa\n",
     "            const_mdata[\"creation-time\"] = \"00:00:00\"\n",
     "            continue\n",
     "        \n",
-    "        if gain_mode and const_name in (\"BadPixelsPC\", \"SlopesPC\", \"BadPixelsFF\", \"SlopesFF\"):\n",
+    "        if gain_mode and const_name in (\n",
+    "            \"BadPixelsPC\", \"SlopesPC\", \"BadPixelsFF\", \"SlopesFF\"\n",
+    "        ):\n",
     "            param_copy = cond_param.copy()\n",
     "            del param_copy[\"gain_mode\"]\n",
     "            condition = getattr(Conditions, cond_type).AGIPD(**param_copy)\n",
@@ -309,7 +324,7 @@
     "            const_mdata[\"file-path\"] = const_dict[const_name][:2]\n",
     "            const_mdata[\"creation-time\"] = None\n",
     "\n",
-    "    return qm, mdata_dict, karabo_da, local_acq_rate, local_max_cells"
+    "    return mdata_dict, karabo_da, local_acq_rate, local_mem_cells"
    ]
   },
   {
@@ -328,18 +343,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# set everything up filewise\n",
-    "mapped_files, _, _, _, _ = tools.map_modules_from_folder(\n",
-    "    str(in_folder), run, path_template, karabo_da, sequences\n",
-    ")\n",
-    "\n",
     "pc_bools = [corr_bools.get(\"rel_gain\"),\n",
     "            corr_bools.get(\"adjust_mg_baseline\"),\n",
     "            corr_bools.get('blc_noise'),\n",
     "            corr_bools.get('blc_hmatch'),\n",
     "            corr_bools.get('blc_stripes'),\n",
     "            melt_snow]\n",
-    "\n",
     "inp = []\n",
     "only_dark = False\n",
     "nodb_with_dark = False\n",
@@ -350,20 +359,13 @@
     "\n",
     "da_to_qm = dict()\n",
     "for module_index, k_da in zip(modules, karabo_da):\n",
-    "    qm = tools.module_index_to_qm(module_index)\n",
-    "    da_to_qm[k_da] = qm\n",
-    "    \n",
+    "    da_to_qm[k_da] = tools.module_index_to_qm(module_index)\n",
     "    if k_da in retrieved_constants:\n",
-    "        print(f\"Constant for {k_da} already in calibration_metadata.yml, won't query again.\")\n",
-    "        continue\n",
-    "    \n",
-    "    if qm in mapped_files and not mapped_files[qm].empty():\n",
-    "        # TODO: make map_modules_from_folder just return list(s)\n",
-    "        qm_files = [Path(mapped_files[qm].get()) for _ in range(mapped_files[qm].qsize())]\n",
-    "    else:\n",
+    "        print(\n",
+    "            f\"Constant for {k_da} already in calibration_metadata.yml, won't query again.\")\n",
     "        continue\n",
     "\n",
-    "    inp.append((qm_files, qm, k_da, module_index))"
+    "    inp.append((k_da, module_index))"
    ]
   },
   {
@@ -382,8 +384,20 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "for qm, md_dict, karabo_da, acq_rate, max_cells in results:\n",
+    "acq_rate_mods = []\n",
+    "mem_cells_mods = []\n",
+    "for md_dict, karabo_da, acq_rate, mem_cells in results:\n",
     "    retrieved_constants[karabo_da] = md_dict\n",
+    "    mem_cells_mods.append(mem_cells)\n",
+    "    acq_rate_mods.append(acq_rate)\n",
+    "\n",
+    "# Validate that mem_cells and acq_rate are the same for all modules.\n",
+    "# TODO: Should a warning be enough?\n",
+    "if len(set(mem_cells_mods)) != 1 or len(set(acq_rate_mods)) != 1:\n",
+    "    print(\n",
+    "        \"WARNING: Number of memory cells or \"\n",
+    "        \"acquisition rate are not identical for all modules.\\n\"\n",
+    "        f\"mem_cells: {mem_cells_mods}.\\nacq_rate: {acq_rate_mods}.\")\n",
     "\n",
     "# check if it is requested not to retrieve any constants from the database\n",
     "if nodb_with_dark:\n",
@@ -394,7 +408,7 @@
     "          ', '.join([tools.module_index_to_qm(x) for x in modules]))\n",
     "    print(f\"Operating conditions are:\")\n",
     "    print(f\"• Bias voltage: {bias_voltage}\")\n",
-    "    print(f\"• Memory cells: {max_cells}\")\n",
+    "    print(f\"• Memory cells: {mem_cells}\")\n",
     "    print(f\"• Acquisition rate: {acq_rate}\")\n",
     "    print(f\"• Gain mode: {gain_mode.name}\")\n",
     "    print(f\"• Gain setting: {gain_setting}\")\n",
@@ -451,9 +465,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.11"
+   "version": "3.8.12"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }
diff --git a/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb b/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
index ff66d2a70934726d0febb03f3a9c11f46300cab3..be96463e0e8c90733e0e8532d1df0a732a714cc1 100644
--- a/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
+++ b/notebooks/AGIPD/Characterize_AGIPD_Gain_Darks_NBC.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# AGIPD Characterize Dark Images #\n",
     "\n",
-    "Author: S. Hauf, Version: 0.1\n",
+    "Author: European XFEL Detector Group, Version: 2.0\n",
     "\n",
     "The following code analyzes a set of dark images taken with the AGIPD detector to deduce detector offsets , noise, bad-pixel maps and thresholding. All four types of constants are evaluated per-pixel and per-memory cell. Data for the detector's three gain stages needs to be present, separated into separate runs.\n",
     "\n",
@@ -21,7 +21,6 @@
    "source": [
     "in_folder = \"/gpfs/exfel/d/raw/CALLAB/202031/p900113\" # path to input data, required\n",
     "out_folder = \"\" # path to output to, required\n",
-    "sequences = [-1] # sequence files to evaluate.\n",
     "modules = [-1]  # list of modules to evaluate, RANGE ALLOWED\n",
     "run_high = 9985 # run number in which high gain data was recorded, required\n",
     "run_med = 9984 # run number in which medium gain data was recorded, required\n",
@@ -30,13 +29,10 @@
     "\n",
     "karabo_id = \"HED_DET_AGIPD500K2G\" # karabo karabo_id\n",
     "karabo_da = ['-1']  # a list of data aggregators names, Default [-1] for selecting all data aggregators\n",
-    "receiver_id = \"{}CH0\" # inset for receiver devices\n",
-    "path_template = 'RAW-R{:04d}-{}-S{:05d}.h5' # the template to use to access data\n",
-    "h5path = '/INSTRUMENT/{}/DET/{}:xtdf/image' # path in the HDF5 file to images\n",
-    "h5path_idx = '/INDEX/{}/DET/{}:xtdf/image' # path in the HDF5 file to images\n",
-    "h5path_ctrl = '/CONTROL/{}/MDL/FPGA_COMP' # path to control information\n",
+    "receiver_template = \"{}CH0\" # inset for receiver devices\n",
+    "instrument_source_template = '{}/DET/{}:xtdf'  # path in the HDF5 file to images\n",
+    "ctrl_source_template = '{}/MDL/FPGA_COMP'  # path to control information\n",
     "karabo_id_control = \"HED_EXP_AGIPD500K2G\" # karabo-id for control device '\n",
-    "karabo_da_control = \"AGIPD500K2G00\" # karabo DA for control infromation\n",
     "\n",
     "use_dir_creation_date = True  # use dir creation date as data production reference date\n",
     "cal_db_interface = \"tcp://max-exfl016:8020\" # the database interface to use\n",
@@ -45,12 +41,12 @@
     "db_output = False # output constants to database\n",
     "\n",
     "mem_cells = 0 # number of memory cells used, set to 0 to automatically infer\n",
-    "bias_voltage = 0 # detector bias voltage\n",
-    "gain_setting = 0.1 # the gain setting, use 0.1 to try to auto-determine\n",
+    "bias_voltage = 0 # bias voltage, set to 0 to use stored value in slow data.\n",
+    "gain_setting = -1  # the gain setting, use -1 to use value stored in slow data.\n",
+    "gain_mode = -1  # gain mode, use -1 to use value stored in slow data.\n",
     "integration_time = -1 # integration time, negative values for auto-detection.\n",
     "acq_rate = 0. # the detector acquisition rate, use 0 to try to auto-determine\n",
     "interlaced = False # assume interlaced data format, for data prior to Dec. 2017\n",
-    "rawversion = 2 # RAW file format version\n",
     "\n",
     "thresholds_offset_sigma = 3. # offset sigma thresholds for offset deduced bad pixels\n",
     "thresholds_offset_hard = [0, 0]  # For setting the same threshold offset for the 3 gains. Left for backcompatability. Default [0, 0] to take the following parameters.\n",
@@ -67,8 +63,9 @@
     "thresholds_noise_hard_mg = [4, 20] # Medium-gain thresholds in absolute ADU terms for offset deduced bad pixels\n",
     "thresholds_noise_hard_lg = [4, 20] # Low-gain thresholds in absolute ADU terms for offset deduced bad pixels\n",
     "\n",
-    "thresholds_gain_sigma = 5. # Gain separation sigma threshold\n",
-    "\n",
+    "thresholds_gain_sigma = 5.  # Gain separation sigma threshold\n",
+    "max_trains = 0  # Maximum number of trains to use for processing dark. Set to 0 to process all available trains.\n",
+    "min_trains = 1  # Miniumum number of trains for processing dark. If raw folder has less than minimum trains processing is stopped.\n",
     "high_res_badpix_3d = False # set this to True if you need high-resolution 3d bad pixel plots. ~7mins extra time for 64 memory cells\n",
     "\n",
     "# This is used if modules is not specified:\n",
@@ -96,28 +93,24 @@
     "import os\n",
     "from collections import OrderedDict\n",
     "from datetime import timedelta\n",
-    "from typing import Tuple\n",
+    "from pathlib import Path\n",
+    "from typing import List, Tuple\n",
     "\n",
     "import dateutil.parser\n",
-    "import h5py\n",
     "import matplotlib\n",
     "import numpy as np\n",
     "import pasha as psh\n",
+    "import psutil\n",
     "import tabulate\n",
     "import yaml\n",
+    "from IPython.display import Latex, Markdown, display\n",
+    "from extra_data import RunDirectory\n",
     "\n",
     "matplotlib.use('agg')\n",
     "\n",
     "import iCalibrationDB\n",
     "import matplotlib.pyplot as plt\n",
-    "from cal_tools.agipdlib import (\n",
-    "    get_acq_rate,\n",
-    "    get_bias_voltage,\n",
-    "    get_gain_mode,\n",
-    "    get_gain_setting,\n",
-    "    get_integration_time,\n",
-    "    get_num_cells,\n",
-    ")\n",
+    "from cal_tools.agipdlib import AgipdCtrl\n",
     "from cal_tools.enums import AgipdGainMode, BadPixels\n",
     "from cal_tools.plotting import (\n",
     "    create_constant_overview,\n",
@@ -137,7 +130,6 @@
     "    save_const_to_h5,\n",
     "    send_to_db,\n",
     ")\n",
-    "from IPython.display import Latex, Markdown, display\n",
     "\n",
     "%matplotlib inline"
    ]
@@ -149,14 +141,19 @@
    "outputs": [],
    "source": [
     "# insert control device if format string (does nothing otherwise)\n",
-    "h5path_ctrl = h5path_ctrl.format(karabo_id_control)\n",
-    "\n",
-    "max_cells = mem_cells\n",
-    "\n",
-    "offset_runs = OrderedDict()\n",
-    "offset_runs[\"high\"] = run_high\n",
-    "offset_runs[\"med\"] = run_med\n",
-    "offset_runs[\"low\"] = run_low\n",
+    "ctrl_src = ctrl_source_template.format(karabo_id_control)\n",
+    "\n",
+    "runs_dict = OrderedDict()\n",
+    "\n",
+    "for gain_idx, (run_name, run_number) in enumerate(zip(\n",
+    "    [\"high\", \"med\", \"low\"],\n",
+    "    [run_high, run_med, run_low]\n",
+    ")):\n",
+    "    runs_dict[run_name] = {\n",
+    "        \"number\": run_number,\n",
+    "        \"gain\": gain_idx,\n",
+    "        \"dc\": RunDirectory(f'{in_folder}/r{run_number:04d}/')\n",
+    "    }\n",
     "\n",
     "creation_time=None\n",
     "if use_dir_creation_date:\n",
@@ -166,6 +163,10 @@
     "\n",
     "run, prop, seq = run_prop_seq_from_path(in_folder)\n",
     "\n",
+    "# Read report path and create file location tuple to add with the injection\n",
+    "file_loc = f\"proposal:{prop} runs:{run_low} {run_med} {run_high}\"\n",
+    "\n",
+    "report = get_report(out_folder)\n",
     "cal_db_interface = get_random_db_interface(cal_db_interface)\n",
     "print(f'Calibration database interface: {cal_db_interface}')\n",
     "\n",
@@ -181,31 +182,18 @@
     "    dinstance = \"AGIPD500K\"\n",
     "    nmods = 8\n",
     "\n",
-    "if sequences == [-1]:\n",
-    "    sequences = None\n",
-    "control_names = [f'{in_folder}/r{r:04d}/RAW-R{r:04d}-{karabo_da_control}-S00000.h5'\n",
-    "                 for r in (run_high, run_med, run_low)]\n",
+    "instrument_src = instrument_source_template.format(karabo_id, receiver_template)\n",
+    "run_numbers = [run_high, run_med, run_low]\n",
     "\n",
-    "if operation_mode not in (\"ADAPTIVE_GAIN\", \"FIXED_GAIN\"):\n",
-    "    print(f\"WARNING: unknown operation_mode \\\"{operation_mode}\\\" parameter set\")\n",
-    "run_gain_modes = [get_gain_mode(fn, h5path_ctrl) for fn in control_names]\n",
-    "if all(gm == AgipdGainMode.ADAPTIVE_GAIN for gm in run_gain_modes):\n",
-    "    fixed_gain_mode = False\n",
-    "    if operation_mode == \"FIXED_GAIN\":\n",
-    "        print(\"WARNING: operation_mode parameter is FIXED_GAIN, slow data indicates adaptive gain\")\n",
-    "elif run_gain_modes == [AgipdGainMode.FIXED_HIGH_GAIN, AgipdGainMode.FIXED_MEDIUM_GAIN, AgipdGainMode.FIXED_LOW_GAIN]:\n",
-    "    if operation_mode == \"ADAPTIVE_GAIN\":\n",
-    "        print(\"WARNING: operation_mode parameter ix ADAPTIVE_GAIN, slow data indicates fixed gain\")\n",
-    "    fixed_gain_mode = True\n",
-    "else:\n",
-    "    print(f'Something is clearly wrong; slow data indicates gain modes {run_gain_modes}')\n",
+    "def create_karabo_da_list(modules):\n",
+    "    return([\"AGIPD{:02d}\".format(i) for i in modules])\n",
     "\n",
-    "if integration_time < 0:\n",
-    "    integration_times = [get_integration_time(fn, h5path_ctrl) for fn in control_names]\n",
-    "    if len(set(integration_times)) > 1:\n",
-    "        print(f'WARNING: integration time is not constant across the specified dark runs')\n",
-    "\n",
-    "integration_time = integration_times[0]\n",
+    "if karabo_da[0] == '-1':\n",
+    "    if modules[0] == -1:\n",
+    "        modules = list(range(nmods))\n",
+    "    karabo_da = create_karabo_da_list(modules)\n",
+    "else:\n",
+    "    modules = [int(x[-2:]) for x in karabo_da]\n",
     "\n",
     "print(f\"Detector in use is {karabo_id}\")\n",
     "print(f\"Instrument {instrument}\")\n",
@@ -218,30 +206,85 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "runs = [run_high, run_med, run_low]\n",
-    "\n",
-    "if gain_setting == 0.1:\n",
-    "    if creation_time.replace(tzinfo=None) < dateutil.parser.parse('2020-01-31'):\n",
-    "        print(\"Set gain-setting to None for runs taken before 2020-01-31\")\n",
-    "        gain_setting = None\n",
+    "# Create out_folder if it doesn't exist.\n",
+    "Path(out_folder).mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "n_files = 0\n",
+    "total_file_sizes = 0\n",
+    "max_trains_list = []\n",
+    "\n",
+    "for run_dict in runs_dict.values():\n",
+    "    missing_modules = []\n",
+    "    image_dc = run_dict[\"dc\"].select(f\"{karabo_id_control}*\", \"*\", require_all=True)\n",
+    "    # This is important in case of no slurm parallelization over modules is done.\n",
+    "    # (e.g. running notebook interactively)\n",
+    "    sources_l = [(f\"{karabo_id_control}*\", \"*\")]\n",
+    "    sources_l += [(instrument_src.format(m), \"*\") for m in modules]\n",
+    "    image_dc = run_dict[\"dc\"].select(sources_l, require_all=True)\n",
+    "    # validate that there are trains and that data sources are\n",
+    "    # present for any of the selected modules.\n",
+    "    if (\n",
+    "        len(image_dc.train_ids) == 0 or\n",
+    "        not np.any([\n",
+    "            karabo_id in s for s in run_dict[\"dc\"].select(sources_l, require_all=True).all_sources])  # noqa\n",
+    "    ):\n",
+    "        raise ValueError(f\"No images to process for run: {run_dict['number']}\")\n",
+    "\n",
+    "    max_trains_list.append(len(image_dc.train_ids))\n",
+    "\n",
+    "    # update run_dc with selected module sources\n",
+    "    run_dict[\"dc\"] = image_dc\n",
+    "\n",
+    "# Update modules and karabo_da lists based on available modules to processes.\n",
+    "modules = [m for m in modules if m not in missing_modules]\n",
+    "karabo_da = create_karabo_da_list(modules)\n",
+    "\n",
+    "# Remodifing run data collections to display actual total files number and size. \n",
+    "for run_dict in runs_dict.values():\n",
+    "    file_sizes = [os.path.getsize(f.filename) / 1e9 for f in run_dict[\"dc\"].deselect(f\"{karabo_id_control}*\").files]\n",
+    "    total_file_sizes += sum(file_sizes)\n",
+    "    n_files += len(file_sizes)\n",
+    "\n",
+    "print(f\"Will process data in a total of {n_files} files ({total_file_sizes:.02f} GB).\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Read and validate the runs control data."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def read_run_conditions(runs_dict: dict):\n",
+    "    agipd_cond = AgipdCtrl(\n",
+    "        run_dc=runs_dict[\"dc\"],\n",
+    "        image_src=instrument_src_mod,\n",
+    "        ctrl_src=ctrl_src,\n",
+    "    )\n",
+    "    cond_dict[\"runs\"].append(runs_dict[\"number\"])\n",
+    "    if acq_rate == 0:\n",
+    "        cond_dict[\"acq_rate\"].append(agipd_cond.get_acq_rate())\n",
+    "    if mem_cells == 0:\n",
+    "        cond_dict[\"mem_cells\"].append(agipd_cond.get_num_cells())\n",
+    "    if gain_setting == -1:    \n",
+    "        cond_dict[\"gain_setting\"].append(\n",
+    "            agipd_cond.get_gain_setting(creation_time))\n",
+    "    if bias_voltage == 0.:\n",
+    "        cond_dict[\"bias_voltage\"].append(\n",
+    "            agipd_cond.get_bias_voltage(karabo_id_control))\n",
+    "    if integration_time == -1:\n",
+    "        cond_dict[\"integration_time\"].append(\n",
+    "            agipd_cond.get_integration_time())\n",
+    "    if gain_mode == -1:\n",
+    "        cond_dict[\"gain_mode\"].append(agipd_cond.get_gain_mode())\n",
     "    else:\n",
-    "        try:\n",
-    "            # extract gain setting and validate that all runs have the same setting\n",
-    "            gsettings = []\n",
-    "            for r in runs:\n",
-    "                control_fname = '{}/r{:04d}/RAW-R{:04d}-{}-S00000.h5'.format(in_folder, r, r,\n",
-    "                                                                             karabo_da_control)\n",
-    "                gsettings.append(get_gain_setting(control_fname, h5path_ctrl))\n",
-    "            if not all(g == gsettings[0] for g in gsettings):\n",
-    "                raise ValueError(f\"Different gain settings for the 3 input runs {gsettings}\")\n",
-    "            gain_setting = gsettings[0]\n",
-    "        except Exception as e:\n",
-    "            print(f'Error while reading gain setting from: \\n{control_fname}')\n",
-    "            print(f'Error: {e}')\n",
-    "            if \"component not found\" in str(e):\n",
-    "                print(\"Gain setting is not found in the control information\")\n",
-    "            print(\"Data will not be processed\")\n",
-    "            sequences = []"
+    "        cond_dict[\"gain_mode\"].append(AgipdGainMode(gain_mode))"
    ]
   },
   {
@@ -250,26 +293,116 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if karabo_da[0] == '-1':\n",
-    "    if modules[0] == -1:\n",
-    "        modules = list(range(nmods))\n",
-    "    karabo_da = [\"AGIPD{:02d}\".format(i) for i in modules]\n",
-    "else:\n",
-    "    modules = [int(x[-2:]) for x in karabo_da]\n",
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
-    "h5path_idx = h5path_idx.format(karabo_id, receiver_id)\n",
-    "\n",
-    "if bias_voltage == 0:\n",
-    "    # Read the bias voltage from files, if recorded.\n",
-    "    # If not available, make use of the historical voltage the detector is running at\n",
-    "    bias_voltage = get_bias_voltage(control_names[0], karabo_id_control)\n",
-    "    bias_voltage = bias_voltage if bias_voltage is not None else 300\n",
+    "def validate_gain_modes(gain_modes: List[AgipdGainMode]):\n",
+    "    # Validate that gain modes are not a mix of adaptive and fixed gain.\n",
+    "    if all(\n",
+    "        gm == AgipdGainMode.ADAPTIVE_GAIN for gm in gain_modes\n",
+    "    ):\n",
+    "        fixed_gain_mode = False\n",
+    "    elif any(\n",
+    "        gm == AgipdGainMode.ADAPTIVE_GAIN for gm in gain_modes\n",
+    "    ):\n",
+    "        raise ValueError(\n",
+    "            f\"ERROR: Given runs {self.read_conditions['run_number']}\"\n",
+    "            \" have a mix of ADAPTIVE and FIXED gain modes: \"\n",
+    "            f\"{self.read_conditions['gain_mode']}.\"\n",
+    "    )\n",
+    "    else:\n",
+    "        fixed_gain_mode = True\n",
+    "    return fixed_gain_mode"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Read slow data from 1st channel only.\n",
+    "# Read all modules in one notebook and validate the conditions across detectors?\n",
+    "# Currently slurm jobs run per one module.\n",
+    "\n",
+    "# TODO: what if first module is not available. Maybe only channel 2 available\n",
+    "instrument_src_mod = instrument_src.format(modules[0])\n",
+    "\n",
+    "cond_dict = dict()\n",
+    "fixed_gain_mode = None\n",
+    "\n",
+    "with multiprocessing.Manager() as manager:\n",
+    "    cond_dict[\"runs\"] = manager.list()\n",
+    "    cond_dict[\"acq_rate\"] = manager.list()\n",
+    "    cond_dict[\"mem_cells\"] = manager.list()\n",
+    "    cond_dict[\"gain_setting\"] = manager.list()\n",
+    "    cond_dict[\"gain_mode\"] = manager.list()\n",
+    "    cond_dict[\"bias_voltage\"] = manager.list()\n",
+    "    cond_dict[\"integration_time\"] = manager.list()\n",
+    "\n",
+    "    with multiprocessing.Pool(processes=len(modules)) as pool:\n",
+    "        pool.starmap(read_run_conditions, zip(runs_dict.values()))\n",
+    "\n",
+    "    for cond, vlist in cond_dict.items():\n",
+    "        if cond == \"runs\":\n",
+    "            continue\n",
+    "        elif cond == \"gain_mode\":\n",
+    "            fixed_gain_mode = validate_gain_modes(cond_dict[\"gain_mode\"])\n",
+    "        if not all(x == vlist[0] for x in vlist):\n",
+    "            # TODO: raise ERROR??\n",
+    "            print(\n",
+    "                f\"WARNING: {cond} is not the same for the runs \"\n",
+    "                f\"{cond_dict['runs']} with values\"\n",
+    "                f\" of {cond_dict[cond]}, respectively.\"\n",
+    "            )\n",
+    "    if cond_dict[\"acq_rate\"]: acq_rate = cond_dict[\"acq_rate\"][0]\n",
+    "    if cond_dict[\"mem_cells\"]: mem_cells = cond_dict[\"mem_cells\"][0]\n",
+    "    if cond_dict[\"gain_setting\"]: gain_setting = cond_dict[\"gain_setting\"][0]\n",
+    "    if cond_dict[\"gain_mode\"]: gain_mode = list(cond_dict[\"gain_mode\"])\n",
+    "    if cond_dict[\"bias_voltage\"]: bias_voltage = cond_dict[\"bias_voltage\"][0]\n",
+    "    if cond_dict[\"integration_time\"]: integration_time = cond_dict[\"integration_time\"][0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Determine the gain operation mode based on the gain_mode stored in control h5file.\n",
+    "if operation_mode not in (\"ADAPTIVE_GAIN\", \"FIXED_GAIN\"):\n",
+    "    print(f\"WARNING: unknown operation_mode \\\"{operation_mode}\\\" parameter set\")\n",
     "\n",
+    "if (\n",
+    "    gain_mode == [\n",
+    "        AgipdGainMode.FIXED_HIGH_GAIN,\n",
+    "        AgipdGainMode.FIXED_MEDIUM_GAIN,\n",
+    "        AgipdGainMode.FIXED_LOW_GAIN\n",
+    "    ] and\n",
+    "    operation_mode == \"ADAPTIVE_GAIN\"\n",
+    "):\n",
+    "    print(\n",
+    "        \"WARNING: operation_mode parameter is ADAPTIVE_GAIN, \"\n",
+    "        \"slow data indicates FIXED_GAIN.\")\n",
+    "elif not fixed_gain_mode and operation_mode == \"FIXED_GAIN\":\n",
+    "    print(\n",
+    "        \"WARNING: operation_mode parameter is FIXED_GAIN, \"\n",
+    "        \"slow data indicates ADAPTIVE_GAIN\")\n",
+    "elif not all(gm == AgipdGainMode.ADAPTIVE_GAIN for gm in gain_mode):\n",
+    "    raise ValueError(\n",
+    "        \"ERROR: Wrong arrangment of given dark runs. \"\n",
+    "        f\"Given runs' gain_modes are {gain_mode} for runs: {runs}.\"\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "print(\"Parameters are:\")\n",
     "print(f\"Proposal: {prop}\")\n",
-    "print(f\"Memory cells: {mem_cells}/{max_cells}\")\n",
-    "print(\"Runs: {}\".format([v for v in offset_runs.values()]))\n",
-    "print(f\"Sequences: {sequences if sequences else 'All'}\")\n",
+    "print(f\"Acquisition rate: {acq_rate}\")\n",
+    "print(f\"Memory cells: {mem_cells}\")\n",
+    "print(f\"Runs: {run_numbers}\")\n",
     "print(f\"Interlaced mode: {interlaced}\")\n",
     "print(f\"Using DB: {db_output}\")\n",
     "print(f\"Input: {in_folder}\")\n",
@@ -315,51 +448,28 @@
     "    ]"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The following lines will create a queue of files which will the be executed module-parallel. Distiguishing between different gains."
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# set everything up filewise\n",
-    "os.makedirs(out_folder, exist_ok=True)\n",
-    "gain_mapped_files, total_files, total_file_size = map_gain_stages(\n",
-    "    in_folder, offset_runs, path_template, karabo_da, sequences\n",
-    ")\n",
-    "print(f\"Will process a total of {total_files} files ({total_file_size:.02f} GB).\")\n",
-    "\n",
-    "inp = []\n",
-    "for gain_index, (gain, qm_file_map) in enumerate(gain_mapped_files.items()):\n",
-    "    gain_input = []\n",
-    "    for module_index in modules:\n",
-    "        qm = module_index_to_qm(module_index)\n",
-    "        if qm not in qm_file_map:\n",
-    "            print(f\"Did not find files for {qm}\")\n",
-    "            continue\n",
-    "        file_queue = qm_file_map[qm]\n",
-    "        while not file_queue.empty():\n",
-    "            filename = file_queue.get()\n",
-    "            # TODO: remove after using EXtra-data to read files\n",
-    "            # and skip empty trains.\n",
-    "            with h5py.File(filename, \"r\") as fin:\n",
-    "                if fin[h5path.format(module_index)+\"/trainId\"].shape[0] != 0:\n",
-    "                    print(f\"Process {filename} for {qm}\")\n",
-    "                    gain_input.append((filename, module_index, gain_index))\n",
-    "                else:\n",
-    "                    print(f\"Do not process {filename} because it is empty.\")\n",
-    "    if not gain_input:\n",
-    "        raise ValueError(\n",
-    "            \"No images to process for run: \"\n",
-    "            f\"{[v for v in offset_runs.values()][gain_index]}\"\n",
-    "        )\n",
-    "    inp += gain_input"
+    "# Check if max_trains can be processed.\n",
+    "\n",
+    "# more relevant if running on multiple modules (i.e. within notebook)\n",
+    "# mem_cells * gains * n_constants * modules * agipd_[x,y]image_size * 2\n",
+    "av_mem = psutil.virtual_memory().available\n",
+    "possible_trains = av_mem // (352 * 3 * 3 * len(modules) * 131072 * 2)\n",
+    "if max_trains == 0:\n",
+    "    max_trains = max(max_trains_list)\n",
+    "if max_trains > possible_trains:\n",
+    "    max_trains = possible_trains\n",
+    "    print(\n",
+    "        f\"WARNING: available memory for processing is { av_mem / 1e9:.02f} GB.\"\n",
+    "        f\" Modifing max_trains to process to {max_trains}\")\n",
+    "\n",
+    "for run_dict in runs_dict.values():\n",
+    "    run_dict[\"dc\"] = run_dict[\"dc\"].select_trains(np.s_[:max_trains])"
    ]
   },
   {
@@ -377,49 +487,31 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# min() only relevant if running on multiple modules (i.e. within notebook)\n",
-    "parallel_num_procs = min(12, total_files)\n",
+    "parallel_num_procs = min(12, len(modules)*3)\n",
     "parallel_num_threads = multiprocessing.cpu_count() // parallel_num_procs\n",
     "print(f\"Will use {parallel_num_procs} processes with {parallel_num_threads} threads each\")\n",
     "\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",
-    "    if max_cells == 0:\n",
-    "        num_cells = get_num_cells(fast_data_filename, karabo_id, channel)\n",
-    "    else:\n",
-    "        num_cells = max_cells\n",
+    "    channel: int, runs_dict: dict,\n",
+    ") -> Tuple[int, int, np.array, np.array, np.array, np.array, np.array]:\n",
     "\n",
-    "    if acq_rate == 0.:\n",
-    "        slow_paths = control_names[gain_index], karabo_id_control\n",
-    "        fast_paths = fast_data_filename, karabo_id, channel\n",
-    "        local_acq_rate = get_acq_rate(fast_paths, slow_paths)\n",
-    "    else:\n",
-    "        local_acq_rate = acq_rate\n",
+    "    # Select the corresponding module channel.\n",
+    "    instrument_src_mod = instrument_src.format(channel)\n",
     "\n",
-    "    local_thresholds_offset_hard = thresholds_offset_hard[gain_index]\n",
-    "    local_thresholds_noise_hard = thresholds_noise_hard[gain_index]\n",
+    "    run_dc = runs_dict[\"dc\"]\n",
+    "    gain_index = runs_dict[\"gain\"]\n",
     "\n",
-    "    h5path_f = h5path.format(channel)\n",
-    "    h5path_idx_f = h5path_idx.format(channel)\n",
+    "    if run_dc[instrument_src_mod, \"image.data\"].shape[0] < min_trains:\n",
+    "        print(\n",
+    "            f\"WARNING: {run_dc.files} have less than \"\n",
+    "            \"minimum trains: {min_trains}.\")\n",
     "\n",
-    "    with h5py.File(fast_data_filename, \"r\") as infile:\n",
-    "        if rawversion == 2:\n",
-    "            count = np.squeeze(infile[f\"{h5path_idx_f}/count\"])\n",
-    "            first = np.squeeze(infile[f\"{h5path_idx_f}/first\"])\n",
-    "            last_index = int(first[count != 0][-1]+count[count != 0][-1])\n",
-    "            first_index = int(first[count != 0][0])\n",
-    "        else:\n",
-    "            status = np.squeeze(infile[f\"{h5path_idx_f}/status\"])\n",
-    "            if np.count_nonzero(status != 0) == 0:\n",
-    "                return\n",
-    "            last = np.squeeze(infile[f\"{h5path_idx_f}/last\"])\n",
-    "            first = np.squeeze(infile[f\"{h5path_idx_f}/first\"])\n",
-    "            last_index = int(last[status != 0][-1]) + 1\n",
-    "            first_index = int(first[status != 0][0])\n",
-    "        im = np.array(infile[f\"{h5path_f}/data\"][first_index:last_index,...])\n",
-    "        cell_ids = np.squeeze(infile[f\"{h5path_f}/cellId\"][first_index:last_index,...])\n",
+    "    # Read module's image and cellId data.\n",
+    "    im = run_dc[instrument_src_mod, \"image.data\"].ndarray()\n",
+    "    cell_ids = np.squeeze(run_dc[instrument_src_mod, \"image.cellId\"].ndarray())\n",
+    "\n",
+    "    local_thresholds_offset_hard = thresholds_offset_hard[gain_index]\n",
+    "    local_thresholds_noise_hard = thresholds_noise_hard[gain_index]    \n",
     "\n",
     "    if interlaced:\n",
     "        if not fixed_gain_mode:\n",
@@ -430,13 +522,12 @@
     "        if not fixed_gain_mode:\n",
     "            ga = im[:, 1, ...]\n",
     "        im = im[:, 0, ...].astype(np.float32)\n",
-    "\n",
     "    im = np.transpose(im)\n",
     "    if not fixed_gain_mode:\n",
     "        ga = np.transpose(ga)\n",
     "\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",
+    "    offset = context.alloc(shape=(im.shape[0], im.shape[1], mem_cells), dtype=np.float64)\n",
     "    noise = context.alloc(like=offset)\n",
     "\n",
     "    if fixed_gain_mode:\n",
@@ -455,7 +546,6 @@
     "            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",
-    "\n",
     "    context.map(process_cell, np.unique(cell_ids))\n",
     "\n",
     "    # bad pixels\n",
@@ -478,7 +568,7 @@
     "    bp[(noise < local_thresholds_noise_hard[0]) | (noise > local_thresholds_noise_hard[1])] |= BadPixels.NOISE_OUT_OF_THRESHOLD\n",
     "    bp[~np.isfinite(noise)] |= BadPixels.OFFSET_NOISE_EVAL_ERROR\n",
     "\n",
-    "    return offset, noise, gains, gains_std, bp, num_cells, local_acq_rate"
+    "    return channel, gain_index, offset, noise, gains, gains_std, bp"
    ]
   },
   {
@@ -488,7 +578,18 @@
    "outputs": [],
    "source": [
     "with multiprocessing.Pool(processes=parallel_num_procs) as pool:\n",
-    "    results = pool.starmap(characterize_module, inp)"
+    "    results = pool.starmap(\n",
+    "        characterize_module, itertools.product(modules, list(runs_dict.values())))\n",
+    "\n",
+    "# mapped values for processing 2 modules example:\n",
+    "# [\n",
+    "#     0, {\"gain\": 0, \"run_number\": <run-high>, \"dc\": <high-dc>},\n",
+    "#     0, {\"gain\": 1, \"run_number\": <run-med>, \"dc\": <med-dc>},\n",
+    "#     0, {\"gain\": 2, \"run_number\": <run-low>, \"dc\": <low-dc>},\n",
+    "#     1, {\"gain\": 0, \"run_number\": <run-high>, \"dc\": <high-dc>},\n",
+    "#     1, {\"gain\": 1, \"run_number\": <run-med>, \"dc\": <med-dc>},\n",
+    "#     1, {\"gain\": 2, \"run_number\": <run-low>, \"dc\": <low-dc>},\n",
+    "# ]"
    ]
   },
   {
@@ -504,13 +605,8 @@
     "    gain_g = OrderedDict()\n",
     "    gainstd_g = OrderedDict()\n",
     "\n",
-    "all_cells = []\n",
-    "all_acq_rate = []\n",
     "\n",
-    "for (_, module_index, gain_index), (offset, noise, gains, gains_std, bp,\n",
-    "                                    thiscell, thisacq) in zip(inp, results):\n",
-    "    all_cells.append(thiscell)\n",
-    "    all_acq_rate.append(thisacq)\n",
+    "for module_index, gain_index, offset, noise, gains, gains_std, bp in results:\n",
     "    qm = module_index_to_qm(module_index)\n",
     "    if qm not in offset_g:\n",
     "        offset_g[qm] = np.zeros((offset.shape[0], offset.shape[1], offset.shape[2], 3))\n",
@@ -525,13 +621,7 @@
     "    badpix_g[qm][..., gain_index] = bp\n",
     "    if not fixed_gain_mode:\n",
     "        gain_g[qm][..., gain_index] = gains\n",
-    "        gainstd_g[qm][..., gain_index] = gains_std\n",
-    "\n",
-    "\n",
-    "max_cells = np.max(all_cells)\n",
-    "print(f\"Using {max_cells} memory cells\")\n",
-    "acq_rate = np.max(all_acq_rate)\n",
-    "print(f\"Using {acq_rate} MHz acquisition rate\")"
+    "        gainstd_g[qm][..., gain_index] = gains_std"
    ]
   },
   {
@@ -593,19 +683,6 @@
     "        res[qm]['ThresholdsDark'] = thresholds_g[qm]"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Read report path and create file location tuple to add with the injection\n",
-    "proposal = list(filter(None, in_folder.strip('/').split('/')))[-2]\n",
-    "file_loc = 'proposal:{} runs:{} {} {}'.format(proposal, run_low, run_med, run_high)\n",
-    "\n",
-    "report = get_report(out_folder)"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -615,7 +692,7 @@
     "# set the operating condition\n",
     "# note: iCalibrationDB only adds gain_mode if it is truthy, so we don't need to handle None\n",
     "condition = iCalibrationDB.Conditions.Dark.AGIPD(\n",
-    "    memory_cells=max_cells,\n",
+    "    memory_cells=mem_cells,\n",
     "    bias_voltage=bias_voltage,\n",
     "    acquisition_rate=acq_rate,\n",
     "    gain_setting=gain_setting,\n",
@@ -638,7 +715,7 @@
     "    constant=iCalibrationDB.CalibrationConstant(),\n",
     "    condition=condition,\n",
     "    cal_db_interface=cal_db_interface,\n",
-    "    snapshot_at=creation_time.isoformat(),\n",
+    "    snapshot_at=creation_time.isoformat() if creation_time else None,\n",
     "    timeout=cal_db_timeout\n",
     ")\n",
     "for module_index, module_da, module_pdu in zip(modules, karabo_da, all_pdus):\n",
@@ -649,6 +726,13 @@
     "    }"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Sending calibration constants to the database."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -673,11 +757,18 @@
     "                                  file_loc, report, creation_time, out_folder)\n",
     "            print(f\"Calibration constant {const} for {qm} is stored locally in {file_loc}.\\n\")\n",
     "\n",
-    "    print(\"Constants parameter conditions are:\\n\")\n",
-    "    print(f\"• memory_cells: {max_cells}\\n• bias_voltage: {bias_voltage}\\n\"\n",
-    "          f\"• acquisition_rate: {acq_rate}\\n• gain_setting: {gain_setting}\\n\"\n",
-    "          f\"• gain_mode: {fixed_gain_mode}\\n• integration_time: {integration_time}\\n\"\n",
-    "          f\"• creation_time: {md.calibration_constant_version.begin_at if md is not None else creation_time}\\n\")"
+    "print(\"Constants parameter conditions are:\\n\")\n",
+    "print(f\"• memory_cells: {mem_cells}\\n• bias_voltage: {bias_voltage}\\n\"\n",
+    "      f\"• acquisition_rate: {acq_rate}\\n• gain_setting: {gain_setting}\\n\"\n",
+    "      f\"• gain_mode: {fixed_gain_mode}\\n• integration_time: {integration_time}\\n\"\n",
+    "      f\"• creation_time: {md.calibration_constant_version.begin_at if md is not None else creation_time}\\n\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Retrieving previous calibration constants for comparison."
    ]
   },
   {
@@ -693,11 +784,6 @@
     "def retrieve_old_constant(qm, const):\n",
     "    dconst = getattr(iCalibrationDB.Constants.AGIPD, const)()\n",
     "\n",
-    "    # This should be used in case of running notebook\n",
-    "    # by a different method other than myMDC which already\n",
-    "    # sends CalCat info.\n",
-    "    # TODO: Set db_module to \"\" by default in the first cell\n",
-    "\n",
     "    data, mdata = get_from_db(\n",
     "        karabo_id=karabo_id,\n",
     "        karabo_da=qm_dict[qm][\"karabo_da\"],\n",
@@ -705,7 +791,7 @@
     "        condition=condition,\n",
     "        empty_constant=None,\n",
     "        cal_db_interface=cal_db_interface,\n",
-    "        creation_time=creation_time-timedelta(seconds=1),\n",
+    "        creation_time=creation_time-timedelta(seconds=1) if creation_time else None,\n",
     "        strategy=\"pdu_prior_in_time\",\n",
     "        verbosity=1,\n",
     "        timeout=cal_db_timeout\n",
@@ -771,7 +857,7 @@
    "source": [
     "cell = 3\n",
     "gain = 0\n",
-    "show_overview(res, cell, gain, infix=\"{}-{}-{}\".format(*offset_runs.values()))"
+    "show_overview(res, cell, gain, infix=\"{}-{}-{}\".format(*run_numbers))"
    ]
   },
   {
@@ -789,7 +875,7 @@
    "source": [
     "cell = 3\n",
     "gain = 1\n",
-    "show_overview(res, cell, gain, infix=\"{}-{}-{}\".format(*offset_runs.values()))"
+    "show_overview(res, cell, gain, infix=\"{}-{}-{}\".format(*run_numbers))"
    ]
   },
   {
@@ -807,7 +893,7 @@
    "source": [
     "cell = 3\n",
     "gain = 2\n",
-    "show_overview(res, cell, gain, infix=\"{}-{}-{}\".format(*offset_runs.values()))"
+    "show_overview(res, cell, gain, infix=\"{}-{}-{}\".format(*run_numbers))"
    ]
   },
   {
@@ -851,6 +937,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "\n",
     "## Aggregate values, and per Cell behaviour ##\n",
     "\n",
     "The following tables and plots give an overview of statistical aggregates for each constant, as well as per cell behavior."
@@ -862,7 +949,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "create_constant_overview(offset_g, \"Offset (ADU)\", max_cells, 4000, 8000,\n",
+    "create_constant_overview(offset_g, \"Offset (ADU)\", mem_cells, 4000, 8000,\n",
     "                         badpixels=[badpix_g, np.nan])"
    ]
   },
@@ -872,7 +959,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "create_constant_overview(noise_g, \"Noise (ADU)\", max_cells, 0, 100,\n",
+    "create_constant_overview(noise_g, \"Noise (ADU)\", mem_cells, 0, 100,\n",
     "                         badpixels=[badpix_g, np.nan])"
    ]
   },
@@ -890,7 +977,7 @@
     "        bp_thresh[mod][...,:2] = con[...,:2]\n",
     "        bp_thresh[mod][...,2:] = con\n",
     "\n",
-    "    create_constant_overview(thresholds_g, \"Threshold (ADU)\", max_cells, 4000, 10000, 5,\n",
+    "    create_constant_overview(thresholds_g, \"Threshold (ADU)\", mem_cells, 4000, 10000, 5,\n",
     "                             badpixels=[bp_thresh, np.nan],\n",
     "                             gmap=['HG-MG Threshold', 'MG-LG Threshold', 'High gain', 'Medium gain', 'low gain'],\n",
     "                             marker=['d','d','','','']\n",
@@ -906,7 +993,7 @@
     "bad_pixel_aggregate_g = OrderedDict()\n",
     "for m, d in badpix_g.items():\n",
     "    bad_pixel_aggregate_g[m] = d.astype(np.bool).astype(np.float)\n",
-    "create_constant_overview(bad_pixel_aggregate_g, \"Bad pixel fraction\", max_cells, 0, 0.10, 3)"
+    "create_constant_overview(bad_pixel_aggregate_g, \"Bad pixel fraction\", mem_cells, 0, 0.10, 3)"
    ]
   },
   {
@@ -1115,7 +1202,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.11"
+   "version": "3.8.12"
   }
  },
  "nbformat": 4,
diff --git a/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb b/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb
index 021e16ec615b20c9ecf79b9570d87bcd0387f2c0..8854d0f5c6d0ca5791829d12ecd84c775bbbd77a 100644
--- a/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb
+++ b/notebooks/AGIPD/Characterize_AGIPD_Gain_FlatFields_Summary.ipynb
@@ -22,13 +22,8 @@
     "run = 449 # runs of image data used to create histograms\n",
     "\n",
     "karabo_id = \"MID_DET_AGIPD1M-1\" # karabo karabo_id\n",
-    "receiver_id = \"{}CH0\" # inset for receiver devices\n",
-    "path_template = 'RAW-R{:04d}-{}-S{:05d}.h5' # the template to use to access data\n",
-    "h5path = 'INSTRUMENT/{}/DET/{}:xtdf/' # path in the HDF5 file to images\n",
-    "h5path_idx = 'INDEX/{}/DET/{}:xtdf/' # path in the HDF5 file to images\n",
-    "h5path_ctrl = '/CONTROL/{}/MDL/FPGA_COMP' # path to control information\n",
-    "karabo_id_control = \"MID_IRU_AGIPD1M1\" # karabo-id for control device\n",
-    "karabo_da_control = 'AGIPD1MCTRL00' # karabo DA for control infromation\n",
+    "ctrl_source_template = '{}/MDL/FPGA_COMP' # path to control information\n",
+    "karabo_id_control = \"MID_EXP_AGIPD1M1\" # karabo-id for control device\n",
     "\n",
     "use_dir_creation_date = True # use the creation data of the input dir for database queries\n",
     "cal_db_interface = \"tcp://max-exfl016:8015#8045\" # the database interface to use\n",
@@ -56,9 +51,9 @@
     "\n",
     "# Detector conditions\n",
     "max_cells = 0 # number of memory cells used, set to 0 to automatically infer\n",
-    "bias_voltage = 300 # Bias voltage\n",
+    "bias_voltage = 0. # Bias voltage\n",
     "acq_rate = 0. # the detector acquisition rate, use 0 to try to auto-determine\n",
-    "gain_setting = 0.1 # the gain setting, use 0.1 to try to auto-determine\n",
+    "gain_setting = -1 # the gain setting, use 0.1 to try to auto-determine\n",
     "photon_energy = 8.05 # photon energy in keV\n",
     "integration_time = -1 # integration time, negative values for auto-detection."
    ]
@@ -80,13 +75,7 @@
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import tabulate\n",
-    "from cal_tools.agipdlib import (\n",
-    "    get_acq_rate,\n",
-    "    get_bias_voltage,\n",
-    "    get_gain_setting,\n",
-    "    get_integration_time,\n",
-    "    get_num_cells,\n",
-    ")\n",
+    "from cal_tools.agipdlib import AgipdCtrl\n",
     "from cal_tools.agipdutils_ff import (\n",
     "    BadPixelsFF,\n",
     "    any_in,\n",
@@ -104,7 +93,7 @@
     "    send_to_db\n",
     ")\n",
     "from dateutil import parser\n",
-    "from extra_data import RunDirectory, stack_detector_data\n",
+    "from extra_data import H5File, RunDirectory, stack_detector_data\n",
     "from extra_geom import AGIPD_1MGeometry, AGIPD_500K2GGeometry\n",
     "from iCalibrationDB import Conditions, Constants, Detectors\n",
     "from iminuit import Minuit\n",
@@ -132,45 +121,38 @@
    "outputs": [],
    "source": [
     "# Get operation conditions\n",
-    "filename = glob.glob(f\"{raw_folder}/r{run:04d}/*-AGIPD[0-1][0-9]-*\")[0]\n",
-    "channel = int(re.findall(r\".*-AGIPD([0-9]+)-.*\", filename)[0])\n",
-    "control_fname = f'{raw_folder}/r{run:04d}/RAW-R{run:04d}-{karabo_da_control}-S00000.h5'\n",
-    "h5path_ctrl = h5path_ctrl.format(karabo_id_control)\n",
-    "\n",
-    "# Evaluate number of memory cells\n",
-    "mem_cells = get_num_cells(filename, karabo_id, channel)\n",
-    "if mem_cells is None:\n",
-    "    raise ValueError(f\"No raw images found in {filename}\")\n",
+    "ctrl_source = ctrl_source_template.format(karabo_id_control)\n",
     "\n",
-    "# Evaluate aquisition rate\n",
-    "fast_paths = (filename, karabo_id, channel)\n",
-    "slow_paths = (control_fname, karabo_id_control)\n",
-    "\n",
-    "if acq_rate == 0.:\n",
-    "    acq_rate = get_acq_rate(fast_paths,slow_paths)\n",
+    "raw_dc = RunDirectory(f'{raw_folder}/r{run:04d}/')\n",
     "\n",
+    "# Read operating conditions from AGIPD00 files\n",
+    "instrument_src_mod = [\n",
+    "    s for s in list(raw_dc.all_sources) if \"0CH\" in s][0]\n",
+    "ctrl_src = [\n",
+    "    s for s in list(raw_dc.all_sources) if ctrl_source in s][0]\n",
     "# Evaluate creation time\n",
     "creation_time = None\n",
     "if use_dir_creation_date:\n",
     "    creation_time = get_dir_creation_date(raw_folder, run)\n",
-    "    \n",
-    "# Evaluate gain setting\n",
-    "if gain_setting == 0.1:\n",
-    "    if creation_time.replace(tzinfo=None) < parser.parse('2020-01-31'):\n",
-    "        print(\"Set gain-setting to None for runs taken before 2020-01-31\")\n",
-    "        gain_setting = None\n",
-    "    else:\n",
-    "        try:\n",
-    "            gain_setting = get_gain_setting(control_fname, h5path_ctrl)\n",
-    "        except Exception as e:\n",
-    "            print(f'Error while reading gain setting from: \\n{control_fname}')\n",
-    "            print(e)\n",
-    "            print(\"Set gain settion to 0\")\n",
-    "            gain_setting = 0\n",
-    "\n",
-    "# Evaluate integration time\n",
-    "if integration_time < 0:\n",
-    "    integration_time = get_integration_time(control_fname, h5path_ctrl)\n",
+    "\n",
+    "agipd_cond = AgipdCtrl(\n",
+    "    run_dc=raw_dc,\n",
+    "    image_src=instrument_src_mod,\n",
+    "    ctrl_src=ctrl_src,\n",
+    "    raise_error=False,  # to be able to process very old data without mosetting value\n",
+    ")\n",
+    "\n",
+    "mem_cells = agipd_cond.get_num_cells()\n",
+    "if mem_cells is None:\n",
+    "    raise ValueError(f\"No raw images found in {raw_dc[instrument_src_mod].files}\")\n",
+    "if acq_rate == 0.:\n",
+    "    acq_rate = agipd_cond.get_acq_rate()\n",
+    "if gain_setting == -1:\n",
+    "    gain_setting = agipd_cond.get_gain_setting(creation_time)\n",
+    "if bias_voltage == 0.:\n",
+    "    bias_voltage = agipd_cond.get_bias_voltage(karabo_id_control)\n",
+    "if integration_time == -1:\n",
+    "    integration_time = agipd_cond.get_integration_time()\n",
     "\n",
     "# Evaluate detector instance for mapping\n",
     "instrument = karabo_id.split(\"_\")[0]\n",
diff --git a/notebooks/FastCCD/Characterize_Darks_NewDAQ_FastCCD_NBC_New_Common_Mode.ipynb b/notebooks/FastCCD/Characterize_Darks_NewDAQ_FastCCD_NBC_New_Common_Mode.ipynb
index a91683d1ba28baf421109730ba4bcc9d60056a00..045b9c3d053dfeb8a744b4d47a0fe524ab41ab8c 100644
--- a/notebooks/FastCCD/Characterize_Darks_NewDAQ_FastCCD_NBC_New_Common_Mode.ipynb
+++ b/notebooks/FastCCD/Characterize_Darks_NewDAQ_FastCCD_NBC_New_Common_Mode.ipynb
@@ -1122,14 +1122,9 @@
     "            parm.lower_deviation = temp_limits\n",
     "            parm.upper_deviation = temp_limits\n",
     "\n",
-    "    # This should be used in case of running notebook \n",
-    "    # by a different method other than myMDC which already\n",
-    "    # sends CalCat info.\n",
-    "    # TODO: Set db_module to \"\" by default in the first cell\n",
-    "    if not db_module:\n",
-    "        db_module = get_pdu_from_db(karabo_id, karabo_da, dconst,\n",
-    "                                    condition, cal_db_interface,\n",
-    "                                    snapshot_at=creation_time)[0]\n",
+    "    db_module = get_pdu_from_db(karabo_id, karabo_da, dconst,\n",
+    "                                condition, cal_db_interface,\n",
+    "                                snapshot_at=creation_time)[0]\n",
     "    \n",
     "    if db_output:\n",
     "        md = send_to_db(db_module, karabo_id, dconst, condition, file_loc, report,\n",
diff --git a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
index cbfae2b5d78c233e0ea9e1ddc393100f64e7b36f..19206fc2b559670d62febcf698286079fe1849be 100644
--- a/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
+++ b/notebooks/Jungfrau/Jungfrau_Gain_Correct_and_Verify_NBC.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# Jungfrau Offline Correction #\n",
     "\n",
-    "Author: European XFEL Detector Group, Version: 0.1\n",
+    "Author: European XFEL Detector Group, Version: 2.0\n",
     "\n",
     "Offline Calibration for the Jungfrau Detector"
    ]
@@ -17,38 +17,44 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "in_folder = \"/gpfs/exfel/exp/CALLAB/202031/p900113/raw\"  # the folder to read data from, required\n",
-    "out_folder =  \"\"  # the folder to output to, required\n",
+    "in_folder = \"/gpfs/exfel/exp/SPB/202130/p900204/raw\"  # the folder to read data from, required\n",
+    "out_folder =  \"/gpfs/exfel/data/scratch/ahmedk/test/remove\"  # the folder to output to, required\n",
+    "run = 112  # run to process, required\n",
     "sequences = [-1]  # sequences to correct, set to [-1] for all, range allowed\n",
-    "run = 9979  # run to process, required\n",
+    "sequences_per_node = 1  # number of sequence files per cluster node if run as slurm job, set to 0 to not run SLURM parallel\n",
     "\n",
+    "# Parameters used to access raw data.\n",
     "karabo_id = \"SPB_IRDA_JF4M\"  # karabo prefix of Jungfrau devices\n",
-    "karabo_da = ['JNGFR01']  # data aggregators\n",
-    "receiver_id = \"JNGFR{:02d}\"  # inset for receiver devices\n",
-    "receiver_control_id = \"CONTROL\"  # inset for control devices\n",
-    "path_template = 'RAW-R{:04d}-{}-S{:05d}.h5'  # template to use for file name\n",
-    "h5path = '/INSTRUMENT/{}/DET/{}:daqOutput/data'  # path in H5 file under which images are located\n",
-    "h5path_run = '/RUN/{}/DET/{}'  # path to run data\n",
-    "h5path_cntrl = '/CONTROL/{}/DET/{}'  # path to control data\n",
+    "karabo_da = ['JNGFR01', 'JNGFR02', 'JNGFR03', 'JNGFR04', 'JNGFR05', 'JNGFR06', 'JNGFR07', 'JNGFR08']  # data aggregators\n",
+    "receiver_template = \"JNGFR{:02d}\"  # Detector receiver template for accessing raw data files. e.g. \"JNGFR{:02d}\"\n",
+    "instrument_source_template = '{}/DET/{}:daqOutput'  # template for source name (filled with karabo_id & receiver_id). e.g. 'SPB_IRDA_JF4M/DET/JNGFR01:daqOutput'\n",
+    "ctrl_source_template = '{}/DET/CONTROL'  # template for control source name (filled with karabo_id_control)\n",
     "karabo_id_control = \"\"  # if control is on a different ID, set to empty string if it is the same a karabo-id\n",
-    "karabo_da_control = \"JNGFRCTRL00\"  # file inset for control data\n",
     "\n",
+    "# Parameters for calibration database.\n",
     "use_dir_creation_date = True  # use the creation data of the input dir for database queries\n",
     "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",
     "overwrite = True  # set to True if existing data should be overwritten\n",
     "relative_gain = True  # do relative gain correction\n",
-    "bias_voltage = 180  # will be overwritten by value in file\n",
-    "sequences_per_node = 5  # number of sequence files per cluster node if run as slurm job, set to 0 to not run SLURM parallel\n",
-    "photon_energy = 9.2  # photon energy in keV\n",
-    "chunk_size_idim = 1  # chunking size of imaging dimension, adjust if user software is sensitive to this.\n",
+    "plt_images = 100  # Number of images to plot after applying selected corrections.\n",
+    "limit_images = 0  # ONLY FOR TESTING. process only first N images, Use 0 to process all.\n",
+    "\n",
+    "# Parameters for retrieving calibration constants\n",
+    "manual_slow_data = False  # if true, use manually entered bias_voltage and integration_time values\n",
     "integration_time = 4.96  # integration time in us, will be overwritten by value in file\n",
     "gain_setting = 0  # 0 for dynamic gain, 1 for dynamic HG0, will be overwritten by value in file\n",
+    "gain_mode = 0  # 0 for runs with dynamic gain setting, 1 for fixgain. It will be overwritten by value in file, if manual_slow_data is set to True.\n",
     "mem_cells = 0  # leave memory cells equal 0, as it is saved in control information starting 2019.\n",
+    "bias_voltage = 180  # will be overwritten by value in file\n",
+    "\n",
+    "# Parameters for plotting\n",
+    "skip_plots = False  # exit after writing corrected files\n",
+    "\n",
+    "# TODO: Remove\n",
     "db_module = [\"\"]  # ID of module in calibration database, this parameter is ignore in the notebook. TODO: remove from calibration_configurations.\n",
-    "manual_slow_data = False  # if true, use manually entered bias_voltage and integration_time values\n",
-    "chunk_size = 0\n",
     "\n",
     "\n",
     "def balance_sequences(in_folder, run, sequences, sequences_per_node, karabo_da):\n",
@@ -62,9 +68,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import copy\n",
     "import multiprocessing\n",
-    "import time\n",
     "import warnings\n",
     "from functools import partial\n",
     "from pathlib import Path\n",
@@ -73,16 +77,23 @@
     "import matplotlib\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pasha as psh\n",
     "import tabulate\n",
+    "from IPython.display import Latex, Markdown, display\n",
+    "from extra_data import H5File, RunDirectory\n",
+    "from matplotlib.colors import LogNorm\n",
+    "\n",
+    "from cal_tools import h5_copy_except\n",
+    "from cal_tools.jungfraulib import JungfrauCtrl\n",
     "from cal_tools.enums import BadPixels\n",
+    "from cal_tools.step_timing import StepTimer\n",
     "from cal_tools.tools import (\n",
     "    get_constant_from_db_and_time,\n",
     "    get_dir_creation_date,\n",
-    "    map_modules_from_folder,\n",
+    "    map_seq_files,\n",
+    "    write_compressed_frames,\n",
     ")\n",
     "from iCalibrationDB import Conditions, Constants\n",
-    "from IPython.display import Latex, display\n",
-    "from matplotlib.colors import LogNorm\n",
     "\n",
     "warnings.filterwarnings('ignore')\n",
     "\n",
@@ -98,22 +109,16 @@
    "source": [
     "in_folder = Path(in_folder)\n",
     "out_folder = Path(out_folder)\n",
-    "ped_dir = in_folder / f'r{run:04d}'\n",
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
+    "run_dc = RunDirectory(in_folder / f'r{run:04d}')\n",
+    "instrument_src = instrument_source_template.format(karabo_id, receiver_template)\n",
     "\n",
     "if out_folder.exists() and not overwrite:\n",
     "    raise AttributeError(\"Output path exists! Exiting\")\n",
     "else:\n",
     "    out_folder.mkdir(parents=True, exist_ok=True)\n",
     "\n",
-    "fp_name_contr = path_template.format(run, karabo_da_control, 0)\n",
-    "fp_path_contr = ped_dir / fp_name_contr\n",
-    "\n",
-    "if sequences[0] == -1:\n",
-    "    sequences = None\n",
-    "\n",
     "print(f\"Run is: {run}\")\n",
-    "print(f\"HDF5 path: {h5path}\")\n",
+    "print(f\"Instrument H5File source: {instrument_src}\")\n",
     "print(f\"Process modules: {karabo_da}\")\n",
     "\n",
     "creation_time = None\n",
@@ -131,17 +136,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def check_memory_cells(file_name, path):\n",
-    "    with h5py.File(file_name, 'r') as f:\n",
-    "        t_stamp = np.array(f[path + '/storageCells/timestamp'])\n",
-    "        st_cells = np.array(f[path + '/storageCells/value'])\n",
-    "        sc_start = np.array(f[path + '/storageCellStart/value'])\n",
-    "\n",
-    "    valid_train = t_stamp > 0\n",
-    "    n_scs = st_cells[valid_train][0] + 1\n",
-    "    sc_s = sc_start[valid_train][0]\n",
+    "# Read available sequence files to correct.\n",
+    "mapped_files, num_seq_files = map_seq_files(\n",
+    "    run_dc, karabo_id, karabo_da, sequences)\n",
     "\n",
-    "    return n_scs, sc_s"
+    "if not len(mapped_files):\n",
+    "    raise IndexError(\n",
+    "        \"No sequence files available to correct for the selected sequences and karabo_da.\")"
    ]
   },
   {
@@ -150,30 +151,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# set everything up filewise\n",
-    "mapped_files, _, total_sequences, _, _ = map_modules_from_folder(\n",
-    "    in_folder, run, path_template, karabo_da, sequences\n",
-    ")\n",
-    "\n",
-    "print(f\"Processing a total of {total_sequences} sequence files\")\n",
+    "print(f\"Processing a total of {num_seq_files} sequence files\")\n",
     "table = []\n",
     "fi = 0\n",
-    "if total_sequences > 0: # create table\n",
-    "    for i, key in enumerate(mapped_files):\n",
-    "        for k, f in enumerate(list(mapped_files[key].queue)):\n",
-    "            if k == 0:\n",
-    "                table.append((fi, karabo_da[i], k, f))\n",
-    "            else:\n",
-    "                table.append((fi, \"\", k,  f))\n",
-    "            fi += 1\n",
-    "    md = display(Latex(tabulate.tabulate(\n",
-    "        table, tablefmt='latex',\n",
-    "        headers=[\"#\", \"module\", \"# module\", \"file\"])))\n",
-    "\n",
-    "# restore the queue\n",
-    "mapped_files, _, total_sequences, _, _ = map_modules_from_folder(\n",
-    "    in_folder, run, path_template, karabo_da, sequences\n",
-    ")"
+    "for kda, sfiles in mapped_files.items():\n",
+    "    for k, f in enumerate(sfiles):\n",
+    "        if k == 0:\n",
+    "            table.append((fi, kda, k, f))\n",
+    "        else:\n",
+    "            table.append((fi, \"\", k,  f))\n",
+    "        fi += 1\n",
+    "md = display(Latex(tabulate.tabulate(\n",
+    "    table, tablefmt='latex',\n",
+    "    headers=[\"#\", \"module\", \"# module\", \"file\"])))"
    ]
   },
   {
@@ -182,45 +172,47 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if not manual_slow_data:\n",
-    "    with h5py.File(fp_path_contr, 'r') as f:\n",
-    "        run_path = h5path_run.format(karabo_id_control, receiver_control_id)\n",
-    "        integration_time = float(f[f'{run_path}/exposureTime/value'][()]*1e6)\n",
-    "        bias_voltage = int(np.squeeze(f[f'{run_path}/vHighVoltage/value'])[0])\n",
-    "        try:\n",
-    "            gain_s = f[f'/RUN/{karabo_id_control}/DET/CONTROL/settings/value'][0].decode()\n",
-    "        except KeyError:\n",
-    "            print(\n",
-    "                \"ERROR: gain_setting is not available for h5 ctrl path \"\n",
-    "                f\"/RUN/{karabo_id_control}/DET/CONTROL/settings/value,\\nfor file: {fp_path_contr}. \\n\"\n",
-    "                \"WARNING: Setting gain_setting to 0, assuming that this is an old run.\\n\"\n",
-    "            )\n",
-    "            gain_s = \"KeyError\"\n",
-    "    gain_setting = 1 if gain_s == 'dynamichg0' else 0  # gain_s == 'dynamicgain'\n",
-    "\n",
-    "control_path = h5path_cntrl.format(karabo_id_control, receiver_control_id)\n",
+    "ctrl_src = ctrl_source_template.format(karabo_id_control)\n",
+    "ctrl_data = JungfrauCtrl(run_dc, ctrl_src)\n",
     "try:\n",
-    "    this_run_mcells, sc_start = check_memory_cells(fp_path_contr, control_path)\n",
+    "    this_run_mcells, sc_start = ctrl_data.get_memory_cells()\n",
+    "\n",
     "    if this_run_mcells == 1:\n",
     "        memory_cells = 1\n",
-    "        print(f'Dark runs in single cell mode\\n storage cell start: {sc_start:02d}')\n",
+    "        print(\"Run is in single cell mode.\\n\"\n",
+    "              f\"Storage cell start: {sc_start:02d}\")\n",
     "    else:\n",
     "        memory_cells = 16\n",
-    "        print(f'Dark runs in burst mode\\n storage cell start: {sc_start:02d}')\n",
-    "except Exception as e:\n",
-    "    if \"Unable to open object\" in str(e):\n",
-    "        if mem_cells==0:\n",
-    "            memory_cells = 1\n",
-    "        else:\n",
-    "            memory_cells = mem_cells\n",
-    "        print(f'Set memory cells to {memory_cells} as it is not saved in control information.')\n",
+    "        print(f\"Run is in burst mode.\\n\"\n",
+    "              f\"Storage cell start: {sc_start:02d}\")\n",
+    "except KeyError as e:\n",
+    "    print(\"WARNING: KeyError while reading number of memory cells.\")\n",
+    "    if mem_cells == 0:\n",
+    "        memory_cells = 1\n",
     "    else:\n",
-    "        print(f\"Error trying to access memory cell from contol information: {e}\")\n",
+    "        memory_cells = mem_cells\n",
+    "    print(f\"WARNING: Set memory cells to {memory_cells}, as \"\n",
+    "          \"it is not saved in control information.\")\n",
+    "\n",
+    "if not manual_slow_data:\n",
+    "    integration_time = ctrl_data.get_integration_time()\n",
+    "    bias_voltage = ctrl_data.get_bias_voltage()\n",
+    "    gain_setting = ctrl_data.get_gain_setting()\n",
+    "    gain_mode = ctrl_data.get_gain_mode()\n",
     "\n",
     "print(f\"Integration time is {integration_time} us\")\n",
-    "print(f\"Gain setting is {gain_setting}\")\n",
+    "print(f\"Gain setting is {gain_setting} (run settings: \"\n",
+    "      f\"{ctrl_data.run_settings.value if ctrl_data.run_settings else ctrl_data.run_settings})\")  # noqa\n",
+    "print(f\"Gain mode is {gain_mode}\")\n",
     "print(f\"Bias voltage is {bias_voltage} V\")\n",
-    "print(f\"Number of memory cells is {memory_cells}\")"
+    "print(f\"Number of memory cells are {memory_cells}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Retrieving calibration constants ###"
    ]
   },
   {
@@ -229,12 +221,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "condition = Conditions.Dark.jungfrau(\n",
-    "    memory_cells=memory_cells,\n",
-    "    bias_voltage=bias_voltage,\n",
-    "    integration_time=integration_time,\n",
-    "    gain_setting=gain_setting,\n",
-    ")\n",
+    "condition = {\n",
+    "    \"memory_cells\": memory_cells,\n",
+    "    \"bias_voltage\": bias_voltage,\n",
+    "    \"integration_time\": integration_time,\n",
+    "    \"gain_setting\": gain_setting,\n",
+    "}\n",
+    "# TODO: Currently there are no gain constants for fixed gain mode.\n",
+    "# This should not be the case in the future.\n",
+    "gain_condition = Conditions.Dark.jungfrau(**condition)\n",
+    "\n",
+    "condition[\"gain_mode\"] = gain_mode\n",
+    "dark_condition = Conditions.Dark.jungfrau(**condition)\n",
+    "\n",
     "\n",
     "def get_constants_for_module(karabo_da: str):\n",
     "    \"\"\" Get calibration constants for given module of Jungfrau\n",
@@ -252,23 +251,28 @@
     "        get_constant_from_db_and_time,\n",
     "        karabo_id=karabo_id,\n",
     "        karabo_da=karabo_da,\n",
-    "        condition=condition,\n",
     "        cal_db_interface=cal_db_interface,\n",
     "        creation_time=creation_time,\n",
     "        timeout=cal_db_timeout,\n",
+    "        print_once=False,\n",
     "    )\n",
     "    offset_map, when[\"Offset\"] = retrieval_function(\n",
-    "        constant=Constants.jungfrau.Offset(), empty_constant=np.zeros((1024, 512, 1, 3))\n",
+    "        condition=dark_condition,\n",
+    "        constant=Constants.jungfrau.Offset(),\n",
+    "        empty_constant=np.zeros((512, 1024, memory_cells, 3))\n",
     "    )\n",
     "    mask, when[\"BadPixelsDark\"] = retrieval_function(\n",
+    "        condition=dark_condition,\n",
     "        constant=Constants.jungfrau.BadPixelsDark(),\n",
-    "        empty_constant=np.zeros((1024, 512, 1, 3)),\n",
+    "        empty_constant=np.zeros((512, 1024, memory_cells, 3), dtype=np.uint32),\n",
     "    )\n",
     "    mask_ff, when[\"BadPixelsFF\"] = retrieval_function(\n",
+    "        condition=gain_condition,\n",
     "        constant=Constants.jungfrau.BadPixelsFF(),\n",
     "        empty_constant=None\n",
     "    )\n",
     "    gain_map, when[\"Gain\"] = retrieval_function(\n",
+    "        condition=gain_condition,\n",
     "        constant=Constants.jungfrau.RelativeGain(),\n",
     "        empty_constant=None\n",
     "    )\n",
@@ -277,26 +281,29 @@
     "    if mask_ff is not None:\n",
     "        mask |= np.moveaxis(mask_ff, 0, 1)\n",
     "\n",
-    "    # move from x,y,cell,gain to cell,x,y,gain\n",
-    "    offset_map = np.squeeze(offset_map)\n",
-    "    mask = np.squeeze(mask)\n",
-    "\n",
     "    if memory_cells > 1:\n",
-    "        offset_map = np.moveaxis(np.moveaxis(offset_map, 0, 2), 0, 2)\n",
-    "        mask = np.moveaxis(np.moveaxis(mask, 0, 2), 0, 2)\n",
+    "        # move from x, y, cell, gain to cell, x, y, gain\n",
+    "        offset_map = np.moveaxis(offset_map, [0, 1], [1, 2])\n",
+    "        mask = np.moveaxis(mask, [0, 1], [1, 2])\n",
+    "    else:\n",
+    "        offset_map = np.squeeze(offset_map)\n",
+    "        mask = np.squeeze(mask)\n",
     "\n",
     "    if gain_map is not None:\n",
     "        if memory_cells > 1:\n",
-    "            gain_map = np.moveaxis(np.moveaxis(gain_map, 0, 2), 0, 1)\n",
+    "            gain_map = np.moveaxis(gain_map, [0, 2], [2, 0])\n",
+    "            # add extra empty cell constant\n",
+    "            b = np.ones(((1,)+gain_map.shape[1:]))\n",
+    "            gain_map = np.concatenate((gain_map, b), axis=0)\n",
     "        else:\n",
-    "            gain_map = np.squeeze(gain_map)\n",
-    "            gain_map = np.moveaxis(gain_map, 1, 0)\n",
+    "            gain_map = np.moveaxis(np.squeeze(gain_map), 1, 0)\n",
     "\n",
     "    return offset_map, mask, gain_map, karabo_da, when\n",
     "\n",
     "with multiprocessing.Pool() as pool:\n",
     "    r = pool.map(get_constants_for_module, karabo_da)\n",
     "\n",
+    "# Print timestamps for the retrieved constants.\n",
     "constants = {}\n",
     "for offset_map, mask, gain_map, k_da, when in r:\n",
     "    print(f'Constants for module {k_da}:')\n",
@@ -316,23 +323,52 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def copy_and_sanitize_non_cal_data(infile, outfile, h5base):\n",
-    "    \"\"\"Copy and sanitize data from `infile` that is not calibrated.\"\"\"\n",
+    "# Correct a chunk of images for offset and gain\n",
+    "def correct_train(wid, index, d):\n",
+    "    d = d.astype(np.float32)  # [cells, x, y]\n",
+    "    g = gain[index]\n",
+    "\n",
+    "    # Jungfrau gains 0[00], 1[01], 3[11]\n",
+    "    g[g==3] = 2\n",
+    "\n",
+    "    # Select memory cells\n",
+    "    if memory_cells > 1:\n",
+    "        \"\"\"\n",
+    "        Even though it is correct to assume that memory cells pattern \n",
+    "        can be the same across all trains (for one correction run\n",
+    "        taken with one acquisition), it is preferred to not assume\n",
+    "        this to account for exceptions that can happen.\n",
+    "        \"\"\"\n",
+    "        m = memcells[index].copy()\n",
+    "        # 255 is a cell value pointing to no cell image data (image of 0 pixels).\n",
+    "        # Corresponding image will be corrected with constant of cell 0. To avoid values of 0.\n",
+    "        # This line is depending on not storing the modified memory cells in the corrected data.\n",
+    "        m[m==255] = 0\n",
+    "\n",
+    "        offset_map_cell = offset_map[m, ...]  # [16 + empty cell, x, y]\n",
+    "        mask_cell = mask[m, ...]\n",
+    "    else:\n",
+    "        offset_map_cell = offset_map\n",
+    "        mask_cell = mask\n",
+    "\n",
+    "    # Offset correction\n",
+    "    offset = np.choose(g, np.moveaxis(offset_map_cell, -1, 0))\n",
     "\n",
-    "    h5base = h5base.lstrip(\"/\")\n",
-    "    dont_copy = [\"adc\", ]\n",
-    "    dont_copy = [f'{h5base}/{dnc}' for dnc in dont_copy]\n",
+    "    d -= offset\n",
     "\n",
-    "    def visitor(k, item):\n",
-    "        if k not in dont_copy:\n",
-    "            if isinstance(item, h5py.Group):\n",
-    "                outfile.create_group(k)\n",
-    "            elif isinstance(item, h5py.Dataset):\n",
-    "                group = str(k).split(\"/\")\n",
-    "                group = \"/\".join(group[:-1])\n",
-    "                infile.copy(k, outfile[group])\n",
+    "    # Gain correction\n",
+    "    if relative_gain:\n",
+    "        if memory_cells > 1:\n",
+    "            gain_map_cell = gain_map[m, ...]\n",
+    "        else:\n",
+    "            gain_map_cell = gain_map\n",
+    "        cal = np.choose(g, np.moveaxis(gain_map_cell, -1, 0))\n",
+    "        d /= cal\n",
+    "\n",
+    "    msk = np.choose(g, np.moveaxis(mask_cell, -1, 0))\n",
     "\n",
-    "    infile.visititems(visitor)"
+    "    data_corr[index, ...] = d\n",
+    "    mask_corr[index, ...] = msk"
    ]
   },
   {
@@ -341,85 +377,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Correct a chunk of images for offset and gain\n",
-    "def correct_chunk(\n",
-    "    offset_map,\n",
-    "    mask,\n",
-    "    gain_map,\n",
-    "    memory_cells,\n",
-    "    relative_gain,\n",
-    "    inp,\n",
-    "):\n",
-    "    fim_data = None\n",
-    "    gim_data = None\n",
-    "    rim_data = None\n",
-    "    msk_data = None\n",
-    "    err = ''\n",
-    "\n",
-    "    try:\n",
-    "        d, g, m, ind, copy_sample = inp\n",
-    "        g[g==3] = 2\n",
-    "\n",
-    "        if copy_sample and ind == 0:\n",
-    "            if memory_cells == 1:\n",
-    "                rim_data = np.squeeze(copy.copy(d))\n",
-    "            else:\n",
-    "                rim_data = np.squeeze(copy.copy(d[:, 0, ...]))\n",
-    "\n",
-    "        # Select memory cells\n",
-    "        if memory_cells > 1:\n",
-    "            m[m>16] = 0\n",
-    "            offset_map_cell = offset_map[m, ...]\n",
-    "            mask_cell = mask[m, ...]\n",
-    "        else:\n",
-    "            offset_map_cell = offset_map\n",
-    "            mask_cell = mask\n",
-    "\n",
-    "        # Offset correction\n",
-    "        offset = np.choose(\n",
-    "            g,\n",
-    "            (offset_map_cell[..., 0],\n",
-    "             offset_map_cell[..., 1],\n",
-    "             offset_map_cell[..., 2]),\n",
-    "        )\n",
-    "        d -= offset\n",
-    "\n",
-    "        # Gain correction\n",
-    "        if relative_gain:\n",
-    "            if memory_cells > 1:\n",
-    "                gain_map_cell = gain_map[m, ...]\n",
-    "            else:\n",
-    "                gain_map_cell = gain_map\n",
-    "            cal = np.choose(\n",
-    "                g,\n",
-    "                (gain_map_cell[..., 0],\n",
-    "                 gain_map_cell[..., 1],\n",
-    "                 gain_map_cell[..., 2]),\n",
-    "            )\n",
-    "            d /= cal      \n",
-    "\n",
-    "        msk = np.choose(\n",
-    "            g,\n",
-    "            (mask_cell[..., 0],\n",
-    "             mask_cell[..., 1],\n",
-    "             mask_cell[..., 2]),\n",
-    "        )\n",
-    "\n",
-    "        # Store sample of data for plotting\n",
-    "        if copy_sample and ind == 0:\n",
-    "            if memory_cells == 1:\n",
-    "                fim_data = np.squeeze(copy.copy(d))\n",
-    "                gim_data = np.squeeze(copy.copy(g))\n",
-    "                msk_data = np.squeeze(copy.copy(msk))\n",
-    "            else:\n",
-    "                fim_data = np.squeeze(copy.copy(d[:, 1, ...]))\n",
-    "                gim_data = np.squeeze(copy.copy(g[:, 1, ...]))\n",
-    "                msk_data = np.squeeze(copy.copy(msk[:, 1, ...]))\n",
-    "\n",
-    "    except Exception as e:\n",
-    "        err = e\n",
+    "step_timer = StepTimer()\n",
     "\n",
-    "    return ind, d, msk, rim_data, fim_data, gim_data, msk_data, err"
+    "n_cpus = multiprocessing.cpu_count()\n",
+    "context = psh.context.ProcessContext(num_workers=n_cpus)\n",
+    "print(f\"Using {n_cpus} workers for correction.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Correcting RAW data ###"
    ]
   },
   {
@@ -433,103 +402,129 @@
     "rim_data = {}\n",
     "msk_data = {}\n",
     "\n",
-    "# For each module, chunks will be processed by pool\n",
-    "pool = multiprocessing.Pool()\n",
     "# Loop over modules\n",
-    "for local_karabo_da, mapped_files_module in zip(karabo_da, mapped_files.values()):\n",
-    "    h5path_f = h5path.format(int(local_karabo_da[-2:]))\n",
-    "    # Loop over sequences for given module\n",
-    "    for sequence_file_number, sequence_file in enumerate(mapped_files_module.queue):\n",
+    "for local_karabo_da, mapped_files_module in mapped_files.items():\n",
+    "    instrument_src_kda = instrument_src.format(int(local_karabo_da[-2:]))\n",
+    "    data_path = \"INSTRUMENT/\"+instrument_src_kda+\"/data\"\n",
+    "\n",
+    "    for sequence_file in mapped_files_module:  # noqa\n",
     "        sequence_file = Path(sequence_file)\n",
+    "        seq_dc = H5File(sequence_file)\n",
+    "\n",
+    "        # Save corrected data in an output file with name\n",
+    "        # of corresponding raw sequence file.\n",
+    "        ofile_name = sequence_file.name.replace(\"RAW\", \"CORR\")\n",
+    "        out_file = out_folder / ofile_name\n",
+    "\n",
+    "        # load shape of data for memory cells, and detector size (imgs, cells, x, y)\n",
+    "        # dshape[0] = number of available images to correct.\n",
+    "        dshape = seq_dc[instrument_src_kda, \"data.adc\"].shape\n",
+    "\n",
+    "        if dshape[0] == 0:\n",
+    "            print(f\"\\t- WARNING: No image data for {ofile_name}: data shape is {dshape}\")\n",
+    "            continue\n",
+    "\n",
+    "        # load number of data available, including trains with empty data.\n",
+    "        n_trains = len(seq_dc.train_ids)\n",
+    "\n",
+    "        n_imgs = dshape[0]\n",
+    "\n",
+    "        # For testing, only correct limit_images\n",
+    "        if limit_images > 0:\n",
+    "            n_imgs = min(n_imgs, limit_images)\n",
+    "\n",
+    "        print(f\"\\nNumber of images to correct: {n_imgs} for {ofile_name}\")\n",
+    "        if n_trains - dshape[0] != 0:\n",
+    "            print(f\"\\t- WARNING: {sequence_file.name} has {n_trains - dshape[0]} \"\n",
+    "                  \"trains with empty data.\")\n",
+    "\n",
+    "        # Just in case if n_imgs is less than the chosen plt_images.\n",
+    "        plt_images = min(plt_images, n_imgs)\n",
+    "\n",
+    "        # load constants from the constants dictionary.\n",
     "        offset_map, mask, gain_map = constants[local_karabo_da]\n",
-    "                                 \n",
-    "        with h5py.File(sequence_file, 'r') as infile:\n",
-    "            # The processed files are saved here in a folder with the run name.\n",
-    "            out_filename = out_folder / sequence_file.name.replace(\"RAW\", \"CORR\")\n",
-    "            print(f'Process file: {sequence_file}, with path {h5path_f}')\n",
-    "            try:\n",
-    "                with h5py.File(out_filename, \"w\") as outfile:\n",
-    "                    copy_and_sanitize_non_cal_data(infile, outfile, h5path_f)\n",
-    "\n",
-    "                    oshape = infile[h5path_f+\"/adc\"].shape\n",
-    "                    print(f'Data shape: {oshape}')\n",
-    "                    if not oshape[0]:\n",
-    "                        raise ValueError(f\"No image data: shape {oshape}\")\n",
-    "                    # Chunk always contains >= 1 complete image\n",
-    "                    chunk_shape = (chunk_size_idim, 1) + oshape[-2:]\n",
-    "\n",
-    "                    ddset = outfile.create_dataset(\n",
-    "                        h5path_f+\"/adc\",\n",
-    "                        oshape,\n",
-    "                        chunks=chunk_shape,\n",
-    "                        dtype=np.float32)\n",
-    "\n",
-    "                    mskset = outfile.create_dataset(\n",
-    "                        h5path_f+\"/mask\",\n",
-    "                        oshape,\n",
-    "                        chunks=chunk_shape,\n",
-    "                        dtype=np.uint32,\n",
-    "                        compression=\"gzip\",\n",
-    "                        compression_opts=1,\n",
-    "                        shuffle=True)\n",
-    "                    # Parallelize over chunks of images\n",
-    "                    inp = []\n",
-    "                    max_ind = oshape[0]\n",
-    "                    ind = 0\n",
-    "\n",
-    "                    # If chunk size is not given maximum 12+1 chunks is expected\n",
-    "                    if chunk_size == 0:\n",
-    "                        chunk_size = max_ind // 12\n",
-    "                        print(f'Chunk size: {chunk_size}')\n",
-    "\n",
-    "                    ts = time.time()\n",
-    "                    while ind<max_ind:\n",
-    "                        d = infile[h5path_f+\"/adc\"][ind:ind+chunk_size, ...].astype(np.float32)\n",
-    "                        g = infile[h5path_f+\"/gain\"][ind:ind+chunk_size, ...]\n",
-    "                        if h5path_f+\"/memoryCell\" in infile:\n",
-    "                            m = infile[h5path_f+\"/memoryCell\"][ind:ind+chunk_size, ...]\n",
-    "                        else:\n",
-    "                            m = None\n",
-    "                        print(f'To process: {d.shape}')\n",
-    "                        inp.append((d, g, m, ind, sequence_file_number==0))\n",
-    "                        ind += chunk_size\n",
-    "\n",
-    "                    print('Preparation time: ', time.time() - ts)\n",
-    "                    ts = time.time()\n",
-    "\n",
-    "                    print(f'Run {len(inp)} processes')\n",
-    "\n",
-    "                    p = partial(\n",
-    "                        correct_chunk,\n",
-    "                        offset_map,\n",
-    "                        mask,\n",
-    "                        gain_map,\n",
-    "                        memory_cells,\n",
-    "                        relative_gain,\n",
-    "                    )\n",
-    "\n",
-    "                    r = pool.map(p, inp)\n",
-    "                    \n",
-    "                    if sequence_file_number == 0:\n",
-    "                        (_,_,_,\n",
-    "                         rim_data[local_karabo_da], fim_data[local_karabo_da],\n",
-    "                         gim_data[local_karabo_da], msk_data[local_karabo_da], _) = r[0]\n",
-    "\n",
-    "                    print('Correction time: ', time.time() - ts)\n",
-    "                    ts = time.time()\n",
-    "\n",
-    "                    for rr in r:\n",
-    "                        ind, cdata, cmask, _, _, _, _, err = rr\n",
-    "                        data_size = cdata.shape[0]\n",
-    "                        ddset[ind:ind+data_size, ...] = cdata\n",
-    "                        mskset[ind:ind+data_size, ...] = cmask\n",
-    "                        if err != '':\n",
-    "                            print(f'Error: {err}')\n",
-    "\n",
-    "                    print('Saving time: ', time.time() - ts)\n",
-    "            except Exception as e:\n",
-    "                print(f\"Error: {e}\")\n",
-    "pool.close()"
+    "\n",
+    "        # Allocate shared arrays.\n",
+    "        data_corr = context.alloc(shape=dshape, dtype=np.float32)\n",
+    "        mask_corr = context.alloc(shape=dshape, dtype=np.uint32)\n",
+    "\n",
+    "        step_timer.start()\n",
+    "        seq_dc = seq_dc.select(\n",
+    "            instrument_src_kda, \"*\", require_all=True).select_trains(np.s_[:n_imgs])\n",
+    "\n",
+    "        data = seq_dc[instrument_src_kda, \"data.adc\"].ndarray()\n",
+    "        gain = seq_dc[instrument_src_kda, \"data.gain\"].ndarray()\n",
+    "        memcells = seq_dc[instrument_src_kda, \"data.memoryCell\"].ndarray()\n",
+    "\n",
+    "        rim_data[local_karabo_da] = data[:plt_images, 0, ...].copy()\n",
+    "\n",
+    "        context.map(correct_train, data)\n",
+    "        step_timer.done_step(f'Correction time.')\n",
+    "\n",
+    "        step_timer.start()\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",
+    "            # Copy RAW non-calibrated sources.\n",
+    "            with h5py.File(sequence_file, 'r') as sfile:\n",
+    "                h5_copy_except.h5_copy_except_paths(\n",
+    "                    sfile, ofile, [f\"{data_path}/adc\"])\n",
+    "\n",
+    "            # Create datasets with the available corrected data\n",
+    "            ddset = ofile.create_dataset(\n",
+    "                f\"{data_path}/adc\",\n",
+    "                data=data_corr,\n",
+    "                chunks=((1,) + dshape[1:]),  # 1 chunk == 1 image\n",
+    "                dtype=np.float32,\n",
+    "            )\n",
+    "\n",
+    "            write_compressed_frames(\n",
+    "                mask_corr,\n",
+    "                ofile,\n",
+    "                dataset_path=f\"{data_path}/mask\",\n",
+    "                comp_threads=n_cpus,\n",
+    "            )\n",
+    "\n",
+    "        step_timer.done_step(f'Saving data time.')\n",
+    "\n",
+    "        # Prepare plotting arrays\n",
+    "        step_timer.start()\n",
+    "\n",
+    "        # TODO: Print out which cell is being selected for plotting\n",
+    "        fim_data[local_karabo_da] = data_corr[:plt_images, 0, ...].copy()\n",
+    "        msk_data[local_karabo_da] = mask_corr[:plt_images, 0, ...].copy()\n",
+    "        gim_data[local_karabo_da] = gain[:plt_images, 0, ...].copy()\n",
+    "\n",
+    "        step_timer.done_step(f'Preparing plotting data of {plt_images} images.')"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Processing time summary ###"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"Total processing time {step_timer.timespan():.01f} s\")\n",
+    "step_timer.print_summary()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if skip_plots:\n",
+    "    print('Skipping plots')\n",
+    "    import sys\n",
+    "    sys.exit(0)"
    ]
   },
   {
@@ -577,7 +572,7 @@
     "    do_2d_plot(\n",
     "        h, (ex, ey),\n",
     "        \"Signal (ADU)\",\n",
-    "        \"Gain Bit Value\",\n",
+    "        \"Gain Bit Value (high gain=0[00], medium gain=1[01], low gain=3[11])\",\n",
     "        f\"Module {mod}\")"
    ]
   },
@@ -602,8 +597,8 @@
     "\n",
     "    im = ax.imshow(\n",
     "        np.mean(rim_data[mod],axis=0),\n",
-    "                vmin=min(0.75*np.median(rim_data[mod][rim_data[mod] > 0]), 2000),\n",
-    "                vmax=max(1.5*np.median(rim_data[mod][rim_data[mod] > 0]), 16000),\n",
+    "        vmin=min(0.75*np.median(rim_data[mod][rim_data[mod] > 0]), 2000),\n",
+    "        vmax=max(1.5*np.median(rim_data[mod][rim_data[mod] > 0]), 16000),\n",
     "        cmap=\"jet\")\n",
     "\n",
     "    ax.set_title(f'Module {mod}')\n",
@@ -800,5 +795,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }
diff --git a/notebooks/Jungfrau/Jungfrau_dark_analysis_all_gains_burst_mode_NBC.ipynb b/notebooks/Jungfrau/Jungfrau_dark_analysis_all_gains_burst_mode_NBC.ipynb
index 1c887202cd64d7e71327dcebfc92d8335e814b7e..356e75e8c3cb178838724491f697259d22c64add 100644
--- a/notebooks/Jungfrau/Jungfrau_dark_analysis_all_gains_burst_mode_NBC.ipynb
+++ b/notebooks/Jungfrau/Jungfrau_dark_analysis_all_gains_burst_mode_NBC.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# Jungfrau Dark Image Characterization #\n",
     "\n",
-    "Version: 0.1, Author: M. Ramilli, S. Hauf\n",
+    "Author: European XFEL Detector Group, Version: 2.0\n",
     "\n",
     "Analyzes Jungfrau dark image data to deduce offset, noise and resulting bad pixel maps"
    ]
@@ -17,43 +17,49 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cluster_profile = 'noDB'  # the ipcluster profile name\n",
     "in_folder = '/gpfs/exfel/exp/SPB/202130/p900204/raw/'  # folder under which runs are located, required\n",
-    "out_folder = '' # path to place reports at, required\n",
+    "out_folder = '/gpfs/exfel/data/scratch/ahmedk/test/remove' # path to place reports at, required\n",
     "run_high = 141 # run number for G0 dark run, required\n",
     "run_med = 142 # run number for G1 dark run, required\n",
     "run_low = 143 # run number for G2 dark run, required\n",
     "\n",
+    "# Parameters used to access raw data.\n",
     "karabo_da = ['JNGFR01', 'JNGFR02','JNGFR03','JNGFR04', 'JNGFR05', 'JNGFR06','JNGFR07','JNGFR08'] # list of data aggregators, which corresponds to different JF modules\n",
-    "karabo_id = \"SPB_IRDA_JF4M\"  # karabo_id (detector identifier) prefix of Jungfrau detector to process.\n",
-    "karabo_id_control = \"\"  # if control is on a different ID, set to empty string if it is the same a karabo-id\n",
-    "receiver_id = 'JNGFR{:02}' # inset for receiver devices\n",
-    "receiver_control_id = \"CONTROL\" # inset for control devices\n",
-    "path_template = 'RAW-R{:04d}-{}-S{{:05d}}.h5'  # template to use for file name, double escape sequence number\n",
-    "h5path = '/INSTRUMENT/{}/DET/{}:daqOutput/data'  # path in H5 file under which images are located\n",
-    "h5path_run = '/RUN/{}/DET/{}' # path to run data\n",
-    "h5path_cntrl = '/CONTROL/{}/DET/{}' # path to control data\n",
-    "karabo_da_control = \"JNGFRCTRL00\" # file inset for control data\n",
+    "karabo_id = 'SPB_IRDA_JF4M'  # karabo_id (detector identifier) prefix of Jungfrau detector to process.\n",
+    "karabo_id_control = ''  # if control is on a different ID, set to empty string if it is the same a karabo-id\n",
+    "receiver_template = 'JNGFR{:02}' # inset for receiver devices\n",
+    "instrument_source_template = '{}/DET/{}:daqOutput'  # template for instrument source name (filled with karabo_id & receiver_id). e.g. 'SPB_IRDA_JF4M/DET/JNGFR01:daqOutput'\n",
+    "ctrl_source_template = '{}/DET/CONTROL'  # template for control source name (filled with karabo_id_control)\n",
     "\n",
+    "# Parameters for calibration database and storing constants.\n",
     "use_dir_creation_date = True  # use dir creation date\n",
     "cal_db_interface = 'tcp://max-exfl016:8016'  # calibrate db interface to connect to\n",
     "cal_db_timeout = 300000 # timeout on caldb requests\n",
     "local_output = True  # output constants locally\n",
     "db_output = False  # output constants to database\n",
     "\n",
-    "integration_time = 1000 # integration time in us, will be overwritten by value in file\n",
-    "gain_setting = 0  # 0 for dynamic, forceswitchg1, forceswitchg2, 1 for dynamichg0, fixedgain1, fixgain2. Will be overwritten by value in file\n",
-    "bias_voltage = 90  # sensor bias voltage in V, will be overwritten by value in file\n",
+    "# Parameters affecting creating dark calibration constants.\n",
     "badpixel_threshold_sigma = 5.  # bad pixels defined by values outside n times this std from median\n",
     "offset_abs_threshold_low = [1000, 10000, 10000]  # absolute bad pixel threshold in terms of offset, lower values\n",
     "offset_abs_threshold_high = [8000, 15000, 15000]  # absolute bad pixel threshold in terms of offset, upper values\n",
-    "chunkSize = 10  # iteration chunk size, needs to match or be less than number of images in a sequence file\n",
-    "imageRange = [0, 500]  # image range in which to evaluate\n",
-    "memoryCells = 16  # number of memory cells\n",
-    "db_module = [\"\"]  # ID of module in calibration database, this parameter is ignored in the notebook. TODO: remove from calibration_configurations.\n",
+    "max_trains = 0  # Maximum trains to process darks. Set to 0 to process all available train images.\n",
+    "min_trains = 1  # Minimum number of trains that should be available to process dark constants. Default 1.\n",
     "manual_slow_data = False  # if true, use manually entered bias_voltage and integration_time values\n",
     "time_limits = 0.025  # to find calibration constants later on, the integration time is allowed to vary by 0.5 us\n",
-    "operation_mode = ''  # Detector operation mode, optional"
+    "\n",
+    "# Parameters to be used for injecting dark calibration constants.\n",
+    "integration_time = 1000 # integration time in us, will be overwritten by value in file\n",
+    "gain_setting = 0  # 0 for dynamic, forceswitchg1, forceswitchg2, 1 for dynamichg0, fixgain1, fixgain2. Will be overwritten by value in file\n",
+    "gain_mode = 0  # 1 if medium and low runs are  fixgain1 and fixgain2, otherwise 0. It will be overwritten by value in file, if manual_slow_data\n",
+    "bias_voltage = 90  # sensor bias voltage in V, will be overwritten by value in file\n",
+    "memory_cells = 16  # number of memory cells\n",
+    "\n",
+    "# TODO: this is used for only Warning check at AGIPD dark.\n",
+    "# Need to rethink if it makes sense to use it here as well.\n",
+    "operation_mode = 'ADAPTIVE_GAIN'  # Detector operation mode, optional\n",
+    "\n",
+    "# TODO: Remove\n",
+    "db_module = [\"\"]  # ID of module in calibration database.  TODO: remove from calibration_configurations."
    ]
   },
   {
@@ -67,29 +73,25 @@
     "import glob\n",
     "import os\n",
     "import warnings\n",
-    "\n",
+    "from pathlib import Path\n",
     "warnings.filterwarnings('ignore')\n",
     "\n",
-    "import h5py\n",
     "import matplotlib\n",
     "import matplotlib.pyplot as plt\n",
+    "import multiprocessing\n",
     "import numpy as np\n",
-    "from h5py import File as h5file\n",
+    "import pasha as psh\n",
+    "from IPython.display import Markdown, display\n",
+    "from extra_data import RunDirectory\n",
     "\n",
     "matplotlib.use('agg')\n",
     "%matplotlib inline\n",
     "\n",
-    "from XFELDetAna.detectors.jungfrau.util import (\n",
-    "    rollout_data,\n",
-    "    sanitize_data_cellid,\n",
-    ")\n",
     "from XFELDetAna.plotting.heatmap import heatmapPlot\n",
     "from XFELDetAna.plotting.histogram import histPlot\n",
-    "from XFELDetAna.detectors.jungfrau import reader as jfreader\n",
-    "from XFELDetAna.detectors.jungfrau.jf_chunk_reader import JFChunkReader\n",
-    "from XFELDetAna.util import env\n",
+    "from cal_tools import jungfraulib, step_timing\n",
     "from cal_tools.ana_tools import save_dict_to_hdf5\n",
-    "from cal_tools.enums import BadPixels\n",
+    "from cal_tools.enums import BadPixels, JungfrauSettings\n",
     "from cal_tools.tools import (\n",
     "    get_dir_creation_date,\n",
     "    get_pdu_from_db,\n",
@@ -98,9 +100,7 @@
     "    save_const_to_h5,\n",
     "    send_to_db,\n",
     ")\n",
-    "from iCalibrationDB import Conditions, Constants\n",
-    "\n",
-    "env.iprofile = cluster_profile"
+    "from iCalibrationDB import Conditions, Constants"
    ]
   },
   {
@@ -109,40 +109,48 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "path_inset = karabo_da[0] # karabo_da is a concurrency parameter\n",
-    "receiver_id = receiver_id.format(int(path_inset[-2:]))\n",
-    "proposal = list(filter(None, in_folder.strip('/').split('/')))[-2]\n",
-    "file_loc = 'proposal:{} runs:{} {} {}'.format(proposal, run_high, run_med, run_low)\n",
-    "\n",
-    "report = get_report(out_folder)\n",
-    "\n",
-    "os.makedirs(out_folder, exist_ok=True)\n",
-    "    \n",
     "# Constants relevant for the analysis\n",
-    "run_nums = [run_high, run_med, run_low] # run number for G0/HG0, G1, G2 \n",
-    "sensorSize = [1024, 512]\n",
-    "blockSize = [1024, 512]\n",
-    "xRange = [0, 0+sensorSize[0]]\n",
-    "yRange = [0, 0+sensorSize[1]]\n",
+    "run_nums = [run_high, run_med, run_low]  # run number for G0/HG0, G1, G2\n",
+    "sensor_size = (1024, 512)\n",
     "gains = [0, 1, 2]\n",
     "\n",
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
+    "fixed_settings = [\n",
+    "    JungfrauSettings.FIX_GAIN_1.value, JungfrauSettings.FIX_GAIN_2.value]  # noqa\n",
+    "dynamic_settings = [\n",
+    "    JungfrauSettings.FORCE_SWITCH_HG1.value, JungfrauSettings.FORCE_SWITCH_HG2.value]  # noqa\n",
     "\n",
     "creation_time = None\n",
     "if use_dir_creation_date:\n",
     "    creation_time = get_dir_creation_date(in_folder, run_high)\n",
-    "    print(\"Using {} as creation time\".format(creation_time))\n",
+    "    print(f\"Using {creation_time} as creation time\")\n",
+    "os.makedirs(out_folder, exist_ok=True)\n",
     "\n",
     "cal_db_interface = get_random_db_interface(cal_db_interface)\n",
-    "print('Calibration database interface: {}'.format(cal_db_interface))\n",
-    "    \n",
-    "offset_abs_threshold = [offset_abs_threshold_low, offset_abs_threshold_high]\n",
+    "print(f'Calibration database interface: {cal_db_interface}')\n",
     "\n",
     "if karabo_id_control == \"\":\n",
-    "    karabo_id_control = karabo_id\n",
+    "    karabo_id_control = karabo_id"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "proposal = list(filter(None, in_folder.strip('/').split('/')))[-2]\n",
+    "file_loc = f\"proposal:{proposal} runs:{run_high} {run_med} {run_low}\"\n",
+    "\n",
+    "report = get_report(out_folder)\n",
     "\n",
-    "print('Path inset ', path_inset)\n",
-    "print('Receiver Id ', receiver_id)"
+    "step_timer = step_timing.StepTimer()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Reading control data"
    ]
   },
   {
@@ -151,17 +159,60 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def check_memoryCells(file_name, path):\n",
-    "    with h5file(file_name, 'r') as f:\n",
-    "        t_stamp = np.array(f[path + '/storageCells/timestamp'])\n",
-    "        st_cells = np.array(f[path + '/storageCells/value'])\n",
-    "        sc_start = np.array(f[path + '/storageCellStart/value'])\n",
-    "        \n",
-    "    valid_train = t_stamp > 0        \n",
-    "    n_scs = st_cells[valid_train][0] + 1\n",
-    "    sc_s = sc_start[valid_train][0]       \n",
-    "        \n",
-    "    return n_scs, sc_s"
+    "step_timer.start()\n",
+    "gain_runs = dict()\n",
+    "\n",
+    "med_low_settings = []\n",
+    "\n",
+    "ctrl_src = ctrl_source_template.format(karabo_id_control)\n",
+    "\n",
+    "for gain, run_n in enumerate(run_nums):\n",
+    "    run_dc = RunDirectory(f\"{in_folder}/r{run_n:04d}/\")\n",
+    "    gain_runs[run_n] = [gain, run_dc]\n",
+    "    ctrl_data = jungfraulib.JungfrauCtrl(run_dc, ctrl_src)\n",
+    "    run_settings = ctrl_data.run_settings.value if ctrl_data.run_settings else ctrl_data.run_settings  # noqa\n",
+    "    # Read control data for the high gain run only.\n",
+    "    if run_n == run_high:\n",
+    "\n",
+    "        run_mcells, sc_start = ctrl_data.get_memory_cells()\n",
+    "\n",
+    "        if not manual_slow_data:\n",
+    "            integration_time = ctrl_data.get_integration_time()\n",
+    "            bias_voltage = ctrl_data.get_bias_voltage()\n",
+    "            gain_setting = ctrl_data.get_gain_setting()\n",
+    "            print(f\"Gain setting is {gain_setting} ({run_settings})\")\n",
+    "            print(f\"Integration time is {integration_time} us\")\n",
+    "            print(f\"Bias voltage is {bias_voltage} V\")\n",
+    "        if run_mcells == 1:\n",
+    "            memory_cells = 1\n",
+    "            print('Dark runs in single cell mode, '\n",
+    "                  f'storage cell start: {sc_start:02d}')\n",
+    "        else:\n",
+    "            memory_cells = 16\n",
+    "            print('Dark runs in burst mode, '\n",
+    "                  f'storage cell start: {sc_start:02d}')\n",
+    "    else:\n",
+    "        gain_mode = ctrl_data.get_gain_mode()\n",
+    "\n",
+    "        med_low_settings.append(run_settings)\n",
+    "\n",
+    "# A transperent workaround for old raw data with wrong/missing medium and low settings\n",
+    "if med_low_settings == [None, None]:\n",
+    "    print(\"WARNING: run.settings is not stored in the data to read. \"\n",
+    "          f\"Hence assuming gain_mode = {gain_mode} for adaptive old data.\")\n",
+    "elif med_low_settings == [\"dynamicgain\", \"forceswitchg1\"]:\n",
+    "    print(f\"WARNING: run.settings for medium and low gain runs are wrong {med_low_settings}. \"\n",
+    "          f\"This is an expected bug for old raw data. Setting gain_mode to {gain_mode}.\")\n",
+    "# Validate that low_med_settings is not a mix of adaptive and fixed settings.\n",
+    "elif not (sorted(med_low_settings) in [fixed_settings, dynamic_settings]):  # noqa\n",
+    "    raise ValueError(\n",
+    "        \"Medium and low run settings are not as expected. \"\n",
+    "        f\"Either {dynamic_settings} or {fixed_settings} are expected.\\n\"\n",
+    "        f\"Got {sorted(med_low_settings)} for both runs, respectively.\")\n",
+    "\n",
+    "print(f\"Gain mode is {gain_mode} ({med_low_settings})\")\n",
+    "\n",
+    "step_timer.done_step(f'Reading control data.')"
    ]
   },
   {
@@ -170,136 +221,117 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "chunkSize = 100\n",
-    "filep_size = 1000\n",
-    "memoryCells = None\n",
+    "# Use only high gain threshold for all gains in case of fixed_gain.\n",
+    "\n",
+    "if gain_mode:  # fixed_gain\n",
+    "    offset_abs_threshold = [[offset_abs_threshold_low[0]]*3, [offset_abs_threshold_high[0]]*3]\n",
+    "else:\n",
+    "    offset_abs_threshold = [offset_abs_threshold_low, offset_abs_threshold_high]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "context = psh.context.ThreadContext(num_workers=multiprocessing.cpu_count())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "All jungfrau runs are taken through one acquisition, except for the forceswitch runs.\n",
+    "While taking non-fixed dark runs, a procedure of multiple acquisitions is used to switch the storage cell indices.\n",
+    "\n",
+    "This is done for medium and low gain dark dynamic runs, only [forceswitchg1, forceswitchg2]:\n",
+    "Switching the cell indices in burst mode is a work around for hardware procedure\n",
+    "deficiency that produces wrong data for dark runs except for the first storage cell.\n",
+    "This is why multiple acquisitions are taken to switch the used storage cells and\n",
+    "acquire data through two cells for each of the 16 cells instead of acquiring darks through all 16 cells.\n",
+    "\"\"\"\n",
+    "\n",
+    "print(f\"Maximum trains to process is set to {max_trains}\")\n",
     "\n",
     "noise_map = dict()\n",
     "offset_map = dict()\n",
+    "bad_pixels_map = dict()\n",
     "\n",
-    "# TODO: parallelize with multiprocessing.\n",
     "for mod in karabo_da:\n",
-    "    for gain, r_n in enumerate(run_nums):\n",
-    "\n",
-    "        print(f\"Gain stage {gain}, run {r_n}\")\n",
-    "        valid_data = []\n",
-    "        valid_cellids = []\n",
-    "\n",
-    "        n_tr = 0\n",
-    "        n_empty_trains = 0\n",
-    "        n_empty_sc = 0\n",
-    "\n",
-    "        ped_dir = \"{}/r{:04d}/\".format(in_folder, r_n)\n",
-    "        fp_name = path_template.format(r_n, karabo_da_control)\n",
-    "        fp_path = '{}/{}'.format(ped_dir, fp_name)\n",
-    "      \n",
-    "        files_pattern = \"{}/*{}*.h5\".format(ped_dir, path_inset)\n",
-    "        n_files = len(glob.glob(files_pattern))\n",
-    "        if n_files == 0:\n",
-    "            raise Exception(f\"No files found matching {files_pattern!r}\")\n",
-    "\n",
-    "        myRange = range(0, n_files)\n",
-    "        control_path = h5path_cntrl.format(karabo_id_control, receiver_control_id)\n",
-    "\n",
-    "        this_run_mcells, sc_start = check_memoryCells(fp_path.format(0).format(myRange[0]), control_path)\n",
-    "\n",
-    "        if mod not in noise_map:\n",
-    "            if not manual_slow_data:\n",
-    "                run_path = h5path_run.format(karabo_id_control, receiver_control_id)\n",
-    "                filename = fp_path.format(0)\n",
-    "                with h5py.File(filename, 'r') as f:\n",
-    "                    integration_time = float(f[f'{run_path}/exposureTime/value'][()]*1e6)\n",
-    "                    bias_voltage = int(np.squeeze(f[f'{run_path}/vHighVoltage/value'])[0])\n",
-    "                    if r_n == run_high:\n",
-    "                        try:\n",
-    "                            gain_s = f[f'/RUN/{karabo_id_control}/DET/CONTROL/settings/value'][0].decode()\n",
-    "                        except KeyError:\n",
-    "                            print(\n",
-    "                                \"ERROR: gain_setting is not available for h5 ctrl path \"\n",
-    "                                f\"/RUN/{karabo_id_control}/DET/CONTROL/settings/value,\\nfor file: {filename}. \\n\"\n",
-    "                                \"WARNING: Setting gain_setting to 0, assuming that this is an old run.\\n\")\n",
-    "                            gain_s = \"KeyError\"\n",
-    "                        gain_setting = 1 if gain_s == \"dynamichg0\"  else 0\n",
-    "                        print(f\"Constants Gain setting is {gain_setting} ({gain_s})\")\n",
-    "\n",
-    "\n",
-    "            print(\"Integration time is {} us\".format(integration_time))\n",
-    "            print(\"Bias voltage is {} V\".format(bias_voltage))\n",
-    "\n",
-    "            if this_run_mcells == 1:\n",
-    "                memoryCells = 1\n",
-    "                print('Dark runs in single cell mode\\n storage cell start: {:02d}'.format(sc_start))\n",
-    "            else:\n",
-    "                memoryCells = 16\n",
-    "                print('Dark runs in burst mode\\n storage cell start: {:02d}'.format(sc_start))\n",
-    "\n",
-    "            noise_map[mod] = np.zeros(sensorSize+[memoryCells, 3])\n",
-    "            offset_map[mod] = np.zeros(sensorSize+[memoryCells, 3])\n",
-    "\n",
-    "        fp_name = path_template.format(r_n, path_inset)\n",
-    "        fp_path = '{}/{}'.format(ped_dir, fp_name)\n",
-    "\n",
-    "        print(\"Reading data from {}\".format(fp_path))\n",
-    "        print(\"Run is: {}\".format(r_n))\n",
-    "        print(\"HDF5 path: {}\".format(h5path))\n",
-    "\n",
-    "        imageRange = [0, filep_size*len(myRange)]\n",
-    "        reader = JFChunkReader(filename = fp_path, readFun = jfreader.readData, size = filep_size, chunkSize = chunkSize,\n",
-    "                               path = h5path, image_range=imageRange, pixels_x = sensorSize[0], pixels_y = sensorSize[1],\n",
-    "                               x_range = xRange, y_range = yRange, imagesPerChunk=chunkSize, filesRange = myRange,\n",
-    "                               memoryCells=this_run_mcells, blockSize=blockSize)\n",
-    "\n",
-    "        for data in reader.readChunks():\n",
-    "\n",
-    "            images = np.array(data[0], dtype=np.float)\n",
-    "            gainmaps = np.array(data[1], dtype=np.uint16)\n",
-    "            trainId = np.array(data[2])\n",
-    "            fr_num = np.array(data[3])\n",
-    "            acelltable = np.array(data[4])\n",
-    "            n_tr += acelltable.shape[-1]\n",
-    "            this_tr = acelltable.shape[-1]\n",
-    "\n",
-    "            idxs = np.nonzero(trainId)[0]\n",
-    "            images = images[..., idxs]\n",
-    "            gainmaps = gainmaps[..., idxs]\n",
-    "            fr_num = fr_num[..., idxs]\n",
-    "            acelltable = acelltable[..., idxs]\n",
-    "\n",
-    "            if memoryCells == 1:\n",
-    "                acelltable -= sc_start\n",
-    "\n",
-    "            n_empty_trains += this_tr - acelltable.shape[-1]\n",
-    "            n_empty_sc += len(acelltable[acelltable > 15])\n",
-    "\n",
-    "            # throwing away all the SC entries except\n",
-    "            # the first for lower gains.\n",
-    "            if gain > 0 and memoryCells == 16: \n",
-    "                acelltable[1:] = 255\n",
-    "\n",
-    "            # makes 4-dim vecs into 3-dim\n",
-    "            # makes 2-dim into 1-dim\n",
-    "            # leaves  1-dim and 3-dim vecs\n",
-    "\n",
-    "            images, gainmaps, acelltable = rollout_data([images, gainmaps, acelltable])\n",
-    "\n",
-    "            # removes entries with cellID 255\n",
-    "            images, gainmaps, acelltable = sanitize_data_cellid([images, gainmaps], acelltable)\n",
-    "            valid_data.append(images)\n",
-    "            valid_cellids.append(acelltable)\n",
-    "\n",
-    "        valid_data = np.concatenate(valid_data, axis=2)\n",
-    "        valid_cellids = np.concatenate(valid_cellids, axis=0)\n",
-    "\n",
-    "        for cell in range(memoryCells):\n",
-    "\n",
-    "            thiscell = valid_data[..., valid_cellids == cell]\n",
-    "            noise_map[mod][..., cell, gain] = np.std(thiscell, axis=2)\n",
-    "            offset_map[mod][..., cell, gain] = np.mean(thiscell, axis=2)\n",
-    "\n",
-    "        print(f'G{gain:01d} dark calibration')\n",
-    "        print(f'Missed {n_empty_trains:d} out of {n_tr:d} trains')\n",
-    "        print(\n",
-    "            f'Lost {n_empty_sc:d} images out of {this_run_mcells*(n_tr-n_empty_trains):d}')  # noqa"
+    "    step_timer.start()\n",
+    "    instrument_src = instrument_source_template.format(\n",
+    "        karabo_id, receiver_template.format(int(mod[-2:])))\n",
+    "\n",
+    "    print(f\"\\n- Instrument data path for {mod} is {instrument_src}.\")\n",
+    "\n",
+    "    offset_map[mod] = context.alloc(shape=(sensor_size+(memory_cells, 3)), fill=0)\n",
+    "    noise_map[mod] = context.alloc(like=offset_map[mod], fill=0)\n",
+    "    bad_pixels_map[mod] = context.alloc(like=offset_map[mod], dtype=np.uint32, fill=0)\n",
+    "\n",
+    "    for run_n, [gain, run_dc] in gain_runs.items():\n",
+    "\n",
+    "        def process_cell(worker_id, array_index, cell_number):\n",
+    "            cell_slice_idx = acelltable == cell_number\n",
+    "            thiscell = images[..., cell_slice_idx]\n",
+    "\n",
+    "            offset_map[mod][..., cell_number, gain] = np.mean(thiscell, axis=2)\n",
+    "            noise_map[mod][..., cell_number, gain] = np.std(thiscell, axis=2)\n",
+    "\n",
+    "            # Check if there are wrong bad gain values.\n",
+    "            # Indicate pixels with wrong gain value across all trains for each cell.\n",
+    "            bad_pixels_map[mod][\n",
+    "                np.average(gain_vals[..., cell_slice_idx], axis=2) != raw_g] |= BadPixels.WRONG_GAIN_VALUE.value\n",
+    "        print(f\"Gain stage {gain}, run {run_n}\")\n",
+    "\n",
+    "        # load shape of data for memory cells, and detector size (imgs, cells, x, y)\n",
+    "        n_imgs = run_dc[instrument_src, \"data.adc\"].shape[0]\n",
+    "        # load number of data available, including trains with empty data.\n",
+    "        n_trains = len(run_dc.train_ids)\n",
+    "        instr_dc = run_dc.select(instrument_src, require_all=True)\n",
+    "        empty_trains = n_trains - n_imgs\n",
+    "        if empty_trains != 0:\n",
+    "            print(\n",
+    "                f\"\\tWARNING: {mod} has {empty_trains} trains with empty data out of {n_trains} trains at \"  # noqa\n",
+    "                f\"{Path(run_dc[instrument_src, 'data.adc'].files[0].filename).parent}.\")\n",
+    "        if max_trains > 0:\n",
+    "            n_imgs = min(n_imgs, max_trains)\n",
+    "        print(f\"Processing {n_imgs} images.\")\n",
+    "        # Select only requested number of images to process darks.\n",
+    "        instr_dc = instr_dc.select_trains(np.s_[:n_imgs])\n",
+    "\n",
+    "        if n_imgs < min_trains:\n",
+    "            raise ValueError(\n",
+    "                f\"Less than {min_trains} trains are available in RAW data.\"\n",
+    "                \" Not enough data to process darks.\")\n",
+    "\n",
+    "        images = np.transpose(\n",
+    "            instr_dc[instrument_src, \"data.adc\"].ndarray(), (3, 2, 1, 0))\n",
+    "        acelltable = np.transpose(instr_dc[instrument_src, \"data.memoryCell\"].ndarray())\n",
+    "        gain_vals = np.transpose(\n",
+    "            instr_dc[instrument_src, \"data.gain\"].ndarray(), (3, 2, 1, 0))\n",
+    "\n",
+    "        # define gain value as saved in raw gain map\n",
+    "        raw_g = 3 if gain == 2 else gain\n",
+    "\n",
+    "        if memory_cells == 1:\n",
+    "            acelltable -= sc_start\n",
+    "        # Only for dynamic medium and low gain runs [forceswitchg1, forceswitchg2] in burst mode.\n",
+    "\n",
+    "        if gain_mode == 0 and gain > 0 and memory_cells == 16:\n",
+    "            # 255 similar to the receiver which uses the 255\n",
+    "            # value to indicate a cell without an image.\n",
+    "            # image shape for forceswitchg1 and forceswitchg2 = (1024, 512, 2, trains)\n",
+    "            # compared to expected shape of (1024, 512, 16, trains) for high gain run.\n",
+    "            acelltable[1:] = 255\n",
+    "\n",
+    "        # Calculate offset and noise maps\n",
+    "        context.map(process_cell, range(memory_cells))\n",
+    "\n",
+    "    step_timer.done_step(f'Creating Offset and noise constants for a module.')"
    ]
   },
   {
@@ -327,14 +359,18 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [],
    "source": [
+    "# TODO: Fix plots arrangment and speed for Jungfrau burst mode.\n",
+    "step_timer.start()\n",
     "for mod in karabo_da:\n",
     "    for g_idx in gains:\n",
-    "        for cell in range(0, memoryCells):\n",
+    "        for cell in range(0, memory_cells):\n",
     "            f_o0 = heatmapPlot(\n",
-    "                np.swapaxes(offset_map[mod][..., cell, g_idx], 0, 1), \n",
+    "                np.swapaxes(offset_map[mod][..., cell, g_idx], 0, 1),\n",
     "                y_label=\"Row\",\n",
     "                x_label=\"Column\",\n",
     "                lut_label=unit,\n",
@@ -380,13 +416,14 @@
     "                histotype='stepfilled',\n",
     "            )\n",
     "\n",
-    "            ax_n0.tick_params(axis='both',which='major', labelsize=15)\n",
+    "            ax_n0.tick_params(axis='both', which='major', labelsize=15)\n",
     "            ax_n0.set_title(\n",
     "                f'Module noise distribution - Cell {cell:02d} - Module {mod}',\n",
     "                fontsize=15)\n",
     "            ax_n0.set_xlabel(\n",
     "                f'RMS noise {g_name[g_idx]} ' + unit, fontsize=15)\n",
-    "            plt.show()"
+    "            plt.show()\n",
+    "step_timer.done_step(f'Plotting offset and noise maps.')"
    ]
   },
   {
@@ -415,19 +452,13 @@
    "outputs": [],
    "source": [
     "def print_bp_entry(bp):\n",
-    "    print(\"{:<30s} {:032b}\".format(bp.name, bp.value))\n",
+    "    print(\"{:<30s} {:032b} -> {}\".format(bp.name, bp.value, int(bp.value)))\n",
     "\n",
     "print_bp_entry(BadPixels.OFFSET_OUT_OF_THRESHOLD)\n",
     "print_bp_entry(BadPixels.NOISE_OUT_OF_THRESHOLD)\n",
-    "print_bp_entry(BadPixels.OFFSET_NOISE_EVAL_ERROR)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "print_bp_entry(BadPixels.OFFSET_NOISE_EVAL_ERROR)\n",
+    "print_bp_entry(BadPixels.WRONG_GAIN_VALUE)\n",
+    "\n",
     "def eval_bpidx(d):\n",
     "\n",
     "    mdn = np.nanmedian(d, axis=(0, 1))[None, None, :, :]\n",
@@ -440,17 +471,18 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [],
    "source": [
-    "bad_pixels_map = dict()\n",
+    "step_timer.start()\n",
     "\n",
     "for mod in karabo_da:\n",
-    "\n",
-    "    bad_pixels_map[mod] = np.zeros(noise_map[mod].shape, np.uint32)\n",
+    "    display(Markdown(f\"### Badpixels for module {mod}:\"))\n",
     "    offset_abs_threshold = np.array(offset_abs_threshold)\n",
     "\n",
-    "    bad_pixels_map[mod][eval_bpidx(offset_map[mod])] = BadPixels.OFFSET_OUT_OF_THRESHOLD.value\n",
+    "    bad_pixels_map[mod][eval_bpidx(offset_map[mod])] |= BadPixels.OFFSET_OUT_OF_THRESHOLD.value\n",
     "\n",
     "    bad_pixels_map[mod][~np.isfinite(offset_map[mod])] |= BadPixels.OFFSET_NOISE_EVAL_ERROR.value\n",
     "\n",
@@ -460,16 +492,18 @@
     "\n",
     "    bad_pixels_map[mod][(offset_map[mod] < offset_abs_threshold[0][None, None, None, :]) | (offset_map[mod] > offset_abs_threshold[1][None, None, None, :])] |= BadPixels.OFFSET_OUT_OF_THRESHOLD.value  # noqa\n",
     "\n",
-    "for g_idx in gains:\n",
-    "    for cell in range(memoryCells):\n",
-    "        bad_pixels = bad_pixels_map[mod][:, :, cell, g_idx]\n",
-    "        fn_0 = heatmapPlot(\n",
-    "            np.swapaxes(bad_pixels, 0, 1),\n",
-    "            y_label=\"Row\",\n",
-    "            x_label=\"Column\",\n",
-    "            lut_label=f\"Badpixels {g_name[g_idx]} [ADCu]\",\n",
-    "            aspect=1.,\n",
-    "            vmin=0, title=f'G{g_idx} Bad pixel map - Cell {cell:02d} - Module {mod}')"
+    "    for g_idx in gains:\n",
+    "        for cell in range(memory_cells):\n",
+    "            bad_pixels = bad_pixels_map[mod][:, :, cell, g_idx]\n",
+    "            fn_0 = heatmapPlot(\n",
+    "                np.swapaxes(bad_pixels, 0, 1),\n",
+    "                y_label=\"Row\",\n",
+    "                x_label=\"Column\",\n",
+    "                lut_label=f\"Badpixels {g_name[g_idx]} [ADCu]\",\n",
+    "                aspect=1.,\n",
+    "                vmin=0, vmax=5,\n",
+    "                title=f'G{g_idx} Bad pixel map - Cell {cell:02d} - Module {mod}')\n",
+    "step_timer.done_step(f'Creating bad pixels constant and plotting it for a module.')"
    ]
   },
   {
@@ -480,10 +514,12 @@
    "source": [
     "# set the operating condition\n",
     "condition = Conditions.Dark.jungfrau(\n",
-    "    memory_cells=memoryCells,\n",
+    "    memory_cells=memory_cells,\n",
     "    bias_voltage=bias_voltage,\n",
     "    integration_time=integration_time,\n",
-    "    gain_setting=gain_setting)\n",
+    "    gain_setting=gain_setting,\n",
+    "    gain_mode=gain_mode,\n",
+    ")\n",
     "\n",
     "db_modules = get_pdu_from_db(\n",
     "    karabo_id=karabo_id,\n",
@@ -494,12 +530,20 @@
     "    snapshot_at=creation_time)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Inject and save calibration constants"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "for mod, db_mod in zip(karabo_da, db_modules):\n",
     "    constants = {\n",
     "        'Offset': np.moveaxis(offset_map[mod], 0, 1),\n",
@@ -548,10 +592,22 @@
     "print(\"Constants parameter conditions are:\\n\")\n",
     "print(\n",
     "    f\"• Bias voltage: {bias_voltage}\\n\"\n",
-    "    f\"• Memory cells: {memoryCells}\\n\"\n",
+    "    f\"• Memory cells: {memory_cells}\\n\"\n",
     "    f\"• Integration time: {integration_time}\\n\"\n",
     "    f\"• Gain setting: {gain_setting}\\n\"\n",
-    "    f\"• Creation time: {md.calibration_constant_version.begin_at if md is not None else creation_time}\\n\")  # noqa"
+    "    f\"• Gain mode: {gain_mode}\\n\"\n",
+    "    f\"• Creation time: {md.calibration_constant_version.begin_at if md is not None else creation_time}\\n\")  # noqa\n",
+    "step_timer.done_step(\"Injecting constants.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"Total processing time {step_timer.timespan():.01f} s\")\n",
+    "step_timer.print_summary()"
    ]
   }
  ],
@@ -571,9 +627,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.11"
+   "version": "3.8.12"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }
diff --git a/notebooks/REMI/REMI_Digitize_and_Transform.ipynb b/notebooks/REMI/REMI_Digitize_and_Transform.ipynb
index 1c78263ce6caffbda31d6f5b0647273b90965966..8d7ce5a454e0517b27cedbb4b18600f2ad371f7e 100644
--- a/notebooks/REMI/REMI_Digitize_and_Transform.ipynb
+++ b/notebooks/REMI/REMI_Digitize_and_Transform.ipynb
@@ -484,8 +484,12 @@
     "        col = (signal_idx % 2) if signal_idx < 6 else np.s_[:]\n",
     "        ax = fig.add_subplot(grid[row, col])\n",
     "        \n",
-    "        pulse_idx = np.where(np.isfinite(edges[:, signal_idx, 0]))[0][0]\n",
-    "        train_idx = np.where(pulse_idx >= pulse_offsets)[0][-1]\n",
+    "        finite_edges = np.isfinite(edges[:, signal_idx, 0])\n",
+    "        if not finite_edges.any():\n",
+    "            continue\n",
+    "            \n",
+    "        pulse_idx = finite_edges.nonzero()[0][0]\n",
+    "        train_idx = (pulse_idx >= pulse_offsets).nonzero()[0][-1]\n",
     "        trigger = triggers[pulse_idx]\n",
     "        \n",
     "        sourcekey = remi.get_channel_sourcekey(\n",
@@ -500,7 +504,7 @@
     "        \n",
     "        pulse_trace = corr_trace[np.s_[trigger['start']:trigger['stop']]]\n",
     "        \n",
-    "        x_time = remi.get_time_calibration() * (np.arange(len(pulse_trace) + trigger['offset']))\n",
+    "        x_time = remi.get_time_calibration() * (np.arange(len(pulse_trace)) + trigger['offset'])\n",
     "        \n",
     "        ax.plot(x_time, pulse_trace, lw=1)\n",
     "        ax.set_xlim(x_time[0], x_time[-1])\n",
@@ -895,9 +899,9 @@
     "dataset_kwargs = {k[8:]: v for k, v in locals().items() if k.startswith('dataset_compression')}\n",
     "\n",
     "metadata = dc.run_metadata()\n",
-    "daq_library_bytes = metadata['daqLibrary'].encode('ascii')\n",
-    "karabo_framework_bytes = metadata['karaboFramework'].encode('ascii')\n",
-    "proposal_number = int(proposal) if proposal else metadata['proposalNumber']\n",
+    "daq_library_bytes = metadata.get('daqLibrary', '0.0').encode('ascii')\n",
+    "karabo_framework_bytes = metadata.get('karaboFramework', '0.0').encode('ascii')\n",
+    "proposal_number = int(proposal) if proposal else metadata.get('proposalNumber', -1)\n",
     "\n",
     "print('Writing sequence files', flush=True, end='')\n",
     "\n",
diff --git a/notebooks/ePix100/Characterize_Darks_ePix100_NBC.ipynb b/notebooks/ePix100/Characterize_Darks_ePix100_NBC.ipynb
index 89ea1f460e2b16e56cbe860f328686722c901af2..0f1c2cd3858339186ecda900cdc05f910c73cc03 100644
--- a/notebooks/ePix100/Characterize_Darks_ePix100_NBC.ipynb
+++ b/notebooks/ePix100/Characterize_Darks_ePix100_NBC.ipynb
@@ -6,9 +6,9 @@
    "source": [
     "# ePix100 Dark Characterization\n",
     "\n",
-    "Author: M. Karnevskiy, Version 1.0\n",
+    "Author: European XFEL Detector Group, Version: 2.0\n",
     "\n",
-    "The following notebook provides dark image analysis of the ePix100 detector.\n",
+    "The following notebook provides dark image analysis and calibration constants of the ePix100 detector.\n",
     "\n",
     "Dark characterization evaluates offset and noise of the detector and gives information about bad pixels. Resulting maps are saved as .h5 files for a latter use and injected to calibration DB."
    ]
@@ -19,33 +19,40 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cluster_profile = \"noDB\" # ipcluster profile to use\n",
     "in_folder = '/gpfs/exfel/exp/HED/202030/p900136/raw' # input folder, required\n",
-    "out_folder = '/gpfs/exfel/data/scratch/ahmedk/test/HED_dark/' # output folder, required\n",
+    "out_folder = '' # output folder, required\n",
     "sequence = 0 # sequence file to use\n",
     "run = 182 # which run to read data from, required\n",
     "\n",
+    "# Parameters for accessing the raw data.\n",
     "karabo_id = \"HED_IA1_EPX100-2\" # karabo karabo_id\n",
     "karabo_da = [\"EPIX02\"]  # data aggregators\n",
-    "receiver_id = \"RECEIVER\" # inset for receiver devices\n",
+    "receiver_template = \"RECEIVER\" # detector receiver template for accessing raw data files\n",
     "path_template = 'RAW-R{:04d}-{}-S{{:05d}}.h5' # the template to use to access data\n",
-    "h5path = '/INSTRUMENT/{}/DET/{}:daqOutput/data/image/pixels' # path in the HDF5 file to images\n",
-    "h5path_t = '/INSTRUMENT/{}/DET/{}:daqOutput/data/backTemp'  # path to find temperature at\n",
-    "h5path_cntrl = '/CONTROL/{}/DET'  # path to control data\n",
+    "instrument_source_template = '{}/DET/{}:daqOutput' # instrument detector data source in h5files\n",
     "\n",
+    "# Parameters for the calibration database.\n",
     "use_dir_creation_date = True\n",
     "cal_db_interface = \"tcp://max-exfl016:8020\" # calibration DB interface to use\n",
     "cal_db_timeout = 300000 # timeout on caldb requests\n",
     "db_output = False # Output constants to the calibration database\n",
     "local_output = True # output constants locally\n",
     "\n",
-    "number_dark_frames = 0 # number of images to be used, if set to 0 all available images are used\n",
-    "temp_limits = 5 # limit for parameter Operational temperature\n",
-    "db_module = 'ePix100_M17' # detector karabo_id\n",
+    "# Conditions used for injected calibration constants.\n",
     "bias_voltage = 200 # bias voltage\n",
     "in_vacuum = False # detector operated in vacuum\n",
     "fix_temperature = 290. # fix temperature to this value\n",
-    "operation_mode = ''  # Detector operation mode, optional"
+    "temp_limits = 5 # limit for parameter Operational temperature\n",
+    "\n",
+    "# Parameters used during selecting raw data trains.\n",
+    "min_trains = 1 # Minimum number of trains that should be available to process dark constants. Default 1.\n",
+    "max_trains = 1000  # Maximum number of trains to use for processing dark constants. Set to 0 to use all available trains.\n",
+    "\n",
+    "# Don't delete! myMDC sends this by default.\n",
+    "operation_mode = ''  # Detector operation mode, optional\n",
+    "\n",
+    "# TODO: delete after removing from calibration_configurations\n",
+    "db_module = ''  # ID of module in calibration database, this parameter is ignore in the notebook. TODO: remove from calibration_configurations."
    ]
   },
   {
@@ -57,14 +64,13 @@
     "import os\n",
     "import warnings\n",
     "\n",
-    "warnings.filterwarnings('ignore')\n",
-    "\n",
-    "import h5py\n",
-    "import matplotlib.pyplot as plt\n",
-    "from IPython.display import Latex, Markdown, display\n",
-    "\n",
-    "%matplotlib inline\n",
     "import numpy as np\n",
+    "import pasha as psh\n",
+    "from extra_data import RunDirectory\n",
+    "\n",
+    "import XFELDetAna.xfelprofiler as xprof\n",
+    "from XFELDetAna import xfelpyanatools as xana\n",
+    "from XFELDetAna.plotting.util import prettyPlotting\n",
     "from cal_tools.tools import (\n",
     "    get_dir_creation_date,\n",
     "    get_pdu_from_db,\n",
@@ -73,32 +79,25 @@
     "    save_const_to_h5,\n",
     "    send_to_db,\n",
     ")\n",
-    "from iCalibrationDB import Conditions, Constants, Detectors, Versions\n",
-    "from iCalibrationDB.detectors import DetectorTypes\n",
-    "from XFELDetAna.util import env\n",
+    "from iCalibrationDB import Conditions, Constants"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%matplotlib inline\n",
     "\n",
-    "env.iprofile = cluster_profile\n",
-    "from XFELDetAna import xfelpyanatools as xana\n",
-    "from XFELDetAna import xfelpycaltools as xcal\n",
-    "from XFELDetAna.detectors.fastccd import readerh5 as fastccdreaderh5\n",
-    "from XFELDetAna.plotting.util import prettyPlotting\n",
+    "warnings.filterwarnings('ignore')\n",
     "\n",
     "prettyPlotting = True\n",
-    "import XFELDetAna.xfelprofiler as xprof\n",
     "\n",
     "profiler = xprof.Profiler()\n",
     "profiler.disable()\n",
-    "from XFELDetAna.xfelreaders import ChunkReader\n",
-    "\n",
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
-    "h5path_t = h5path_t.format(karabo_id, receiver_id)\n",
-    "h5path_cntrl = h5path_cntrl.format(karabo_id)\n",
     "\n",
-    "def nImagesOrLimit(nImages, limit):\n",
-    "    if limit == 0:\n",
-    "        return nImages\n",
-    "    else:\n",
-    "        return min(nImages, limit)"
+    "instrument_src = instrument_source_template.format(karabo_id, receiver_template)"
    ]
   },
   {
@@ -110,19 +109,16 @@
     "# Read report path and create file location tuple to add with the injection\n",
     "proposal = list(filter(None, in_folder.strip('/').split('/')))[-2]\n",
     "file_loc = f'proposal:{proposal} runs:{run}'\n",
-    "\n",
+    "pixels_x = 708\n",
+    "pixels_y = 768\n",
     "report = get_report(out_folder)\n",
     "\n",
-    "x = 708  # rows of the xPix100\n",
-    "y = 768  # columns of the xPix100\n",
-    "\n",
     "ped_dir = os.path.join(in_folder, f\"r{run:04d}\")\n",
     "fp_name = path_template.format(run, karabo_da[0]).format(sequence)\n",
-    "filename = os.path.join(ped_dir, fp_name)\n",
+    "run_dir = RunDirectory(ped_dir)\n",
     "\n",
-    "print(f\"Reading data from: {filename}\\n\")\n",
     "print(f\"Run number: {run}\")\n",
-    "print(f\"HDF5 path: {h5path}\")\n",
+    "print(f\"Instrument H5File source: {instrument_src}\")\n",
     "if use_dir_creation_date:\n",
     "    creation_time = get_dir_creation_date(in_folder, run)\n",
     "    print(f\"Using {creation_time.isoformat()} as creation time\")\n",
@@ -135,60 +131,58 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sensorSize = [x, y]\n",
-    "chunkSize = 100  #Number of images to read per chunk\n",
-    "\n",
-    "#Sensor area will be analysed according to blocksize\n",
-    "blockSize = [sensorSize[0] // 2, sensorSize[1] // 2]\n",
-    "xcal.defaultBlockSize = blockSize\n",
-    "cpuCores = 4  #Specifies the number of running cpu cores\n",
-    "memoryCells = 1  #No mamery cells\n",
-    "\n",
-    "#Specifies total number of images to proceed\n",
-    "nImages = fastccdreaderh5.getDataSize(filename, h5path)[0]\n",
-    "nImages = nImagesOrLimit(nImages, number_dark_frames)\n",
-    "print(\"\\nNumber of dark images to analyze: \", nImages)\n",
-    "run_parallel = False\n",
-    "\n",
-    "with h5py.File(filename, 'r') as f:\n",
-    "    integration_time = int(f[os.path.join(h5path_cntrl, \"CONTROL\",\"expTime\", \"value\")][0])\n",
-    "    temperature = np.mean(f[h5path_t])/100.\n",
+    "pixel_data = (instrument_src, \"data.image.pixels\")\n",
+    "\n",
+    "# Specifies total number of images to proceed\n",
+    "n_trains = run_dir.get_data_counts(*pixel_data).shape[0]\n",
+    "\n",
+    "# Modify n_trains to process based on given maximum\n",
+    "# and minimun number of trains.\n",
+    "if max_trains:\n",
+    "    n_trains = min(max_trains, n_trains)\n",
+    "\n",
+    "if n_trains < min_trains:\n",
+    "    raise ValueError(\n",
+    "        f\"Less than {min_trains} trains are available in RAW data.\"\n",
+    "         \" Not enough data to process darks.\")\n",
+    "\n",
+    "print(f\"Number of dark images to analyze: {n_trains}.\")\n",
+    "\n",
+    "integration_time = int(run_dir.get_array(\n",
+    "    f\"{karabo_id}/DET/CONTROL\",\n",
+    "    \"expTime.value\")[0])\n",
+    "temperature = np.mean(run_dir.get_array(\n",
+    "    instrument_src,\n",
+    "    \"data.backTemp\").values) / 100.\n",
+    "\n",
+    "if fix_temperature != 0:\n",
+    "    temperature_k = fix_temperature\n",
+    "    print(\"Temperature is fixed!\")\n",
+    "else:\n",
     "    temperature_k = temperature + 273.15\n",
-    "    if fix_temperature != 0:\n",
-    "        temperature_k = fix_temperature\n",
-    "        print(\"Temperature is fixed!\")\n",
-    "    print(f\"Bias voltage is {bias_voltage} V\")\n",
-    "    print(f\"Detector integration time is set to {integration_time}\")\n",
-    "    print(f\"Mean temperature was {temperature:0.2f} °C / {temperature_k:0.2f} K\")\n",
-    "    print(f\"Operated in vacuum: {in_vacuum} \")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "reader = ChunkReader(filename, fastccdreaderh5.readData,\n",
-    "                     nImages, chunkSize,\n",
-    "                     path=h5path,\n",
-    "                     pixels_x=sensorSize[0],\n",
-    "                     pixels_y=sensorSize[1], )"
+    "\n",
+    "print(f\"Bias voltage is {bias_voltage} V\")\n",
+    "print(f\"Detector integration time is set to {integration_time}\")\n",
+    "print(f\"Mean temperature was {temperature:0.2f} °C / {temperature_k:0.2f} K\")\n",
+    "print(f\"Operated in vacuum: {in_vacuum} \")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [],
    "source": [
-    "noiseCal = xcal.NoiseCalculator(sensorSize, memoryCells,\n",
-    "                                cores=cpuCores, blockSize=blockSize,\n",
-    "                                parallel=run_parallel)\n",
-    "histCalRaw = xcal.HistogramCalculator(sensorSize, bins=1000,\n",
-    "                                      range=[0, 10000], parallel=False,\n",
-    "                                      memoryCells=memoryCells,\n",
-    "                                      cores=cpuCores, blockSize=blockSize)"
+    "data_dc = run_dir.select(\n",
+    "    *pixel_data, require_all=True).select_trains(np.s_[:n_trains])\n",
+    "print(f\"Reading data from: {data_dc.files}\\n\")\n",
+    "\n",
+    "data = data_dc[pixel_data].ndarray()\n",
+    "\n",
+    "noise_data = np.std(data, axis=0)\n",
+    "offset_data = np.mean(data, axis=0)"
    ]
   },
   {
@@ -197,18 +191,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "for data in reader.readChunks():\n",
-    "    dx = np.count_nonzero(data, axis=(0, 1))\n",
-    "    data = data[:, :, dx != 0]\n",
-    "    histCalRaw.fill(data)\n",
-    "    noiseCal.fill(data) #Fill calculators with data\n",
-    "    \n",
     "constant_maps = {}\n",
-    "constant_maps['Offset'] = noiseCal.getOffset()  #Produce offset map\n",
-    "constant_maps['Noise'] = noiseCal.get()  #Produce noise map\n",
-    "\n",
-    "noiseCal.reset()  #Reset noise calculator\n",
-    "print(\"Initial maps were created\")"
+    "constant_maps['Offset'] = offset_data[..., np.newaxis]\n",
+    "constant_maps['Noise'] = noise_data[..., np.newaxis]\n",
+    "print(\"Initial constant maps are created\")"
    ]
   },
   {
@@ -251,31 +237,26 @@
     "fig = xana.heatmapPlot(constant_maps['Offset'][:, :, 0],\n",
     "                       x_label='Columns', y_label='Rows',\n",
     "                       lut_label='Offset (ADU)',\n",
-    "                       x_range=(0, y),\n",
-    "                       y_range=(0, x), vmin=1000, vmax=4000)\n",
+    "                       x_range=(0, pixels_y),\n",
+    "                       y_range=(0, pixels_x), vmin=1000, vmax=4000)\n",
     "\n",
     "fig = xana.heatmapPlot(constant_maps['Noise'][:, :, 0],\n",
     "                       x_label='Columns', y_label='Rows',\n",
     "                       lut_label='Noise (ADU)',\n",
-    "                       x_range=(0, y),\n",
-    "                       y_range=(0, x), vmax=2 * np.mean(constant_maps['Noise']))"
+    "                       x_range=(0, pixels_y),\n",
+    "                       y_range=(0, pixels_x), vmax=2 * np.mean(constant_maps['Noise']))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Save constants to DB\n",
     "dclass=\"ePix100\"\n",
     "md = None\n",
     "\n",
-    "# If PDU(db_module) is not given with input parameters \n",
-    "# retrieve the connected physical detector unit\n",
-    "\n",
     "for const_name in constant_maps.keys():\n",
     "    det = getattr(Constants, dclass)\n",
     "    const = getattr(det, const_name)()\n",
@@ -291,14 +272,9 @@
     "            parm.lower_deviation = temp_limits\n",
     "            parm.upper_deviation = temp_limits\n",
     "\n",
-    "    # This should be used in case of running notebook \n",
-    "    # by a different method other than myMDC which already\n",
-    "    # sends CalCat info.\n",
-    "    # TODO: Set db_module to \"\" by default in the first cell\n",
-    "    if not db_module:\n",
-    "        db_module = get_pdu_from_db(karabo_id, karabo_da, const,\n",
-    "                                    condition, cal_db_interface,\n",
-    "                                    snapshot_at=creation_time)[0]\n",
+    "    db_module = get_pdu_from_db(karabo_id, karabo_da, const,\n",
+    "                                condition, cal_db_interface,\n",
+    "                                snapshot_at=creation_time)[0]\n",
     "\n",
     "    if db_output:\n",
     "        md = send_to_db(db_module, karabo_id, const, condition,\n",
@@ -317,13 +293,6 @@
     "      f\"• Temperature: {temperature_k}\\n• In Vacuum: {in_vacuum}\\n\"\n",
     "      f\"• Creation time: {md.calibration_constant_version.begin_at if md is not None else creation_time}\\n\")"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
@@ -342,7 +311,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.8.11"
   },
   "latex_envs": {
    "LaTeX_envs_menu_present": true,
diff --git a/notebooks/ePix100/Correction_ePix100_NBC.ipynb b/notebooks/ePix100/Correction_ePix100_NBC.ipynb
index c241f8e270cdc3bbaffc18c13211915bdd562771..8c5ec9fc0ce8b4d43782ca6aed0c2f63e63c9666 100644
--- a/notebooks/ePix100/Correction_ePix100_NBC.ipynb
+++ b/notebooks/ePix100/Correction_ePix100_NBC.ipynb
@@ -4,11 +4,11 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# ePIX Data Correction\n",
+    "# ePix100 Data Correction\n",
     "\n",
-    "Authors: Q. Tian S. Hauf M. Cascella, Version 1.0\n",
+    "Author: European XFEL Detector Group, Version: 2.0\n",
     "\n",
-    "The following notebook provides Offset correction of images acquired with the ePix100 detector."
+    "The following notebook provides data correction of images acquired with the ePix100 detector."
    ]
   },
   {
@@ -17,44 +17,49 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cluster_profile = \"noDB\"  # ipcluster profile to use\n",
-    "in_folder = \"/gpfs/exfel/exp/MID/202121/p002929/raw\"  # input folder, required\n",
+    "in_folder = \"/gpfs/exfel/exp/CALLAB/202031/p900113/raw\"  # input folder, required\n",
     "out_folder = \"\"  # output folder, required\n",
     "sequences = [-1]  # sequences to correct, set to -1 for all, range allowed\n",
-    "run = 126  # which run to read data from, required\n",
+    "sequences_per_node = 1  # number of sequence files per cluster node if run as slurm job, set to 0 to not run SLURM parallel\n",
+    "run = 9988  # which run to read data from, required\n",
     "\n",
+    "# Parameters for accessing the raw data.\n",
     "karabo_id = \"MID_EXP_EPIX-1\"  # karabo karabo_id\n",
     "karabo_da = \"EPIX01\"  # data aggregators\n",
-    "db_module = \"ePix100_M15\"  # module id in the database\n",
-    "receiver_id = \"RECEIVER\"  # inset for receiver devices\n",
+    "db_module = \"\"  # module id in the database\n",
+    "receiver_template = \"RECEIVER\"  # detector receiver template for accessing raw data files\n",
     "path_template = 'RAW-R{:04d}-{}-S{{:05d}}.h5'  # the template to use to access data\n",
-    "h5path = '/INSTRUMENT/{}/DET/{}:daqOutput/data/image'  # path in the HDF5 file to images\n",
-    "h5path_t = '/INSTRUMENT/{}/DET/{}:daqOutput/data/backTemp'  # path to find temperature at\n",
-    "h5path_cntrl = '/CONTROL/{}/DET'  # path to control data\n",
+    "instrument_source_template = '{}/DET/{}:daqOutput'  # instrument detector data source in h5files\n",
+    "\n",
+    "# Parameters affecting writing corrected data.\n",
+    "chunk_size_idim = 1  # H5 chunking size of output data\n",
+    "overwrite = True  # overwrite output folder\n",
+    "\n",
+    "# Only for testing\n",
+    "limit_images = 0  # ONLY FOR TESTING. process only first N images, 0 - process all.\n",
     "\n",
+    "# Parameters for the calibration database.\n",
     "use_dir_creation_date = True  # date constants injected before directory creation time\n",
     "cal_db_interface = \"tcp://max-exfl016:8015#8025\"  # calibration DB interface to use\n",
     "cal_db_timeout = 300000  # timeout on caldb requests\n",
     "\n",
-    "cpuCores = 4  # Specifies the number of running cpu cores\n",
-    "chunk_size_idim = 1  # H5 chunking size of output data\n",
-    "overwrite = True  # overwrite output folder\n",
-    "limit_images = 0  # process only first N images, 0 - process all\n",
-    "sequences_per_node = 1  # number of sequence files per cluster node if run as slurm job, set to 0 to not run SLURM parallel\n",
-    "\n",
+    "# Conditions for retrieving calibration constants.\n",
     "bias_voltage = 200  # bias voltage\n",
     "in_vacuum = False  # detector operated in vacuum\n",
     "fix_temperature = 290.  # fix temperature to this value\n",
+    "temp_deviations = 5.  # temperature deviation for the constant operating conditions\n",
     "gain_photon_energy = 9.0  # Photon energy used for gain calibration\n",
     "photon_energy = 0.  # Photon energy to calibrate in number of photons, 0 for calibration in keV\n",
     "\n",
+    "# Flags to select type of applied corrections.\n",
     "pattern_classification = True  # do clustering.\n",
     "relative_gain = True  # Apply relative gain correction.\n",
     "absolute_gain = True  # Apply absolute gain correction (implies relative gain).\n",
     "common_mode = True  # Apply common mode correction.\n",
-    "cm_min_frac = 0.25 # No CM correction is performed if after masking the ratio of good pixels falls below this \n",
-    "cm_noise_sigma = 5. # CM correction noise standard deviation\n",
     "\n",
+    "# Parameters affecting applied correction.\n",
+    "cm_min_frac = 0.25  # No CM correction is performed if after masking the ratio of good pixels falls below this \n",
+    "cm_noise_sigma = 5.  # CM correction noise standard deviation\n",
     "split_evt_primary_threshold = 7.  # primary threshold for split event correction\n",
     "split_evt_secondary_threshold = 5.  # secondary threshold for split event correction\n",
     "split_evt_mip_threshold = 1000.  # minimum ionizing particle threshold\n",
@@ -72,23 +77,27 @@
    "outputs": [],
    "source": [
     "import tabulate\n",
+    "import traceback\n",
     "import warnings\n",
     "\n",
     "import h5py\n",
+    "import pasha as psh\n",
+    "import multiprocessing\n",
     "import numpy as np\n",
     "import matplotlib.pyplot as plt\n",
-    "from IPython.display import Latex, display\n",
+    "from IPython.display import Latex, Markdown, display\n",
+    "from extra_data import RunDirectory, H5File\n",
     "from pathlib import Path\n",
     "\n",
-    "import XFELDetAna.xfelprofiler as xprof\n",
     "from XFELDetAna import xfelpyanatools as xana\n",
     "from XFELDetAna import xfelpycaltools as xcal\n",
     "from XFELDetAna.plotting.util import prettyPlotting\n",
-    "from XFELDetAna.util import env\n",
+    "from cal_tools import h5_copy_except\n",
     "from cal_tools.tools import (\n",
     "    get_constant_from_db,\n",
     "    get_dir_creation_date,\n",
     ")\n",
+    "from cal_tools.step_timing import StepTimer\n",
     "from iCalibrationDB import (\n",
     "    Conditions,\n",
     "    Constants,\n",
@@ -98,10 +107,6 @@
     "\n",
     "prettyPlotting = True\n",
     "\n",
-    "profiler = xprof.Profiler()\n",
-    "profiler.disable()\n",
-    "env.iprofile = cluster_profile\n",
-    "\n",
     "%matplotlib inline"
    ]
   },
@@ -111,20 +116,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "instrument_src = instrument_source_template.format(karabo_id, receiver_template)\n",
+    "\n",
     "if absolute_gain:\n",
-    "    relative_gain = True"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
-    "h5path_t = h5path_t.format(karabo_id, receiver_id)\n",
-    "h5path_cntrl = h5path_cntrl.format(karabo_id)\n",
-    "plot_unit = 'ADU'"
+    "    relative_gain = True\n",
+    "\n",
+    "plot_unit = 'ADU'\n"
    ]
   },
   {
@@ -138,12 +135,14 @@
     "\n",
     "in_folder = Path(in_folder)\n",
     "ped_dir = in_folder / f\"r{run:04d}\"\n",
+    "run_dc = RunDirectory(ped_dir)\n",
+    "\n",
     "fp_name = path_template.format(run, karabo_da)\n",
     "\n",
     "print(f\"Reading from: {ped_dir / fp_name}\")\n",
     "print(f\"Run is: {run}\")\n",
-    "print(f\"HDF5 path: {h5path}\")\n",
-    "print(f\"Data is output to: {out_folder}\")\n",
+    "print(f\"Instrument H5File source: {instrument_src}\")\n",
+    "print(f\"Data corrected files are stored at: {out_folder}\")\n",
     "\n",
     "creation_time = None\n",
     "if use_dir_creation_date:\n",
@@ -158,37 +157,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sensorSize = [x, y]\n",
-    "chunkSize = 100  # Number of images to read per chunk\n",
-    "# Sensor area will be analysed according to blocksize\n",
-    "blockSize = [sensorSize[0]//2, sensorSize[1]//2]\n",
-    "xcal.defaultBlockSize = blockSize\n",
-    "memoryCells = 1  # ePIX has no memory cells\n",
-    "run_parallel = True\n",
-    "\n",
-    "# Read slow data from the first available sequence file.\n",
-    "filename = ped_dir / fp_name.format(\n",
-    "    sequences[0] if sequences[0] != -1 else 0)\n",
-    "with h5py.File(filename, 'r') as f:\n",
-    "    integration_time = int(\n",
-    "        f[f\"{h5path_cntrl}/CONTROL/expTime/value\"][0])\n",
-    "    temperature = np.mean(f[h5path_t]) / 100.\n",
-    "    temperature_k = temperature + 273.15\n",
-    "    if fix_temperature != 0:\n",
-    "        temperature_k = fix_temperature\n",
-    "        print(\"Temperature is fixed!\")\n",
-    "    print(f\"Bias voltage is {bias_voltage} V\")\n",
-    "    print(f\"Detector integration time is set to {integration_time}\")\n",
-    "    print(\n",
-    "        f\"Mean temperature was {temperature:0.2f} °C \"\n",
-    "        f\"/ {temperature_k:0.2f} K at beginning of run\"\n",
-    "    )\n",
-    "    print(f\"Operated in vacuum: {in_vacuum} \")\n",
+    "seq_files = [Path(f.filename) for f in run_dc.select(f\"*{karabo_id}*\").files]\n",
     "\n",
-    "out_folder = Path(out_folder)\n",
-    "if out_folder.is_dir() and not overwrite:\n",
-    "    raise AttributeError(\"Output path exists! Exiting\")\n",
-    "out_folder.mkdir(parents=True, exist_ok=True)"
+    "# If a set of sequences requested to correct,\n",
+    "# adapt seq_files list.\n",
+    "if sequences != [-1]:\n",
+    "    seq_files = [f for f in seq_files if any(f.match(f\"*-S{s:05d}.h5\") for s in sequences)]\n",
+    "\n",
+    "if not len(seq_files):\n",
+    "    raise IndexError(\"No sequence files available for the selected sequences.\")\n",
+    "\n",
+    "print(f\"Processing a total of {len(seq_files)} sequence files\")"
    ]
   },
   {
@@ -197,17 +176,52 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Glob the right *.h5 fast data files.\n",
-    "seq_files = sorted(ped_dir.glob(f\"*{karabo_da}*.h5\"))\n",
+    "step_timer = StepTimer()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "step_timer.start()\n",
     "\n",
-    "# If a set of sequences requested to correct,\n",
-    "# adapt seq_files list.\n",
-    "if sequences != [-1]:\n",
-    "    seq_files = [f for f in seq_files\n",
-    "                 if any(\n",
-    "                     f.match(f\"*-S{s:05d}.h5\") for s in sequences)]\n",
+    "sensorSize = [x, y]\n",
+    "# Sensor area will be analysed according to blocksize\n",
+    "blockSize = [sensorSize[0]//2, sensorSize[1]//2]\n",
+    "xcal.defaultBlockSize = blockSize\n",
+    "memoryCells = 1  # ePIX has no memory cells\n",
+    "run_parallel = False\n",
+    "\n",
+    "# Read control data.\n",
+    "integration_time = int(run_dc.get_run_value(\n",
+    "    f\"{karabo_id}/DET/CONTROL\",\n",
+    "    \"expTime.value\"))\n",
+    "temperature = np.mean(run_dc.get_array(\n",
+    "    f\"{karabo_id}/DET/{receiver_template}:daqOutput\",\n",
+    "    f\"data.backTemp\").values) / 100.\n",
+    "\n",
+    "if fix_temperature != 0:\n",
+    "    temperature_k = fix_temperature\n",
+    "    print(\"Temperature is fixed!\")\n",
+    "else:\n",
+    "    temperature_k = temperature + 273.15\n",
     "\n",
-    "print(f\"Processing a total of {len(seq_files)} sequence files\")"
+    "print(f\"Bias voltage is {bias_voltage} V\")\n",
+    "print(f\"Detector integration time is set to {integration_time}\")\n",
+    "print(\n",
+    "    f\"Mean temperature was {temperature:0.2f} °C \"\n",
+    "    f\"/ {temperature_k:0.2f} K at beginning of the run.\"\n",
+    ")\n",
+    "print(f\"Operated in vacuum: {in_vacuum} \")\n",
+    "\n",
+    "out_folder = Path(out_folder)\n",
+    "if out_folder.is_dir() and not overwrite:\n",
+    "    raise AttributeError(\n",
+    "        \"Output path exists! Exiting as overwriting corrected data is prohibited.\")\n",
+    "out_folder.mkdir(parents=True, exist_ok=True)\n",
+    "step_timer.done_step(f'Reading control parameters.')\n"
    ]
   },
   {
@@ -231,19 +245,17 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
+    "## Retrieving calibration constants\n",
+    "\n",
     "As a first step, dark maps have to be loaded."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "temp_limits = 5.\n",
-    "\n",
     "cond_dict = {\n",
     "    \"bias_voltage\": bias_voltage,\n",
     "    \"integration_time\": integration_time,\n",
@@ -274,8 +286,8 @@
     "    # TODO: Fix this logic.\n",
     "    for parm in condition.parameters:\n",
     "        if parm.name == \"Sensor Temperature\":\n",
-    "            parm.lower_deviation = temp_limits\n",
-    "            parm.upper_deviation = temp_limits\n",
+    "            parm.lower_deviation = temp_deviations\n",
+    "            parm.upper_deviation = temp_deviations\n",
     "\n",
     "    const_data[cname] = get_constant_from_db(\n",
     "        karabo_id=karabo_id,\n",
@@ -295,8 +307,7 @@
     "        \"No gain correction will be applied\"\n",
     "    )\n",
     "    relative_gain = False\n",
-    "    absolute_gain = False\n",
-    "    plot_unit = 'ADU'"
+    "    absolute_gain = False\n"
    ]
   },
   {
@@ -305,12 +316,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Initializing some parameters.\n",
+    "hscale = 1\n",
+    "stats = True\n",
     "hrange = np.array([-50, 1000])\n",
     "nbins = hrange[1] - hrange[0]\n",
-    "hscale = 1\n",
-    "\n",
-    "commonModeBlockSize = [x//2, y//2]\n",
-    "stats = True"
+    "commonModeBlockSize = [x//2, y//2]"
    ]
   },
   {
@@ -319,23 +330,23 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "offsetCorrection = xcal.OffsetCorrection(\n",
+    "histCalOffsetCor = xcal.HistogramCalculator(\n",
     "    sensorSize,\n",
-    "    const_data[\"Offset\"],\n",
+    "    bins=nbins,\n",
+    "    range=hrange,\n",
+    "    parallel=run_parallel,\n",
     "    nCells=memoryCells,\n",
-    "    cores=cpuCores,\n",
-    "    gains=None,\n",
-    "    blockSize=blockSize,\n",
-    "    parallel=run_parallel\n",
+    "    blockSize=blockSize\n",
     ")\n",
     "\n",
-    "histCalOffsetCor = xcal.HistogramCalculator(\n",
+    "\n",
+    "# *****************Histogram Calculators****************** #\n",
+    "histCalCor = xcal.HistogramCalculator(\n",
     "    sensorSize,\n",
-    "    bins=nbins,\n",
-    "    range=hrange,\n",
+    "    bins=1050,\n",
+    "    range=[-50, 1000],\n",
     "    parallel=run_parallel,\n",
     "    nCells=memoryCells,\n",
-    "    cores=cpuCores,\n",
     "    blockSize=blockSize\n",
     ")"
    ]
@@ -353,10 +364,8 @@
     "        range=hrange,\n",
     "        parallel=run_parallel,\n",
     "        nCells=memoryCells,\n",
-    "        cores=cpuCores,\n",
     "        blockSize=blockSize,\n",
     "    )\n",
-    "\n",
     "    cmCorrectionB = xcal.CommonModeCorrection(\n",
     "        shape=sensorSize,\n",
     "        blockSize=commonModeBlockSize, \n",
@@ -364,6 +373,7 @@
     "        nCells=memoryCells, \n",
     "        noiseMap=const_data['Noise'],\n",
     "        runParallel=run_parallel,\n",
+    "        parallel=run_parallel,\n",
     "        stats=stats,\n",
     "        minFrac=cm_min_frac,\n",
     "        noiseSigma=cm_noise_sigma,\n",
@@ -375,6 +385,7 @@
     "        nCells=memoryCells, \n",
     "        noiseMap=const_data['Noise'],\n",
     "        runParallel=run_parallel,\n",
+    "        parallel=run_parallel,\n",
     "        stats=stats,\n",
     "        minFrac=cm_min_frac,\n",
     "        noiseSigma=cm_noise_sigma,\n",
@@ -386,10 +397,11 @@
     "        nCells=memoryCells, \n",
     "        noiseMap=const_data['Noise'],\n",
     "        runParallel=run_parallel,\n",
+    "        parallel=run_parallel,\n",
     "        stats=stats,\n",
     "        minFrac=cm_min_frac,\n",
     "        noiseSigma=cm_noise_sigma,\n",
-    "    )\n"
+    "    )"
    ]
   },
   {
@@ -411,7 +423,6 @@
     "        gain_cnst/const_data[\"RelativeGain\"][..., None],\n",
     "        nCells=memoryCells,\n",
     "        parallel=run_parallel,\n",
-    "        cores=cpuCores,\n",
     "        blockSize=blockSize,\n",
     "        gains=None,\n",
     "    )\n",
@@ -422,7 +433,6 @@
     "        range=hrange,\n",
     "        parallel=run_parallel,\n",
     "        nCells=memoryCells,\n",
-    "        cores=cpuCores,\n",
     "        blockSize=blockSize\n",
     "    )\n",
     "    \n",
@@ -433,7 +443,6 @@
     "            range=hrange*hscale,\n",
     "            parallel=run_parallel,\n",
     "            nCells=memoryCells,\n",
-    "            cores=cpuCores,\n",
     "            blockSize=blockSize\n",
     "        )"
    ]
@@ -453,19 +462,16 @@
     "        split_evt_mip_threshold,\n",
     "        tagFirstSingles=0,\n",
     "        nCells=memoryCells,\n",
-    "        cores=cpuCores,\n",
     "        allowElongated=False,\n",
     "        blockSize=[x, y],\n",
-    "        runParallel=run_parallel,\n",
+    "        parallel=run_parallel,\n",
     "    )\n",
-    "\n",
     "    histCalSECor = xcal.HistogramCalculator(\n",
     "        sensorSize,\n",
     "        bins=nbins,\n",
     "        range=hrange,\n",
     "        parallel=run_parallel,\n",
     "        nCells=memoryCells,\n",
-    "        cores=cpuCores,\n",
     "        blockSize=blockSize,\n",
     "    )\n",
     "    histCalGainCorSingles = xcal.HistogramCalculator(\n",
@@ -474,7 +480,6 @@
     "        range=hrange*hscale,\n",
     "        parallel=run_parallel,\n",
     "        nCells=memoryCells,\n",
-    "        cores=cpuCores,\n",
     "        blockSize=blockSize\n",
     "    )"
    ]
@@ -483,7 +488,89 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Applying corrections"
+    "## Applying corrections"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def correct_train(wid, index, tid, d):\n",
+    "\n",
+    "    d = d[pixel_data[0]][pixel_data[1]][..., np.newaxis].astype(np.float32)\n",
+    "    d = np.compress(\n",
+    "        np.any(d > 0, axis=(0, 1)), d, axis=2)\n",
+    "    \n",
+    "    # Offset correction.\n",
+    "    d -= const_data[\"Offset\"]\n",
+    "\n",
+    "    histCalOffsetCor.fill(d)\n",
+    "    # Common Mode correction.\n",
+    "    if common_mode:\n",
+    "        # Block CM\n",
+    "        d = cmCorrectionB.correct(d)\n",
+    "        # Row CM\n",
+    "        d = cmCorrectionR.correct(d)\n",
+    "        # COL CM\n",
+    "        d = cmCorrectionC.correct(d)\n",
+    "        histCalCMCor.fill(d)\n",
+    "\n",
+    "    # relative gain correction.\n",
+    "    if relative_gain:\n",
+    "        d = gainCorrection.correct(d)\n",
+    "        histCalRelGainCor.fill(d)\n",
+    "\n",
+    "    data[index, ...] = np.squeeze(d)\n",
+    "\n",
+    "    \"\"\"The gain correction is currently applying\n",
+    "    an absolute correction (not a relative correction\n",
+    "    as the implied by the name);\n",
+    "    it changes the scale (the unit of measurement)\n",
+    "    of the data from ADU to either keV or n_of_photons.\n",
+    "    But the pattern classification relies on comparing\n",
+    "    data with the noise map, which is still in ADU.\n",
+    "\n",
+    "    The best solution is to do a relative gain\n",
+    "    correction first and apply the global absolute\n",
+    "    gain to the data at the end, after clustering.\n",
+    "    \"\"\"\n",
+    "\n",
+    "    if pattern_classification:\n",
+    "\n",
+    "        d_clu, patterns = patternClassifier.classify(d)\n",
+    "\n",
+    "        d_clu[d_clu < (split_evt_primary_threshold*const_data[\"Noise\"])] = 0\n",
+    "\n",
+    "        data_patterns[index, ...] = np.squeeze(patterns)\n",
+    "\n",
+    "        data_clu[index, ...] = np.squeeze(d_clu)\n",
+    "\n",
+    "        d_clu[patterns != 100] = np.nan\n",
+    "\n",
+    "        histCalSECor.fill(d_clu)\n",
+    "\n",
+    "    # absolute gain correction\n",
+    "    # changes data from ADU to keV (or n. of photons)\n",
+    "    if absolute_gain:\n",
+    "\n",
+    "        d = d * gain_cnst\n",
+    "        if photon_energy > 0:\n",
+    "            d /= photon_energy\n",
+    "        histCalAbsGainCor.fill(d)\n",
+    "\n",
+    "        if pattern_classification:\n",
+    "            # Modify pattern classification.\n",
+    "            d_clu = d_clu * gain_cnst\n",
+    "            if photon_energy > 0:\n",
+    "                d_clu /= photon_energy\n",
+    "\n",
+    "            histCalGainCorSingles.fill(d_clu)\n",
+    "\n",
+    "            data_clu[index, ...] = np.squeeze(d_clu)\n",
+    "\n",
+    "    histCalCor.fill(d)"
    ]
   },
   {
@@ -492,28 +579,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def copy_and_sanitize_non_cal_data(\n",
-    "    infile: h5py,\n",
-    "    outfile: h5py,\n",
-    "    h5base: str\n",
-    "):\n",
-    "    \"\"\" Copy and sanitize data in `infile`\n",
-    "    that is not touched by `correctEPIX`. \"\"\"\n",
-    "\n",
-    "    if h5base.startswith(\"/\"):\n",
-    "        h5base = h5base[1:]\n",
-    "    dont_copy = [h5base+\"/pixels\"]\n",
-    "\n",
-    "    def visitor(k, item):\n",
-    "        if k not in dont_copy:\n",
-    "            if isinstance(item, h5py.Group):\n",
-    "                outfile.create_group(k)\n",
-    "            elif isinstance(item, h5py.Dataset):\n",
-    "                group = str(k).split(\"/\")\n",
-    "                group = \"/\".join(group[:-1])\n",
-    "                infile.copy(k, outfile[group])\n",
-    "\n",
-    "    infile.visititems(visitor)"
+    "pixel_data = (instrument_src, \"data.image.pixels\")\n",
+    "\n",
+    "# 10 is a number chosen after testing 1 ... 71 parallel threads\n",
+    "context = psh.context.ThreadContext(num_workers=10)"
    ]
   },
   {
@@ -523,85 +592,76 @@
    "outputs": [],
    "source": [
     "for f in seq_files:\n",
-    "    data = None\n",
+    "\n",
+    "    seq_dc = H5File(f)\n",
+    "\n",
+    "    n_imgs = seq_dc.get_data_counts(*pixel_data).shape[0]\n",
+    "\n",
+    "    # Data shape in seq_dc excluding trains with empty images. \n",
+    "    dshape = seq_dc[pixel_data].shape\n",
+    "    dataset_chunk = ((chunk_size_idim,) + dshape[1:])  # e.g. (1, pixels_x, pixels_y) \n",
+    "\n",
+    "    if n_imgs - dshape[0] != 0:\n",
+    "        print(f\"- WARNING: {f} has {n_imgs - dshape[0]} trains with empty data.\")\n",
+    "\n",
+    "    # This parameter is only used for testing.\n",
+    "    if limit_images > 0:\n",
+    "        n_imgs = min(n_imgs, limit_images)\n",
+    "\n",
+    "    data = context.alloc(shape=dshape, dtype=np.float32)\n",
+    "\n",
+    "    if pattern_classification:\n",
+    "        data_clu = context.alloc(shape=dshape, dtype=np.float32)\n",
+    "        data_patterns = context.alloc(shape=dshape, dtype=np.int32)\n",
+    "\n",
+    "    step_timer.start()\n",
+    "\n",
+    "    context.map(\n",
+    "        correct_train, seq_dc.select(\n",
+    "            *pixel_data, require_all=True).select_trains(np.s_[:n_imgs])\n",
+    "    )\n",
+    "    step_timer.done_step(f'Correcting {n_imgs} trains.')\n",
+    "\n",
+    "    # Store detector h5 information in the corrected file\n",
+    "    # and deselect data to correct and store later.\n",
+    "    step_timer.start()\n",
+    "\n",
     "    out_file = out_folder / f.name.replace(\"RAW\", \"CORR\")\n",
-    "    with h5py.File(f, \"r\") as infile, h5py.File(out_file, \"w\") as ofile:  # noqa\n",
-    "        try:\n",
-    "            copy_and_sanitize_non_cal_data(infile, ofile, h5path)\n",
-    "            data = infile[h5path+\"/pixels\"][()]\n",
-    "            data = np.compress(\n",
-    "                np.any(data > 0, axis=(1, 2)), data, axis=0)\n",
-    "            if limit_images > 0:\n",
-    "                data = data[:limit_images, ...]\n",
-    "\n",
-    "            oshape = data.shape\n",
-    "            data = np.moveaxis(data, 0, 2)\n",
-    "            ddset = ofile.create_dataset(\n",
-    "                h5path+\"/pixels\",\n",
-    "                oshape,\n",
-    "                chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
+    "    data_path = \"INSTRUMENT/\"+instrument_src+\"/data/image\"\n",
+    "    pixels_path = f\"{data_path}/pixels\"\n",
+    "    \n",
+    "    # First copy all raw data source to the corrected file,\n",
+    "    # while excluding the raw data image /data/image/pixels.\n",
+    "    with h5py.File(out_file, 'w') as ofile:\n",
+    "        # Copy RAW non-calibrated sources.\n",
+    "        with h5py.File(f, 'r') as sfile:\n",
+    "            h5_copy_except.h5_copy_except_paths(\n",
+    "                sfile, ofile,\n",
+    "                [pixels_path])\n",
+    "\n",
+    "        # Create dataset in CORR h5 file and add corrected images.\n",
+    "        dataset = ofile.create_dataset(\n",
+    "            pixels_path,\n",
+    "            data=data,\n",
+    "            chunks=dataset_chunk,\n",
+    "            dtype=np.float32)\n",
+    "\n",
+    "        if pattern_classification:\n",
+    "            # Save /data/image//pixels_classified in corrected file.\n",
+    "            datasetc = ofile.create_dataset(\n",
+    "                f\"{data_path}/pixels_classified\",\n",
+    "                data=data_clu,\n",
+    "                chunks=dataset_chunk,\n",
     "                dtype=np.float32)\n",
     "\n",
-    "            # Offset correction.\n",
-    "            data = offsetCorrection.correct(data.astype(np.float32))\n",
-    "            histCalOffsetCor.fill(data)\n",
-    "            \n",
-    "            # Common Mode correction.\n",
-    "            if common_mode:\n",
-    "                # Block CM\n",
-    "                data = cmCorrectionB.correct(data)\n",
-    "                # Row CM\n",
-    "                data = cmCorrectionR.correct(data)\n",
-    "                # COL CM\n",
-    "                data = cmCorrectionC.correct(data)\n",
-    "                histCalCMCor.fill(data)\n",
-    "\n",
-    "            # relative gain correction.\n",
-    "            if relative_gain:\n",
-    "                data = gainCorrection.correct(data.astype(np.float32))\n",
-    "                histCalRelGainCor.fill(data)\n",
-    "\n",
-    "            ddset[...] = np.moveaxis(data, 2, 0)\n",
-    "\n",
-    "            if pattern_classification:\n",
-    "                ddsetc = ofile.create_dataset(\n",
-    "                    h5path+\"/pixels_classified\",\n",
-    "                    oshape,\n",
-    "                    chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                    dtype=np.float32, compression=\"gzip\")\n",
-    "\n",
-    "                ddsetp = ofile.create_dataset(\n",
-    "                    h5path+\"/patterns\",\n",
-    "                    oshape,\n",
-    "                    chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                    dtype=np.int32, compression=\"gzip\")\n",
-    "\n",
-    "                data_clu, patterns = patternClassifier.classify(data)\n",
-    "\n",
-    "                data_clu[data_clu < (split_evt_primary_threshold*const_data[\"Noise\"])] = 0  # noqa\n",
-    "                ddsetc[...] = np.moveaxis(data_clu, 2, 0)\n",
-    "                ddsetp[...] = np.moveaxis(patterns, 2, 0)\n",
-    "\n",
-    "                data_clu[patterns != 100] = np.nan\n",
-    "                histCalSECor.fill(data_clu)\n",
-    "\n",
-    "            # absolute gain correction\n",
-    "            # changes data from ADU to keV (or n. of photons)\n",
-    "            if absolute_gain:\n",
-    "                data = data * gain_cnst\n",
-    "                if photon_energy > 0:\n",
-    "                    data /= photon_energy\n",
-    "                histCalAbsGainCor.fill(data)\n",
-    "\n",
-    "                if pattern_classification:\n",
-    "                    data_clu = data_clu *gain_cnst\n",
-    "                    if photon_energy > 0:\n",
-    "                        data_clu /= photon_energy\n",
-    "                    ddsetc[...] = np.moveaxis(data_clu, 2, 0)\n",
-    "                    histCalGainCorSingles.fill(data_clu)\n",
-    "\n",
-    "        except Exception as e:\n",
-    "            print(f\"ERROR applying common mode correction for {f}: {e}\")"
+    "            # Save /data/image//patterns in corrected file.\n",
+    "            datasetp = ofile.create_dataset(\n",
+    "                f\"{data_path}/patterns\",\n",
+    "                data=data_patterns,\n",
+    "                chunks=dataset_chunk,\n",
+    "                dtype=np.int32)\n",
+    "\n",
+    "        step_timer.done_step('Storing data.')"
    ]
   },
   {
@@ -610,7 +670,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ho, eo, co, so = histCalOffsetCor.get()\n",
+    "ho, eo, co, so = histCalCor.get()\n",
     "\n",
     "d = [{\n",
     "    'x': co,\n",
@@ -619,9 +679,21 @@
     "    'drawstyle': 'steps-mid',\n",
     "    'errorstyle': 'bars',\n",
     "    'errorcoarsing': 2,\n",
-    "    'label': 'Offset corr.'\n",
+    "    'label': 'Total corr.'\n",
     "}]\n",
     "\n",
+    "ho, eo, co, so = histCalOffsetCor.get()\n",
+    "\n",
+    "d.append({\n",
+    "    'x': co,\n",
+    "    'y': ho,\n",
+    "    'y_err': np.sqrt(ho[:]),\n",
+    "    'drawstyle': 'steps-mid',\n",
+    "    'errorstyle': 'bars',\n",
+    "    'errorcoarsing': 2,\n",
+    "    'label': 'Offset corr.'\n",
+    "})\n",
+    "\n",
     "if common_mode:\n",
     "    ho, eo, co, so = histCalCMCor.get()\n",
     "    d.append({\n",
@@ -715,7 +787,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Mean Image of last Sequence ##"
+    "## Mean Image of the corrected data"
    ]
   },
   {
@@ -724,20 +796,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "fig = xana.heatmapPlot(\n",
-    "    np.nanmedian(data, axis=2),\n",
+    "    np.nanmedian(data, axis=0),\n",
     "    x_label='Columns', y_label='Rows',\n",
     "    lut_label=f'Signal ({plot_unit})',\n",
     "    x_range=(0, y),\n",
     "    y_range=(0, x),\n",
-    "    vmin=-50, vmax=50)"
+    "    vmin=-50, vmax=50)\n",
+    "step_timer.done_step(f'Plotting mean image of {data.shape[0]} trains.')"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Single Shot of last Sequence ##"
+    "## Single Shot of the corrected data"
    ]
   },
   {
@@ -746,13 +820,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "fig = xana.heatmapPlot(\n",
-    "    data[..., 0],\n",
+    "    data[0, ...],\n",
     "    x_label='Columns', y_label='Rows',\n",
     "    lut_label=f'Signal ({plot_unit})',\n",
     "    x_range=(0, y),\n",
     "    y_range=(0, x),\n",
-    "    vmin=-50, vmax=50)"
+    "    vmin=-50, vmax=50)\n",
+    "step_timer.done_step(f'Plotting single shot of corrected data.')"
    ]
   }
  ],
@@ -793,5 +869,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }
diff --git a/notebooks/generic/overallmodules_Darks_Summary_NBC.ipynb b/notebooks/generic/overallmodules_Darks_Summary_NBC.ipynb
index 6fbfb3d357a1bbf6fc04212144c1ee7b2227efcb..e204cf3f605f20c7342638395402119b7f8e7864 100644
--- a/notebooks/generic/overallmodules_Darks_Summary_NBC.ipynb
+++ b/notebooks/generic/overallmodules_Darks_Summary_NBC.ipynb
@@ -6,8 +6,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#Author: K. Ahmed, M. Karnevsky, Version: 0.1\n",
-    "#The following is a summary for the processing of dark images and calibration constants production.\n",
+    "# Author: European XFEL Detector Group, Version: 1.0\n",
+    "\n",
+    "#  Summary for processed of dark calibration constants and a comparison with previous injected constants.\n",
     "\n",
     "out_folder = \"/gpfs/exfel/data/scratch/ahmedk/test/fixed_gain/SPB_summary_fix2\" # path to output to, required\n",
     "karabo_id = \"SPB_DET_AGIPD1M-1\" # detector instance\n",
@@ -25,7 +26,6 @@
     "import os\n",
     "import warnings\n",
     "from collections import OrderedDict\n",
-    "from datetime import datetime\n",
     "from pathlib import Path\n",
     "\n",
     "warnings.filterwarnings('ignore')\n",
@@ -40,7 +40,6 @@
     "\n",
     "matplotlib.use(\"agg\")\n",
     "import matplotlib.gridspec as gridspec\n",
-    "import matplotlib.patches as patches\n",
     "import matplotlib.pyplot as plt\n",
     "\n",
     "%matplotlib inline\n",
@@ -49,8 +48,6 @@
     "from cal_tools.ana_tools import get_range\n",
     "from cal_tools.plotting import show_processed_modules\n",
     "from cal_tools.tools import CalibrationMetadata, module_index_to_qm\n",
-    "from iCalibrationDB import Detectors\n",
-    "from XFELDetAna.plotting.heatmap import heatmapPlot\n",
     "from XFELDetAna.plotting.simpleplot import simplePlot"
    ]
   },
diff --git a/notebooks/pnCCD/Characterize_pnCCD_Dark_NBC.ipynb b/notebooks/pnCCD/Characterize_pnCCD_Dark_NBC.ipynb
index b23ef32bfcfde18bba57315e5f93fe120e7b07df..f3e8d5cf8d7bdb397b27b50d376eedab3058c9ff 100644
--- a/notebooks/pnCCD/Characterize_pnCCD_Dark_NBC.ipynb
+++ b/notebooks/pnCCD/Characterize_pnCCD_Dark_NBC.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# pnCCD Dark Characterization\n",
     "\n",
-    "Author: DET Group, modified by Kiana Setoodehnia, Version: 4.0 (December 2020)\n",
+    "Author: DET Group, modified by Kiana Setoodehnia, Version: 5.0\n",
     "\n",
     "The following notebook provides dark image analysis of the pnCCD detector. Dark characterization evaluates offset and noise of the detector and gives information about bad pixels. \n",
     "\n",
@@ -22,31 +22,21 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:54:38.999974Z",
-     "start_time": "2018-12-06T10:54:38.983406Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "cluster_profile = \"noDB\"  # ipcluster profile to use\n",
     "in_folder = \"/gpfs/exfel/exp/SQS/202031/p900166/raw\"  # input folder, required\n",
-    "out_folder = '/gpfs/exfel/data/scratch/setoodeh'  # output folder, required\n",
-    "sequence = 0  # sequence file to use\n",
+    "out_folder = '/gpfs/exfel/data/scratch/ahmedk/test/remove'  # output folder, required\n",
     "run = 339 # which run to read data from, required\n",
     "\n",
-    "# Data files parameters:\n",
-    "db_module = \"pnCCD_M205_M206\" # the device name for pnCCD detector\n",
+    "# Data files parameters.\n",
     "karabo_da = ['PNCCD01'] # data aggregators\n",
-    "karabo_da_control = \"PNCCD02\" # file inset for control data\n",
     "karabo_id = \"SQS_NQS_PNCCD1MP\" # karabo prefix of PNCCD devices\n",
     "receiver_id = \"PNCCD_FMT-0\" # inset for receiver devices\n",
     "path_template = 'RAW-R{:04d}-{}-S{{:05d}}.h5'  # the template to use to access data\n",
-    "h5path = '/INSTRUMENT/{}/CAL/{}:output/data/image/' # path in the HDF5 file the data is at\n",
-    "h5path_ctrl = '/CONTROL/{}/CTRL/TCTRL'\n",
+    "instrument_source_template = '{}/CAL/{}:output'  # data source path in h5file. Template filled with karabo_id and receiver_id\n",
     "\n",
-    "# Database access parameters:\n",
+    "# Database access parameters.\n",
     "use_dir_creation_date = True  # use dir creation date as data production reference date\n",
     "cal_db_interface = \"tcp://max-exfl016:8021\"  # calibration DB interface to use\n",
     "cal_db_timeout = 300000 # timeout on caldb requests\n",
@@ -54,51 +44,51 @@
     "local_output = True # if True, the notebook saves dark constants locally\n",
     "creation_time = \"\" # To overwrite the measured creation_time. Required Format: YYYY-MM-DD HR:MN:SC.00 e.g. 2019-07-04 11:02:41.00\n",
     "\n",
-    "number_dark_frames = 0  # number of images to be used, if set to 0 all available images are used\n",
-    "chunkSize = 100 # number of images to read per chunk\n",
+    "# Parameters used for injecting calibration constants.\n",
     "fix_temperature_top = 0.  # fix temperature of top pnCCD sensor in K. Set to 0, to use the value from slow data\n",
     "fix_temperature_bot = 0.  # fix temperature of bottom pnCCD sensor in K. Set to 0, to use the value from slow data\n",
-    "gain = 0.1  # the detector's gain setting. Set to 0.1 to use the value from the slow data\n",
+    "temp_limits = 5  # temperature limits in which calibration parameters are considered equal\n",
+    "gain = -1  # the detector's gain setting. Set to -1 to use the value from the slow data.\n",
     "bias_voltage = 0. # the detector's bias voltage. set to 0. to use the value from slow data.\n",
-    "integration_time = 70  # detector's integration time\n",
+    "integration_time = 70  # detector's integration time.\n",
+    "\n",
+    "# Parameters affecting dark constants calculation.\n",
     "commonModeAxis = 0 # axis along which common mode will be calculated (0: along rows, 1: along columns)\n",
     "commonModeBlockSize = [512, 512] # size of the detector in pixels for common mode calculations\n",
     "sigmaNoise = 10.  # pixels whose signal value exceeds sigmaNoise*noise will be considered as cosmics and are masked\n",
     "bad_pixel_offset_sigma = 4.  # any pixel whose offset beyond this standard deviations is a bad pixel\n",
     "bad_pixel_noise_sigma = 4.  # any pixel whose noise beyond this standard deviations is a bad pixel\n",
-    "temp_limits = 5  # temperature limits in which calibration parameters are considered equal\n",
+    "max_trains = 500  # Maximum number of trains to use for dark processing. Set to 0 to process all trains.\n",
+    "min_trains = 1  # Minimum number of trains to proceed with dark processing.\n",
+    "\n",
+    "# Don't delete. myMDC sends this parameter by default.\n",
+    "operation_mode = ''  # Detector operation mode, optional\n",
     "\n",
-    "run_parallel = True # for parallel computation\n",
-    "cpuCores = 40 # specifies the number of running cpu cores\n",
-    "operation_mode = ''  # Detector operation mode, optional"
+    "# TODO: Remove unused parameter\n",
+    "db_module = \"\" # the device name for pnCCD detector"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:54:39.467334Z",
-     "start_time": "2018-12-06T10:54:39.427784Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "import copy\n",
     "import datetime\n",
     "import os\n",
     "import warnings\n",
     "\n",
     "warnings.filterwarnings('ignore')\n",
     "\n",
-    "import h5py\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pasha as psh\n",
+    "from extra_data import RunDirectory\n",
     "\n",
     "%matplotlib inline\n",
-    "import XFELDetAna.xfelprofiler as xprof\n",
+    "from cal_tools import pnccdlib\n",
     "from cal_tools.enums import BadPixels\n",
-    "from cal_tools.pnccdlib import VALID_GAINS, extract_slow_data\n",
+    "from cal_tools.step_timing import StepTimer\n",
     "from cal_tools.tools import (\n",
     "    get_dir_creation_date,\n",
     "    get_pdu_from_db,\n",
@@ -107,41 +97,13 @@
     "    save_const_to_h5,\n",
     "    send_to_db,\n",
     ")\n",
-    "from iCalibrationDB import Conditions, Constants, Detectors, Versions\n",
+    "from iCalibrationDB import Conditions, Constants\n",
     "from iCalibrationDB.detectors import DetectorTypes\n",
     "from IPython.display import Markdown, display\n",
     "from prettytable import PrettyTable\n",
     "\n",
-    "profiler = xprof.Profiler()\n",
-    "profiler.disable()\n",
-    "from XFELDetAna.util import env\n",
-    "\n",
-    "env.iprofile = cluster_profile\n",
     "from XFELDetAna import xfelpyanatools as xana\n",
-    "from XFELDetAna import xfelpycaltools as xcal\n",
-    "from XFELDetAna.plotting.util import prettyPlotting\n",
-    "\n",
-    "prettyPlotting=True\n",
-    "from XFELDetAna.detectors.fastccd import readerh5 as fastccdreaderh5\n",
-    "from XFELDetAna.xfelreaders import ChunkReader"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:54:39.467334Z",
-     "start_time": "2018-12-06T10:54:39.427784Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "def nImagesOrLimit(nImages, limit):\n",
-    "    if limit == 0:\n",
-    "        return nImages\n",
-    "    else:\n",
-    "        return min(nImages, limit)"
+    "from XFELDetAna import xfelpycaltools as xcal"
    ]
   },
   {
@@ -162,12 +124,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:54:40.058101Z",
-     "start_time": "2018-12-06T10:54:40.042615Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "# Calibration Database Settings, and Some Initial Run Parameters & Paths:\n",
@@ -177,11 +134,8 @@
     "pixels_y = 1024 # number of columns of the pnCCD \n",
     "print(f\"pnCCD size is: {pixels_x}x{pixels_y} pixels.\")\n",
     "\n",
-    "ped_dir = \"{}/r{:04d}\".format(in_folder, run)\n",
-    "fp_name = path_template.format(run, karabo_da[0])\n",
-    "fp_path = '{}/{}'.format(ped_dir, fp_name)\n",
-    "filename = fp_path.format(sequence)\n",
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
+    "instrument_src = instrument_source_template.format(\n",
+    "    karabo_id, receiver_id)\n",
     "\n",
     "# Output Folder Creation:\n",
     "os.makedirs(out_folder, exist_ok=True)\n",
@@ -206,8 +160,7 @@
     "cal_db_interface = get_random_db_interface(cal_db_interface)\n",
     "print(f'Calibration database interface: {cal_db_interface}')\n",
     "print(f\"Sending constants to the calibration database: {db_output}\")\n",
-    "print(f\"HDF5 path to data: {h5path}\")\n",
-    "print(f\"Reading data from: {filename}\")\n",
+    "print(f\"Instrument H5File source: {instrument_src}\")\n",
     "print(f\"Run number: {run}\")"
    ]
   },
@@ -217,38 +170,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Extracting slow data:\n",
-    "if karabo_da_control:\n",
-    "    ctrl_fname = os.path.join(ped_dir, path_template.format(run, karabo_da_control)).format(sequence)\n",
-    "    ctrl_path = h5path_ctrl.format(karabo_id)\n",
-    "    mdl_ctrl_path = f\"/CONTROL/{karabo_id}/MDL/\"\n",
-    "\n",
-    "    (bias_voltage, gain,\n",
-    "     fix_temperature_top,\n",
-    "     fix_temperature_bot) = extract_slow_data(karabo_id, karabo_da_control, ctrl_fname, ctrl_path,\n",
-    "                                              mdl_ctrl_path, bias_voltage, gain,\n",
-    "                                              fix_temperature_top, fix_temperature_bot)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:54:40.555804Z",
-     "start_time": "2018-12-06T10:54:40.452978Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "# Reading Parameters such as Detector Bias, Gain, etc. from the Data:\n",
-    "memoryCells = 1 # pnCCD has 1 memory cell\n",
-    "sensorSize = [pixels_x, pixels_y]\n",
-    "blockSize = [sensorSize[0]//2, sensorSize[1]//2]# sensor area will be analysed according to blocksize\n",
-    "xcal.defaultBlockSize = blockSize\n",
-    "nImages = fastccdreaderh5.getDataSize(filename, h5path)[0] # specifies total number of images to proceed\n",
-    "nImages = nImagesOrLimit(nImages, number_dark_frames)\n",
-    "profile = False"
+    "step_timer = StepTimer()"
    ]
   },
   {
@@ -257,50 +179,55 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
+    "\n",
+    "run_dc = RunDirectory(f\"{in_folder}/r{run:04d}\")\n",
+    "\n",
+    "# extract control data\n",
+    "ctrl_data = pnccdlib.PnccdCtrl(run_dc, karabo_id)\n",
+    "\n",
+    "if bias_voltage == 0:\n",
+    "    bias_voltage = ctrl_data.get_bias_voltage()\n",
+    "if gain == -1:\n",
+    "    gain = ctrl_data.get_gain()\n",
+    "if fix_temperature_top == 0:\n",
+    "    fix_temperature_top = ctrl_data.get_fix_temperature_top()\n",
+    "if fix_temperature_bot == 0:\n",
+    "    fix_temperature_bot = ctrl_data.get_fix_temperature_bot()\n",
+    "\n",
     "# Printing the Parameters Read from the Data File:\n",
     "display(Markdown('### Detector Parameters'))\n",
-    "print(f\"Bias voltage is {bias_voltage:0.2f} V.\")\n",
-    "print(f\"Detector gain is set to {gain}.\")\n",
+    "print(f\"Bias voltage is {bias_voltage:0.1f} V.\")\n",
+    "print(f\"Detector gain is set to 1/{int(gain)}.\")\n",
     "print(f\"Detector integration time is set to {integration_time} ms\")\n",
     "print(f\"Top pnCCD sensor is at temperature of {fix_temperature_top:0.2f} K\")\n",
     "print(f\"Bottom pnCCD sensor is at temperature of {fix_temperature_bot:0.2f} K\")\n",
-    "print(\"Number of dark images to analyze:\", nImages) "
+    "step_timer.done_step(f'Reading control data.')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:54:41.584031Z",
-     "start_time": "2018-12-06T10:54:41.578462Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "# Reading Files in Chunks:\n",
+    "# Reading Parameters such as Detector Bias, Gain, etc. from the Data:\n",
+    "sensorSize = [pixels_x, pixels_y]\n",
+    "# sensor area will be analysed according to blocksize\n",
+    "blockSize = [sensorSize[0]//2, sensorSize[1]//2]\n",
+    "xcal.defaultBlockSize = blockSize\n",
     "\n",
-    "# Chunk reader returns an iterator to access the data in the file within the ranges:\n",
-    "reader = ChunkReader(filename, fastccdreaderh5.readData, nImages, chunkSize, path=h5path, \n",
-    "                     pixels_x=pixels_x, pixels_y=pixels_y)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:54:41.899511Z",
-     "start_time": "2018-12-06T10:54:41.864816Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "# Calculators:\n",
+    "pixel_data = (instrument_src, \"data.image\")\n",
+    "data_dc = run_dc.select(*pixel_data, require_all=True)\n",
+    "n_trains = data_dc[pixel_data].shape[0]\n",
     "\n",
-    "# noiseCal is a noise map calculator, which internally also produces a per-pixel mean map, i.e., an offset map:\n",
-    "noiseCal = xcal.NoiseCalculator(sensorSize, memoryCells, cores=cpuCores, blockSize=blockSize,\n",
-    "                                runParallel=run_parallel)"
+    "if max_trains != 0:\n",
+    "    n_trains = min(n_trains, max_trains)\n",
+    "if n_trains < min_trains:\n",
+    "    raise ValueError(\n",
+    "        f\"Files {data_dc.files} consists of less than\"\n",
+    "        f\" the required number of {min_trains} trains to proceed with \"\n",
+    "        \"dark processing.\")"
    ]
   },
   {
@@ -315,55 +242,16 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:55:21.238009Z",
-     "start_time": "2018-12-06T10:54:54.586435Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "counter1 = 1 # To count how many \"if data.shape[2] >= chunkSize\" instances are there.\n",
-    "counter2 = 0 # To count how many \"if data.shape[2] < chunkSize\" instances are there.\n",
-    "chunkSize_new = 0 # See below\n",
-    "\n",
-    "for data in reader.readChunks():\n",
-    "    data = data.astype(np.float32)\n",
-    "    dx = np.count_nonzero(data, axis=(0, 1)) \n",
-    "    data = data[:,:,dx != 0] # Getting rid of empty frames\n",
-    "    # Some sequences may have less than 500 frames in them. To find out how many images there are, we will \n",
-    "    # temporarily change chunkSize to be the same as whatever number of frames the last chunk of data has:\n",
-    "    if data.shape[2] < chunkSize:\n",
-    "        chunkSize_new = data.shape[2]\n",
-    "        print(\"Number of images are less than chunkSize. chunkSize is temporarily changed to {}.\"\n",
-    "              .format(chunkSize_new))\n",
-    "        images = images + chunkSize_new\n",
-    "        counter2 += 1 \n",
-    "    else:\n",
-    "        images = counter1*chunkSize + counter2*chunkSize_new\n",
-    "        counter1 += 1\n",
-    "\n",
-    "    noiseCal.fill(data) # Filling the histogram calculator with data\n",
-    "    chunkSize = 100 # resetting the chunkSize to its default value for the next sequence or data-chunk\n",
-    "\n",
-    "print('A total number of {} images are processed.'.format(images))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:55:21.238009Z",
-     "start_time": "2018-12-06T10:54:54.586435Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "offsetMap = noiseCal.getOffset() # Producing offset map\n",
-    "noiseMap = noiseCal.get() # Producing noise map\n",
-    "noiseCal.reset() # Resetting noise calculator\n",
-    "print(\"Initial maps are created.\")"
+    "step_timer.start()\n",
+    "data = data_dc.select_trains(np.s_[:n_trains])[pixel_data].ndarray()\n",
+    "data = data.astype(np.float32)\n",
+    "\n",
+    "noiseMap = np.std(data, axis=0, dtype=np.float64)\n",
+    "offsetMap = np.mean(data, axis=0, dtype=np.float64)\n",
+    "step_timer.done_step(f'Initial maps are created from {n_trains} trains.')"
    ]
   },
   {
@@ -373,7 +261,7 @@
    "outputs": [],
    "source": [
     "# pnCCD valid gains are 1, 1/4, 1/16, 1/64, 1/256, 1/340 and 1/512:\n",
-    "gain_k = [k for k, v in VALID_GAINS.items() if v == gain][0]\n",
+    "gain_k = [k for k, v in pnccdlib.VALID_GAINS.items() if v == gain][0]\n",
     "if gain_k == 'a' or gain_k == 'b':\n",
     "    xrange = (0, 200) # x-axis range for the noise histogram plots\n",
     "    bins = 2000 # number of bins for Histogram Calculators \n",
@@ -381,7 +269,7 @@
     "else:\n",
     "    xrange = (0, 20)\n",
     "    bins = 100\n",
-    "    bin_range = [-50, 50] "
+    "    bin_range = [-50, 50]"
    ]
   },
   {
@@ -396,17 +284,11 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:56:20.686534Z",
-     "start_time": "2018-12-06T10:56:11.721829Z"
-    },
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "#************** HISTOGRAMS *******************#\n",
-    "\n",
+    "step_timer.start()\n",
     "fig = plt.figure(figsize=(12,4))\n",
     "ax = fig.add_subplot(121)\n",
     "xana.histPlot(ax, offsetMap.flatten(), bins=2000, plot_errors=False)\n",
@@ -425,15 +307,16 @@
     "\n",
     "#************** HEAT MAPS *******************#\n",
     "\n",
-    "fig = xana.heatmapPlot(offsetMap[:,:,0], x_label='Column Number', y_label='Row Number',  aspect=1,\n",
+    "fig = xana.heatmapPlot(offsetMap, x_label='Column Number', y_label='Row Number', aspect=1,\n",
     "                       x_range=(0, pixels_y), y_range=(0, pixels_x), vmin=6000, vmax=15000, \n",
     "                       lut_label='Offset (ADU)', panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)',\n",
     "                       title = 'Offset Map')\n",
     "\n",
-    "fig = xana.heatmapPlot(noiseMap[:,:,0], x_label='Column Number', y_label='Row Number', aspect=1,\n",
+    "fig = xana.heatmapPlot(noiseMap, x_label='Column Number', y_label='Row Number', aspect=1,\n",
     "                       lut_label='Uncorrected Noise (ADU)', x_range=(0, pixels_y),\n",
     "                       y_range=(0, pixels_x), vmax=2*np.mean(noiseMap), panel_x_label='Row Stat (ADU)', \n",
-    "                       panel_y_label='Column Stat (ADU)', title = 'Uncorrected NoiseMap')"
+    "                       panel_y_label='Column Stat (ADU)', title = 'Uncorrected NoiseMap')\n",
+    "step_timer.done_step(f'Offset and Noise Maps before Common Mode Correction.')"
    ]
   },
   {
@@ -469,7 +352,7 @@
     "bad_pixels[:,0] = EDGE_PIXELS \n",
     "bad_pixels[:,1023] = EDGE_PIXELS  \n",
     "\n",
-    "fig = xana.heatmapPlot(np.log2(bad_pixels[:, :, 0]),\n",
+    "fig = xana.heatmapPlot(np.log2(bad_pixels),\n",
     "                       x_label='Columns', y_label='Rows', lut_label='Bad Pixel Value (ADU)',\n",
     "                       x_range=(0, pixels_y), y_range=(0, pixels_x), vmin=0, vmax=32,\n",
     "                       panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)', \n",
@@ -482,11 +365,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Offset Correction:\n",
-    "\n",
-    "offsetCorrection = xcal.OffsetCorrection(sensorSize, offsetMap, nCells = memoryCells, cores=cpuCores, gains=None,\n",
-    "                                         runParallel=run_parallel, blockSize=blockSize)\n",
-    "\n",
     "# Common Mode Correction:\n",
     "# In this method, the median of all (assuming stride = 1) pixels that are read out at the same time along a column \n",
     "# is subtracted from the signal in each pixel in that column. Here, the signals in the pixels refer to the value in \n",
@@ -495,12 +373,15 @@
     "# in a column per quadrant is smaller than the value set for minFrac parameter for a particular column, that column \n",
     "# will be ignored for calculation of common mode values and that column is not corrected for common mode.\n",
     "# minFrac = 0 means no column is ignored except those containing nan values (bad pixels):\n",
-    "cmCorrection = xcal.CommonModeCorrection(sensorSize,\n",
-    "                                         commonModeBlockSize,\n",
-    "                                         commonModeAxis, parallel=run_parallel, dType=np.float32, stride=1,\n",
-    "                                         noiseMap=noiseMap.astype(np.float32), minFrac=0)\n",
-    "\n",
-    "cmCorrection.debug()"
+    "cmCorrection = xcal.CommonModeCorrection(\n",
+    "    sensorSize,\n",
+    "    commonModeBlockSize,\n",
+    "    commonModeAxis,\n",
+    "    parallel=False,\n",
+    "    dType=np.float32, stride=1,\n",
+    "    noiseMap=noiseMap[..., np.newaxis],  # np.float32\n",
+    "    minFrac=0,\n",
+    ")"
    ]
   },
   {
@@ -514,11 +395,11 @@
     "# negative bin to ensure the offset and common mode corrections actually move the signal to zero:\n",
     "\n",
     "# For offset corrected data:\n",
-    "histCalCorrected = xcal.HistogramCalculator(sensorSize, bins=bins, range=bin_range, memoryCells=memoryCells,\n",
-    "                                            cores=cpuCores, gains=None, blockSize=blockSize)\n",
+    "histCalCorrected = xcal.HistogramCalculator(sensorSize, bins=bins, range=bin_range, memoryCells=1,\n",
+    "                                            parallel=False, gains=None, blockSize=blockSize)\n",
     "# For common mode corrected data:\n",
-    "histCalCMCorrected = xcal.HistogramCalculator(sensorSize, bins=bins, range=bin_range, memoryCells=memoryCells,\n",
-    "                                              cores=cpuCores, gains=None, blockSize=blockSize)"
+    "histCalCMCorrected = xcal.HistogramCalculator(sensorSize, bins=bins, range=bin_range, memoryCells=1,\n",
+    "                                              parallel=False, gains=None, blockSize=blockSize)"
    ]
   },
   {
@@ -533,41 +414,33 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "tags": []
+   },
    "outputs": [],
    "source": [
-    "counter1 = 1 # To count how many \"if data.shape[2] >= chunkSize\" instances are there.\n",
-    "counter2 = 0 # To count how many \"if data.shape[2] < chunkSize\" instances are there.\n",
-    "chunkSize_new = 0 # See below\n",
-    "\n",
-    "for data in reader.readChunks():\n",
-    "    data = data.astype(np.float32)\n",
-    "    dx = np.count_nonzero(data, axis=(0, 1))\n",
-    "    data = data[:,:,dx != 0]\n",
-    "    \n",
-    "    # Some sequences may have less than 500 frames in them. To find out how many images there are, we will \n",
-    "    # temporarily change chunkSize to be the same as whatever number of frames the last chunk of data has:\n",
-    "    if data.shape[2] < chunkSize:\n",
-    "        chunkSize_new = data.shape[2]\n",
-    "        print(\"Number of images are less than chunkSize. chunkSize is temporarily changed to {}.\"\n",
-    "              .format(chunkSize_new))\n",
-    "        images = images + chunkSize_new\n",
-    "        counter2 += 1 \n",
-    "    else:\n",
-    "        images = counter1*chunkSize + counter2*chunkSize_new\n",
-    "        counter1 += 1\n",
-    "        \n",
-    "    data -= offsetMap.data # Offset correction\n",
-    "    offset_corr_data = copy.copy(data) # I am copying this so that I can have access to it in the table below\n",
-    "    histCalCorrected.fill(data)\n",
-    "    cellTable=np.zeros(data.shape[2], np.int32) # Common mode correction\n",
-    "    data = cmCorrection.correct(data.astype(np.float32), cellTable=cellTable) \n",
-    "    histCalCMCorrected.fill(data)\n",
-    "    noiseCal.fill(data)  # Filling calculators with data\n",
-    "    chunkSize = 100 # resetting the chunkSize to its default value for the next sequence or data-chunk\n",
-    "\n",
-    "print('A total number of {} images are processed.'.format(images))\n",
-    "print(\"Offset and common mode corrections are applied.\")"
+    "step_timer.start()\n",
+    "def correct_image(wid, idx, d):\n",
+    "    d -= offsetMap\n",
+    "    offset_corr_data[idx, ...] = d\n",
+    "\n",
+    "    d = np.squeeze(cmCorrection.correct(\n",
+    "        d[..., np.newaxis],\n",
+    "        cellTable=np.zeros((1,), np.int32)\n",
+    "    ))\n",
+    "    corr_data[idx, ...] = d\n",
+    "\n",
+    "context = psh.context.ThreadContext(num_workers=10)\n",
+    "\n",
+    "data_shape = (n_trains, sensorSize[0], sensorSize[1])\n",
+    "\n",
+    "corr_data = context.alloc(shape=data_shape, dtype=np.float32)\n",
+    "offset_corr_data = context.alloc(shape=data_shape, dtype=np.float32)\n",
+    "\n",
+    "context.map(correct_image, data.copy())\n",
+    "histCalCorrected.fill(offset_corr_data)\n",
+    "histCalCMCorrected.fill(corr_data)\n",
+    "step_timer.done_step(f'Offset and common mode corrections are applied.')"
    ]
   },
   {
@@ -576,7 +449,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "noiseMapCM = noiseCal.get() # Producing common mode corrected noise map\n",
+    "noiseMapCM = np.std(corr_data, axis=0, dtype=np.float64) # Producing common mode corrected noise map\n",
     "ho, eo, co, so = histCalCorrected.get()\n",
     "hCM, eCM, cCM, sCM = histCalCMCorrected.get()"
    ]
@@ -588,10 +461,10 @@
    "outputs": [],
    "source": [
     "# We are copying these so that we can replot them later after the calculators are reset:\n",
-    "ho_second_trial = copy.copy(ho)\n",
-    "co_second_trial = copy.copy(co)\n",
-    "hCM_second_trial = copy.copy(hCM)\n",
-    "cCM_second_trial = copy.copy(cCM)"
+    "ho_second_trial = ho.copy()\n",
+    "co_second_trial = co.copy()\n",
+    "hCM_second_trial = hCM.copy()\n",
+    "cCM_second_trial = cCM.copy()"
    ]
   },
   {
@@ -609,6 +482,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "do = [{'x': co,\n",
     "     'y': ho,\n",
     "     'drawstyle': 'steps-post',\n",
@@ -630,13 +504,14 @@
     "t0.title = \"Comparison of the First Round of Corrections - Bad Pixels Not Excluded\"\n",
     "t0.field_names = [\"Dark Pedestal After Offset Correction\", \"Dark Pedestal After Offset and Common Mode Corrections\"]\n",
     "t0.add_row([\"Mean: {:0.3f} ADU\".format(np.mean(offset_corr_data)), \"Mean: {:0.3f} ADU\"\n",
-    "            .format(np.mean(data))])\n",
+    "            .format(np.mean(corr_data))])\n",
     "t0.add_row([\"Median: {:0.3f} ADU\".format(np.median(offset_corr_data)), \"Median: {:0.3f} ADU\"\n",
-    "            .format(np.median(data))])\n",
+    "            .format(np.median(corr_data))])\n",
     "t0.add_row([\"Standard Deviation: {:0.3f} ADU\".format(np.std(offset_corr_data)), \n",
     "            \"Standard Deviation: {:0.3f} ADU\"\n",
-    "            .format(np.std(data))])\n",
-    "print(t0,'\\n')"
+    "            .format(np.std(corr_data))])\n",
+    "print(t0,'\\n')\n",
+    "step_timer.done_step(\"Plotting time\")"
    ]
   },
   {
@@ -655,6 +530,7 @@
    "outputs": [],
    "source": [
     "#***** HISTOGRAM OF CORRECTED NOISE MAP *******#\n",
+    "step_timer.start()\n",
     "fig = plt.figure(figsize=(12,4))\n",
     "ax = fig.add_subplot(122)\n",
     "xana.histPlot(ax, noiseMapCM.flatten(), bins=2000, plot_errors=False)\n",
@@ -663,13 +539,14 @@
     "t = ax.set_title(\"Histogram of the Noise Map After Offset and \\n Common Mode Corrections\")\n",
     "t = ax.set_xlim(xrange)\n",
     "\n",
-    "fig = xana.heatmapPlot(noiseMapCM[:, :, 0],\n",
+    "fig = xana.heatmapPlot(noiseMapCM,\n",
     "                       x_label='Columns', y_label='Rows',\n",
     "                       lut_label='Corrected Noise (ADU)',\n",
     "                       x_range=(0, pixels_y),\n",
     "                       y_range=(0, pixels_x), vmax=2*np.mean(noiseMap), panel_x_label='Row Stat (ADU)', \n",
     "                       panel_y_label='Column Stat (ADU)', \n",
-    "                       title='Noise Map After Offset and Common Mode Corrections')"
+    "                       title='Noise Map After Offset and Common Mode Corrections')\n",
+    "step_timer.done_step(\"Plotting time\")"
    ]
   },
   {
@@ -679,7 +556,6 @@
    "outputs": [],
    "source": [
     "# Resetting the calculators:\n",
-    "noiseCal.reset() # resetting noise calculator\n",
     "histCalCorrected.reset() # resetting histogram calculator\n",
     "histCalCMCorrected.reset() # resetting histogram calculator"
    ]
@@ -699,15 +575,25 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "mnoffset = np.nanmedian(offsetMap)\n",
-    "stdoffset = np.nanstd(offsetMap)\n",
+    "stdoffset = np.nanstd(offsetMap, dtype=np.float64)\n",
     "bad_pixels[(offsetMap < mnoffset-bad_pixel_offset_sigma*stdoffset) |\n",
     "           (offsetMap > mnoffset+bad_pixel_offset_sigma*stdoffset)] = BadPixels.OFFSET_OUT_OF_THRESHOLD.value\n",
     "\n",
     "mnnoise = np.nanmedian(noiseMapCM)\n",
-    "stdnoise = np.nanstd(noiseMapCM)\n",
+    "stdnoise = np.nanstd(noiseMapCM, dtype=np.float64)\n",
     "bad_pixels[(noiseMapCM < mnnoise-bad_pixel_noise_sigma*stdnoise) |\n",
-    "           (noiseMapCM > mnnoise+bad_pixel_noise_sigma*stdnoise)] = BadPixels.NOISE_OUT_OF_THRESHOLD.value"
+    "           (noiseMapCM > mnnoise+bad_pixel_noise_sigma*stdnoise)] = BadPixels.NOISE_OUT_OF_THRESHOLD.value\n",
+    "\n",
+    "# Each pnCCD sensor has 22 rows and 60 columns cut in the middle of it out by a laser:\n",
+    "\n",
+    "hole_mask = np.zeros(bad_pixels.shape, np.uint32) \n",
+    "hole_mask[489:533, 481:541] = BadPixels.NON_SENSITIVE.value\n",
+    "\n",
+    "# Assigning this masked area as bad pixels:\n",
+    "bad_pixels = np.bitwise_or(bad_pixels, hole_mask)\n",
+    "step_timer.done_step(\"Calculating BadPixelMap.\")"
    ]
   },
   {
@@ -720,25 +606,17 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "# Each pnCCD sensor has 22 rows and 60 columns cut in the middle of it out by a laser:\n",
-    "\n",
-    "hole_mask = np.zeros(bad_pixels.shape, np.uint32) \n",
-    "hole_mask[489:533,481:541,:] = BadPixels.NON_SENSITIVE.value\n",
-    "\n",
-    "# Assigning this masked area as bad pixels:\n",
-    "bad_pixels = np.bitwise_or(bad_pixels, hole_mask)\n",
-    "\n",
-    "fig = xana.heatmapPlot(np.log2(bad_pixels[:, :, 0]),\n",
+    "step_timer.start()\n",
+    "fig = xana.heatmapPlot(np.log2(bad_pixels),\n",
     "                       x_label='Columns', y_label='Rows',\n",
     "                       lut_label='Bad Pixel Value (ADU)',\n",
     "                       x_range=(0, pixels_y),\n",
     "                       y_range=(0, pixels_x), vmin=0, vmax=32, panel_x_label='Row Stat (ADU)', \n",
-    "                       panel_y_label='Column Stat (ADU)', title = 'Second Bad Pixels Map with the Hole in Center')"
+    "                       panel_y_label='Column Stat (ADU)', title = 'Second Bad Pixels Map with the Hole in Center')\n",
+    "step_timer.done_step(\"Plotting time\")"
    ]
   },
   {
@@ -756,8 +634,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "event_threshold = sigmaNoise*np.median(noiseMapCM) # for exclusion of possible cosmic ray events\n",
-    "noiseCal.setBadPixelMask(bad_pixels != 0) # setting bad pixels map for the noise calculator"
+    "event_threshold = sigmaNoise * np.median(noiseMapCM)  # for exclusion of possible cosmic ray events"
    ]
   },
   {
@@ -766,11 +643,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cmCorrection = xcal.CommonModeCorrection(sensorSize,\n",
-    "                                         commonModeBlockSize,\n",
-    "                                         commonModeAxis, parallel=run_parallel, dType=np.float32, stride=1,\n",
-    "                                         noiseMap=noiseMapCM.astype(np.float32), minFrac=0)\n",
-    "cmCorrection.debug()"
+    "cmCorrection = xcal.CommonModeCorrection(\n",
+    "    sensorSize,\n",
+    "    commonModeBlockSize,\n",
+    "    commonModeAxis,\n",
+    "    parallel=False,\n",
+    "    dType=np.float32, stride=1,\n",
+    "    noiseMap=noiseMapCM[..., np.newaxis], minFrac=0,\n",
+    ")"
    ]
   },
   {
@@ -779,46 +659,33 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "counter1 = 1 # To count how many \"if data.shape[2] >= chunkSize\" instances are there.\n",
-    "counter2 = 0 # To count how many \"if data.shape[2] < chunkSize\" instances are there.\n",
-    "chunkSize_new = 0 # See below\n",
-    "\n",
-    "for data in reader.readChunks():\n",
-    "    data = data.astype(np.float32)\n",
-    "    dx = np.count_nonzero(data, axis=(0, 1))\n",
-    "    data = data[:,:,dx != 0]\n",
-    "    data_mask = np.repeat(bad_pixels, data.shape[2], axis=2) # Converting bad_pixels to the same shape as the data\n",
-    "    data[data_mask != 0] = np.nan # masking data for bad pixels and equating the values to np.nan\n",
-    "    \n",
-    "    # Some sequences may have less than 500 frames in them. To find out how many images there are, we will \n",
-    "    # temporarily change chunkSize to be the same as whatever number of frames the last chunk of data has:\n",
-    "    if data.shape[2] < chunkSize:\n",
-    "        chunkSize_new = data.shape[2]\n",
-    "        print(\"Number of images are less than chunkSize. chunkSize is temporarily changed to {}.\"\n",
-    "              .format(chunkSize_new))\n",
-    "        images = images + chunkSize_new\n",
-    "        counter2 += 1 \n",
-    "    else:\n",
-    "        images = counter1*chunkSize + counter2*chunkSize_new\n",
-    "        counter1 += 1\n",
-    "    \n",
-    "    data_copy = offsetCorrection.correct(copy.copy(data))\n",
-    "    cellTable=np.zeros(data_copy.shape[2], np.int32)\n",
-    "    data_copy = cmCorrection.correct(data_copy.astype(np.float32), cellTable=cellTable)\n",
-    "    data[data_copy > event_threshold] = np.nan # discarding events caused by cosmic rays\n",
-    "    #data = np.ma.MaskedArray(data, np.isnan(data), fill_value=np.nan) # masking cosmics,default fill_value = 1e+20\n",
+    "step_timer.start()\n",
+    "def correct_image(wid, idx, d):\n",
     "    \n",
-    "    data -= offsetMap.data # Offset correction\n",
-    "    offset_corr_data2 = copy.copy(data) # I am copying this so that I can have access to it in the table below\n",
-    "    histCalCorrected.fill(data)\n",
-    "    cellTable=np.zeros(data.shape[2], np.int32) # Common mode correction\n",
-    "    data = cmCorrection.correct(data.astype(np.float32), cellTable=cellTable) \n",
-    "    histCalCMCorrected.fill(data)\n",
-    "    noiseCal.fill(data)  # Filling calculators with data\n",
-    "    chunkSize = 100 # resetting the chunkSize to its default value for the next sequence or data-chunk\n",
-    "\n",
-    "print(\"Final iteration is Performed.\")\n",
-    "print('A total number of {} images are processed.'.format(images))"
+    "    d[bad_pixels != 0] = np.nan # masking data for bad pixels and equating the values to np.nan\n",
+    "    d_off = d - offsetMap\n",
+    "    d_off_cm = np.squeeze(cmCorrection.correct(\n",
+    "        d_off[..., np.newaxis], cellTable=np.zeros((1,), dtype=np.float32)))\n",
+    "\n",
+    "    # discarding events caused by cosmic rays\n",
+    "    d_off[d_off_cm > event_threshold] = np.nan\n",
+    "    d_off_cm[d_off_cm > event_threshold] = np.nan  \n",
+    "\n",
+    "    offset_corr_data[idx, ...] = d_off  # Will be used for plotting\n",
+    "    corr_data[idx, ...] = np.squeeze(d_off_cm)\n",
+    "\n",
+    "context = psh.context.ThreadContext(num_workers=10)\n",
+    "\n",
+    "corr_data = context.alloc(shape=data_shape, dtype=np.float32)\n",
+    "offset_corr_data = context.alloc(shape=data_shape, dtype=np.float32)\n",
+    "\n",
+    "context.map(correct_image, data.copy())\n",
+    "\n",
+    "histCalCorrected.fill(offset_corr_data)\n",
+    "histCalCMCorrected.fill(corr_data)\n",
+    "\n",
+    "print(f\"A total number of {n_trains} images are processed.\")\n",
+    "step_timer.done_step(\"Final iteration is Performed.\")"
    ]
   },
   {
@@ -827,26 +694,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "noiseMapCM_2nd = noiseCal.get().filled(np.nan) # the masked pixels are filled with nans\n",
+    "noiseMapCM_2nd = np.nanstd(corr_data, axis=0, dtype=np.float64)\n",
     "ho2, eo2, co2, so2 = histCalCorrected.get()\n",
     "hCM2, eCM2, cCM2, sCM2 = histCalCMCorrected.get()"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Plots of the Final Results\n",
-    "\n",
-    "The following plot and table compare the offset and common mode corrected signal with and without the bad pixels."
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "do_Final = [{'x': co_second_trial,\n",
     "     'y': ho_second_trial,\n",
     "     'drawstyle': 'steps-post',\n",
@@ -858,7 +717,7 @@
     "     'drawstyle': 'steps-post',\n",
     "     'color': 'red',\n",
     "     'ecolor': 'crimson',\n",
-    "     'label': 'Offset and Common Mode Corrected Signal (BP Incl.)' \n",
+    "     'label': 'Offset and Common Mode Corrected Signal (BP Incl.)'\n",
     "     },\n",
     "    {'x': co2,\n",
     "     'y': ho2,\n",
@@ -873,21 +732,22 @@
     "     'label': 'Offset and Common Mode Corrected Signal (BP Excl.)'\n",
     "     }]\n",
     "\n",
-    "fig = xana.simplePlot(do_Final, figsize='2col', aspect=1, x_label = 'ADU', \n",
-    "                      y_label=\"Counts (logarithmic scale)\", y_log=True, x_range = bin_range, \n",
+    "fig = xana.simplePlot(do_Final, figsize='2col', aspect=1, x_label = 'ADU',\n",
+    "                      y_label=\"Counts (logarithmic scale)\", y_log=True, x_range = bin_range,\n",
     "                      y_range = (0.02, 1e8),\n",
     "                      legend='top-right-frame-1col', title = 'Comparison of Corrections')\n",
     "\n",
     "t0 = PrettyTable()\n",
     "t0.title = \"Comparison of the Second Round of Corrections - Bad Pixels Excluded\"\n",
     "t0.field_names = [\"Dark Pedestal After Offset Correction\", \"Dark Pedestal After Offset and Common Mode Corrections\"]\n",
-    "t0.add_row([\"Mean: {:0.3f} ADU\".format(np.nanmean(offset_corr_data2)), \"Mean: {:0.3f} ADU\"\n",
-    "            .format(np.nanmean(data))])\n",
-    "t0.add_row([\"Median: {:0.3f} ADU\".format(np.nanmedian(offset_corr_data2)), \"Median: {:0.3f} ADU\"\n",
-    "            .format(np.nanmedian(data))])\n",
-    "t0.add_row([\"Standard Deviation: {:0.3f} ADU\".format(np.nanstd(offset_corr_data2)), \n",
-    "            \"Standard Deviation: {:0.3f} ADU\".format(np.nanstd(data))])\n",
-    "print(t0,'\\n')"
+    "t0.add_row([\"Mean: {:0.3f} ADU\".format(np.nanmean(offset_corr_data)), \"Mean: {:0.3f} ADU\"\n",
+    "            .format(np.nanmean(corr_data))])\n",
+    "t0.add_row([\"Median: {:0.3f} ADU\".format(np.nanmedian(offset_corr_data)), \"Median: {:0.3f} ADU\"\n",
+    "            .format(np.nanmedian(corr_data))])\n",
+    "t0.add_row([\"Standard Deviation: {:0.3f} ADU\".format(np.nanstd(offset_corr_data)),\n",
+    "            \"Standard Deviation: {:0.3f} ADU\".format(np.nanstd(corr_data))])\n",
+    "print(t0,'\\n')\n",
+    "step_timer.done_step(\"Plotting time\")"
    ]
   },
   {
@@ -906,7 +766,8 @@
    "outputs": [],
    "source": [
     "#***** HISTOGRAMS OF NOISE MAPS *******#\n",
-    "fig = plt.figure(figsize=(12,4))\n",
+    "step_timer.start()\n",
+    "fig = plt.figure(figsize=(12, 4))\n",
     "ax = fig.add_subplot(121)\n",
     "xana.histPlot(ax, noiseMapCM_2nd.flatten(), bins=2000, plot_errors=False)\n",
     "t = ax.set_xlabel(\"ADU per 2000 bins\")\n",
@@ -914,11 +775,12 @@
     "t = ax.set_title(\"Histogram of the Noise Map After Offset and Common Mode \\n Corrections and exclusion of Bad Pixels\")\n",
     "t = ax.set_xlim(xrange)\n",
     "\n",
-    "fig = xana.heatmapPlot(noiseMapCM_2nd[:,:,0], aspect=1, x_label='Column Number', y_label='Row Number',\n",
+    "fig = xana.heatmapPlot(noiseMapCM_2nd, aspect=1, x_label='Column Number', y_label='Row Number',\n",
     "                       lut_label='Final Corrected Noise (ADU)', x_range=(0, pixels_y), y_range=(0, pixels_x),\n",
     "                       vmax=2*np.mean(noiseMap),\n",
     "                       title = 'Final Offset and Common Mode Corrected Noise Map \\n (Bad Pixels Excluded)', \n",
-    "                       panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)')"
+    "                       panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)')\n",
+    "step_timer.done_step(\"Plotting time\")"
    ]
   },
   {
@@ -933,11 +795,10 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "mnoffset = np.nanmedian(offsetMap)\n",
     "stdoffset = np.nanstd(offsetMap)\n",
     "bad_pixels[(offsetMap < mnoffset-bad_pixel_offset_sigma*stdoffset) |\n",
@@ -950,12 +811,13 @@
     "\n",
     "bad_pixels = np.bitwise_or(bad_pixels, hole_mask)\n",
     "\n",
-    "fig = xana.heatmapPlot(np.log2(bad_pixels[:, :, 0]),\n",
+    "fig = xana.heatmapPlot(np.log2(bad_pixels),\n",
     "                       x_label='Columns', y_label='Rows',\n",
     "                       lut_label='Bad Pixel Value (ADU)',\n",
     "                       x_range=(0, pixels_y),\n",
     "                       y_range=(0, pixels_x), vmin=0, vmax=32, panel_x_label='Row Stat (ADU)', \n",
-    "                       panel_y_label='Column Stat (ADU)', title = 'Final Bad Pixels Map')"
+    "                       panel_y_label='Column Stat (ADU)', title = 'Final Bad Pixels Map')\n",
+    "step_timer.done_step(\"Plotting time\")"
    ]
   },
   {
@@ -983,47 +845,39 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T10:56:22.741284Z",
-     "start_time": "2018-12-06T10:56:20.688393Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
+    "step_timer.start()\n",
     "constant_maps = {\n",
-    "    'Offset': offsetMap,\n",
-    "    'Noise': noiseMapCM_2nd,\n",
-    "    'BadPixelsDark': bad_pixels\n",
+    "    'Offset': offsetMap[..., np.newaxis],\n",
+    "    'Noise': noiseMapCM_2nd[..., np.newaxis],\n",
+    "    'BadPixelsDark': bad_pixels[..., np.newaxis],\n",
     "}\n",
     "md = None\n",
+    "\n",
+    "det = Constants.CCD(DetectorTypes.pnCCD)\n",
+    "\n",
+    "# set the operating condition\n",
+    "condition = Conditions.Dark.CCD(bias_voltage=bias_voltage,\n",
+    "                                integration_time=integration_time,\n",
+    "                                gain_setting=gain,\n",
+    "                                temperature=fix_temperature_top,\n",
+    "                                pixels_x=pixels_x,\n",
+    "                                pixels_y=pixels_y)\n",
+    "for parm in condition.parameters:\n",
+    "    if parm.name == \"Sensor Temperature\":\n",
+    "        parm.lower_deviation = temp_limits\n",
+    "        parm.upper_deviation = temp_limits\n",
+    "\n",
+    "db_module = get_pdu_from_db(karabo_id, karabo_da, getattr(det, 'Offset')(),\n",
+    "                            condition, cal_db_interface,\n",
+    "                            snapshot_at=creation_time)[0]\n",
+    "\n",
     "for const_name in constant_maps.keys():\n",
-    "    det = Constants.CCD(DetectorTypes.pnCCD)\n",
     "    const = getattr(det, const_name)()\n",
     "    const.data = constant_maps[const_name].data\n",
     "\n",
-    "    # set the operating condition\n",
-    "    condition = Conditions.Dark.CCD(bias_voltage=bias_voltage,\n",
-    "                                    integration_time=integration_time,\n",
-    "                                    gain_setting=gain,\n",
-    "                                    temperature=fix_temperature_top,\n",
-    "                                    pixels_x=pixels_x,\n",
-    "                                    pixels_y=pixels_y)\n",
-    "\n",
-    "    for parm in condition.parameters:\n",
-    "        if parm.name == \"Sensor Temperature\":\n",
-    "            parm.lower_deviation = temp_limits\n",
-    "            parm.upper_deviation = temp_limits\n",
-    "\n",
-    "    # This should be used in case of running notebook \n",
-    "    # by a different method other than myMDC which already\n",
-    "    # sends CalCat info.\n",
-    "    # TODO: Set db_module to \"\" by default in the first cell\n",
-    "    if not db_module:\n",
-    "        db_module = get_pdu_from_db(karabo_id, karabo_da, const,\n",
-    "                                    condition, cal_db_interface,\n",
-    "                                    snapshot_at=creation_time)[0]\n",
-    "\n",
     "    if db_output:\n",
     "        md = send_to_db(db_module, karabo_id, const, condition,\n",
     "                        file_loc=file_loc, report_path=report,\n",
@@ -1039,7 +893,8 @@
     "print(\"Constants parameter conditions are:\\n\")\n",
     "print(f\"• bias_voltage: {bias_voltage}\\n• gain_setting: {gain}\\n\"\n",
     "      f\"• top_temperature: {fix_temperature_top}\\n• integration_time: {integration_time}\\n\"\n",
-    "      f\"• creation_time: {md.calibration_constant_version.begin_at if md is not None else creation_time}\\n\")"
+    "      f\"• creation_time: {md.calibration_constant_version.begin_at if md is not None else creation_time}\\n\")\n",
+    "step_timer.done_step(\"Storing calibration constants.\")"
    ]
   },
   {
@@ -1047,7 +902,10 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "print(f\"Total processing time {step_timer.timespan():.01f} s\")\n",
+    "step_timer.print_summary()"
+   ]
   }
  ],
  "metadata": {
@@ -1066,7 +924,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.8.11"
   },
   "latex_envs": {
    "LaTeX_envs_menu_present": true,
@@ -1087,5 +945,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }
diff --git a/notebooks/pnCCD/Correct_pnCCD_NBC.ipynb b/notebooks/pnCCD/Correct_pnCCD_NBC.ipynb
index db8da2b619ff4062322f6782a48894ef4d124f4c..f7bc8a91781e608289f94dd99dd377f4bff0bd23 100644
--- a/notebooks/pnCCD/Correct_pnCCD_NBC.ipynb
+++ b/notebooks/pnCCD/Correct_pnCCD_NBC.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# pnCCD Data Correction #\n",
     "\n",
-    "Authors: DET Group, Modified by Kiana Setoodehnia on December 2020 - Version 4.0\n",
+    "Authors: DET Group, Modified by Kiana Setoodehnia - Version 5.0\n",
     "\n",
     "The following notebook provides offset, common mode, relative gain, split events and pattern classification corrections of images acquired with the pnCCD. This notebook *does not* yet correct for charge transfer inefficiency."
    ]
@@ -14,48 +14,32 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T15:54:23.218849Z",
-     "start_time": "2018-12-06T15:54:23.166497Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "in_folder = \"/gpfs/exfel/exp/SQS/202031/p900166/raw\" # input folder\n",
-    "out_folder = '/gpfs/exfel/data/scratch/setoodeh/Test' # output folder\n",
-    "run = 347 # which run to read data from\n",
-    "sequences = [-1] # sequences to correct, set to -1 for all, range allowed\n",
+    "in_folder = \"/gpfs/exfel/exp/SQS/202031/p900166/raw\"  # input folder\n",
+    "out_folder = \"/gpfs/exfel/data/scratch/ahmedk/test/remove\"  # output folder\n",
+    "run = 347  # which run to read data from\n",
+    "sequences = [-1]  # sequences to correct, set to -1 for all, range allowed\n",
+    "sequences_per_node = 1  # number of sequences running on the same slurm node.\n",
     "\n",
-    "db_module = \"pnCCD_M205_M206\"\n",
     "karabo_da = 'PNCCD01' # data aggregators\n",
-    "karabo_da_control = \"PNCCD02\" # file inset for control data\n",
     "karabo_id = \"SQS_NQS_PNCCD1MP\" # karabo prefix of PNCCD devices\n",
     "receiver_id = \"PNCCD_FMT-0\" # inset for receiver devices\n",
-    "path_template = 'RAW-R{:04d}-{}-S{{:05d}}.h5' # the template to use to access data\n",
-    "path_template_seqs = \"{}/r{:04d}/*PNCCD01-S*.h5\"\n",
-    "h5path = '/INSTRUMENT/{}/CAL/{}:output/data/' # path to data in the HDF5 file \n",
-    "h5path_ctrl = '/CONTROL/{}/CTRL/TCTRL'\n",
-    "\n",
-    "overwrite = True # keep this as True to not overwrite the output \n",
-    "use_dir_creation_date = True # To obtain creation time of the run\n",
-    "number_dark_frames = 0 # number of images to be used, if set to 0 all available images are used\n",
-    "chunk_size_idim = 1 # H5 chunking size of output data\n",
-    "cluster_profile = \"noDB\" # ipcluster profile to use\n",
-    "\n",
-    "cpuCores = 40\n",
+    "path_template = 'RAW-R{:04d}-{}-S{:05d}.h5'  # the template to use to access data\n",
+    "instrument_source_template = '{}/CAL/{}:output'  # template for data source name, will be filled with karabo_id and receiver_id.\n",
+    "\n",
+    "# Parameters affecting data correction.\n",
+    "commonModeAxis = 0 # axis along which common mode will be calculated, 0 = row, and 1 = column\n",
     "commonModeBlockSize = [512, 512] # size of the detector in pixels for common mode calculations\n",
-    "commonModeAxis = 0 # axis along which common mode will be calculated, 0 = row, and 1 = column \n",
     "split_evt_primary_threshold = 4. # primary threshold for split event classification in terms of n sigma noise\n",
     "split_evt_secondary_threshold = 3. # secondary threshold for split event classification in terms of n sigma noise\n",
-    "sequences_per_node = 1\n",
-    "limit_images = 0\n",
-    "seq_num = 0  # sequence number for which the last plot at the end of the notebook is plotted\n",
+    "saturated_threshold = 32000.  # full well capacity in ADU\n",
     "\n",
-    "# pnCCD parameters:\n",
+    "# Conditions for retrieving calibration constants\n",
     "fix_temperature_top = 0. # fix temperature for top sensor in K, set to 0. to use value from slow data.\n",
     "fix_temperature_bot = 0. # fix temperature for bottom senspr in K, set to 0. to use value from slow data.\n",
-    "gain = 0.1 # the detector's gain setting, It is later read from file and this value is overwritten\n",
+    "gain = -1  # the detector's gain setting. Set to -1 to use the value from the slow data.\n",
     "bias_voltage = 0. # the detector's bias voltage. set to 0. to use value from slow data.\n",
     "integration_time = 70  # detector's integration time\n",
     "photon_energy = 1.6 # Al fluorescence in keV\n",
@@ -63,14 +47,27 @@
     "cal_db_interface = \"tcp://max-exfl016:8015\" # calibration DB interface to use\n",
     "cal_db_timeout = 300000 # timeout on caldb requests\n",
     "creation_time = \"\" # To overwrite the measured creation_time. Required Format: YYYY-MM-DD HR:MN:SC.ms e.g. 2019-07-04 11:02:41.00\n",
+    "use_dir_creation_date = True  # To obtain creation time of the run\n",
     "\n",
+    "# Booleans for selecting corrections to apply.\n",
     "only_offset = False # Only, apply offset.\n",
     "common_mode = True # Apply common mode correction\n",
     "relgain = True # Apply relative gain correction\n",
-    "cti = False # Apply charge transfer inefficiency correction (not implemented, yet)\n",
-    "do_pattern_classification = True # classify split events\n",
+    "pattern_classification = True  # classify split events\n",
+    "\n",
+    "# parameters affecting stored output data.\n",
+    "chunk_size_idim = 1  # H5 chunking size of output data\n",
+    "overwrite = True  # keep this as True to not overwrite the output \n",
+    "# ONLY FOR TESTING\n",
+    "limit_images = 0  # this parameter is used for limiting number of images to correct from a sequence file. ONLY FOR TESTING.\n",
+    "\n",
+    "# TODO: REMOVE\n",
+    "db_module = \"\"\n",
     "\n",
-    "saturated_threshold = 32000. # full well capacity in ADU"
+    "\n",
+    "def balance_sequences(in_folder, run, sequences, sequences_per_node, karabo_da):\n",
+    "    from xfel_calibrate.calibrate import balance_sequences as bs\n",
+    "    return bs(in_folder, run, sequences, sequences_per_node, karabo_da)"
    ]
   },
   {
@@ -87,66 +84,46 @@
     "# Dont apply any corrections if only_offset is requested.\n",
     "if not only_offset:\n",
     "    corr_bools[\"relgain\"] = relgain\n",
-    "    corr_bools[\"cti\"] = cti\n",
     "    corr_bools[\"common_mode\"] = common_mode\n",
-    "    corr_bools[\"pattern_class\"] = do_pattern_classification"
+    "    corr_bools[\"pattern_class\"] = pattern_classification"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T15:54:23.455376Z",
-     "start_time": "2018-12-06T15:54:23.413579Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "import copy\n",
     "import datetime\n",
-    "import glob\n",
     "import os\n",
-    "import time\n",
-    "import traceback\n",
     "import warnings\n",
-    "from datetime import timedelta\n",
+    "from pathlib import Path\n",
     "from typing import Tuple\n",
-    "\n",
     "warnings.filterwarnings('ignore')\n",
     "\n",
     "import h5py\n",
     "import matplotlib.pyplot as plt\n",
-    "from iminuit import Minuit\n",
+    "import numpy as np\n",
+    "import pasha as psh\n",
     "from IPython.display import Markdown, display\n",
+    "from extra_data import H5File, RunDirectory\n",
+    "from prettytable import PrettyTable\n",
     "\n",
     "%matplotlib inline\n",
-    "import numpy as np\n",
-    "import XFELDetAna.xfelprofiler as xprof\n",
-    "from cal_tools.pnccdlib import VALID_GAINS, extract_slow_data\n",
+    "\n",
+    "from XFELDetAna import xfelpyanatools as xana\n",
+    "from XFELDetAna import xfelpycaltools as xcal\n",
+    "from cal_tools import pnccdlib\n",
     "from cal_tools.tools import (\n",
     "    get_constant_from_db_and_time,\n",
     "    get_dir_creation_date,\n",
     "    get_random_db_interface,\n",
+    "    map_modules_from_folder,\n",
     ")\n",
-    "from iCalibrationDB import Conditions, ConstantMetaData, Constants, Detectors, Versions\n",
-    "from iCalibrationDB.detectors import DetectorTypes\n",
-    "from prettytable import PrettyTable\n",
-    "\n",
-    "profiler = xprof.Profiler()\n",
-    "profiler.disable()\n",
-    "from XFELDetAna.util import env\n",
-    "\n",
-    "env.iprofile = cluster_profile\n",
-    "from XFELDetAna import xfelpyanatools as xana\n",
-    "from XFELDetAna import xfelpycaltools as xcal\n",
-    "from XFELDetAna.plotting.util import prettyPlotting\n",
-    "\n",
-    "prettyPlotting=True\n",
-    "\n",
-    "\n",
-    "if sequences[0] == -1:\n",
-    "    sequences = None"
+    "from cal_tools.step_timing import StepTimer\n",
+    "from cal_tools import h5_copy_except\n",
+    "from iCalibrationDB import Conditions, ConstantMetaData, Constants\n",
+    "from iCalibrationDB.detectors import DetectorTypes"
    ]
   },
   {
@@ -156,23 +133,21 @@
    "outputs": [],
    "source": [
     "# Calibration Database Settings, and Some Initial Run Parameters & Paths:\n",
-    "\n",
     "display(Markdown('### Initial Settings and Paths'))\n",
-    "pixels_x = 1024 # rows of pnCCD in pixels \n",
-    "pixels_y = 1024 # columns of pnCCD in pixels\n",
+    "\n",
+    "# Sensor size and block size definitions (important for common mode and other corrections):\n",
+    "pixels_x = 1024  # rows of pnCCD in pixels \n",
+    "pixels_y = 1024  # columns of pnCCD in pixels\n",
+    "sensorSize = [pixels_x, pixels_y]\n",
+    "# For xcal.HistogramCalculators.\n",
+    "blockSize = [pixels_x//2, pixels_y//2]  # sensor area will be analysed according to blockSize.\n",
+    "\n",
     "print(f\"pnCCD size is: {pixels_x}x{pixels_y} pixels.\")\n",
     "print(f'Calibration database interface selected: {cal_db_interface}')\n",
     "\n",
-    "proposal = list(filter(None, in_folder.strip('/').split('/')))[-2]\n",
-    "file_loc =f'Proposal: {proposal}, Run: {run}'\n",
-    "print(f'Proposal: {proposal}, Run: {run}')\n",
-    "\n",
     "# Paths to the data:\n",
-    "ped_dir = \"{}/r{:04d}\".format(in_folder, run)\n",
-    "fp_name = path_template.format(run, karabo_da)\n",
-    "fp_path = '{}/{}'.format(ped_dir, fp_name)\n",
-    "h5path = h5path.format(karabo_id, receiver_id)\n",
-    "print(\"HDF5 path to data: {}\\n\".format(h5path))\n",
+    "instrument_src = instrument_source_template.format(karabo_id, receiver_id)\n",
+    "print(f\"Instrument H5File source: {instrument_src}\\n\")\n",
     "\n",
     "# Run's creation time:\n",
     "if creation_time:\n",
@@ -189,7 +164,6 @@
     "if not creation_time and use_dir_creation_date:\n",
     "    creation_time = get_dir_creation_date(in_folder, run)\n",
     "\n",
-    "\n",
     "print(f\"Creation time: {creation_time}\")"
    ]
   },
@@ -199,26 +173,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Reading all sequences of the run:\n",
-    "file_list = []\n",
-    "total_sequences = 0\n",
-    "fsequences = []\n",
-    "\n",
-    "if sequences is None:\n",
-    "    file_list = glob.glob(fp_path.format(0).replace('00000', '*'))\n",
-    "    file_list = sorted(file_list, key = lambda x: (len (x), x))\n",
-    "    total_sequences = len(file_list)\n",
-    "    fsequences = range(total_sequences)\n",
-    "else:\n",
-    "    for seq in sequences:\n",
-    "        abs_entry = fp_path.format(seq)\n",
-    "        if os.path.isfile(abs_entry):\n",
-    "            file_list.append(abs_entry)\n",
-    "            total_sequences += 1\n",
-    "            fsequences.append(seq)\n",
-    "\n",
-    "sequences = fsequences    \n",
-    "print(f\"This run has a total number of {total_sequences} sequences.\\n\")"
+    "step_timer = StepTimer()"
    ]
   },
   {
@@ -227,17 +182,30 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# extract slow data\n",
-    "if karabo_da_control:\n",
-    "    ctrl_fname = os.path.join(ped_dir, path_template.format(run, karabo_da_control)).format(sequences[0])\n",
-    "    ctrl_path = h5path_ctrl.format(karabo_id)\n",
-    "    mdl_ctrl_path = f\"/CONTROL/{karabo_id}/MDL/\"\n",
-    "\n",
-    "    (bias_voltage, gain,\n",
-    "     fix_temperature_top,\n",
-    "     fix_temperature_bot) = extract_slow_data(karabo_id, karabo_da_control, ctrl_fname, ctrl_path,\n",
-    "                                              mdl_ctrl_path, bias_voltage, gain,\n",
-    "                                              fix_temperature_top, fix_temperature_bot)"
+    "run_dc = RunDirectory(Path(in_folder) / f\"r{run:04d}\")\n",
+    "ctrl_data = pnccdlib.PnccdCtrl(run_dc, karabo_id)\n",
+    "\n",
+    "# extract control data\n",
+    "step_timer.start()\n",
+    "\n",
+    "if bias_voltage == 0.:\n",
+    "    bias_voltage = ctrl_data.get_bias_voltage()\n",
+    "if gain == -1:\n",
+    "    gain = ctrl_data.get_gain()\n",
+    "if fix_temperature_top == 0:\n",
+    "    fix_temperature_top = ctrl_data.get_fix_temperature_top()\n",
+    "if fix_temperature_bot == 0:\n",
+    "    fix_temperature_bot = ctrl_data.get_fix_temperature_bot()\n",
+    "\n",
+    "step_timer.done_step(\"Reading control parameters.\")\n",
+    "\n",
+    "# Printing the Parameters Read from the Data File:\n",
+    "display(Markdown('### Detector Parameters'))\n",
+    "print(f\"Bias voltage is {bias_voltage:0.1f} V.\")\n",
+    "print(f\"Detector gain is set to 1/{int(gain)}.\")\n",
+    "print(f\"Detector integration time is set to {integration_time} ms\")\n",
+    "print(f\"Top pnCCD sensor is at temperature of {fix_temperature_top:0.2f} K\")\n",
+    "print(f\"Bottom pnCCD sensor is at temperature of {fix_temperature_bot:0.2f} K\")"
    ]
   },
   {
@@ -246,14 +214,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Printing the Parameters Read from the Data File:\n",
-    "\n",
-    "display(Markdown('### Detector Parameters'))\n",
-    "print(f\"Bias voltage is {bias_voltage:0.1f} V.\")\n",
-    "print(f\"Detector gain is set to 1/{int(gain)}.\")\n",
-    "print(f\"Detector integration time is set to {integration_time} ms\")\n",
-    "print(f\"Top pnCCD sensor is at temperature of {fix_temperature_top:0.2f} K\")\n",
-    "print(f\"Bottom pnCCD sensor is at temperature of {fix_temperature_bot:0.2f} K\")"
+    "seq_files = []\n",
+    "for f in run_dc.select(instrument_src).files:\n",
+    "    fpath = Path(f.filename)\n",
+    "    if fpath.match(f\"*{karabo_da}*.h5\"):\n",
+    "        seq_files.append(fpath)\n",
+    "if sequences != [-1]:\n",
+    "    seq_files = sorted([f for f in seq_files if any(f.match(f\"*-S{s:05d}.h5\") for s in sequences)])\n",
+    "print(f\"Processing a total of {len(seq_files)} sequence files:\")\n",
+    "print(*seq_files, sep='\\n')"
    ]
   },
   {
@@ -262,7 +231,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "gain_k = [k for k, v in VALID_GAINS.items() if v == gain][0]\n",
+    "gain_k = [k for k, v in pnccdlib.VALID_GAINS.items() if v == gain][0]\n",
     "if gain_k == 'a':\n",
     "    split_evt_mip_threshold = 1000. # MIP threshold in ADU for event classification (10 times average noise)\n",
     "\n",
@@ -343,45 +312,9 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "display(Markdown('### List of Files to be Processed'))\n",
-    "print(\"Reading data from the following files:\")\n",
-    "for i in range(len(file_list)):\n",
-    "    print(file_list[i])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T15:54:23.913269Z",
-     "start_time": "2018-12-06T15:54:23.868910Z"
-    }
-   },
-   "outputs": [],
-   "source": [
-    "# Sensor size and block size definitions (important for common mode and other corrections):\n",
-    "\n",
-    "sensorSize = [pixels_x, pixels_y]\n",
-    "blockSize = [sensorSize[0]//2, sensorSize[1]//2] # sensor area will be analysed according to blockSize\n",
-    "xcal.defaultBlockSize = blockSize # for xcal.HistogramCalculators \n",
-    "memoryCells = 1 # pnCCD has 1 memory cell"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T15:54:23.913269Z",
-     "start_time": "2018-12-06T15:54:23.868910Z"
-    }
-   },
-   "outputs": [],
    "source": [
     "# Output Folder Creation:\n",
-    "os.makedirs(out_folder, exist_ok=True)"
+    "os.makedirs(out_folder, exist_ok=True if overwrite else False)"
    ]
   },
   {
@@ -443,16 +376,21 @@
     "    return constants"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Retrieving calibration constants"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "display(Markdown('### Dark Data Retrieval'))\n",
-    "\n",
+    "display(Markdown('### Dark constants retrieval'))\n",
+    "step_timer.start()\n",
     "db_parms = cal_db_interface, cal_db_timeout\n",
     "\n",
     "constants = get_dark(db_parms, bias_voltage, gain, integration_time,\n",
@@ -474,7 +412,8 @@
     "                       lut_label='Bad Pixel Value (ADU)', \n",
     "                       aspect=1, x_range=(0, pixels_y), y_range=(0, pixels_x), \n",
     "                       panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)', \n",
-    "                       title = 'Dark Bad Pixels Map')"
+    "                       title = 'Dark Bad Pixels Map')\n",
+    "step_timer.done_step(\"Dark constants retrieval\")"
    ]
   },
   {
@@ -484,7 +423,8 @@
    "outputs": [],
    "source": [
     "if corr_bools.get('relgain'):\n",
-    "    display(Markdown('We will now retrieve the relative gain map from the calibration database.'))\n",
+    "    step_timer.start()\n",
+    "    display(Markdown('### Relative gain constant retrieval'))\n",
     "    metadata = ConstantMetaData()\n",
     "    relgain = Constants.CCD(DetectorTypes.pnCCD).RelativeGain()\n",
     "    metadata.calibration_constant = relgain\n",
@@ -514,7 +454,8 @@
     "                           panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)', \n",
     "                           panel_top_low_lim = 0.5, panel_top_high_lim = 1.5, panel_side_low_lim = 0.5, \n",
     "                           panel_side_high_lim = 1.5, \n",
-    "                           title = f'Relative Gain Map for pnCCD (Gain = 1/{int(gain)})')"
+    "                           title = f'Relative Gain Map for pnCCD (Gain = 1/{int(gain)})')\n",
+    "    step_timer.done_step(\"Relative gain constant retrieval\")"
    ]
   },
   {
@@ -524,7 +465,6 @@
    "outputs": [],
    "source": [
     "#************************ Calculators ************************#\n",
-    "\n",
     "if corr_bools.get('common_mode'):\n",
     "    # Common Mode Correction Calculator:\n",
     "    cmCorrection = xcal.CommonModeCorrection([pixels_x, pixels_y],\n",
@@ -532,7 +472,6 @@
     "                                             commonModeAxis,\n",
     "                                             parallel=False, dType=np.float32, stride=1,\n",
     "                                             noiseMap=constants[\"Noise\"].astype(np.float32), minFrac=0.25)\n",
-    "    cmCorrection.debug()\n",
     "\n",
     "if corr_bools.get('pattern_class'):\n",
     "    # Pattern Classifier Calculator:\n",
@@ -543,11 +482,10 @@
     "                                                 split_evt_secondary_threshold,\n",
     "                                                 split_evt_mip_threshold,\n",
     "                                                 tagFirstSingles=3, # track along y-axis, left to right (see \n",
-    "                                                 nCells=memoryCells, # split_event.py file in pydetlib/lib/src/\n",
-    "                                                 cores=cpuCores,     # XFELDetAna/algorithms)\n",
-    "                                                 allowElongated=False,\n",
+    "                                                 nCells=1, # split_event.py file in pydetlib/lib/src/\n",
+    "                                                 allowElongated=False, # XFELDetAna/algorithms)\n",
     "                                                 blockSize=[pixels_x, pixels_y//2],\n",
-    "                                                 runParallel=True)\n",
+    "                                                 parallel=False)\n",
     "\n",
     "    # Right Hemisphere:\n",
     "    patternClassifierRH = xcal.PatternClassifier([pixels_x, pixels_y//2],\n",
@@ -556,31 +494,25 @@
     "                                                 split_evt_secondary_threshold,\n",
     "                                                 split_evt_mip_threshold,\n",
     "                                                 tagFirstSingles=4, # track along y-axis, right to left\n",
-    "                                                 nCells=memoryCells,\n",
-    "                                                 cores=cpuCores,\n",
+    "                                                 nCells=1,\n",
     "                                                 allowElongated=False,\n",
     "                                                 blockSize=[pixels_x, pixels_y//2],\n",
-    "                                                 runParallel=True)\n",
+    "                                                 parallel=False)\n",
     "\n",
-    "    patternClassifierLH._imagesPerChunk = 500\n",
-    "    patternClassifierRH._imagesPerChunk = 500\n",
-    "    patternClassifierLH.debug()\n",
-    "    patternClassifierRH.debug()\n",
+    "    patternClassifierLH._imagesPerChunk = 1\n",
+    "    patternClassifierRH._imagesPerChunk = 1\n",
     "\n",
+    "    patternClassifierLH._noisemap = constants[\"Noise\"][:, :pixels_x//2]\n",
+    "    patternClassifierRH._noisemap = constants[\"Noise\"][:, pixels_x//2:]\n",
     "    # Setting bad pixels:\n",
-    "    patternClassifierLH.setBadPixelMask(constants[\"BadPixelsDark\"][:, :pixels_y//2] != 0)\n",
-    "    patternClassifierRH.setBadPixelMask(constants[\"BadPixelsDark\"][:, pixels_y//2:] != 0)"
+    "    patternClassifierLH.setBadPixelMask(constants[\"BadPixelsDark\"][:, :pixels_x//2] != 0)\n",
+    "    patternClassifierRH.setBadPixelMask(constants[\"BadPixelsDark\"][:, pixels_x//2:] != 0)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T15:54:28.771629Z",
-     "start_time": "2018-12-06T15:54:28.346051Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "#***************** Histogram Calculators ******************#\n",
@@ -588,91 +520,136 @@
     "histCalRaw = xcal.HistogramCalculator(sensorSize, \n",
     "                                      bins=bins, \n",
     "                                      range=bin_range,\n",
-    "                                      nCells=memoryCells, \n",
-    "                                      cores=cpuCores,\n",
-    "                                      blockSize=blockSize) \n",
-    "histCalRaw.debug()\n",
+    "                                      nCells=1, \n",
+    "                                      parallel=False,\n",
+    "                                      blockSize=blockSize)\n",
     "# Will contain offset corrected data:\n",
     "histCalOffsetCor = xcal.HistogramCalculator(sensorSize, \n",
     "                                            bins=bins, \n",
     "                                            range=bin_range,\n",
-    "                                            nCells=memoryCells, \n",
-    "                                            cores=cpuCores,\n",
+    "                                            nCells=1, \n",
+    "                                            parallel=False,\n",
     "                                            blockSize=blockSize)\n",
-    "histCalOffsetCor.debug()\n",
     "if corr_bools.get('common_mode'):\n",
     "    # Will contain common mode corrected data:\n",
-    "    histCalCommonModeCor = xcal.HistogramCalculator(sensorSize, \n",
-    "                                                    bins=bins, \n",
+    "    histCalCommonModeCor = xcal.HistogramCalculator(sensorSize,\n",
+    "                                                    bins=bins,\n",
     "                                                    range=bin_range,\n",
-    "                                                    nCells=memoryCells, \n",
-    "                                                    cores=cpuCores,\n",
+    "                                                    nCells=1, \n",
+    "                                                    parallel=False,\n",
     "                                                    blockSize=blockSize)\n",
-    "    histCalCommonModeCor.debug()\n",
-    "    \n",
     "if corr_bools.get('pattern_class'):\n",
     "    # Will contain split events pattern data:\n",
-    "    histCalPcorr = xcal.HistogramCalculator(sensorSize, \n",
-    "                                            bins=bins, \n",
+    "    histCalPcorr = xcal.HistogramCalculator(sensorSize,\n",
+    "                                            bins=bins,\n",
     "                                            range=bin_range,\n",
-    "                                            nCells=memoryCells, \n",
-    "                                            cores=cpuCores,\n",
+    "                                            nCells=1, \n",
+    "                                            parallel=False,\n",
     "                                            blockSize=blockSize)\n",
-    "    histCalPcorr.debug()\n",
     "    # Will contain singles events data:\n",
-    "    histCalPcorrS = xcal.HistogramCalculator(sensorSize, \n",
-    "                                             bins=bins, \n",
+    "    histCalPcorrS = xcal.HistogramCalculator(sensorSize,\n",
+    "                                             bins=bins,\n",
     "                                             range=bin_range,\n",
-    "                                             nCells=memoryCells, \n",
-    "                                             cores=cpuCores,\n",
+    "                                             nCells=1, \n",
+    "                                             parallel=False,\n",
     "                                             blockSize=blockSize)\n",
-    "    histCalPcorrS.debug()\n",
     "if corr_bools.get('relgain'):\n",
     "    # Will contain gain corrected data:\n",
-    "    histCalGainCor = xcal.HistogramCalculator(sensorSize, \n",
-    "                                              bins=bins, \n",
+    "    histCalGainCor = xcal.HistogramCalculator(sensorSize,\n",
+    "                                              bins=bins,\n",
     "                                              range=bin_range,\n",
-    "                                              nCells=memoryCells, \n",
-    "                                              cores=cpuCores,\n",
-    "                                              blockSize=blockSize)\n",
-    "    histCalGainCor.debug()"
+    "                                              nCells=1, \n",
+    "                                              parallel=False,\n",
+    "                                              blockSize=blockSize)"
    ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {
+    "tags": []
+   },
+   "source": [
+    "## Applying corrections to the raw data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
    "source": [
-    "Applying offset and common mode corrections to the raw data"
+    "def offset_correction(wid, index, d):\n",
+    "    \"\"\"offset correction.\n",
+    "    Equating bad pixels' values to np.nan,\n",
+    "    so that the pattern classifier ignores them:\n",
+    "    \"\"\"\n",
+    "    d = d.copy()\n",
+    "\n",
+    "    # TODO: To clear this up. Is it on purpose to save corrected data with nans?\n",
+    "    d[bpix != 0] = np.nan\n",
+    "    d -= offset  # offset correction\n",
+    "\n",
+    "    # TODO: to clear this up. why save the badpixels map in the corrected data?\n",
+    "    bpix_data[index, ...] = bpix\n",
+    "    data[index, ...] = d\n",
+    "\n",
+    "def common_mode(wid, index, d):\n",
+    "    \"\"\"common-mode correction.\n",
+    "    Discarding events caused by saturated pixels:\n",
+    "    \"\"\"\n",
+    "    d = np.squeeze(cmCorrection.correct(d, cellTable=np.zeros(pixels_y, np.int32)))\n",
+    "    # we equate these values to np.nan so that the pattern classifier ignores them:\n",
+    "    d[d >= saturated_threshold] = np.nan\n",
+    "    data[index, ...] = d\n",
+    "\n",
+    "\n",
+    "def gain_correction(wid, index, d):\n",
+    "    \"\"\"relative gain correction.\"\"\"\n",
+    "    d /= relativegain  \n",
+    "    data[index, ...] = d\n",
+    "\n",
+    "\n",
+    "def pattern_classification_correction(wid, index, d):\n",
+    "    \"\"\"pattern classification correction.\n",
+    "    data set to save split event corrected images\n",
+    "    \n",
+    "    The calculation of the cluster map:]\n",
+    "    Dividing the data into left and right hemispheres:\n",
+    "    \"\"\"\n",
+    "\n",
+    "    # pattern classification on corrected data\n",
+    "    dataLH, patternsLH = patternClassifierLH.classify(d[:, :pixels_x//2])\n",
+    "    dataRH, patternsRH = patternClassifierRH.classify(d[:, pixels_x//2:])\n",
+    "\n",
+    "    d[:, :pixels_x//2] = np.squeeze(dataLH)\n",
+    "    d[:, pixels_x//2:] = np.squeeze(dataRH)\n",
+    "\n",
+    "    patterns = np.zeros(d.shape, patternsLH.dtype)\n",
+    "    patterns[:, :pixels_x//2] = np.squeeze(patternsLH)\n",
+    "    patterns[:, pixels_x//2:] = np.squeeze(patternsRH)\n",
+    "    d[d < split_evt_primary_threshold*noise] = 0\n",
+    "    data[index, ...] = d\n",
+    "    ptrn_data[index, ...] = patterns\n",
+    "    d[patterns != 100] = np.nan  # Discard doubles, triples, quadruple, clusters, first singles\n",
+    "    filtered_data[index, ...] = d"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T16:08:53.551111Z",
-     "start_time": "2018-12-06T16:08:53.531064Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "def copy_and_sanitize_non_cal_data(infile: str, outfile: str, h5base: str):\n",
-    "    '''This function reads the .h5 data and writes the corrected .h5 data.'''\n",
-    "    if h5base.startswith(\"/\"):\n",
-    "        h5base = h5base[1:]\n",
-    "    dont_copy = ['image']\n",
-    "    dont_copy = [h5base+\"/{}\".format(do) for do in dont_copy]\n",
-    "\n",
-    "    def visitor(k, item):\n",
-    "        if k not in dont_copy:\n",
-    "            if isinstance(item, h5py.Group):\n",
-    "                outfile.create_group(k)\n",
-    "            elif isinstance(item, h5py.Dataset):\n",
-    "                group = str(k).split(\"/\")\n",
-    "                group = \"/\".join(group[:-1])\n",
-    "                infile.copy(k, outfile[group])\n",
-    "                \n",
-    "    infile.visititems(visitor)"
+    "# 10 is a number chosen after testing 1 ... 71 parallel threads for a node with 72 cpus.\n",
+    "parallel_num_threads = 10\n",
+    "context = psh.context.ThreadContext(num_workers=parallel_num_threads)\n",
+    "\n",
+    "data_path = \"INSTRUMENT/\"+instrument_src+\"/data/\"\n",
+    "\n",
+    "offset = np.squeeze(constants[\"Offset\"])\n",
+    "noise = np.squeeze(constants[\"Noise\"])\n",
+    "bpix = np.squeeze(constants[\"BadPixelsDark\"])\n",
+    "relativegain = constants.get(\"RelativeGain\")"
    ]
   },
   {
@@ -681,172 +658,170 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Data corrections and event classifications happen here. Also, the corrected data are written to datasets:\n",
+    "def write_datasets(corr_arrays, ofile):\n",
+    "    \"\"\"\n",
+    "    Creating datasets first then adding data.\n",
+    "    To have metadata together available at the start of the file,\n",
+    "    so it's quick to see what the file contains\n",
+    "    \"\"\"\n",
+    "    comp_fields = [\"gain\", \"patterns\", \"pixels_classified\"]\n",
+    "    img_grp = ofile[data_path]\n",
+    "\n",
+    "    for field, data in corr_arrays.items():\n",
+    "        kw = dict(chunks=(chunk_size_idim, pixels_x, pixels_y))\n",
+    "        if field in comp_fields:\n",
+    "            kw[\"compression\"] = \"gzip\"\n",
+    "\n",
+    "        img_grp.create_dataset(\n",
+    "            field, shape=data.shape, dtype=data.dtype, **kw)\n",
+    "\n",
+    "    for field, data in corr_arrays.items():\n",
+    "        img_grp[field][:] = data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Data corrections and event classifications happen here.\n",
+    "# Also, the corrected data are written to datasets:\n",
+    "for seq_n, seq_f in enumerate(seq_files):\n",
+    "    f_dc = H5File(seq_f)\n",
+    "    out_file = f\"{out_folder}/{seq_f.name}\".replace(\"RAW\", \"CORR\")\n",
     "\n",
-    "# Initialize 5 numpy array of zeros with the shape of (1024, 1024, 0)\n",
-    "uncor, off_cor, g_cor, cm_cor, final_cor = np.zeros((5, 1024, 1024, 0), dtype=np.float32)\n",
+    "    step_timer.start()\n",
     "\n",
-    "offsetMap = np.squeeze(constants[\"Offset\"])\n",
-    "noiseMap = np.squeeze(constants[\"Noise\"])\n",
-    "badPixelMap = np.squeeze(constants[\"BadPixelsDark\"])\n",
+    "    dshape = f_dc[instrument_src, \"data.image\"].shape\n",
+    "    n_trains = dshape[0]\n",
     "\n",
-    "if corr_bools.get('relgain'):\n",
-    "    relGain = constants[\"RelativeGain\"]\n",
+    "    # If you want to analyze only a certain number of frames\n",
+    "    # instead of all available good frames.\n",
+    "    if limit_images > 0:\n",
+    "        n_trains = min(n_trains, limit_images)\n",
+    "    data_shape = (n_trains, dshape[1], dshape[2])\n",
     "\n",
-    "for k, f in enumerate(file_list):\n",
-    "    with h5py.File(f, 'r') as infile:\n",
-    "        out_fileb = \"{}/{}\".format(out_folder, f.split(\"/\")[-1])\n",
-    "        out_file = out_fileb.replace(\"RAW\", \"CORR\")\n",
+    "    print(f\"Correcting file: {seq_f} of shape {data_shape}.\")\n",
     "\n",
-    "        data = None\n",
-    "        noise = None\n",
-    "        \n",
-    "        try:\n",
-    "            with h5py.File(out_file, \"w\") as ofile:\n",
-    "                copy_and_sanitize_non_cal_data(infile, ofile, h5path)\n",
-    "                data = infile[h5path+\"/image\"][()]\n",
-    "                # Getting rid of empty frames:\n",
-    "                nzidx = np.count_nonzero(data, axis=(1, 2))\n",
-    "                data = data[nzidx != 0, ...]\n",
-    "                \n",
-    "                # If you want to analyze only a certain number of frames instead of all available good frames: \n",
-    "                if limit_images > 0:\n",
-    "                    data = data[:limit_images,...]\n",
-    "                    \n",
-    "                # used for array shapes in the corrected data sets that we create and save in this loop:\n",
-    "                oshape = data.shape\n",
-    "                \n",
-    "                data = np.moveaxis(data, 0, 2)\n",
-    "                \n",
-    "                # data set to save offset corrected images:\n",
-    "                ddset = ofile.create_dataset(h5path+\"/pixels\",\n",
-    "                                             oshape,\n",
-    "                                             chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                                             dtype=np.float32)\n",
-    "                # data set to create bad pixels:\n",
-    "                ddsetm = ofile.create_dataset(h5path+\"/mask\",\n",
-    "                                             oshape,\n",
-    "                                             chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                                             dtype=np.uint32, compression=\"gzip\")\n",
-    "                \n",
-    "                data = data.astype(np.float32) \n",
-    "                \n",
-    "                # Creating maps for correction usage:\n",
-    "                offset = np.repeat(offsetMap[...,None], data.shape[2], axis=2)\n",
-    "                noise = np.repeat(noiseMap[...,None], data.shape[2], axis=2)\n",
-    "                bpix = np.repeat(badPixelMap[...,None], data.shape[2], axis=2)\n",
-    "                if corr_bools.get('relgain'):\n",
-    "                    rg = np.repeat(relGain[:,:,None], data.shape[2], axis=2) # rg = relative gain\n",
-    "                \n",
-    "                # non-corrected images for first sequence only:\n",
-    "                if k == 0:\n",
-    "                    uncor = np.append(uncor, data, axis=2)\n",
-    "                        \n",
-    "                histCalRaw.fill(data) # filling histogram with raw uncorrected data\n",
-    "                \n",
-    "                # equating bad pixels' values to np.nan so that the pattern classifier ignores them:\n",
-    "                data[bpix != 0] = np.nan\n",
-    "            \n",
-    "                data -= offset # offset correction  \n",
-    "                \n",
-    "                # Offset corrected images for first sequence only:\n",
-    "                if k == 0:\n",
-    "                    off_cor = np.append(off_cor, data, axis=2)\n",
-    "                histCalOffsetCor.fill(data) # filling histogram with offset corrected data\n",
-    "\n",
-    "                ddset[...] = np.moveaxis(data, 2, 0)\n",
-    "                ddsetm[...] = np.moveaxis(bpix, 2, 0)\n",
-    "                ofile.flush()\n",
-    "\n",
-    "                # cm: common mode\n",
-    "                if corr_bools.get('common_mode'):\n",
-    "                    \n",
-    "                    # data set to save common mode corrected images:\n",
-    "                    ddsetcm = ofile.create_dataset(h5path+\"/pixels_cm\",\n",
-    "                                                   oshape,\n",
-    "                                                   chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                                                   dtype=np.float32)\n",
-    "                    \n",
-    "                    # common mode correction:\n",
-    "                    data = cmCorrection.correct(data.astype(np.float32),  # common mode correction\n",
-    "                                                cellTable=np.zeros(data.shape[2], np.int32)) \n",
-    "                    \n",
-    "                    # discarding events caused by saturated pixels:\n",
-    "                    # we equate these values to np.nan so that the pattern classifier ignores them:\n",
-    "                    data[data >= saturated_threshold] = np.nan \n",
-    "                    histCalCommonModeCor.fill(data) # filling histogram with common mode corrected data\n",
-    "                    # common mode corrected images for first sequence only:\n",
-    "                    if k == 0:\n",
-    "                        cm_cor = np.append(cm_cor, data, axis=2)\n",
-    "                    \n",
-    "                    ddsetcm[...] = np.moveaxis(data, 2, 0)\n",
-    "                    \n",
-    "                if corr_bools.get('relgain'):\n",
-    "                    # data set to save gain corrected images:\n",
-    "                    ddsetg = ofile.create_dataset(h5path+\"/gain\",\n",
-    "                                                  oshape,\n",
-    "                                                  chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                                                  dtype=np.float32, compression=\"gzip\")\n",
-    "                \n",
-    "                    data /= rg  # relative gain correction \n",
-    "                    histCalGainCor.fill(data) # filling histogram with gain corrected data\n",
-    "                    # gain corrected images for first sequence only:\n",
-    "                    if k == 0:\n",
-    "                        g_cor = np.append(g_cor, data, axis=2)\n",
-    "                    ddsetg[...] = np.moveaxis(rg, 2, 0).astype(np.float32)\n",
-    "\n",
-    "                if corr_bools.get('pattern_class'):\n",
-    "                    # data set to save split event corrected images:\n",
-    "                    # c: classifications, p: even patterns\n",
-    "                    ddsetc = ofile.create_dataset(h5path+\"/pixels_classified\",\n",
-    "                                                  oshape,\n",
-    "                                                  chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                                                  dtype=np.float32, compression=\"gzip\")\n",
-    "                    # data set to save different valid patterns:\n",
-    "                    ddsetp = ofile.create_dataset(h5path+\"/patterns\",\n",
-    "                                                 oshape,\n",
-    "                                                 chunks=(chunk_size_idim, oshape[1], oshape[2]),\n",
-    "                                                 dtype=np.int32, compression=\"gzip\")\n",
-    "                    \n",
-    "                    # The calculation of the cluster map:\n",
-    "                    patternClassifierLH._noisemap = noise[:, :pixels_x//2,  :]\n",
-    "                    patternClassifierRH._noisemap = noise[:, pixels_x//2:,  :]\n",
-    "\n",
-    "                    # Dividing the data into left and right hemispheres:\n",
-    "                    dataLH = data[:, :pixels_x//2, :]\n",
-    "                    dataRH = data[:, pixels_x//2:, :]\n",
-    "                    \n",
-    "                    # pattern classification on corrected data\n",
-    "                    dataLH, patternsLH = patternClassifierLH.classify(dataLH) \n",
-    "                    dataRH, patternsRH = patternClassifierRH.classify(dataRH)\n",
-    "\n",
-    "                    data[:, :pixels_x//2, :] = dataLH\n",
-    "                    data[:, pixels_x//2:, :] = dataRH\n",
-    "\n",
-    "                    patterns = np.zeros(data.shape, patternsLH.dtype)\n",
-    "                    patterns[:, :pixels_x//2, :] = patternsLH\n",
-    "                    patterns[:, pixels_x//2:, :] = patternsRH\n",
-    "\n",
-    "                    data[data < split_evt_primary_threshold*noise] = 0\n",
-    "                    ddsetc[...] = np.moveaxis(data, 2, 0)\n",
-    "                    ddsetp[...] = np.moveaxis(patterns, 2, 0)\n",
-    "\n",
-    "                    histCalPcorr.fill(data) # filling histogram with split events corrected data\n",
-    "                    data[patterns != 100] = np.nan # Discard doubles, triples, quadruple, clusters, first singles\n",
-    "                    histCalPcorrS.fill(data) # filling histogram with singles events data\n",
-    "                     \n",
-    "                    # split event corrected images for first sequence only (also these events are only \n",
-    "                    # singles events):\n",
-    "                    if k == 0:\n",
-    "                        final_cor = np.append(final_cor, data, axis=2)\n",
-    "                    \n",
-    "        except Exception as e:\n",
-    "            print(f\"Couldn't calibrate data in {f}: {e}\\n\")\n",
+    "    data_dc = f_dc.select(instrument_src, \"data.image\", require_all=True).select_trains(np.s_[:n_trains])  # noqa\n",
+    "\n",
+    "    raw_data = data_dc[instrument_src, \"data.image\"].ndarray().astype(np.float32)\n",
+    "\n",
+    "    if seq_n == 0:\n",
+    "        raw_plt = raw_data.copy()  # plot first sequence only\n",
+    "\n",
+    "    step_timer.start()\n",
+    "\n",
+    "    # Allocating shared arrays for data arrays for each correction stage.\n",
+    "    data = context.alloc(shape=data_shape, dtype=np.float32)\n",
+    "    bpix_data = context.alloc(shape=data_shape, dtype=np.uint32)\n",
+    "    histCalRaw.fill(raw_data)  # filling histogram with raw uncorrected data\n",
+    "    \n",
+    "    # Applying offset correction\n",
+    "    context.map(offset_correction, raw_data)\n",
+    "    histCalOffsetCor.fill(data)  # filling histogram with offset corrected data\n",
+    "\n",
+    "    if seq_n == 0:\n",
+    "        off_data = data.copy()  # plot first sequence only\n",
+    "\n",
+    "\n",
+    "    corr_arrays = {\n",
+    "        \"pixels\": data.copy(),\n",
+    "        \"mask\": bpix_data, \n",
+    "    }\n",
+    "    step_timer.done_step(f'offset correction.')\n",
+    "\n",
+    "    if corr_bools.get('common_mode'):\n",
+    "        step_timer.start()\n",
+    "\n",
+    "        # Applying common mode correction\n",
+    "        context.map(common_mode, data)\n",
+    "        if seq_n == 0:\n",
+    "            cm_data = data.copy()  # plot first sequence only\n",
+    "        corr_arrays[\"pixels_cm\"] = data.copy()\n",
+    "        histCalCommonModeCor.fill(data)  # filling histogram with common mode corrected data\n",
+    "\n",
+    "        step_timer.done_step(f'common-mode correction.')\n",
+    "\n",
+    "    if corr_bools.get('relgain'):\n",
+    "\n",
+    "        step_timer.start()\n",
+    "\n",
+    "        # Applying gain correction\n",
+    "        context.map(gain_correction, data)\n",
+    "        if seq_n == 0:\n",
+    "            rg_data = data.copy()  # plot first sequence only\n",
+    "        # TODO: Why storing a repeated constant for each image in corrected files.\n",
+    "        corr_arrays[\"gain\"] = np.repeat(relativegain[np.newaxis, ...], n_trains, axis=0).astype(np.float32)  # noqa\n",
+    "        histCalGainCor.fill(data)  # filling histogram with gain corrected data\n",
+    "        step_timer.done_step(f'gain correction.')\n",
+    "\n",
+    "    if corr_bools.get('pattern_class'):\n",
+    "        step_timer.start()\n",
+    "\n",
+    "        ptrn_data = context.alloc(shape=data_shape, dtype=np.int32)\n",
+    "        filtered_data = context.alloc(shape=data_shape, dtype=np.int32)\n",
+    "        # Applying pattern classification correction\n",
+    "        # Even thougth data is indeed of dtype np.float32,\n",
+    "        # not specifiying this again screw with the data quality.\n",
+    "        context.map(pattern_classification_correction, data.astype(np.float32))\n",
+    "\n",
+    "        if seq_n == 0:\n",
+    "            cls_data = data.copy()  # plot first sequence only\n",
+    "        # split event corrected images plotted for first sequence only\n",
+    "        # (also these events are only singles events):\n",
+    "        corr_arrays[\"pixels_classified\"] = data.copy()\n",
+    "        corr_arrays[\"patterns\"] = ptrn_data\n",
+    "\n",
+    "        histCalPcorr.fill(data)  # filling histogram with split events corrected data\n",
+    "        # filling histogram with corr data after discarding doubles, triples, quadruple, clusters, and first singles\n",
+    "        histCalPcorrS.fill(filtered_data)\n",
+    "        step_timer.done_step(f'pattern classification correction.')\n",
+    "\n",
+    "    step_timer.start()\n",
+    "\n",
+    "    # Storing corrected data sources.\n",
+    "    with h5py.File(out_file, 'w') as ofile:\n",
+    "        # Copy RAW non-calibrated sources.\n",
+    "        with h5py.File(seq_f, 'r') as sfile:\n",
+    "            h5_copy_except.h5_copy_except_paths(\n",
+    "                sfile, ofile, [],\n",
+    "            )\n",
+    "        # TODO: to clear this up: why save corrected data in data/pixels rather than data/image.\n",
+    "        write_datasets(corr_arrays, ofile)\n",
+    "\n",
+    "    step_timer.done_step(f'Storing data.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"In addition to offset correction, the following corrections were performed:\")\n",
+    "for k, v in corr_bools.items():\n",
+    "    if v:\n",
+    "        print(\"  -\", k.upper())\n",
     "\n",
+    "print(f\"Total processing time {step_timer.timespan():.01f} s\")\n",
+    "step_timer.print_summary()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "print(\"In addition to offset correction, the following corrections were performed:\")\n",
     "for k, v in corr_bools.items():\n",
     "    if v:\n",
-    "        print(k)"
+    "        print(\"  -\", k.upper())\n",
+    "\n",
+    "print(f\"Total processing time {step_timer.timespan():.01f} s\")\n",
+    "step_timer.print_summary()"
    ]
   },
   {
@@ -862,7 +837,7 @@
     "# if you use histCalRaw.get(cumulatative = True) and so on, the cumulatative = True turns the counts array such as \n",
     "# RawHistVals and so on into a 1D array instead of keeping the original shape:\n",
     "\n",
-    "RawHistVals, _, RawHistMids, _ = histCalRaw.get() \n",
+    "RawHistVals, _, RawHistMids, _ = histCalRaw.get()\n",
     "off_cor_HistVals, _, off_cor_HistMids, _ = histCalOffsetCor.get()\n",
     "if corr_bools.get('common_mode'):\n",
     "    cm_cor_HistVals, _, cm_HistMids, _ = histCalCommonModeCor.get()\n",
@@ -880,7 +855,7 @@
    "outputs": [],
    "source": [
     "# Saving intermediate data to disk:\n",
-    "\n",
+    "step_timer.start()\n",
     "np.savez(os.path.join(out_folder, 'Raw_Events.npz'), RawHistMids, RawHistVals)\n",
     "np.savez(os.path.join(out_folder, 'Offset_Corrected_Events.npz'), off_cor_HistMids, off_cor_HistVals)\n",
     "if corr_bools.get('common_mode'):\n",
@@ -890,7 +865,7 @@
     "if corr_bools.get('pattern_class'):\n",
     "    np.savez(os.path.join(out_folder, 'Split_Events_Corrected_Events.npz'), split_HistMids, split_HistVals)\n",
     "    np.savez(os.path.join(out_folder, 'Singles_Events.npz'), singles_HistMids, singles_HistVals)\n",
-    "\n",
+    "step_timer.done_step(f'Saving intermediate data to disk.')\n",
     "print(\"Various spectra are saved to disk in the form of histograms. Please check {}\".format(out_folder))"
    ]
   },
@@ -901,6 +876,7 @@
    "outputs": [],
    "source": [
     "display(Markdown('### Raw vs. Corrected Spectra'))\n",
+    "step_timer.start()\n",
     "\n",
     "figure = [{'x': RawHistMids,\n",
     "      'y': RawHistVals,\n",
@@ -937,8 +913,8 @@
     "                   'errorstyle': 'bars',\n",
     "                   'errorcoarsing': 2,\n",
     "                   'label': 'Gain Corrected'})\n",
-    "    \n",
-    "if corr_bools.get('pattern_class'):    \n",
+    "\n",
+    "if corr_bools.get('pattern_class'):\n",
     "    figure.extend([{'x': split_HistMids,\n",
     "                   'y': split_HistVals,\n",
     "                   'y_err': np.sqrt(split_HistVals[:]),\n",
@@ -957,7 +933,8 @@
     "                  }])\n",
     "fig = xana.simplePlot(figure, aspect=1, x_label='ADU', y_label='Number of Occurrences', figsize='2col',\n",
     "                      y_log=True, x_range=bin_range, title = '1 ADU per bin is used.',\n",
-    "                      legend='top-right-frame-1col')"
+    "                      legend='top-right-frame-1col')\n",
+    "step_timer.done_step('Plotting')"
    ]
   },
   {
@@ -969,7 +946,7 @@
     "# This function plots pattern statistics:\n",
     "\n",
     "def classification_plot(patternStats, hemisphere):\n",
-    "    \n",
+    "\n",
     "    print(\"****************** {} HEMISPHERE ******************\\n\"\n",
     "          .format(hemisphere))\n",
     "    fig = plt.figure(figsize=(15, 15))\n",
@@ -994,7 +971,6 @@
     "\n",
     "    smaps = [\"singlemap\", \"firstsinglemap\", \"clustermap\"]\n",
     "    for i, m in enumerate(smaps):\n",
-    "\n",
     "        ax = fig.add_subplot(4, 4, 2+i)\n",
     "        pmap = ax.imshow(patternStats[m], interpolation=\"nearest\", vmax=2*np.nanmedian(patternStats[m]))\n",
     "        ax.set_title(m)\n",
@@ -1003,7 +979,6 @@
     "    mmaps = [\"doublemap\", \"triplemap\", \"quadmap\"]\n",
     "    k = 0\n",
     "    for i, m in enumerate(mmaps):\n",
-    "\n",
     "        for j in range(4):\n",
     "            ax = fig.add_subplot(4, 4, 2+len(smaps)+k)\n",
     "            pmap = ax.imshow(patternStats[m][j], interpolation=\"nearest\", vmax=2*np.median(patternStats[m][j]))\n",
@@ -1045,6 +1020,7 @@
     "display(Markdown('### Classification Results - Tabulated Statistics'))\n",
     "\n",
     "if corr_bools.get('pattern_class'):\n",
+    "    step_timer.start()\n",
     "    t0 = PrettyTable()\n",
     "    t0.title = \"Total Number of Counts after All Corrections\"\n",
     "    t0.field_names = [\"Hemisphere\", \"Singles\", \"First-Singles\", \"Clusters\"]\n",
@@ -1063,18 +1039,14 @@
     "                patternStatsRH['triples'][2], patternStatsLH['quads'][2], patternStatsRH['quads'][2]])\n",
     "    t1.add_row([3, patternStatsLH['doubles'][3], patternStatsRH['doubles'][3], patternStatsLH['triples'][3], \n",
     "                patternStatsRH['triples'][3], patternStatsLH['quads'][3], patternStatsRH['quads'][3]])\n",
-    "    print(t1)"
+    "    print(t1)\n",
+    "    step_timer.done_step('Classification Results - Tabulated Statistics')"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T16:10:56.190150Z",
-     "start_time": "2018-12-06T16:10:56.177570Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "if corr_bools.get('pattern_class'):\n",
@@ -1109,17 +1081,13 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T16:10:56.203219Z",
-     "start_time": "2018-12-06T16:10:56.191509Z"
-    }
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
     "display(Markdown('### Classification Results - Pie Charts'))\n",
     "\n",
     "if corr_bools.get('pattern_class'):\n",
+    "    step_timer.start()\n",
     "    fig = plt.figure(figsize=(12, 7))\n",
     "    ax = fig.add_subplot(1, 2, 1)\n",
     "    labels = ['Singles', 'Doubles', 'Triples', 'Quads']\n",
@@ -1131,7 +1099,8 @@
     "    pie = ax.pie(reloccurRH, labels=labels, autopct='%1.1f%%', shadow=True)\n",
     "    ax.set_title(\"Pattern Occurrence in RH\")\n",
     "    # Set aspect ratio to be equal so that pie is drawn as a circle.\n",
-    "    a = ax.axis('equal')"
+    "    a = ax.axis('equal')\n",
+    "    step_timer.done_step('Classification Results - Pie Charts')"
    ]
   },
   {
@@ -1144,24 +1113,20 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T16:10:56.212586Z",
-     "start_time": "2018-12-06T16:10:56.204731Z"
-    },
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "uncor_mean_im = np.nanmean(uncor, axis=2)\n",
-    "offset_mean_im = np.nanmean(off_cor, axis=2)\n",
+    "step_timer.start()\n",
+    "\n",
+    "uncor_mean_im = np.nanmean(raw_data, axis=0)\n",
+    "offset_mean_im = np.nanmean(off_data, axis=0)\n",
     "\n",
     "if corr_bools.get('common_mode'):\n",
-    "    cm_mean_im = np.nanmean(cm_cor, axis=2)\n",
+    "    cm_mean_im = np.nanmean(cm_data, axis=0)\n",
     "if corr_bools.get('relgain'):\n",
-    "    gain_mean_im = np.nanmean(g_cor, axis=2)\n",
+    "    gain_mean_im = np.nanmean(rg_data, axis=0)\n",
     "if corr_bools.get('pattern_class'):\n",
-    "    mean_im_cc = np.nanmean(final_cor, axis=2)\n",
+    "    mean_im_cc = np.nanmean(cls_data, axis=0)\n",
     "\n",
     "fig = xana.heatmapPlot(uncor_mean_im, x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1, \n",
     "                       x_range=(0, pixels_y), y_range=(0, pixels_x),  \n",
@@ -1189,7 +1154,8 @@
     "if corr_bools.get('pattern_class'):\n",
     "    fig = xana.heatmapPlot(mean_im_cc, x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1,\n",
     "                           x_range=(0, pixels_y), y_range=(0, pixels_x), vmin=0, vmax= 18000,\n",
-    "                           title = 'Image of Single Events Averaged over Frames in the First Sequence')"
+    "                           title = 'Image of Single Events Averaged over Frames in the First Sequence')\n",
+    "step_timer.done_step(\"Plotting\")"
    ]
   },
   {
@@ -1202,44 +1168,40 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2018-12-06T16:11:08.317130Z",
-     "start_time": "2018-12-06T16:11:05.788655Z"
-    },
-    "scrolled": false
-   },
+   "metadata": {},
    "outputs": [],
    "source": [
-    "fig = xana.heatmapPlot(uncor[:,:,0], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1, \n",
+    "step_timer.start()\n",
+    "fig = xana.heatmapPlot(raw_data[0, :, :], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1, \n",
     "                       x_range=(0, pixels_y), y_range=(0, pixels_x), \n",
     "                       panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)',  \n",
     "                       title = 'Uncorrected Image (First Frame of the First Sequence)')\n",
     "\n",
-    "fig = xana.heatmapPlot(off_cor[:,:,0], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1, \n",
+    "fig = xana.heatmapPlot(off_data[0, :, :], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1, \n",
     "                       x_range=(0, pixels_y), y_range=(0, pixels_x),\n",
     "                       panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)',   \n",
     "                       title = 'Offset Corrected Image (First Frame of the First Sequence)')\n",
     "\n",
     "if corr_bools.get('common_mode'):\n",
-    "        fig = xana.heatmapPlot(cm_cor[:,:,2], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', \n",
+    "        fig = xana.heatmapPlot(cm_data[0, :, :], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', \n",
     "                               aspect=1, \n",
     "                               x_range=(0, pixels_y), y_range=(0, pixels_x), \n",
     "                               panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)', \n",
     "                               title = 'Common Mode Corrected Image (First Frame of the First Sequence)')\n",
     "        \n",
     "if corr_bools.get('relgain'):\n",
-    "        fig = xana.heatmapPlot(g_cor[:,:,0], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', \n",
+    "        fig = xana.heatmapPlot(rg_data[0, :, :], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', \n",
     "                               aspect=1, \n",
     "                               x_range=(0, pixels_y), y_range=(0, pixels_x), \n",
     "                               panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)', \n",
     "                               title = 'Gain Corrected Image (First Frame of the First Sequence)')\n",
     "\n",
     "if corr_bools.get('pattern_class'):    \n",
-    "    fig = xana.heatmapPlot(final_cor[:,:,0], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1,\n",
+    "    fig = xana.heatmapPlot(cls_data[0, :, :], x_label='Columns', y_label='Rows', lut_label='Signal (ADU)', aspect=1,\n",
     "                           x_range=(0, pixels_y), y_range=(0, pixels_x), \n",
     "                           panel_x_label='Row Stat (ADU)', panel_y_label='Column Stat (ADU)',  \n",
-    "                           title = 'Image of Single Events (First Frame of the First Sequence)')"
+    "                           title = 'Image of Single Events (First Frame of the First Sequence)')\n",
+    "step_timer.done_step(\"Plotting\")"
    ]
   },
   {
@@ -1282,37 +1244,35 @@
     "    doubles = []\n",
     "    triples = []\n",
     "    quads =  []\n",
-    "\n",
-    "    with h5py.File(\"{}/CORR-R{:04d}-PNCCD01-S{:05d}.h5\".format(out_folder, run, sequences[seq_num]), 'r') as infile:\n",
-    "            data = infile[h5path+\"/pixels_classified\"][()].astype(np.float32) # classifications\n",
-    "            patterns = infile[h5path+\"/patterns\"][()].astype(np.float32) # event patterns\n",
-    "            \n",
-    "            # events' patterns indices are as follows: 100 (singles), 101 (first singles), 200 - 203 (doubles),\n",
-    "            # 300 - 303 (triples), and 400 - 403 (quadruples). Note that for the last three types of patterns, \n",
-    "            # there are left, right, up, and down indices.\n",
-    "\n",
-    "            # Separating the events:\n",
-    "            # Singles and First Singles:\n",
-    "            for s in range(100, 102):\n",
-    "                single = copy.copy(data[...])\n",
-    "                single[patterns != s] = np.nan\n",
-    "                singles.append(single)\n",
-    "            \n",
-    "\n",
-    "            for d in range(200, 204):\n",
-    "                double = copy.copy(data[...])\n",
-    "                double[patterns != d] = np.nan\n",
-    "                doubles.append(double)\n",
-    "\n",
-    "            for t in range(300, 304):\n",
-    "                triple = copy.copy(data[...])\n",
-    "                triple[patterns != t] = np.nan\n",
-    "                triples.append(triple)  \n",
-    "\n",
-    "            for q in range(400, 404):\n",
-    "                quad = copy.copy(data[...])\n",
-    "                quad[patterns != q] = np.nan\n",
-    "                quads.append(quad)"
+    "    with H5File(f\"{out_folder}/{seq_files[0].name.replace('RAW', 'CORR')}\") as dc:  # noqa\n",
+    "        data = dc[instrument_src, \"data.pixels_classified\"].ndarray()\n",
+    "        patterns = dc[instrument_src, \"data.patterns\"].ndarray()\n",
+    "    # events' patterns indices are as follows: 100 (singles), 101 (first singles), 200 - 203 (doubles),\n",
+    "    # 300 - 303 (triples), and 400 - 403 (quadruples). Note that for the last three types of patterns, \n",
+    "    # there are left, right, up, and down indices.\n",
+    "\n",
+    "    # Separating the events:\n",
+    "    # Singles and First Singles:\n",
+    "    for s in range(100, 102):\n",
+    "        single = data.copy()\n",
+    "        single[patterns != s] = np.nan\n",
+    "        singles.append(single)\n",
+    "\n",
+    "\n",
+    "    for d in range(200, 204):\n",
+    "        double = data.copy()\n",
+    "        double[patterns != d] = np.nan\n",
+    "        doubles.append(double)\n",
+    "\n",
+    "    for t in range(300, 304):\n",
+    "        triple = data.copy()\n",
+    "        triple[patterns != t] = np.nan\n",
+    "        triples.append(triple)  \n",
+    "\n",
+    "    for q in range(400, 404):\n",
+    "        quad = data.copy()\n",
+    "        quad[patterns != q] = np.nan\n",
+    "        quads.append(quad)"
    ]
   },
   {
@@ -1322,13 +1282,13 @@
    "outputs": [],
    "source": [
     "if corr_bools.get('pattern_class'):\n",
+    "    step_timer.start()\n",
     "    hA = 0\n",
     "    h = 0\n",
     "    for single in singles:\n",
     "        hs, e = np.histogram(single.flatten(), bins=event_bins, range=b_range) # h: histogram counts, e: bin edges\n",
     "        h += hs\n",
     "        hA += hs # hA: counts all events (see below)\n",
-    "\n",
     "   \n",
     "    # bin edges array has one extra element => need to plot from 0 to the one before the last element to have the \n",
     "    # same size as h-array => in what follows, we use e[:-1] (-1 means one before the last element)\n",
@@ -1367,7 +1327,8 @@
     "    ax.step(e[:-1], h, color='purple', label='Events Splitting on Quadruple Pixels')\n",
     "\n",
     "    ax.step(e[:-1], hA, color='grey', label='All Valid Events')\n",
-    "    l = ax.legend()"
+    "    l = ax.legend()\n",
+    "    step_timer.done_step(\"Plotting\")"
    ]
   },
   {
@@ -1375,7 +1336,10 @@
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "print(f\"Total processing time {step_timer.timespan():.01f} s\")\n",
+    "step_timer.print_summary()"
+   ]
   }
  ],
  "metadata": {
@@ -1394,7 +1358,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.6.7"
+   "version": "3.8.11"
   },
   "latex_envs": {
    "LaTeX_envs_menu_present": true,
@@ -1415,5 +1379,5 @@
   }
  },
  "nbformat": 4,
- "nbformat_minor": 1
+ "nbformat_minor": 4
 }
diff --git a/setup.py b/setup.py
index d3a00fe1b44b261c37d356cf6b56a3d7838dc8a6..e73c52388c666bcf62a753bf39d580c60c70b61e 100644
--- a/setup.py
+++ b/setup.py
@@ -5,8 +5,8 @@ from subprocess import check_output
 import numpy
 from Cython.Build import cythonize
 from Cython.Distutils import build_ext
-from setuptools.extension import Extension
 from setuptools import find_packages, setup
+from setuptools.extension import Extension
 
 from src.xfel_calibrate.notebooks import notebooks
 
@@ -42,56 +42,59 @@ for ctypes in notebooks.values():
 data_files = list(filter(None, data_files))  # Get rid of `None` entries
 
 install_requires = [
-    "Cython==0.29.21",
-    "Jinja2==2.11.2",
-    "astcheck==0.2.5",
-    "astsearch==0.2.0",
-    "cfelpyutils==1.0.1",
-    "dill==0.3.0",
-    "docutils==0.17.1",
-    "dynaconf==3.1.4",
-    "extra_data==1.4.1",
-    "extra_geom==1.6.0",
-    "gitpython==3.1.0",
-    "h5py==3.5.0",
-    "iminuit==1.3.8",
-    "ipykernel==5.1.4",
-    "ipyparallel==6.2.4",
-    "ipython==7.12.0",
-    "ipython_genutils==0.2.0",
-    "jupyter-core==4.6.1",
-    "jupyter_client==6.1.7",
-    "jupyter_console==6.1.0",
-    "kafka-python==2.0.2",
-    "karabo_data==0.7.0",
-    "lxml==4.5.0",
-    "matplotlib==3.4.2",
-    "metadata_client==3.0.8",
-    "nbclient==0.5.1",
-    "nbconvert==5.6.1",
-    "nbformat==5.0.7",
-    "nbparameterise==0.5",
-    "notebook==6.1.5",
-    "numpy==1.20.3",
-    "pasha==0.1.0",
-    "prettytable==0.7.2",
-    "princess==0.5",
-    "pypandoc==1.4",
-    "python-dateutil==2.8.1",
-    "pyyaml==5.3",
-    "pyzmq==19.0.0",
-    "requests==2.22.0",
-    "scikit-learn==0.22.2.post1",
-    "scipy==1.7.0",
-    "sharedmem==0.3.8",
-    "sphinx==1.8.5",
-    "tabulate==0.8.6",
-    "traitlets==4.3.3",
+        "Cython==0.29.21",
+        "Jinja2==2.11.2",
+        "markupsafe==2.0.1",
+        "astcheck==0.2.5",
+        "astsearch==0.2.0",
+        "cfelpyutils==1.0.1",
+        "dill==0.3.0",
+        "docutils==0.17.1",
+        "dynaconf==3.1.4",
+        "extra_data==1.9.1",
+        "extra_geom==1.6.0",
+        "gitpython==3.1.0",
+        "h5py==3.5.0",
+        "iminuit==1.3.8",
+        "ipykernel==5.1.4",
+        "ipyparallel==6.2.4",
+        "ipython==7.12.0",
+        "ipython_genutils==0.2.0",
+        "jupyter-core==4.6.1",
+        "jupyter_client==6.1.7",
+        "jupyter_console==6.1.0",
+        "kafka-python==2.0.2",
+        "karabo_data==0.7.0",
+        "lxml==4.5.0",
+        "markupsafe==2.0.1",
+        "matplotlib==3.4.2",
+        "metadata_client==3.0.8",
+        "nbclient==0.5.1",
+        "nbconvert==5.6.1",
+        "nbformat==5.0.7",
+        "nbparameterise==0.5",
+        "notebook==6.1.5",
+        "numpy==1.20.3",
+        "pasha==0.1.0",
+        "prettytable==0.7.2",
+        "princess==0.5",
+        "psutil==5.9.0",
+        "pypandoc==1.4",
+        "python-dateutil==2.8.1",
+        "pyyaml==5.3",
+        "pyzmq==19.0.0",
+        "requests==2.22.0",
+        "scikit-learn==0.22.2.post1",
+        "scipy==1.7.0",
+        "sharedmem==0.3.8",
+        "sphinx==1.8.5",
+        "tabulate==0.8.6",
+        "traitlets==4.3.3",
 ]
 
 if "readthedocs.org" not in sys.executable:
     install_requires += [
-        "iCalibrationDB @ git+ssh://git@git.xfel.eu:10022/detectors/cal_db_interactive.git@2.0.9",  # noqa
+        "iCalibrationDB @ git+ssh://git@git.xfel.eu:10022/detectors/cal_db_interactive.git@2.1.0",  # noqa
         "XFELDetectorAnalysis @ git+ssh://git@git.xfel.eu:10022/karaboDevices/pyDetLib.git@2.7.0",  # noqa
     ]
 
diff --git a/src/cal_tools/agipdlib.py b/src/cal_tools/agipdlib.py
index feb27962279b1f2047fa433417d8e5be1559278d..a39ca0ea0ab304fad6d35c8f2be7b25bd0193317 100644
--- a/src/cal_tools/agipdlib.py
+++ b/src/cal_tools/agipdlib.py
@@ -1,14 +1,14 @@
 import os
 import posixpath
-import traceback
 import zlib
 from multiprocessing.pool import ThreadPool
-from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
 import h5py
 import numpy as np
 import sharedmem
+from dateutil import parser
+from extra_data import DataCollection, H5File, by_id, components
 from iCalibrationDB import Conditions, Constants
 
 from cal_tools import agipdalgs as calgs
@@ -26,10 +26,35 @@ from cal_tools.enums import AgipdGainMode, BadPixels, SnowResolution
 from cal_tools.h5_copy_except import h5_copy_except_paths
 from cal_tools.tools import get_constant_from_db_and_time
 
+class AgipdCtrl:
+    def __init__(
+        self,
+        run_dc: DataCollection,
+        image_src: str,
+        ctrl_src: str,
+        raise_error: bool = True,
+    ):
+        """
+        Initialize AgipdCondition class to read all required AGIPD parameters.
 
-def get_num_cells(fname, loc, module):
-    with h5py.File(fname, "r") as f:
-        cells = f[f"INSTRUMENT/{loc}/DET/{module}CH0:xtdf/image/cellId"][()]
+        :param image_src: H5 source for image data.
+        :param ctrl_src: H5 source for control (slow) data.
+        """
+        self.run_dc = run_dc
+        self.image_src = image_src
+        self.ctrl_src = ctrl_src
+
+        self.raise_error = raise_error
+
+    def get_num_cells(self) -> Optional[int]:
+        """
+        :return mem_cells: Number of memory cells.
+                          return None, if no data available.
+        """
+        cells = np.squeeze(
+            self.run_dc[
+                self.image_src, "image.cellId"].drop_empty_trains().ndarray()
+        )
         if cells.shape[0] == 0:
             return None
         maxcell = np.max(cells)
@@ -37,91 +62,77 @@ def get_num_cells(fname, loc, module):
         dists = [abs(o - maxcell) for o in options]
         return options[np.argmin(dists)]
 
+    def get_acq_rate(self) -> Optional[float]:
+        """Get the acquisition rate from said detector module.
 
-def get_acq_rate(fast_paths: Tuple[str, str, int],
-                 slow_paths: Optional[Tuple[str, str]] = ('', '')
-                 ) -> Optional[float]:
-    """Get the acquisition rate from said detector module.
+        If the data is available from the middlelayer FPGA_COMP device,
+        then it is retrieved from there.
+        If not, the rate is calculated from two different pulses time.
 
-    If the data is available from the middlelayer FPGA_COMP device, then it is
-    retrieved from there. If not, the rate is calculated from two different
-    pulses time.
+        The first entry is deliberately not used, as the detector just began
+        operating, and it might have skipped a train.
 
-    The first entry is deliberatly not used, as the detector just began
-    operating, and it might have skipped a train.
+        :return acq_rate: the acquisition rate.
+                          return None, if not available.
+        """
+        # Attempt to look for acquisition rate in slow data
+        rep_rate_src = (
+            self.ctrl_src, "bunchStructure.repetitionRate.value")
+
+        if (
+            rep_rate_src[0] in self.run_dc.all_sources and
+            rep_rate_src[1] in self.run_dc.keys_for_source(rep_rate_src[0])
+        ):
+            # The acquisition rate value is stored in a 1D array of type
+            # float.
+            # It is desired to loose precision here because the usage is
+            # about bucketing the rate for managing meta-data.
+
+            return round(float(self.run_dc[rep_rate_src].as_single_value()), 1)
+
+        train_pulses = np.squeeze(
+            self.run_dc[
+                self.image_src, "image.pulseId"
+            ].drop_empty_trains().train_from_index(0)[1]
+        )
 
-    :param slow_paths: in which file and h5 path to look for slow data.
-                       The first string is the filename with complete path,
-                       the second string is the key `karabo_id_control`
+        # Compute acquisition rate from fast data
+        diff = train_pulses[1] - train_pulses[0]
+        options = {8: 0.5, 4: 1.1, 2: 2.2, 1: 4.5}
+        return options.get(diff, None)
 
-    :param fast_paths: in which module file and h5 path to look for pulses.
-                       The first string is the filename with complete path,
-                       the second string is the module device name `karabo_id`,
-                       the third parameter is the module number, used to
-                       navigate through the h5 file structure.
+    def get_gain_setting(
+        self,
+        creation_time: "datetime.datetime",
+    ) -> Optional[int]:
+        """Retrieve Gain setting.
 
-    :return acq_rate: the acquisition rate.
-                      If not found in either files, return None.
-    """
-    # Attempt to look for acquisition rate in slow data
-    slow_data_file, karabo_id_control = slow_paths
-    slow_data_file = Path(slow_data_file)
-    if slow_data_file.is_file():
-        slow_data_path = f'CONTROL/{karabo_id_control}/MDL/FPGA_COMP/bunchStructure/repetitionRate/value'  # noqa
-        with h5py.File(slow_data_file, "r") as fin:
-            if slow_data_path in fin:
-                # The acquisition rate value is stored in a 1D array of type
-                # float. Use the 3rd value, arbitrarily chosen.
-                # It is desired to loose precision here because the usage is
-                # about bucketing the rate for managing meta-data.
-                return round(float(fin[slow_data_path][3]), 1)
-
-    # Compute acquisition rate from fast data
-    fast_data_file, karabo_id, module = fast_paths
-    fast_data_file = Path(fast_data_file)
-    if fast_data_file.is_file():
-        fast_data_path = f'INSTRUMENT/{karabo_id}/DET/{module}CH0:xtdf/image/pulseId'  # noqa
-        with h5py.File(fast_data_file, "r") as fin:
-            if fast_data_path in fin:
-                # pulses is of shape (NNNN, 1), of type uint8.
-                # Squeeze out the data, and subtract the 3rd entry from the 2nd
-                # to get a rate.
-                pulses = np.squeeze(fin[fast_data_path][1:3])
-                diff = pulses[1] - pulses[0]
-                options = {8: 0.5, 4: 1.1, 2: 2.2, 1: 4.5}
-                return options.get(diff, None)
-
-
-def get_gain_setting(fname: str, h5path_ctrl: str) -> int:
-    """Retrieve Gain setting.
-
-    If the data is available from the middlelayer FPGA_COMP device, then it is
-    retrieved from there.
-    If not, the setting is calculated off `setupr` and `patternTypeIndex`
-
-    gain-setting 1: setupr@dark=8, setupr@slopespc=40
-    gain-setting 0: setupr@dark=0, setupr@slopespc=32
-
-    patternTypeIndex 1: High-gain
-    patternTypeIndex 2: Medium-gain
-    patternTypeIndex 3: Low-gain
-    patternTypeIndex 4: SlopesPC
-
-    :param fname: path to file with control information
-    :param h5path_ctrl: path to control information inside the file
-    :return: gain setting
-    """
-    gain_path = f'{h5path_ctrl}/gain/value'
-    with h5py.File(fname, "r") as fin:
-        if gain_path in fin:
-            return fin[gain_path][0]
+        If the data is available from the middlelayer FPGA_COMP device,
+        then it is retrieved from there.
+        If not, the setting is calculated off `setupr` and `patternTypeIndex`
 
-        # Get the index at which the train is not zero.
-        train_id = fin["INDEX/trainId"][()]
-        idx = np.nonzero(train_id)[0][0]
+        gain-setting 1: setupr@dark=8, setupr@slopespc=40
+        gain-setting 0: setupr@dark=0, setupr@slopespc=32
 
-        setupr = fin[f'{h5path_ctrl}/setupr/value'][idx]
-        pattern_type_idx = fin[f'{h5path_ctrl}/patternTypeIndex/value'][idx]
+        patternTypeIndex 1: High-gain
+        patternTypeIndex 2: Medium-gain
+        patternTypeIndex 3: Low-gain
+        patternTypeIndex 4: SlopesPC
+
+        :return: gain setting.
+                 return 0, if not available.
+        """
+        # TODO: remove after fixing get_possible_conditions
+        if creation_time and creation_time.replace(tzinfo=None) < parser.parse('2020-01-31'):
+            print("Set gain-setting to None for runs taken before 2020-01-31")
+            return
+
+        if "gain.value" in self.run_dc.keys_for_source(self.ctrl_src):
+            return self.run_dc[self.ctrl_src, "gain.value"].as_single_value()
+
+        setupr = self.run_dc[self.ctrl_src, "setupr.value"].as_single_value()
+        pattern_type_idx = self.run_dc[
+            self.ctrl_src, "patternTypeIndex.value"].as_single_value()
 
         if (setupr == 0 and pattern_type_idx < 4) or (
                 setupr == 32 and pattern_type_idx == 4):
@@ -130,73 +141,133 @@ def get_gain_setting(fname: str, h5path_ctrl: str) -> int:
                 setupr == 40 and pattern_type_idx == 4):
             return 1
         else:
-            raise ValueError('Could not derive gain setting from setupr and patternTypeIndex')  # noqa
+            if self.raise_error:
+                raise ValueError(
+                    "Could not derive gain setting from"
+                    " setupr and patternTypeIndex"
+                )
 
+            print(
+                "WARNING: gain_setting is not available "
+                f"at source {self.ctrl_src}.\nSet gain_setting to 0.")
+            # TODO: why return 0 and not None?
+            return 0
 
-def get_gain_mode(fname: str, h5path_ctrl: str) -> AgipdGainMode:
-    """Returns the gain mode (adaptive or fixed) from slow data"""
+    def get_gain_mode(self) -> AgipdGainMode:
+        """Returns the gain mode (adaptive or fixed) from slow data"""
 
-    h5path_run = h5path_ctrl.replace("CONTROL/", "RUN/", 1)
-    h5path_gainmode = f'{h5path_run}/gainModeIndex/value'
-    with h5py.File(fname, "r") as fd:
-        if h5path_gainmode in fd:
-            return AgipdGainMode(fd[h5path_gainmode][0])
-    return AgipdGainMode.ADAPTIVE_GAIN
+        if (
+            self.ctrl_src in self.run_dc.all_sources and
+            "gainModeIndex.value" in self.run_dc.keys_for_source(
+                self.ctrl_src)
+        ):
+            return AgipdGainMode(int(
+                self.run_dc.get_run_value(
+                    self.ctrl_src, "gainModeIndex.value")))
 
+        return AgipdGainMode.ADAPTIVE_GAIN
 
-def get_bias_voltage(fname: str, karabo_id_control: str,
-                     module: Optional[int] = 0) -> int:
-    """Read the voltage information from the FPGA device of module 0.
+    def get_bias_voltage(
+        self,
+        karabo_id_control: str,
+        module: Optional[int] = 0
+    ) -> int:
+        """Read the voltage information from the FPGA device of module 0.
 
-    Different modules may operate at different voltages. In practice, they all
-    operate at the same voltage. As such, it is okay to read a single module's
-    value.
+        Different modules may operate at different voltages.
+        In practice, they all operate at the same voltage.
+        As such, it is okay to read a single module's value.
 
-    This value is read from slow data.
+        If the FPGA module source is not available, 300 will be returned.
+        300 is the default bias_voltage value before adding it to slow data.
 
-    If the file cannot be accessed, an OSError will be raised.
-    If the hdf5 path cannot be accessed, None will be returned.
+        :param karabo_id_control: The detector karabo id, for the control device.
+        :param module: defaults to module 0
+        :return: voltage, a uint16
+        """
+        voltage_src = (
+            f"{karabo_id_control}/FPGA/M_{module}",
+            "highVoltage.actual.value")
 
-    :param fname: path to slow data file with control information
-    :param karabo_id: The detector Karabo id, for creating the hdf5 path
-    :param module: defaults to module 0
-    :return: voltage, a uint16
-    """
-    voltage_path = f'/CONTROL/{karabo_id_control}/FPGA/M_{module}/highVoltage/actual/value'  # noqa
-    with h5py.File(fname, "r") as fin:
-        if voltage_path in fin:
-            return fin[voltage_path][0]
+        if (
+            voltage_src[0] in self.run_dc.all_sources and
+            voltage_src[1] in self.run_dc.keys_for_source(voltage_src[0])
+        ):
 
+            return self.run_dc[voltage_src].as_single_value(atol=1, reduce_by='max')
+        else:
+            print(
+                "WARNING: Unable to read bias_voltage from"
+                f" {voltage_src[0]}/{voltage_src[1].replace('.','/')} "
+                "Returning 300 as default bias voltage value."
+            )
+            return 300
 
-def get_integration_time(fname: str, h5path_ctrl: str) -> int:
-    """Read integration time from the FPGA device.
+    def get_integration_time(self) -> int:
+        """Read integration time from the FPGA device.
 
-    The integration time is specified as an integer number of clock
-    cycles each spanning ~9ns. The default (and legacy) value is 12.
+        The integration time is specified as an integer number of clock
+        cycles each spanning ~9ns. The default (and legacy) value is 12.
 
-    :param fname: path to file with control information
-    :param h5path_ctrl: path to control information inside the file
-    :return: integration time
-    """
-    h5path_run = h5path_ctrl.replace('CONTROL/', 'RUN/', 1)
-    h5path_time = f'{h5path_run}/integrationTime/value'
-    with h5py.File(fname, 'r') as fd:
-        if h5path_time in fd:
-            return int(fd[h5path_time][0])
-    return 12
+        :return: integration time
+        """
+        if (
+            self.ctrl_src in self.run_dc.all_sources and
+            'integrationTime.value' in self.run_dc.keys_for_source(
+                self.ctrl_src)
+        ):
+            return int(self.run_dc.get_run_value(
+                self.ctrl_src, 'integrationTime.value'))
+
+        return 12
+
+
+class CellSelection:
+    """Selection of detector memory cells (abstract class)"""
+    row_size = 32
+    ncell_max = 352
+    CM_NONE = 0
+    CM_PRESEL = 1
+    CM_FINSEL = 2
+
+    def get_cells_on_trains(
+        self, trains_sel: List[int], cm: int = 0
+    ) -> np.array:
+        """Returns mask of cells selected for processing
+
+        :param train_sel: list of a train ids selected for processing
+        :param cm: flag indicates the final selection or interim selection
+            for common-mode correction
+        """
+        raise NotImplementedError
+
+    def msg(self):
+        """Return log message on initialization"""
+        raise NotImplementedError
+
+    @staticmethod
+    def _sel_for_cm(flag, flag_cm, cm):
+        if cm == CellSelection.CM_NONE:
+            return flag
+        elif cm == CellSelection.CM_PRESEL:
+            return flag_cm
+        elif cm == CellSelection.CM_FINSEL:
+            return flag[flag_cm]
+        else:
+            raise ValueError("param 'cm' takes only 0,1,2")
 
 
 class AgipdCorrections:
 
     def __init__(
         self,
-        max_cells,
-        max_pulses,
-        h5_data_path="INSTRUMENT/SPB_DET_AGIPD1M-1/DET/{}CH0:xtdf/",
-        h5_index_path="INDEX/SPB_DET_AGIPD1M-1/DET/{}CH0:xtdf/",
+        max_cells: int,
+        cell_sel: CellSelection,
+        h5_data_path: str = "SPB_DET_AGIPD1M-1/DET/{}CH0:xtdf/",
+        h5_index_path: str = "SPB_DET_AGIPD1M-1/DET/{}CH0:xtdf/",
         corr_bools: Optional[dict] = None,
         gain_mode: AgipdGainMode = AgipdGainMode.ADAPTIVE_GAIN,
-        comp_threads=1,
+        comp_threads: int = 1,
         train_ids: Optional[np.ndarray] = None
     ):
         """
@@ -204,8 +275,8 @@ class AgipdCorrections:
 
         :param max_cells: maximum number of memory cells to handle, e.g. if
                          calibration constants only exist for a subset of cells
-        :param max_pulses: a range list of pulse indices used for
-            calibration and histogram. [start, end, step]
+        :param cell_sel: the CellSelection indicates cells selected for
+            calibration
         :param h5_data_path: path in HDF5 file which is prefixed to the
             image/data section
         :param h5_index_path: path in HDF5 file which is prefixed to the
@@ -218,7 +289,8 @@ class AgipdCorrections:
         The following example shows a typical use case:
         .. code-block:: python
 
-            agipd_corr = AgipdCorrections(max_cells, max_pulses,
+            cell_sel = CellRange(max_pulses, max_cells)
+            agipd_corr = AgipdCorrections(max_cells, cell_sel,
                                           h5_data_path=h5path,
                                           h5_index_path=h5path_idx,
                                           corr_bools=corr_bools)
@@ -256,8 +328,7 @@ class AgipdCorrections:
         self.comp_threads = comp_threads
         self.train_ids = np.array(train_ids) if train_ids is not None else None
 
-        self.start, self.last, self.step = self._validate_selected_pulses(
-            max_pulses, max_cells)
+        self.cell_sel = cell_sel
 
         # Correction parameters
         self.baseline_corr_noise_threshold = -1000
@@ -325,57 +396,63 @@ class AgipdCorrections:
         """
         module_idx = int(file_name.split('/')[-1].split('-')[2][-2:])
         agipd_base = self.h5_data_path.format(module_idx)
-        idx_base = self.h5_index_path.format(module_idx)
         data_dict = self.shared_dict[i_proc]
         data_dict['moduleIdx'][0] = module_idx
-        try:
-            f = h5py.File(file_name, "r")
-
-            (_, first_index, last_index,
-             _, valid_indices) = self.get_valid_image_idx(idx_base, f)
-
-            if len(valid_indices) == 0:
-                # If there's not a single valid index, exit early.
-                data_dict['nImg'][0] = 0
-                return 0
-
-            group = f[agipd_base]['image']
-            allcells = np.squeeze(group['cellId'])
-            allpulses = np.squeeze(group['pulseId'])
-
-            firange = self.gen_valid_range(first_index, last_index,
-                                           self.max_cells, allcells,
-                                           allpulses, valid_indices,
-                                           apply_sel_pulses)
-
-            if firange is None:
-                # gen_valid_range() returns None if there are no cells
-                # to correct, exit early.
-                data_dict['nImg'][0] = 0
-                return 0
-
-            n_img = firange.shape[0]
-            data_dict['nImg'][0] = n_img
-            if np.all(np.diff(firange) == 1):
-                # if firange consists of contiguous indices
-                # convert firange from fancy indexing to slicing
-                firange = slice(firange[0], firange[-1]+1)
-                raw_data = group['data'][firange]
-            else:
-                # Avoid very slow performance using fancing indexing,
-                # if firange consists of non-contiguous indices.
-                raw_data = group['data'][:][firange]
-            data_dict['data'][:n_img] = raw_data[:, 0]
-            data_dict['rawgain'][:n_img] = raw_data[:, 1]
-            data_dict['cellId'][:n_img] = allcells[firange]
-            data_dict['pulseId'][:n_img] = allpulses[firange]
-            data_dict['trainId'][:n_img] = np.squeeze(group['trainId'][:][firange])  # noqa
-        except Exception as e:
-            print(f'Error during reading data from file {file_name}: {e}')
-            print(f'Error traceback: {traceback.format_exc()}')
-            n_img = 0
+
+        h5_dc = H5File(file_name)
+
+        # Exclude trains without data.
+        im_dc = h5_dc.select(agipd_base, "image.*", require_all=True)
+
+        valid_train_ids = self.get_valid_image_idx(
+            im_dc[agipd_base, "image.trainId"])
+
+        if len(valid_train_ids) == 0:
+            # If there's not a single valid train, exit early.
+            print(f"WARNING: No valid trains for {im_dc.files} to process.")
             data_dict['nImg'][0] = 0
+            return 0
+
+        # store valid trains in shared memory
+        # valid_train_ids = train_ids[valid]
+        n_valid_trains = len(valid_train_ids)
+        data_dict["n_valid_trains"][0] = n_valid_trains
+        data_dict["valid_trains"][:n_valid_trains] = valid_train_ids
+
+        # get cell selection for the images in this file
+        cm = ( self.cell_sel.CM_NONE if apply_sel_pulses
+                else self.cell_sel.CM_PRESEL )
+
+        img_selected = self.cell_sel.get_cells_on_trains(
+            valid_train_ids, cm=cm)
+        data_dict["cm_presel"][0] = (cm == self.cell_sel.CM_PRESEL)
+
+        # Exclude non_valid trains from the selected data collection.
+        im_dc = im_dc.select_trains(by_id(valid_train_ids))
+
+        if "AGIPD500K" in agipd_base:
+            agipd_comp = components.AGIPD500K(im_dc)
+        else:
+            agipd_comp = components.AGIPD1M(im_dc)
 
+        kw = {
+            "unstack_pulses": False,
+            "pulses": np.nonzero(img_selected),
+        }
+
+        # [n_modules, n_imgs, 2, x, y]
+        raw_data = agipd_comp.get_array("image.data", **kw)[0]
+        n_img = raw_data.shape[0]
+
+        data_dict['nImg'][0] = n_img
+        data_dict['data'][:n_img] = raw_data[:, 0]
+        data_dict['rawgain'][:n_img] = raw_data[:, 1]
+        data_dict['cellId'][:n_img] = agipd_comp.get_array(
+            "image.cellId", **kw)[0]
+        data_dict['pulseId'][:n_img] = agipd_comp.get_array(
+            "image.pulseId", **kw)[0]
+        data_dict['trainId'][:n_img] = agipd_comp.get_array(
+            "image.trainId", **kw)[0]
         return n_img
 
     def write_file(self, i_proc, file_name, ofile_name):
@@ -388,7 +465,7 @@ class AgipdCorrections:
         """
 
         module_idx = int(file_name.split('/')[-1].split('-')[2][-2:])
-        agipd_base = self.h5_data_path.format(module_idx)
+        agipd_base = f'INSTRUMENT/{self.h5_data_path}/'.format(module_idx)
         idx_base = self.h5_index_path.format(module_idx)
         data_path = f'{agipd_base}/image'
         data_dict = self.shared_dict[i_proc]
@@ -490,6 +567,8 @@ class AgipdCorrections:
         fraction = self.cm_dark_fraction
         n_itr = self.cm_n_itr
         n_img = self.shared_dict[i_proc]['nImg'][0]
+        if n_img == 0:
+            return
         cell_id = self.shared_dict[i_proc]['cellId'][:n_img]
         train_id = self.shared_dict[i_proc]['trainId'][:n_img]
         cell_ids = cell_id[train_id == train_id[0]]
@@ -770,97 +849,55 @@ class AgipdCorrections:
         # Copy the data across into the existing shared-memory array
         mask[...] = msk[...]
 
-    def get_valid_image_idx(
-        self, idx_base: str, infile: str, raw_format_version: int = 2
-    ):
-        """Return the indices of valid data"""
-        if raw_format_version == 2:
-            idxtrains = np.squeeze(infile['/INDEX/trainId'])
-
-            # Check against train ID filter list, if any
-            if self.train_ids is not None:
-                valid = np.in1d(idxtrains, self.train_ids)
+    def get_valid_image_idx(self, im_dc: DataCollection) -> list:  # noqa
+        """Return a list of valid train ids.
 
-                if not valid.any():
-                    # Shortcut to avoid any further loading.
-                    return valid, 0, 0, idxtrains, np.zeros(0, dtype=np.int32)
-            else:
-                valid = np.ones_like(idxtrains, dtype=bool)
-
-            # Load count and offsets and filter for non-emtpy trains.
-            count = np.squeeze(infile[idx_base + "image/count"])
-            first = np.squeeze(infile[idx_base + "image/first"])
-            valid &= count != 0
-
-            # Validate that train indices values fall
-            # between medianTrain +- 1e4
-            medianTrain = np.median(idxtrains)
-            lowok = (idxtrains > medianTrain - 1e4)
-            highok = (idxtrains < medianTrain + 1e4)
-            valid &= lowok & highok
-
-            if not valid.any():
-                # Shortcut if no valid trains are left.
-                return valid, 0, 0, idxtrains, np.zeros(0, dtype=np.int32)
-
-            # Last index = last valid train + max. number of memory cells
-            last_index = int(first[valid][-1] + count[valid][-1])
-            first_index = int(first[valid][0])
-            # do actual validity filtering:
-            validc, validf = count[valid], first[valid]
-
-            # Creating an array of validated indices.
-            # If all indices were validated this array will be the same,
-            # as what is stored at /DET/image/trainId
-            valid_indices = np.concatenate(
-                [
-                    np.arange(validf[i], validf[i] + validc[i])
-                    for i in range(validf.size)
-                ],
-                axis=0,
-            )
-            valid_indices = np.squeeze(valid_indices).astype(np.int32)
-
-        elif raw_format_version == 1:
-            status = np.squeeze(infile[idx_base + "image/status"])
-            if np.count_nonzero(status != 0) == 0:
-                raise IOError(f"File {infile} has no valid counts")
-            last = np.squeeze(infile[idx_base + "image/last"])
-            first = np.squeeze(infile[idx_base + "image/first"])
-            valid = status != 0
-            last_index = int(last[status != 0][-1]) + 1
-            first_index = int(first[status != 0][0])
-
-            idxtrains = np.squeeze(infile["/INDEX/trainId"])
-            medianTrain = np.nanmedian(idxtrains)
-            lowok = (idxtrains > medianTrain - 1e4)
-            highok = (idxtrains < medianTrain + 1e4)
-            valid &= lowok & highok
-            valid_indices = None
+        Exclude non-valid train ids from past or future.
+        """
+        dc_trains = im_dc.train_ids
+        if len(dc_trains) == 0:
+            return 0
+        # Check against train ID filter list, if any
+        if self.train_ids is not None:
+            valid = np.in1d(dc_trains, self.train_ids)
         else:
-            raise AttributeError(
-                f"Not a known raw format version: {raw_format_version}")
+            valid = np.ones_like(dc_trains, dtype=bool)
 
-        return (valid, first_index, last_index, idxtrains, valid_indices)
+        # Train indices are of type=f32
+        # Validate that train indices values fall
+        # between medianTrain +- 1e4
+        medianTrain = np.nanmedian(dc_trains)
+        lowok = (dc_trains > medianTrain - 1e4)
+        highok = (dc_trains < medianTrain + 1e4)
+        valid &= lowok & highok
+
+        # exclude non valid trains
+        valid_trains = valid * dc_trains
+
+        return valid_trains[valid_trains!=0]
 
     def apply_selected_pulses(self, i_proc: int) -> int:
         """Select sharedmem data indices to correct based on selected
         pulses indices.
 
-
         :param i_proc: the index of sharedmem for a given file/module
         :return n_img: number of images to correct
         """
-
         data_dict = self.shared_dict[i_proc]
         n_img = data_dict['nImg'][0]
 
-        allpulses = data_dict['pulseId'][:n_img]
+        if not data_dict["cm_presel"][0]:
+            return n_img
+
+        ntrains = data_dict["n_valid_trains"][0]
+        train_ids = data_dict["valid_trains"][:ntrains]
 
         # Initializing can_calibrate array
-        can_calibrate = self.choose_selected_pulses(
-            allpulses,
-            can_calibrate=np.ones(shape=(len(allpulses),), dtype=np.bool))
+        can_calibrate = self.cell_sel.get_cells_on_trains(
+            train_ids, cm=self.cell_sel.CM_FINSEL
+        )
+        if np.all(can_calibrate):
+            return n_img
 
         # Only select data corresponding to selected pulses
         # and overwrite data in shared-memory leaving
@@ -884,143 +921,6 @@ class AgipdCorrections:
 
         return n_img
 
-    def _validate_selected_pulses(
-        self, max_pulses: List[int],
-        max_cells: int,
-    ) -> List[int]:
-        """Validate the selected pulses given from the notebook
-
-        Validate the selected range of pulses to correct raw data
-        of at least one image.
-
-        1) A pulse indices within one train can't be greater
-        than the operating memory cells.
-        2) Validate the order of the given raneg of pulses.
-        3) Raise value error if generate list of pulses is empty.
-
-        :param max_pulses: a list of at most 3 elements defining the
-        range of pulses to calibrate.
-        :param max_cells: operating memory cells.
-
-        :return adjusted_range: An adjusted range of pulse indices to correct.
-        """
-
-        # Validate selected pulses range:
-        # 1) A pulseId can't be greater than the operating memory cells.
-        pulses_range = [max_cells if p > max_cells else p for p in max_pulses]
-
-        if pulses_range != max_pulses:
-            print(
-                "WARNING: \"max_pulses\" list has been modified from "
-                f"{max_pulses} to {pulses_range}. As the number of "
-                "operating memory cells are less than the selected "
-                "maximum pulse."
-            )
-
-        if len(pulses_range) == 1:
-            adjusted_range = (0, pulses_range[0], 1)
-        elif len(pulses_range) == 2:
-            adjusted_range = (pulses_range[0], pulses_range[1], 1)
-        elif len(pulses_range) == 3:
-            adjusted_range = tuple(pulses_range)
-        else:
-            raise ValueError(
-                "ERROR: Wrong length for the list of pulses indices range. "
-                "Please check the given range for \"max_pulses\":"
-                f"{max_pulses}. \"max_pulses\" needs to be a list of "
-                "3 elements, [start, last, step]")
-
-        if adjusted_range[0] > adjusted_range[1]:
-            raise ValueError(
-                "ERROR: Pulse range start is greater than range end. "
-                "Please check the given range for \"max_pulses\":"
-                f"{max_pulses}. \"max_pulses\" needs to be a list of "
-                "3 elements, [start, last, step]")
-
-        if not np.all([isinstance(p, int) for p in max_pulses]):
-            raise TypeError(
-                    "ERROR: \"max_pulses\" elements needs to be integers:"
-                    f" {max_pulses}.")
-
-        print(
-            "A range of pulse indices is selected to correct:"
-            f" {pulses_range}")
-
-        return adjusted_range
-
-    def choose_selected_pulses(self, allpulses: np.array,
-                               can_calibrate: np.array) -> np.array:
-        """
-        Choose given selected pulse from pulseId array of
-        raw data. The selected pulses range is validated then
-        used to add a booleans in can_calibrate and guide the
-        later appliance.
-
-
-        :param allpulses: array of pulseIds from raw data
-        :param can_calibrate: array of booleans for indices to calibrate
-        :return can_calibrate: array of booleans to calibrate depending on
-                               selected pulses
-        """
-
-        # Check interesection between array of booleans and
-        # array of pulses to calibrate.
-        can_calibrate &= (
-            (allpulses >= allpulses[self.start]) &
-            (allpulses <= allpulses[self.last-1]) &
-            (((allpulses - allpulses[self.start]) % allpulses[self.step]) == 0)  # noqa
-        )
-        return can_calibrate
-
-    def gen_valid_range(self, first_index: int, last_index: int,
-                        max_cells: int, allcells: np.array,
-                        allpulses: np.array,
-                        valid_indices: Optional[np.array] = None,
-                        apply_sel_pulses: Optional[bool] = True
-                        ) -> np.array:
-        """ Validate the arrays of image.cellId and image.pulseId
-        to check presence of data and to avoid empty trains.
-
-        selected pulses range given from the AGIPD correction notebook
-        is taken into account if apply_sel_pulses is True
-
-        :param first_index: first index of image data
-        :param last_index: last index of image data
-        :param max_cells: number of memory cells to correct
-        :param allcells: array of image.cellsIds of raw data
-        :param allpulses: array of image.pulseIds of raw data
-        :param valid_indices: validated indices of image.data
-        :param apply_sel_pulses: A flag for applying selected pulses
-                                 after validation for correction
-        :return firange: An array of validated image.data
-                         indices to correct
-        # TODO: Ignore rows (32 pulse) of empty pulses even if
-        common-mode is selected
-        """
-
-        if valid_indices is not None:
-            allcells = allcells[valid_indices]
-            allpulses = allpulses[valid_indices]
-        else:
-            allcells = allcells[first_index:last_index]
-            allpulses = allpulses[first_index:last_index]
-
-        can_calibrate = (allcells < max_cells)
-
-        if not np.any(can_calibrate):
-            return
-
-        if apply_sel_pulses:
-            can_calibrate = self.choose_selected_pulses(
-                allpulses, can_calibrate=can_calibrate)
-        if valid_indices is None:
-            firange = np.arange(first_index, last_index)
-        else:
-            firange = valid_indices
-        firange = firange[can_calibrate]
-
-        return firange
-
     def copy_and_sanitize_non_cal_data(self, infile, outfile, agipd_base,
                                        idx_base, trains):
         """ Copy and sanitize data in `infile` that is not touched by
@@ -1466,7 +1366,6 @@ class AgipdCorrections:
         :param shape: Shape of expected data (nImg, x, y)
         :param n_cores_files: Number of files, handled in parallel
         """
-
         self.shared_dict = []
         for i in range(n_cores_files):
             self.shared_dict.append({})
@@ -1485,3 +1384,200 @@ class AgipdCorrections:
             self.shared_dict[i]["raw_data"] = sharedmem.empty(shape, dtype="f4")  # noqa
             self.shared_dict[i]["rel_corr"] = sharedmem.empty(shape, dtype="f4")  # noqa
             self.shared_dict[i]["t0_rgain"] = sharedmem.empty(shape, dtype="u2")  # noqa
+            # Valid trains
+            self.shared_dict[i]["cm_presel"] = sharedmem.empty(1, dtype="b")
+            self.shared_dict[i]["n_valid_trains"] = sharedmem.empty(1, dtype="i4")  # noqa
+            self.shared_dict[i]["valid_trains"] = sharedmem.empty(1024, dtype="u8")  # noqa
+
+
+def validate_selected_pulses(
+    max_pulses: List[int], max_cells: int
+) -> List[int]:
+    """Validate the selected pulses given from the notebook
+
+    Validate the selected range of pulses to correct raw data
+    of at least one image.
+
+    1) A pulse indices within one train can't be greater
+    than the operating memory cells.
+    2) Validate the order of the given raneg of pulses.
+    3) Raise value error if generate list of pulses is empty.
+
+    :param max_pulses: a list of at most 3 elements defining the
+    range of pulses to calibrate.
+    :param max_cells: operating memory cells.
+
+    :return adjusted_range: An adjusted range of pulse indices to correct.
+    """
+
+    # Validate selected pulses range:
+    # 1) A pulseId can't be greater than the operating memory cells.
+    pulses_range = [max_cells if p > max_cells else p for p in max_pulses]
+
+    if pulses_range != max_pulses:
+        print(
+            "WARNING: \"max_pulses\" list has been modified from "
+            f"{max_pulses} to {pulses_range}. As the number of "
+            "operating memory cells are less than the selected "
+            "maximum pulse."
+        )
+
+    if len(pulses_range) == 1:
+        adjusted_range = (0, pulses_range[0], 1)
+    elif len(pulses_range) == 2:
+        adjusted_range = (pulses_range[0], pulses_range[1], 1)
+    elif len(pulses_range) == 3:
+        adjusted_range = tuple(pulses_range)
+    else:
+        raise ValueError(
+            "ERROR: Wrong length for the list of pulses indices range. "
+            "Please check the given range for \"max_pulses\":"
+            f"{max_pulses}. \"max_pulses\" needs to be a list of "
+            "3 elements, [start, last, step]")
+
+    if adjusted_range[0] > adjusted_range[1]:
+        raise ValueError(
+            "ERROR: Pulse range start is greater than range end. "
+            "Please check the given range for \"max_pulses\":"
+            f"{max_pulses}. \"max_pulses\" needs to be a list of "
+            "3 elements, [start, last, step]")
+
+    if not np.all([isinstance(p, int) for p in max_pulses]):
+        raise TypeError(
+            "ERROR: \"max_pulses\" elements needs to be integers:"
+            f" {max_pulses}.")
+
+    print(
+        "A range of pulse indices is selected to correct:"
+        f" {pulses_range}")
+
+    return adjusted_range
+
+
+class CellRange(CellSelection):
+    """Selection of detector memory cells given as a range"""
+
+    def __init__(self, crange: List[int], max_cells: int):
+        """Initialize range selection
+
+        :param crange: range parameters of selected cells,
+            list up to 3 elements
+        :param max_cells: number of exposed cells
+        """
+        self.max_cells = max_cells
+        self.crange = validate_selected_pulses(crange, max_cells)
+        self.flag = np.zeros(self.max_cells, bool)
+        self.flag_cm = np.zeros(self.ncell_max, bool)
+        self.flag[slice(*crange)] = True
+        self.flag_cm[:self.max_cells] = self.flag
+        self.flag_cm = (self.flag_cm.reshape(-1, self.row_size).any(1)
+                        .repeat(self.row_size)[:self.max_cells])
+
+    def msg(self):
+        return (
+            f"Use range selection with crange={self.crange}, "
+            f"max_cells={self.max_cells}\n"
+            f"Frames per train: {self.flag.sum()}"
+        )
+
+    def get_cells_on_trains(
+        self, train_sel: List[int], cm: int = 0
+    ) -> np.array:
+        return np.tile(self._sel_for_cm(self.flag, self.flag_cm, cm),
+                       len(train_sel))
+
+class LitFrameSelection(CellSelection):
+    """Selection of detector memery cells indicated as lit frames
+    by the AgipdLitFrameFinder
+    """
+    def __init__(self, dev: str, dc: DataCollection, train_ids: List[int],
+                 crange: Optional[List[int]] = None,
+                 energy_threshold: Optional[float] = None):
+        """Initialize lit frame selection
+
+        :param dev: AgipdLitFrameFinder device name
+        :param dc: EXtra-data DataCollection of a run
+        :param train_ids: the list of selected trains
+        :param crange: range parameters of selected cells,
+            list up to 3 elements
+        """
+        # read AgipdLitFrameFinder data
+        self.dev = dev
+        self.crange = validate_selected_pulses(crange, self.ncell_max)
+        self.ethr = energy_threshold
+        intr_src = dev + ':output'
+        nfrm = dc[intr_src, 'data.nFrame'].ndarray()
+        litfrm_train_ids = dc[intr_src, 'data.trainId'].ndarray()
+        litfrm = dc[intr_src, 'data.nPulsePerFrame'].ndarray() > 0
+        if (energy_threshold != -1000 and
+                'data.energyPerFrame' in dc.keys_for_source(intr_src)):
+
+            litfrm &= (dc[intr_src, 'data.energyPerFrame'].ndarray()
+                       > energy_threshold)
+
+        # apply range selection
+        if crange is None:
+            cellsel = np.ones(self.ncell_max, bool)
+        else:
+            cellsel = np.zeros(self.ncell_max, bool)
+            cellsel[slice(*crange)] = True
+
+        # update train selection, removing trains without litframe data
+        if train_ids is None:
+            train_ids = np.unique(litfrm_train_ids)
+            ntrain = len(train_ids)
+        else:
+            ntrain_orig = len(train_ids)
+            train_ids, _, litfrm_train_idx = np.intersect1d(
+                train_ids, litfrm_train_ids, return_indices=True
+            )
+            litfrm = litfrm[litfrm_train_idx]
+            nfrm = nfrm[litfrm_train_idx]
+            ntrain = len(train_ids)
+            if ntrain != ntrain_orig:
+                print(f"Lit frame data missed for {ntrain_orig - ntrain}. "
+                      "Skip them.")
+
+        # initiate instance attributes
+        nimg = np.sum(nfrm)
+        self.train_ids = train_ids
+        self.count = nfrm
+        self.first = np.zeros(ntrain, int)
+        self.flag = np.zeros(nimg, bool)
+        self.flag_cm = np.zeros(nimg, bool)
+        self.min_sel = nfrm.max()
+        self.max_sel = 0
+        # compress frame pattern
+        i0 = 0
+        for i in range(ntrain):
+            n = nfrm[i]
+            iN = i0 + n
+            trnflg = litfrm[i] & cellsel
+            trnflg[n:] = False
+            self.first[i] = i0
+            self.flag[i0:iN] = trnflg[:n]
+            self.flag_cm[i0:iN] = (trnflg.reshape(-1, self.row_size).any(1)
+                                   .repeat(self.row_size)[:n])
+            nlit = trnflg[:n].sum()
+            self.min_sel = min(self.min_sel, nlit)
+            self.max_sel = max(self.max_sel, nlit)
+            i0 = iN
+
+    def msg(self):
+        srng = (f"{self.min_sel}" if self.min_sel == self.max_sel
+                else f"{self.min_sel}~{self.max_sel}")
+        return (
+            f"Use lit frame selection from {self.dev}, crange={self.crange}\n"
+            f"Frames per train: {srng}"
+        )
+
+    def get_cells_on_trains(
+        self, train_sel: List[int], cm: int = 0
+    ) -> np.array:
+        train_idx = np.flatnonzero(np.in1d(self.train_ids, train_sel))
+        i0 = self.first[train_idx]
+        iN = i0 + self.count[train_idx]
+        idx = np.concatenate(
+            [np.arange(i0[i], iN[i], dtype=int) for i in range(train_idx.size)]
+        )
+        return self._sel_for_cm(self.flag[idx], self.flag_cm[idx], cm)
diff --git a/src/cal_tools/enums.py b/src/cal_tools/enums.py
index e02d59994925aa5530d265ff944839fb160c9613..a9bfc0182f4972c41259992b1e31cd50b04fba7d 100644
--- a/src/cal_tools/enums.py
+++ b/src/cal_tools/enums.py
@@ -5,28 +5,28 @@ class BadPixels(IntFlag):
     """ The European XFEL Bad Pixel Encoding
     """
 
-    OFFSET_OUT_OF_THRESHOLD  = 0b000000000000000000001 # bit 1
-    NOISE_OUT_OF_THRESHOLD   = 0b000000000000000000010 # bit 2
-    OFFSET_NOISE_EVAL_ERROR  = 0b000000000000000000100 # bit 3
-    NO_DARK_DATA             = 0b000000000000000001000 # bit 4
-    CI_GAIN_OF_OF_THRESHOLD  = 0b000000000000000010000 # bit 5
-    CI_LINEAR_DEVIATION      = 0b000000000000000100000 # bit 6
-    CI_EVAL_ERROR            = 0b000000000000001000000 # bit 7
-    FF_GAIN_EVAL_ERROR       = 0b000000000000010000000 # bit 8
-    FF_GAIN_DEVIATION        = 0b000000000000100000000 # bit 9
-    FF_NO_ENTRIES            = 0b000000000001000000000 # bit 10
-    CI2_EVAL_ERROR           = 0b000000000010000000000 # bit 11
-    VALUE_IS_NAN             = 0b000000000100000000000 # bit 12
-    VALUE_OUT_OF_RANGE       = 0b000000001000000000000 # bit 13
-    GAIN_THRESHOLDING_ERROR  = 0b000000010000000000000 # bit 14
-    DATA_STD_IS_ZERO         = 0b000000100000000000000 # bit 15
-    ASIC_STD_BELOW_NOISE     = 0b000001000000000000000 # bit 16
-    INTERPOLATED             = 0b000010000000000000000 # bit 17
-    NOISY_ADC                = 0b000100000000000000000 # bit 18
-    OVERSCAN                 = 0b001000000000000000000 # bit 19
-    NON_SENSITIVE            = 0b010000000000000000000 # bit 20
-    NON_LIN_RESPONSE_REGION  = 0b100000000000000000000 # bit 21
-
+    OFFSET_OUT_OF_THRESHOLD  = 1 << 0 # bit 1
+    NOISE_OUT_OF_THRESHOLD   = 1 << 1 # bit 2
+    OFFSET_NOISE_EVAL_ERROR  = 1 << 2 # bit 3
+    NO_DARK_DATA             = 1 << 3 # bit 4
+    CI_GAIN_OF_OF_THRESHOLD  = 1 << 4 # bit 5
+    CI_LINEAR_DEVIATION      = 1 << 5 # bit 6
+    CI_EVAL_ERROR            = 1 << 6 # bit 7
+    FF_GAIN_EVAL_ERROR       = 1 << 7 # bit 8
+    FF_GAIN_DEVIATION        = 1 << 8 # bit 9
+    FF_NO_ENTRIES            = 1 << 9 # bit 10
+    CI2_EVAL_ERROR           = 1 << 10 # bit 11
+    VALUE_IS_NAN             = 1 << 11 # bit 12
+    VALUE_OUT_OF_RANGE       = 1 << 12 # bit 13
+    GAIN_THRESHOLDING_ERROR  = 1 << 13 # bit 14
+    DATA_STD_IS_ZERO         = 1 << 14 # bit 15
+    ASIC_STD_BELOW_NOISE     = 1 << 15 # bit 16
+    INTERPOLATED             = 1 << 16 # bit 17
+    NOISY_ADC                = 1 << 17 # bit 18
+    OVERSCAN                 = 1 << 18 # bit 19
+    NON_SENSITIVE            = 1 << 19 # bit 20
+    NON_LIN_RESPONSE_REGION  = 1 << 20 # bit 21
+    WRONG_GAIN_VALUE         = 1 << 21 # bit 22
 
 class SnowResolution(Enum):
     """ An Enum specifying how to resolve snowy pixels
@@ -42,3 +42,13 @@ class AgipdGainMode(IntEnum):
     FIXED_HIGH_GAIN = 1
     FIXED_MEDIUM_GAIN = 2
     FIXED_LOW_GAIN = 3
+
+
+class JungfrauSettings(Enum):
+    """ Jungfrau run gain settings."""
+    DYNAMIC_GAIN = "dynamicgain"
+    DYNAMIC_GAIN_HG0 = "dynamichg0"
+    FIX_GAIN_1 = "fixgain1"
+    FIX_GAIN_2 = "fixgain2"
+    FORCE_SWITCH_HG1 = "forceswitchg1"
+    FORCE_SWITCH_HG2 = "forceswitchg2"
diff --git a/src/cal_tools/jungfraulib.py b/src/cal_tools/jungfraulib.py
new file mode 100644
index 0000000000000000000000000000000000000000..402888dfa58068ccd173800834dfd165c206fcdc
--- /dev/null
+++ b/src/cal_tools/jungfraulib.py
@@ -0,0 +1,62 @@
+from typing import Tuple
+
+from cal_tools.enums import JungfrauSettings
+
+
+def _get_settings(run_dc, ctrl_src) -> str:
+    try:
+        return(JungfrauSettings(
+            run_dc.get_run_value(ctrl_src, "settings.value")))
+    except KeyError:
+        print(
+            "WARNING: \'settings.value\' key is not available at "
+            f"{run_dc.select(ctrl_src).files[0].filename},\n")
+        return
+
+
+class JungfrauCtrl():
+    def __init__(
+        self,
+        run_dc: "extra_data.DataCollection",  # noqa
+        ctrl_src: str,
+    ):
+
+        """Read slow data from RUN source.
+        :param run_dir: EXtra-data RunDirectory DataCollection object.
+        :param ctrl_src: Control source name for accessing run slow data.
+        """
+        self.run_dc = run_dc
+        self.ctrl_src = ctrl_src
+        self.run_settings = _get_settings(run_dc, ctrl_src)
+
+    def get_memory_cells(self) -> Tuple[int, int]:
+        n_storage_cells = int(self.run_dc.get_run_value(
+            self.ctrl_src, "storageCells.value")) + 1
+        storage_cell_st = int(self.run_dc.get_run_value(
+            self.ctrl_src, "storageCellStart.value"))
+        return n_storage_cells, storage_cell_st
+
+    def get_bias_voltage(self) -> int:
+        return(int(self.run_dc.get_run_value(
+            self.ctrl_src, "vHighVoltage.value")[0]))
+
+    def get_integration_time(self) -> float:
+        return(float(self.run_dc.get_run_value(
+            self.ctrl_src, "exposureTime.value")) * 1e6)
+
+    def get_gain_setting(self) -> int:
+        if self.run_settings == JungfrauSettings.DYNAMIC_GAIN_HG0:
+            gain_setting = 1
+        else:  # JungfrauSettings.DYNAMIC_GAIN
+            gain_setting = 0
+            if self.run_settings != JungfrauSettings.DYNAMIC_GAIN:
+                print(
+                    "WARNING: Setting gain_setting to 0, "
+                    "assuming that this is an old run.\n")
+        return gain_setting
+
+    def get_gain_mode(self) -> int:
+        if self.run_settings in [JungfrauSettings.FIX_GAIN_1, JungfrauSettings.FIX_GAIN_2]:  # noqa
+            return 1
+        else:
+            return 0
diff --git a/src/cal_tools/pnccdlib.py b/src/cal_tools/pnccdlib.py
index 785d289c359ffe6d3e62aa15bb6e41de56d4145d..98594427ed23b03224c856294d800235a200436a 100644
--- a/src/cal_tools/pnccdlib.py
+++ b/src/cal_tools/pnccdlib.py
@@ -1,56 +1,40 @@
-# Extracting slow data:
-import os
-import traceback
-from typing import Tuple
-
-import h5py
-
 VALID_GAINS = {
-    "a" : 1.0,
-    "b" : 4.0,
-    "c" : 16.0,
-    "d" : 64.0,
-    "e" : 256.0,
-    "f" : 340.0,
-    "g" : 512.0
+    "a": 1.0,
+    "b": 4.0,
+    "c": 16.0,
+    "d": 64.0,
+    "e": 256.0,
+    "f": 340.0,
+    "g": 512.0
 }
 
 
-def extract_slow_data(karabo_id: str, karabo_da_control: str,
-                      ctrl_fname: str, ctrl_path: str,
-                      mdl_ctrl_path: str,
-                      bias_voltage: float, gain: float,
-                      fix_temperature_top :float,
-                      fix_temperature_bot :float,
-                      ) -> Tuple[float, float, float, float]:
-    """
-    Extract slow data from given control paths.
-    """
-    try:
-        with h5py.File(ctrl_fname, "r") as f:
-            if bias_voltage == 0.:
-                bias_voltage = abs(f[os.path.join(mdl_ctrl_path,
-                                                  "DAQ_MPOD/u0voltage/value")][0])  # noqa
-            if gain == 0.1:
-                gain = f[os.path.join(mdl_ctrl_path,
-                                      "DAQ_GAIN/pNCCDGain/value")][0]
-            if fix_temperature_top == 0.:
-                fix_temperature_top = f[os.path.join(ctrl_path,
-                                                     "inputA/krdg/value")][0]
-            if fix_temperature_bot == 0.:
-                fix_temperature_bot = f[os.path.join(ctrl_path,
-                                                     "inputB/krdg/value")][0]
-    except IOError:
-        print("Error during reading slow data,"
-              " please check the given h5 path for the control parameters")
-        traceback.print_exc(limit=1)
-        print("bias voltage control h5path:",
-              os.path.join(mdl_ctrl_path, "DAQ_MPOD/u0voltage/value"))
-        print("gain control h5path:",
-              os.path.join(mdl_ctrl_path, "DAQ_GAIN/pNCCDGain/value"))
-        print("fix_temperature_top control h5path:",
-              os.path.join(ctrl_path, "inputA/krdg/value"))
-        print("fix_temperature_bot control h5path:",
-              os.path.join(ctrl_path, "inputB/krdg/value"))
+class PnccdCtrl():
+    def __init__(
+        self,
+        run_dc: "extra_data.DataCollection",  # noqa
+        karabo_id: str
+    ):
+        """ Extract control data from given control paths.
+        :param run_dc: run Extra-data data collection.
+        :param karabo_id: Detector identifier name.
+        """
+        self.run_dc = run_dc
+        self.ctrl_src = f"{karabo_id}/CTRL/TCTRL"
+        self.mdl_src_temp = f"{karabo_id}/MDL/{{}}"
+
+    def get_bias_voltage(self):
+        return(
+            abs(self.run_dc.get_run_value(
+                self.mdl_src_temp.format("DAQ_MPOD"), "u0voltage.value")))
+
+    def get_gain(self):
+        return(
+            self.run_dc.get_run_value(
+                self.mdl_src_temp.format("DAQ_GAIN"), "pNCCDGain.value"))
+
+    def get_fix_temperature_top(self):
+        return self.run_dc.get_run_value(self.ctrl_src, "inputA.krdg.value")
 
-    return bias_voltage, gain, fix_temperature_top, fix_temperature_bot
+    def get_fix_temperature_bot(self):
+        return self.run_dc.get_run_value(self.ctrl_src, "inputB.krdg.value")
diff --git a/src/cal_tools/step_timing.py b/src/cal_tools/step_timing.py
index 576f5284493108d923d40fb0565669229f5b689e..790021f0568f5e20d6111904f01c15847cb8fe0b 100644
--- a/src/cal_tools/step_timing.py
+++ b/src/cal_tools/step_timing.py
@@ -33,4 +33,4 @@ class StepTimer:
         """Show mean & std for each step"""
         for step, data in self.steps.items():
             data = np.asarray(data)
-            print(f'{step}: {data.mean():.01f} +- {data.std():.02f}s')
+            print(f'{step}: {data.mean():.01f} +- {data.std():.02f} s')
diff --git a/src/cal_tools/tools.py b/src/cal_tools/tools.py
index af1c02c04d3cb248727e4cf47a1efec577c4ac12..71fdba355ccbe047d958069b6d5567aeb437274a 100644
--- a/src/cal_tools/tools.py
+++ b/src/cal_tools/tools.py
@@ -1,8 +1,10 @@
 import datetime
 import json
 import re
+import zlib
 from collections import OrderedDict
 from glob import glob
+from multiprocessing.pool import ThreadPool
 from os import environ, listdir, path
 from os.path import isfile
 from pathlib import Path
@@ -56,6 +58,43 @@ def run_prop_seq_from_path(filename):
     return run, proposal, sequence
 
 
+def map_seq_files(
+    run_dc: "extra_data.DataCollection",
+    karabo_id: str,
+    karabo_das: List[str],
+    sequences: Optional[List[int]] = None,
+) -> Tuple[dict, int]:
+
+    """Using a DataCollection from extra-data to read
+    available sequence files.
+
+    Returns:
+        Dict: with karabo_das keys and the corresponding sequence files.
+        Int: for number of all sequence files for all karabo_das to process.
+    """
+
+    if sequences == [-1]:
+        sequences = None
+    if sequences is not None:
+        sequences = set(int(seq) for seq in sequences)
+
+    seq_fn_pat = re.compile(r".*-(?P<da>.*?)-S(?P<seq>.*?)\.h5")
+
+    mapped_files = {kda: [] for kda in karabo_das}
+    total_files = 0
+    for fn in run_dc.select(f"*{karabo_id}*").files:
+        fn = Path(fn.filename)
+        if (match := seq_fn_pat.match(fn.name)) is not None:
+            da = match.group("da")
+            if da in mapped_files and (
+                sequences is None or int(match.group("seq")) in sequences
+            ):
+                mapped_files[da].append(fn)
+                total_files += 1
+
+    return mapped_files, total_files
+
+
 def map_modules_from_folder(in_folder, run, path_template, karabo_da,
                             sequences=None):
     """
@@ -664,7 +703,12 @@ def send_to_db(db_module: str, karabo_id: str, constant, condition,
                     raise
             except Exception as e:
                 # TODO: refactor to use custom exception class
-                if "has already been taken" in str(e):
+                # Refactor error message for re-injecting an
+                # identical CCV to the database.
+                if all(s in str(e) for s in [
+                    "Error creating calibration_constant_version",
+                    "has already been taken",
+                ]):
                     print(
                         f"WARNING: {constant.name} for {db_module}"
                         " has already been injected with the same "
@@ -742,14 +786,19 @@ class CalibrationMetadata(dict):
     def __init__(self, output_dir: Union[Path, str], *args, new=False):
         dict.__init__(self, args)
         self._yaml_fn = Path(output_dir) / "calibration_metadata.yml"
-        if (not new) and self._yaml_fn.exists():
-            with self._yaml_fn.open("r") as fd:
-                data = yaml.safe_load(fd)
-            if isinstance(data, dict):
-                self.update(data)
+        if self._yaml_fn.exists():
+            if new:
+                # TODO: update after resolving this discussion
+                # https://git.xfel.eu/detectors/pycalibration/-/merge_requests/624  # noqa
+                self.save()
             else:
-                print(f"Warning: existing {self._yaml_fn} is malformed, "
-                      "will be overwritten")
+                with self._yaml_fn.open("r") as fd:
+                    data = yaml.safe_load(fd)
+                if isinstance(data, dict):
+                    self.update(data)
+                else:
+                    print(f"Warning: existing {self._yaml_fn} is malformed, "
+                           "will be overwritten")
 
     def save(self):
         with self._yaml_fn.open("w") as fd:
@@ -758,3 +807,41 @@ class CalibrationMetadata(dict):
     def save_copy(self, copy_dir: Path):
         with (copy_dir / self._yaml_fn.name).open("w") as fd:
             yaml.safe_dump(dict(self), fd)
+
+
+def write_compressed_frames(
+        arr: np.ndarray,
+        ofile: h5py.File,
+        dataset_path: str,
+        comp_threads: int = 1):
+    """Compress gain/mask frames in multiple threads, and save their data
+
+    This is significantly faster than letting HDF5 do the compression
+    in a single thread.
+    """
+
+    def _compress_frame(idx):
+        # Equivalent to the HDF5 'shuffle' filter: transpose bytes for better
+        # compression.
+        shuffled = np.ascontiguousarray(
+            arr[idx].view(np.uint8).reshape((-1, arr.itemsize)).transpose()
+        )
+        return idx, zlib.compress(shuffled, level=1)
+
+    # gain/mask compressed with gzip level 1, but not
+    # checksummed as we would have to implement this.
+    dataset = ofile.create_dataset(
+        dataset_path,
+        shape=arr.shape,
+        chunks=((1,) + arr.shape[1:]),
+        compression="gzip",
+        compression_opts=1,
+        shuffle=True,
+        dtype=arr.dtype,
+    )
+
+    with ThreadPool(comp_threads) as pool:
+        for i, compressed in pool.imap(_compress_frame, range(len(arr))):
+            # Each frame is 1 complete chunk
+            chunk_start = (i,) + (0,) * (dataset.ndim - 1)
+            dataset.id.write_direct_chunk(chunk_start, compressed)
diff --git a/src/xfel_calibrate/calibrate.py b/src/xfel_calibrate/calibrate.py
index b8735e7c2c154625dd2c23d9f010fa871999572a..b0cc47655858d49218bd484da032f1426d376932 100755
--- a/src/xfel_calibrate/calibrate.py
+++ b/src/xfel_calibrate/calibrate.py
@@ -35,13 +35,8 @@ from .nb_args import (
 from .settings import (
     default_report_path,
     finalize_time_limit,
-    free_nodes_cmd,
     launcher_command,
-    max_reserved,
-    preempt_nodes_cmd,
     python_path,
-    reservation,
-    reservation_char,
     sprof,
     temp_path,
     try_report_to_output,
@@ -225,7 +220,8 @@ def create_finalize_script(fmt_args, temp_path, job_list) -> str:
     return f_name
 
 
-def run_finalize(fmt_args, temp_path, job_list, sequential=False):
+def run_finalize(
+    fmt_args, temp_path, job_list, sequential=False, partition="exfel"):
     finalize_script = create_finalize_script(fmt_args, temp_path, job_list)
 
     cmd = []
@@ -238,7 +234,7 @@ def run_finalize(fmt_args, temp_path, job_list, sequential=False):
             '--open-mode=append',  # So we can see if it's preempted & requeued
             '--job-name', 'xfel-cal-finalize',
             '--time', finalize_time_limit,
-            '--partition', 'upex-middle',
+            '--partition', partition,
             "--dependency=afterany:" + ":".join(str(j) for j in job_list),
         ]
         print(" ".join(cmd))
@@ -276,36 +272,22 @@ def save_executed_command(run_tmp_path, version, argv):
 class SlurmOptions:
     def __init__(
             self, job_name=None, nice=None, mem=None, partition=None, reservation=None,
-            priority_group=2,
     ):
         self.job_name = job_name or 'xfel_calibrate'
         self.nice = nice
         self.mem = mem
         self.partition = partition
         self.reservation = reservation
-        self.priority_group = priority_group
 
     def get_partition_or_reservation(self) -> List[str]:
         """Return sbatch arguments to use a partition or reservation
 
         --reservation and --slurm-partition options have precedence.
-        Otherwise, if --priority is <=1, it will use a configured reservation
-        depending on how many nodes are currently free.
+        Otherwise, a default partition is used.
         """
         if self.reservation:
             return ['--reservation', self.reservation]
-        elif self.partition:
-            return ['--partition', self.partition]
-        elif self.priority_group <= 1:
-            auto_resvn = reservation_char if self.priority_group <= 0 else reservation
-            # Use a reservation if there aren't many general nodes available to us
-            free = int(check_output(free_nodes_cmd, shell=True).decode('utf8'))
-            preempt = int(check_output(preempt_nodes_cmd, shell=True).decode('utf8'))
-            if free + preempt < max_reserved:
-                return ['--reservation', auto_resvn]
-
-        # Fallback to using the configured partition (default: exfel)
-        return ['--partition', sprof]
+        return ['--partition', self.partition or sprof]
 
     def get_launcher_command(self, temp_path, after_ok=(), after_any=()) -> List[str]:
         """
@@ -610,7 +592,7 @@ def run(argv=None):
 
     title = title.rstrip()
 
-    run_uuid = f"t{datetime.now().strftime('%y%m%d_%H%M%S')}"
+    run_uuid = f"t{datetime.now().strftime('%y%m%d_%H%M%S.%f')}"
 
     # check if concurrency parameter is given and we run concurrently
     if concurrency_par is not None and not any(
@@ -837,7 +819,6 @@ def run(argv=None):
             mem=args['slurm_mem'],
             reservation=args['reservation'],
             partition=args['slurm_partition'],
-            priority_group=args['priority'],
         ))
         errors = False
 
@@ -849,7 +830,7 @@ def run(argv=None):
                 'report_to': report_to,
                 'in_folder': folder,
                 'request_time': request_time,
-                'submission_time': submission_time
+                'submission_time': submission_time,
                 }
 
     joblist.append(run_finalize(
@@ -857,6 +838,7 @@ def run(argv=None):
         temp_path=run_tmp_path,
         job_list=joblist,
         sequential=args["no_cluster_job"],
+        partition=args["slurm_partition"] or "exfel",
     ))
 
     if any(j is not None for j in joblist):
diff --git a/src/xfel_calibrate/nb_args.py b/src/xfel_calibrate/nb_args.py
index 90a86de20bd3aedd404fb7047e39560114d892da..e768d9f63a39b59865c28f4c5b0d6b39582edeaa 100644
--- a/src/xfel_calibrate/nb_args.py
+++ b/src/xfel_calibrate/nb_args.py
@@ -77,10 +77,6 @@ def make_initial_parser(**kwargs):
         "retrieved-constants will be copied to use for a new correction."
     ))
 
-    parser.add_argument('--priority', type=int, default=2,
-                        help="Priority of batch jobs. If priority<=1, reserved"
-                             " nodes become available.")
-
     parser.add_argument('--vector-figs', action="store_true", default=False,
                         help="Use vector graphics for figures in the report.")
 
@@ -428,7 +424,7 @@ def parse_argv_and_load_nb(argv) -> Tuple[Dict, NBDetails]:
               'location is sufficient for your '
               'needs.'.format(', '.join(not_reproducible_args)))
 
-        if not args['not_reproducible']:
+        if not args.not_reproducible:
             # If not explicitly specified that reproducibility may be
             # broken, remind the user and exit.
             print('To proceed, you can explicitly allow reproducibility to '
diff --git a/src/xfel_calibrate/notebooks.py b/src/xfel_calibrate/notebooks.py
index 4df3d7453d92e366867a0c96756195a381a47228..b53304395f83b40ab5c3cd9d0145b8aa2b274dd9 100644
--- a/src/xfel_calibrate/notebooks.py
+++ b/src/xfel_calibrate/notebooks.py
@@ -111,8 +111,9 @@ notebooks = {
         },
         "CORRECT": {
             "notebook": "notebooks/pnCCD/Correct_pnCCD_NBC.ipynb",
-            "concurrency": {"parameter": None,
-                            "default concurrency": None,
+            "concurrency": {"parameter": "sequences",
+                            "default concurrency": [-1],
+                            "use function": "balance_sequences",
                             "cluster cores": 32},
         },
     },
diff --git a/src/xfel_calibrate/repeat.py b/src/xfel_calibrate/repeat.py
index 95d41b1714c1ade7aba3c09b935a8fbe8a69ece0..12d4b30057022f657d9d44faaba7a121d5d8dcb0 100644
--- a/src/xfel_calibrate/repeat.py
+++ b/src/xfel_calibrate/repeat.py
@@ -93,6 +93,7 @@ def main(argv=None):
         temp_path=working_dir,
         job_list=joblist,
         sequential=args.no_cluster_job,
+        partition=args.slurm_partition
     ))
 
     if any(j is not None for j in joblist):
diff --git a/src/xfel_calibrate/settings.py b/src/xfel_calibrate/settings.py
index 989cd1022a47f85cb4d72af7b05754f416973ec8..15067c0a6ff9e211f1ca3d16280edc9fa524f2d3 100644
--- a/src/xfel_calibrate/settings.py
+++ b/src/xfel_calibrate/settings.py
@@ -21,11 +21,6 @@ sprof = os.environ.get("XFELCALSLURM", "exfel")  # TODO: https://git.xfel.eu/git
 launcher_command = "sbatch -t 24:00:00 --requeue --output {temp_path}/slurm-%j.out --parsable"
 free_nodes_cmd = "sinfo -p exfel -t idle -N --noheader | wc -l"
 preempt_nodes_cmd = "squeue -p all,grid --noheader | grep max-exfl | egrep -v 'max-exfl18[3-8]|max-exfl100|max-exflg' | wc -l"
-max_reserved = 8
-# "xcal" reservation value removed as ITDM
-# is giving xcal priority by default.
-reservation = ""
-reservation_char = "darks"
 
 # Time limit for the finalize job (creates PDF report & moves files)
 finalize_time_limit = "30:00"
diff --git a/tests/test_cal_tools.py b/tests/test_cal_tools.py
index 026e7512beaf9d0a5a6e455e8cd757b862b0462b..102f613b7b9f3b9bdb032f5ebe5ab218a8d1e098 100644
--- a/tests/test_cal_tools.py
+++ b/tests/test_cal_tools.py
@@ -7,7 +7,7 @@ import pytest
 import zmq
 from iCalibrationDB import Conditions, ConstantMetaData, Constants
 
-from cal_tools.agipdlib import AgipdCorrections
+from cal_tools.agipdlib import AgipdCorrections, CellRange
 from cal_tools.plotting import show_processed_modules
 from cal_tools.tools import (
     get_dir_creation_date,
@@ -344,7 +344,7 @@ def test_initialize_from_db():
 
     agipd_corr = AgipdCorrections(
         max_cells=MEM_CELLS,
-        max_pulses=[0, 500, 1])
+        cell_sel=CellRange([0, 500, 1], MEM_CELLS))
 
     agipd_corr.allocate_constants(
         modules=[0],
diff --git a/tests/test_webservice.py b/tests/test_webservice.py
index 6e1a335a634462ce548c05c2358de62609597ff9..b53b17f78c959ced708f1a4461da0819847f0c4a 100644
--- a/tests/test_webservice.py
+++ b/tests/test_webservice.py
@@ -13,6 +13,7 @@ from webservice.webservice import (  # noqa: import not at top of file
     parse_config,
     run_action,
     wait_on_transfer,
+    get_slurm_partition,
 )
 
 
@@ -151,3 +152,28 @@ async def test_run_action(mode, cmd, retcode, expected):
     webservice.webservice.run_proc_async = mock_run_proc_async
     ret = await run_action(job_db, cmd, mode, 1, 1, 1)
     assert ret.lower().startswith(expected)
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    'proposal_number, action, mock_proposal_status, expected_result',
+    [
+        (42, 'correct', 'A', 'upex-middle'),  # active
+        ('42', 'dark', 'R', 'upex-high'),  # active
+        (404, 'correct', 'FR', 'exfel'),  # finished and reviewed
+        (404, 'dark', 'CLO', 'exfel'),  # closed
+    ],
+)
+async def test_get_slurm_partition(proposal_number,
+                                   action,
+                                   mock_proposal_status,
+                                   expected_result):
+
+    response = mock.Mock()
+    response.status_code = 200
+    response.json = mock.Mock(return_value={'flg_beamtime_status': mock_proposal_status})
+    client = mock.Mock()
+    client.get_proposal_by_number_api = mock.Mock(
+        return_value=response)
+
+    ret = await get_slurm_partition(client, action, proposal_number)
+    assert ret == expected_result
diff --git a/tests/test_xfel_calibrate/conftest.py b/tests/test_xfel_calibrate/conftest.py
index 95bc80f2e7fb0e71bbae94921df10514bc26c658..7533d2f278397fc234f2c780aebdfc217622ade2 100644
--- a/tests/test_xfel_calibrate/conftest.py
+++ b/tests/test_xfel_calibrate/conftest.py
@@ -35,8 +35,6 @@ class FakeProcessCalibrate(FakeProcess):
         """
         super().__init__()
         #  Fake calls to slurm
-        self.register_subprocess(settings.free_nodes_cmd, stdout=["1"])
-        self.register_subprocess(settings.preempt_nodes_cmd, stdout=["1"])
         self.register_subprocess(
             ["sbatch", self.any()], stdout=["000000"]
         )
diff --git a/webservice/config/webservice.yaml b/webservice/config/webservice.yaml
index 26d74e07c8bc2643b9f179afe7c2deec48e8946d..7a764b0fad1807ae1cd3dcd72a41c1962a47e32e 100644
--- a/webservice/config/webservice.yaml
+++ b/webservice/config/webservice.yaml
@@ -35,7 +35,7 @@ correct:
     cmd : >-
         python -m xfel_calibrate.calibrate {detector} CORRECT
         --slurm-scheduling {sched_prio}
-        --slurm-partition upex-middle
+        --slurm-partition {partition}
         --request-time {request_time}
         --slurm-name {action}_{instrument}_{detector}_{cycle}_p{proposal}_{runs}
         --report-to /gpfs/exfel/exp/{instrument}/{cycle}/p{proposal}/usr/Reports/{runs}/{det_instance}_{action}_{proposal}_{runs}_{time_stamp}
@@ -50,7 +50,7 @@ dark:
         python -m xfel_calibrate.calibrate {detector} DARK
         --concurrency-par karabo_da
         --slurm-scheduling {sched_prio}
-        --slurm-partition upex-high
+        --slurm-partition {partition}
         --request-time {request_time}
         --slurm-name {action}_{instrument}_{detector}_{cycle}_p{proposal}_{runs}
         --report-to /gpfs/exfel/d/cal/caldb_store/xfel/reports/{instrument}/{det_instance}/{action}/{action}_{proposal}_{runs}_{time_stamp}
diff --git a/webservice/serve_overview.py b/webservice/serve_overview.py
index 919d48c153aa102b3dc0dfed40e5b0d5b7f8bd6a..ded99ef5897692c578cd24974e8564d1ebc94810 100644
--- a/webservice/serve_overview.py
+++ b/webservice/serve_overview.py
@@ -2,10 +2,10 @@ import argparse
 import glob
 import os
 import sqlite3
-from collections import OrderedDict
 from datetime import datetime, timezone
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
+from shutil import copyfileobj
 from subprocess import check_output
 from typing import Optional
 
@@ -20,34 +20,11 @@ except:
     from config import serve_overview as config
 
 
-
-class LimitedSizeDict(OrderedDict):
-    def __init__(self, *args, **kwds):
-        self.size_limit = kwds.pop("size_limit", None)
-        OrderedDict.__init__(self, *args, **kwds)
-        self._check_size_limit()
-
-    def __setitem__(self, key, value):
-        OrderedDict.__setitem__(self, key, value)
-        self._check_size_limit()
-
-    def _check_size_limit(self):
-        if self.size_limit is not None:
-            while len(self) > self.size_limit:
-                self.popitem(last=False)
-
-
-pdf_queue = LimitedSizeDict(size_limit=50)
-
-
 # HTTPRequestHandler class
 class RequestHandler(BaseHTTPRequestHandler):
     conf_was_init = False
 
     def init_config(self):
-
-        global cal_config
-
         self.total_jobs_cmd = config["shell-commands"]["total-jobs"]
         self.tail_log_cmd = config["shell-commands"]["tail-log"]
         self.cat_log_cmd = config["shell-commands"]["cat-log"]
@@ -57,98 +34,94 @@ class RequestHandler(BaseHTTPRequestHandler):
         for template, tfile in config["templates"].items():
             with open(tfile, "r") as tf:
                 self.templates[template] = tf.read()
-        global pdf_queue
-        self.pdf_queue = pdf_queue
+
         self.conf_was_init = True
 
+    def serve_css(self):
+        """Serve /serve_overview.css"""
+        self.send_response(200)
+        self.send_header('Content-type', 'text/css')
+        self.end_headers()
+        self.wfile.write(self.templates["css"].encode('utf-8'))
+
+    def serve_dark_overview(self):
+        # Send headers
+        self.send_response(200)
+        self.send_header('Content-type', 'text/html')
+        self.end_headers()
+        host = config["server-config"]["host"]
+        port = config["server-config"]["port"]
+        reports = {}
+        for instrument, detectors in cal_config['dark'].items():
+            reports[instrument] = {}
+            for detector in detectors:
+                tmpl = f'/gpfs/exfel/d/cal/caldb_store/xfel/reports/{instrument}/{detector}/dark/*pdf'
+                files = glob.glob(tmpl)
+                files.sort(key=os.path.getmtime, reverse=True)
+                file_info = []
+                for i, file in enumerate(files):
+                    if (len(file_info) % 2) == 0:
+                        bgcolor = 'EEEEEE'
+                    else:
+                        bgcolor = 'FFFFFF'
+                    time = os.stat(file).st_mtime
+                    d_time = datetime.fromtimestamp(time).replace(
+                        tzinfo=timezone.utc)
+                    s_time = d_time.strftime('%y-%m-%d %H:%M')
+                    file_info.append([file, s_time, bgcolor])
+
+                reports[instrument][detector] = file_info
+
+        tmpl = Template(self.templates["dark-overview"])
+        message = tmpl.render(reports=reports, host=host, port=port)
+
+        self.wfile.write(bytes(message, "utf8"))
+
+    def serve_file_from_gpfs(self):
+        """Serve a file from a path starting with /gpfs"""
+        if self.path.endswith(".html"):
+            mimetype = 'text/html'
+        elif self.path.endswith(".jpg"):
+            mimetype = 'image/jpg'
+        elif self.path.endswith(".gif"):
+            mimetype = 'image/gif'
+        elif self.path.endswith(".png"):
+            mimetype = 'image/png'
+        elif self.path.endswith(".pdf"):
+            mimetype = 'application/pdf'
+        elif self.path.endswith(".js"):
+            mimetype = 'application/javascript'
+        elif self.path.endswith(".css"):
+            mimetype = 'text/css'
+        else:
+            return self.send_error(404)
+
+        if os.path.isfile(self.path):
+            self.send_response(200)
+            self.send_header('Content-Length', str(os.stat(self.path).st_size))
+            self.send_header('Content-type', mimetype)
+            self.end_headers()
+            with open(self.path, "rb") as f:
+                copyfileobj(f, self.wfile)
+        else:
+            self.send_error(404)
+
     def do_GET(self):
 
         if not self.conf_was_init:
             self.init_config()
-        # Send response status code
-        self.send_response(200)
+
         if "/serve_overview.css" in self.path:
-            self.send_header('Content-type', 'text/css')
-            self.end_headers()
-            for s in self.templates["css"].split("\n"):
-                self.wfile.write(s.encode())
-            return
-
-        if "pdf?" in self.path:
-            puuid = self.path.split("?")[1]
-            fpath = self.pdf_queue.get(puuid, None)
-            if fpath is None:
-                return
-            self.send_header('Content-type', 'application/pdf')
-            self.end_headers()
-            with open(fpath, "rb") as f:
-                self.wfile.write(f.read())
-            return
+            return self.serve_css()
 
         if "dark?" in self.path:
-            # Send headers
-            self.send_header('Content-type', 'text/html')
-            self.end_headers()
-            host = config["server-config"]["host"]
-            port = config["server-config"]["port"]
-            reports = {}
-            for instrument, detectors in cal_config['dark'].items():
-                reports[instrument] = {}
-                for detector in detectors:
-                    tmpl = f'/gpfs/exfel/d/cal/caldb_store/xfel/reports/{instrument}/{detector}/dark/*pdf'
-                    files = glob.glob(tmpl)
-                    files.sort(key=os.path.getmtime, reverse=True)
-                    file_info = []
-                    for i, file in enumerate(files):
-                        if (len(file_info) % 2) == 0:
-                            bgcolor = 'EEEEEE'
-                        else:
-                            bgcolor = 'FFFFFF'
-                        time = os.stat(file).st_mtime
-                        d_time = datetime.fromtimestamp(time).replace(
-                            tzinfo=timezone.utc)
-                        s_time = d_time.strftime('%y-%m-%d %H:%M')
-                        file_info.append([file, s_time, bgcolor])
-
-                    reports[instrument][detector] = file_info
-
-            tmpl = Template(self.templates["dark-overview"])
-            message = tmpl.render(reports=reports, host=host, port=port)
-
-            self.wfile.write(bytes(message, "utf8"))
-            return
+            return self.serve_dark_overview()
 
         if "/gpfs" in self.path:
-            sendReply = False
-            if self.path.endswith(".html"):
-                mimetype = 'text/html'
-                sendReply = True
-            if self.path.endswith(".jpg"):
-                mimetype = 'image/jpg'
-                sendReply = True
-            if self.path.endswith(".gif"):
-                mimetype = 'image/gif'
-                sendReply = True
-            if self.path.endswith(".png"):
-                mimetype = 'image/png'
-                sendReply = True
-            if self.path.endswith(".pdf"):
-                mimetype = 'application/pdf'
-                sendReply = True
-            if self.path.endswith(".js"):
-                mimetype = 'application/javascript'
-                sendReply = True
-            if self.path.endswith(".css"):
-                mimetype = 'text/css'
-                sendReply = True
-
-            if sendReply == True and os.path.isfile(self.path):
-                with open(self.path, "rb") as f:
-                    self.send_header('Content-type', mimetype)
-                    self.end_headers()
-                    self.wfile.write(f.read())
-            return
+            return self.serve_file_from_gpfs()
 
+        # Send response status code
+        self.send_response(200)
 
         # Send headers
         self.send_header('Content-type', 'text/html')
diff --git a/webservice/webservice.py b/webservice/webservice.py
index 9244559681aef83cba0218ea8389d4e17158a81a..aa1274969d75b9d46060472d49412abf9738c73c 100644
--- a/webservice/webservice.py
+++ b/webservice/webservice.py
@@ -19,8 +19,9 @@ from datetime import datetime, timezone
 from pathlib import Path
 from subprocess import PIPE, run
 from threading import Thread
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
+import requests
 import yaml
 import zmq
 import zmq.asyncio
@@ -462,36 +463,6 @@ def update_job_db(config):
         time.sleep(time_interval)
 
 
-# Do not copy over files of big detectors when they are not being
-# corrected.
-copy_blocklist_pattern = re.compile(
-    r'\S*RAW-R\d{4}-(AGIPD|LPD|DSSC)\d{2}-S\d{5}.h5$')
-
-
-async def copy_untouched_files(files_set):
-    """ Copy those files which are not touched by the calibration
-    to the output directory.
-
-    :param file_list: The list of files to copy
-
-    Copying is done via an asyncio subprocess call
-    """
-    if files_set:
-        Path(next(iter(files_set)).replace("raw", "proc")).parent.mkdir(
-            parents=True, exist_ok=True)
-    else:
-        return
-
-    for f in files_set:
-        if copy_blocklist_pattern.match(f):
-            continue
-
-        of = f.replace("raw", "proc").replace("RAW", "CORR")
-        cmd = ["rsync", "-av", f, of]
-        await asyncio.subprocess.create_subprocess_shell(" ".join(cmd))
-        logging.info(f"Copying {f} to {of}")
-
-
 async def run_action(job_db, cmd, mode, proposal, run, rid) -> str:
     """Run action command (CORRECT or DARK).
 
@@ -657,6 +628,46 @@ def check_files(in_folder: str,
             files_exists = False
     return files_exists
 
+async def get_slurm_partition(mdc: MetadataClient,
+                            action: str,
+                            proposal_number: Union[int, str]) -> str:
+    """Check MyMDC for the proposal status and select the appropriate slurm
+       partition.
+
+       The partition is either upex-high (for darks) or upex-middle (for
+       corrections) if the proposal is 'R'eady or 'A'ctive.
+       In other cases, the jobs default to the exfel partition.
+
+       :param mdc: an authenticated MyMDC client
+       :param action: the type of action ('correct' or 'dark')
+       :param proposal_number: the proposal number
+       :return: 'exfel' on closed proposals,
+                'upex-high' if dark request on active proposals,
+                'upex-middle' if correct request on active proposals.
+    """
+    # See
+    # https://git.xfel.eu/ITDM/metadata_catalog/-/blob/develop/app/models/proposal.rb
+    # for possible proposals states.
+
+    loop = get_event_loop()
+    response = await shield(loop.run_in_executor(None,
+                                                 mdc.get_proposal_by_number_api,
+                                                 proposal_number))
+    if response.status_code != 200:
+        logging.error(f'Failed to check MDC for proposal "{proposal_number}" '
+                      'status. ASSUMING CLOSED')
+        logging.error(Errors.MDC_RESPONSE.format(response))
+
+    partition = 'exfel'
+    status = response.json().get('flg_beamtime_status', 'whoopsie')
+
+    if status in ['R', 'A']:
+        partition = 'upex-high' if action == 'dark' else 'upex-middle'
+
+    logging.debug(f"Using {partition} for {proposal_number} because {status}")
+
+    return partition
+
 
 async def update_darks_paths(mdc: MetadataClient, rid: int, in_path: str,
                              out_path: str, report_path: str):
@@ -728,6 +739,42 @@ async def update_mdc_status(mdc: MetadataClient, action: str,
         logging.error(Errors.MDC_RESPONSE.format(response))
 
 
+def _orca_passthrough(
+    proposal_number = None,
+    runs = None,
+    action = None,
+    route = "execute",
+    **kwargs,
+):
+    """ Passes through requests received by the webservice to Orca for use during
+    the transition to the new webservice
+
+    Due to network restrictions on Maxwell, this sends post requests to localhost on
+    port 42751 (a random port number), so either Orca should be running locally or a
+    ssh tunnel should be set up with `ssh -L 42751:localhost:42751 TARGET-SERVER -N &`
+    """
+    try:
+        base_url = "http://localhost"
+        port = "42751"
+
+        args = []
+
+        args.append(f"action={action}") if action else None
+        args.append(f"runs={','.join(str(r) for r in runs)}") if runs else None
+
+        args.extend([f"{k}={v}" for k, v in kwargs.items()])
+
+        url = f"{base_url}:{port}/{route}/{proposal_number}?{'&'.join(filter(None, args))}"
+    except Exception as e:
+        logging.warning("error building orca passthrough request", exc_info=True)
+        return None
+
+    try:
+        requests.post(url)
+    except Exception as e:
+        logging.error(f"orca post request error for url {url}", exc_info=True)
+
+
 class ActionsServer:
     def __init__(self, config, mode):
         self.config = config
@@ -838,8 +885,8 @@ class ActionsServer:
         :param instrument: is the instrument
         :param cycle: is the facility cycle
         :param proposal: is the proposal id
-        :param runnr: is the run number in integer form, e.g. without leading
-                     "r"
+        :param runnr: the run number in integer form, i.e. without leading "r"
+        :param priority: unused, retained for compatibility
 
         This will trigger a correction process to be launched for that run in
         the given cycle and proposal.
@@ -851,6 +898,12 @@ class ActionsServer:
             proposal = self._normalise_proposal_num(proposal)
             pconf_full = self.load_proposal_config(cycle, proposal)
 
+            _orca_passthrough(
+                proposal_number=proposal,
+                runs=[runnr],
+                route="execute",
+            )
+
             data_conf = pconf_full['data-mapping']
             if instrument in pconf_full['correct']:
                 pconf = pconf_full['correct'][instrument]
@@ -893,17 +946,13 @@ class ActionsServer:
 
                 # Prepare configs for all detectors in run
                 fl = glob.glob(f"{rpath}/*.h5")
-                corr_file_list = set()
-                copy_file_list = set(fl)
                 detectors = {}
                 for karabo_id in pconf:
                     dconfig = data_conf[karabo_id]
+
                     # check for files according to mapping in raw run dir.
                     if any(y in x for x in fl
                            for y in dconfig['karabo-da']):
-                        for karabo_da in dconfig['karabo-da']:
-                            tfl = glob.glob(f"{rpath}/*{karabo_da}*.h5")
-                            corr_file_list = corr_file_list.union(set(tfl))
                         thisconf = copy.copy(dconfig)
                         if isinstance(pconf[karabo_id], dict):
                             thisconf.update(copy.copy(pconf[karabo_id]))
@@ -911,18 +960,22 @@ class ActionsServer:
                         thisconf["out-folder"] = out_folder
                         thisconf["karabo-id"] = karabo_id
                         thisconf["run"] = runnr
-                        if priority:
-                            thisconf["priority"] = str(priority)
 
                         detectors[karabo_id] = thisconf
-                copy_file_set = copy_file_list.difference(corr_file_list)
-                asyncio.ensure_future(copy_untouched_files(copy_file_set))
+
             except Exception as corr_e:
                 logging.error("Error during correction", exc_info=corr_e)
                 await update_mdc_status(self.mdc, 'correct', rid,
                                         Errors.REQUEST_FAILED)
                 return
 
+            for karabo_id in list(detectors.keys()):
+                # Check for any detectors for which corrections are
+                # disabled manually.
+                if detectors[karabo_id].pop('disable-correct', False):
+                    logging.warning(f'Skipping disabled detector {karabo_id}')
+                    del detectors[karabo_id]
+
             if len(detectors) == 0:
                 msg = Errors.NOTHING_TO_DO.format(rpath)
                 logging.warning(msg)
@@ -974,6 +1027,14 @@ class ActionsServer:
 
             proposal = self._normalise_proposal_num(proposal)
 
+            _orca_passthrough(
+                proposal_number=proposal,
+                runs=runs,
+                action="dark",
+                route="execute",
+                karabo_id=karabo_id,
+            )
+
             pconf_full = self.load_proposal_config(cycle, proposal)
 
             data_conf = pconf_full['data-mapping']
@@ -1017,6 +1078,9 @@ class ActionsServer:
             triple = any(det in karabo_id for det in
                          ["LPD", "AGIPD", "JUNGFRAU", "JF", "JNGFR"])
 
+            # This fails silently if the hardcoded strings above are
+            # ever changed (triple = False) but the underlying notebook
+            # still expects run-high/run-med/run-low.
             if triple and len(runs) == 1:
                 runs_dict = {'run-high': runs[0],
                              'run-med': '0',
@@ -1032,6 +1096,9 @@ class ActionsServer:
             # is not migrated, thus skipping some validation here.
             thisconf = copy.copy(data_conf[karabo_id])
 
+            # Pop internal key to avoid propagation to xfel-calibrate.
+            thisconf.pop('disable-correct', None)
+
             if (karabo_id in pconf
                     and isinstance(pconf[karabo_id], dict)):
                 thisconf.update(copy.copy(pconf[karabo_id]))
@@ -1140,6 +1207,9 @@ class ActionsServer:
     ) -> Tuple[str, List[str]]:
         report = []
         ret = []
+
+        partition = await get_slurm_partition(self.mdc, action, proposal)
+
         # run xfel_calibrate
         for karabo_id, dconfig in detectors.items():
             detector = dconfig['detector-type']
@@ -1147,6 +1217,7 @@ class ActionsServer:
             cmd = self.config[action]['cmd'].format(
                 detector=detector,
                 sched_prio=str(self.config[action]['sched-prio']),
+                partition=partition,
                 action=action, instrument=instrument,
                 cycle=cycle, proposal=proposal,
                 runs="_".join([f"r{r}" for r in run_nrs]),