From 4b8c684c88030d90e26416321d9be16f1c10fc68 Mon Sep 17 00:00:00 2001
From: Egor Sobolev <egor.sobolev@xfel.eu>
Date: Thu, 23 Mar 2023 19:35:33 +0100
Subject: [PATCH] Select frames using cellId instead of position in the
 lit-frame finder mode

---
 src/cal_tools/agipdlib.py | 34 +++++++++++-----------------------
 1 file changed, 11 insertions(+), 23 deletions(-)

diff --git a/src/cal_tools/agipdlib.py b/src/cal_tools/agipdlib.py
index c7191edcf..87e898cbf 100644
--- a/src/cal_tools/agipdlib.py
+++ b/src/cal_tools/agipdlib.py
@@ -322,17 +322,6 @@ class CellSelection:
         """
         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:
 
@@ -520,7 +509,6 @@ class AgipdCorrections:
         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
-        #data_dict["nimg_in_trains"][:n_valid_trains] = nimg_in_trains
 
         if "AGIPD500K" in agipd_base:
             agipd_comp = components.AGIPD500K(im_dc)
@@ -1013,9 +1001,10 @@ class AgipdCorrections:
         cellid = data_dict["cellId"][:n_img]
 
         # Initializing can_calibrate array
-        can_calibrate, _ = self.cell_sel.get_cells_on_trains(
+        can_calibrate, nimg_in_trains = self.cell_sel.get_cells_on_trains(
             train_ids, nimg_in_trains, cellid, cm=self.cell_sel.CM_FINSEL
         )
+        data_dict["nimg_in_trains"][:ntrains] = nimg_in_trains
         if np.all(can_calibrate):
             return n_img
 
@@ -1674,14 +1663,11 @@ class LitFrameSelection(CellSelection):
         self.use_super_selection = use_super_selection
 
         if use_super_selection == 'off':
-            self.cm_sel_type = SelType.ROW
-            self.final_sel_type = SelType.CELL
+            self.sel_type = [SelType.CELL, SelType.ROW, SelType.CELL]
         elif use_super_selection == 'cm':
-            self.cm_sel_type = SelType.SUPER_ROW
-            self.final_sel_type = SelType.CELL
+            self.sel_type = [SelType.CELL, SelType.SUPER_ROW, SelType.CELL]
         elif use_super_selection == 'final':
-            self.cm_sel_type = SelType.SUPER_ROW
-            self.final_sel_type = SelType.SUPER_CELL
+            self.sel_type = [SelType.SUPER_CELL, SelType.SUPER_ROW, SelType.SUPER_CELL]
         else:
             raise ValueError("param 'use_super_selection' takes only "
                              "'off', 'cm' or 'final'")
@@ -1720,11 +1706,13 @@ class LitFrameSelection(CellSelection):
         self, train_sel: np.ndarray, nfrm: np.ndarray,
         cellid: np.ndarray, cm: int = 0
     ) -> np.array:
+        if cm < 0 or cm > 2:
+            raise ValueError("param 'cm' takes only 0,1,2")
+
+        sel, = self._sel.litframes_on_trains(
+            train_sel, nfrm, cellid, [self.sel_type[cm]])
 
-        cell_flags, cm_flags = self._sel.litframes_on_trains(
-            train_sel, nfrm, [self.final_sel_type, self.cm_sel_type])
-        sel = self._sel_for_cm(cell_flags, cm_flags, cm)
-        return sel, nfrm
+        return sel
 
     def filter_trains(self, train_sel: np.ndarray):
         return self._sel.filter_trains(train_sel, drop_empty=True)
-- 
GitLab