CookieboxCalibration: added flag to parallelize over etofs (#513)

* Expose flag to parallelize over tofs: good for machine with large number of cores only.
* Updated change log.

---------

Co-authored-by: Danilo Ferreira de Lima <danilo.enoque.ferreira.de.lima@xfel.de>

Author: daniloefl

docs/changelog.md

@@ -120,6 +120,8 @@ Changed:
   parallelization to the `AdqRawChannel` class (!509).
 - [TOFAnalogResponse][extra.applications.TOFAnalogResponse] averages analog pulses using thresholding
   from `AdqRawChannel` to clean up the data (!511).
+- [CookieboxCalibration][extra.applications.CookieboxCalibration] use
+  processes to parallelize over eTOFs (!513).
 
 ## [2025.1]
 

src/extra/applications/cookiebox.py

@@ -475,7 +475,8 @@ def setup(self,
               scan: Scan,
               xgm: XGM,
               tof_response: Dict[int, TOFAnalogResponse]=None,
-              parallel=None
+              parallel=False,
+              parallel_over_tofs=None,
               ):
         """
         Derive calibrations.
@@ -493,6 +494,7 @@ def setup(self,
                For example: `XGM(run, "SQS_DIAG1_XGMD/XGM/DOOCS")`
           tof_response: The response function object for deconvolution if that is desired.
           parallel: Whether to paralellize data reading.
+          parallel_over_tofs: Whether to parallelize over eTOFs.
         """
         # base properties
         self._run = run
@@ -524,7 +526,7 @@ def setup(self,
         self.update_metadata()
 
         # find RoI if needed
-        self.update_roi(parallel)
+        self.update_roi(parallel, parallel_over_tofs)
 
         # find where the peaks are per energy in each Tof
         self.update_fit_result()
@@ -688,7 +690,7 @@ def update_tof_settings(self):
             self.kwargs_adq[tof_id]["name"] = self._tof[tof_id].name
         self.mask = {tof_id: True for tof_id in self.kwargs_adq.keys()}
 
-    def update_roi(self, parallel=None):
+    def update_roi(self, parallel=False, parallel_over_tofs=16):
         """
         Given calibrated data, apply a selection and find RoI if needed.
 
@@ -697,7 +699,7 @@ def update_roi(self, parallel=None):
         """
         # average data for each energy slice
         logging.info("Reading calibration data ... (this takes a while)")
-        self.select_calibration_data(parallel)
+        self.select_calibration_data(parallel, parallel_over_tofs=parallel_over_tofs)
         # find RoI if needed
         for tof_id in self.kwargs_adq.keys():
             if (self.auger_start_roi[tof_id] is None
@@ -728,7 +730,7 @@ def update_calibration(self):
     def fast_response_correction(self, x, tof_id):
         return self._tof_response[tof_id].apply(x.fillna(0.0), method="nn_matrix", n_iter=100, nonneg=True)
 
-    def select_calibration_data(self, parallel=None):
+    def select_calibration_data(self, parallel=False, parallel_over_tofs=None):
         """
         Select data for calibration.
         """
@@ -751,15 +753,27 @@ def select_calibration_data(self, parallel=None):
                      correction_fn=correction_fn,
                      parallel=parallel,
                      )
-        itr_gen = list(itertools.product(tof_ids, energy_ids))
-        data_gen = map(fn, itr_gen)
-        # organize it all in a numpy array
-        for (d, x), (tof_id, energy_id) in zip(data_gen, itr_gen):
-            data[tof_id] += [d]
-            mean_xgm[tof_id] += [x]
-        for tof_id in tof_ids:
-            data[tof_id] = np.stack(data[tof_id], axis=0)
-            mean_xgm[tof_id] = np.stack(mean_xgm[tof_id], axis=0)
+        if parallel_over_tofs is not None:
+            with ProcessPoolExecutor(max_workers=parallel_over_tofs) as p:
+                itr_gen = list(itertools.product(tof_ids, energy_ids))
+                data_gen = p.map(fn, itr_gen)
+                # organize it all in a numpy array
+                for (d, x), (tof_id, energy_id) in zip(data_gen, itr_gen):
+                    data[tof_id] += [d]
+                    mean_xgm[tof_id] += [x]
+                for tof_id in tof_ids:
+                    data[tof_id] = np.stack(data[tof_id], axis=0)
+                    mean_xgm[tof_id] = np.stack(mean_xgm[tof_id], axis=0)
+        else:
+            itr_gen = list(itertools.product(tof_ids, energy_ids))
+            data_gen = map(fn, itr_gen)
+            # organize it all in a numpy array
+            for (d, x), (tof_id, energy_id) in zip(data_gen, itr_gen):
+                data[tof_id] += [d]
+                mean_xgm[tof_id] += [x]
+            for tof_id in tof_ids:
+                data[tof_id] = np.stack(data[tof_id], axis=0)
+                mean_xgm[tof_id] = np.stack(mean_xgm[tof_id], axis=0)
 
         self.calibration_data = data
         self.calibration_mean_xgm = mean_xgm

tests/test_applications_cookiebox.py

@@ -441,6 +441,58 @@ def test_no_parallel(mock_sqs_etof_calibration_run, tmp_path, mock_etof_mono_ene
         # check how well it matches
         assert np.allclose(ts, ts_true, rtol=1e-2, atol=1e-2)
 
+# tests data reading parallelizing over etofs
+def test_parallel_tofs(mock_sqs_etof_calibration_run, tmp_path, mock_etof_mono_energies, mock_etof_calibration_constants):
+    # same as above, but tests only if a crash happens in `calc_mean`
+    # somehow parallelization means that `calc_mean` is not shown in the coverage
+    pulse_timing = 'SQS_RR_UTC/TSYS/TIMESERVER'
+    monochromator_energy = 'SA3_XTD10_MONO/MDL/PHOTON_ENERGY'
+    digitizer = 'SQS_DIGITIZER_UTC4/ADC/1:network'
+    digitizer_control = 'SQS_DIGITIZER_UTC4/ADC/1'
+    pulse_energy = 'SQS_DIAG1_XGMD/XGM/DOOCS'
+    mock_sqs_etof_calibration_run = mock_sqs_etof_calibration_run.select([pulse_timing,
+                                  digitizer, digitizer_control,
+                                  pulse_energy, f"{pulse_energy}:output",
+                                  monochromator_energy], require_all=True).select_trains(np.s_[10:])
+    channel_name = "1_A"
+    tof_ids = [0]
+    tof_channel = {}
+    tof_channel[0] = AdqRawChannel(mock_sqs_etof_calibration_run,
+                                   channel_name,
+                                   digitizer=digitizer,
+                                   first_pulse_offset=1000)
+    scan = Scan(mock_sqs_etof_calibration_run[monochromator_energy, "actualEnergy"], resolution=2)
+    energy_axis = np.linspace(965, 1070, 160)
+    xgm = XGM(mock_sqs_etof_calibration_run, pulse_energy)
+    cal = CookieboxCalibration(
+                    auger_start_roi=1,
+                    start_roi=75,
+                    stop_roi=320,
+    )
+    cal.setup(run=mock_sqs_etof_calibration_run, energy_axis=energy_axis, tof_settings=tof_channel,
+              xgm=xgm,
+              scan=scan,
+              parallel=False,
+              parallel_over_tofs=2,
+              )
+
+    correct_energies = np.unique(mock_etof_mono_energies)
+    correct_constants = np.array(mock_etof_calibration_constants)
+    for tof_id in tof_ids:
+        assert np.allclose(cal.tof_fit_result[tof_id].energy, correct_energies, rtol=1e-2, atol=1e-2)
+
+        energy = correct_energies
+
+        # get calibration curve
+        c, e0, t0 = cal.model_params[tof_id]
+        ts = t0 + np.sqrt(c/(energy - e0))
+
+        c_true, e0_true, t0_true = correct_constants
+        ts_true = t0_true + np.sqrt(c_true/(energy - e0_true))
+
+        # check how well it matches
+        assert np.allclose(ts, ts_true, rtol=1e-2, atol=1e-2)
+
 def test_deconvolve(mock_sqs_etof_calibration_run, tmp_path):
     # use mock data and do the same as before, but with deconvolution
     # it should improve resolution, but lead to the same calibration constants