From f51aefb8797d3ac05cd90a0b2ffd629555be9c08 Mon Sep 17 00:00:00 2001
From: Thomas Kluyver <thomas@kluyver.me.uk>
Date: Fri, 28 May 2021 14:46:25 +0100
Subject: [PATCH] Configure compression threads from the correction notebook
---
notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb | 6 ++++--
src/cal_tools/agipdlib.py | 7 +++++--
2 files changed, 9 insertions(+), 4 deletions(-)
diff --git a/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb b/notebooks/AGIPD/AGIPD_Correct_and_Verify.ipynb
index 59ee32e21..95c88f91a 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 e72710f0d..8f8896ab5 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)
--
GitLab