diff --git a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb index 59ee32e21946eed5f28607d69a564994c342dbf6..95c88f91ab9d5317f6fb5e1991898f1da336a906 100644 --- a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb +++ b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb @@ -101,6 +101,7 @@ "outputs": [], "source": [ "import itertools\n", + "import os\n", "import math\n", "import multiprocessing\n", "import re\n", @@ -422,7 +423,8 @@ " h5_data_path=h5path,\n", " h5_index_path=h5path_idx,\n", " corr_bools=corr_bools,\n", - " gain_mode=gain_mode\n", + " gain_mode=gain_mode,\n", + " comp_threads=os.cpu_count() / n_cores_files,\n", ")\n", "\n", "agipd_corr.baseline_corr_noise_threshold = -blc_noise_threshold\n", @@ -1150,4 +1152,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/src/cal_tools/agipdlib.py b/src/cal_tools/agipdlib.py index e72710f0dfe11ebb64eefcf818b22c16ce911412..8f8896ab5e013d9e8cef8d2ec9b1fe4dc8857fbf 100644 --- a/src/cal_tools/agipdlib.py +++ b/src/cal_tools/agipdlib.py @@ -175,6 +175,7 @@ class AgipdCorrections: h5_index_path="INDEX/SPB_DET_AGIPD1M-1/DET/{}CH0:xtdf/", corr_bools: Optional[dict] = None, gain_mode: AgipdGainMode = AgipdGainMode.ADAPTIVE_GAIN, + comp_threads=1, ): """ Initialize an AgipdCorrections Class @@ -189,6 +190,7 @@ class AgipdCorrections: index section :param corr_bools: A dict with all of the correction booleans requested or available + :param comp_threads: Number of threads to use for compressing gain/mask The following example shows a typical use case: .. code-block:: python @@ -231,6 +233,7 @@ class AgipdCorrections: if not (len(max_pulses) == 1 and max_pulses[0] == 0) else max_pulses # noqa self.max_cells = max_cells self.gain_mode = gain_mode + self.comp_threads = comp_threads # Correction parameters self.baseline_corr_noise_threshold = -1000 @@ -402,7 +405,7 @@ class AgipdCorrections: else: image_grp[field][:] = data_dict[field][:n_img] - def _write_compressed_frames(self, dataset, arr, threads=1): + def _write_compressed_frames(self, dataset, arr): def _compress_frame(i): # Equivalent to the HDF5 'shuffle' filter: transpose bytes for better # compression. @@ -411,7 +414,7 @@ class AgipdCorrections: ) return i, zlib.compress(shuffled, level=1) - with ThreadPool(threads) as pool: + with ThreadPool(self.comp_threads) as pool: for i, compressed in pool.imap(_compress_frame, range(len(arr))): # Each frame is 1 complete chunk dataset.id.write_direct_chunk((i, 0, 0), compressed)