Fix frame selection for the arrays in the shared memory

src/cal_tools/agipdlib.py

@@ -29,7 +29,6 @@ from cal_tools.h5_copy_except import h5_copy_except_paths
 from cal_tools.tools import get_from_db
 
 
-
 class AgipdCtrl:
     def __init__(
         self,
@@ -262,7 +261,6 @@ class AgipdCtrl:
                 raise ValueError(error)
             return default_voltage
 
-
     def get_integration_time(self) -> int:
         """Read integration time from the FPGA device.
 
@@ -997,27 +995,62 @@ class AgipdCorrections:
         if np.all(can_calibrate):
             return n_img
 
+        # Get selected number of images based on
+        # selected pulses to correct
+        n_img_sel = np.count_nonzero(can_calibrate)
+
         # Only select data corresponding to selected pulses
         # and overwrite data in shared-memory leaving
         # the required indices to correct
-        data = data_dict['data'][:n_img][can_calibrate]
-        rawgain = data_dict['rawgain'][:n_img][can_calibrate]
-        cellId = data_dict['cellId'][:n_img][can_calibrate]
-        pulseId = data_dict['pulseId'][:n_img][can_calibrate]
-        trainId = data_dict['trainId'][:n_img][can_calibrate]
+        # - data
+        arr = data_dict['data'][:n_img][can_calibrate]
+        data_dict['data'][:n_img_sel] = arr
+        # - rawgain`
+        arr = data_dict['rawgain'][:n_img][can_calibrate]
+        data_dict['rawgain'][:n_img_sel] = arr
+        # - cellId
+        arr = data_dict['cellId'][:n_img][can_calibrate]
+        data_dict['cellId'][:n_img_sel] = arr
+        # - pulseId
+        arr = data_dict['pulseId'][:n_img][can_calibrate]
+        data_dict['pulseId'][:n_img_sel] = arr
+        # - trainId
+        arr = data_dict['trainId'][:n_img][can_calibrate]
+        data_dict['trainId'][:n_img_sel] = arr
+
+        # arrays generated after reading and initial selection
+        # - gain
+        arr = data_dict['gain'][:n_img][can_calibrate]
+        data_dict['gain'][:n_img_sel] = arr
+
+        is_adaptive = self.gain_mode is AgipdGainMode.ADAPTIVE_GAIN
+        melt_snow = self.corr_bools.get("melt_snow")
+        if (is_adaptive and melt_snow):
+            # - t0_rgain
+            arr = data_dict['t0_rgain'][:n_img][can_calibrate]
+            data_dict['t0_rgain'][:n_img_sel] = arr
+            # - raw_gain
+            arr = data_dict['raw_data'][:n_img][can_calibrate]
+            data_dict['raw_data'][:n_img_sel] = arr
 
-        # Overwrite selected number of images based on
-        # selected pulses to correct
-        n_img = np.count_nonzero(can_calibrate)
+        # if baseline correction was not requested
+        # `msk` and `rel_corr` will still be empty shared_mem arrays
+        if any(self.blc_bools):
+            # - blShift
+            arr = data_dict['blShift'][:n_img][can_calibrate]
+            data_dict['blShift'][:n_img_sel] = arr
+            # - msk
+            arr = data_dict['msk'][:n_img][can_calibrate]
+            data_dict['msk'][:n_img_sel] = arr
+            # - rel_corr
+            if hasattr(self, "rel_gain"):
+                arr = data_dict['rel_corr'][:n_img][can_calibrate]
+                data_dict['rel_corr'][:n_img_sel] = arr
 
-        data_dict['nImg'][0] = n_img
-        data_dict['data'][: n_img] = data
-        data_dict['rawgain'][:n_img] = rawgain
-        data_dict['cellId'][:n_img] = cellId
-        data_dict['pulseId'][:n_img] = pulseId
-        data_dict['trainId'][:n_img] = trainId
+        # Overwrite number of images
+        data_dict['nImg'][0] = n_img_sel
 
-        return n_img
+        return n_img_sel
 
     def copy_and_sanitize_non_cal_data(self, infile, outfile, agipd_base,
                                        idx_base, trains):