Optimize retrieving mono event ids.

ctapipe_io_magic/__init__.py

@@ -1393,59 +1393,25 @@ class MarsRun:
         mono_ids['M1'] = []
         mono_ids['M2'] = []
 
-        n_m1_events = len(self.event_data['M1']['stereo_event_number'])
-        n_m2_events = len(self.event_data['M2']['stereo_event_number'])
-
         if not self.is_mc:
+            m1_data = self.event_data['M1']['stereo_event_number'][np.where(self.event_data['M1']['trigger_pattern'] == DATA_TRIGGER_PATTERN)]
+            m2_data = self.event_data['M2']['stereo_event_number'][np.where(self.event_data['M2']['trigger_pattern'] == DATA_TRIGGER_PATTERN)]
 
-            m1_data_condition = self.event_data['M1']['trigger_pattern'] == DATA_TRIGGER_PATTERN
-            m2_data_condition = self.event_data['M2']['trigger_pattern'] == DATA_TRIGGER_PATTERN
-
-            for m1_id in range(0, n_m1_events):
-                if m1_data_condition[m1_id]:
-                    m2_stereo_condition = (self.event_data['M2']['stereo_event_number'] ==
-                                           self.event_data['M1']['stereo_event_number'][m1_id])
-
-                    m12_match = np.where(
-                        m2_data_condition & m2_stereo_condition)
+            stereo_numbers = np.intersect1d(m1_data, m2_data)
 
-                    if len(m12_match[0]) == 0:
-                        mono_ids['M1'].append(m1_id)
-
-            for m2_id in range(0, n_m2_events):
-                if m2_data_condition[m2_id]:
-                    m1_stereo_condition = (self.event_data['M1']['stereo_event_number'] ==
-                                           self.event_data['M2']['stereo_event_number'][m2_id])
-
-                    m12_match = np.where(
-                        m1_data_condition & m1_stereo_condition)
+            m1_ids = np.searchsorted(self.event_data['M1']['stereo_event_number'], stereo_numbers)
+            m2_ids = np.searchsorted(self.event_data['M2']['stereo_event_number'], stereo_numbers)
 
-                    if len(m12_match[0]) == 0:
-                        mono_ids['M2'].append(m2_id)
+            mono_ids['M1'] = m1_ids
+            mono_ids['M2'] = m2_ids
         else:
-            m1_data_condition = self.event_data['M1']['trigger_pattern'] == MC_TRIGGER_PATTERN
-            m2_data_condition = self.event_data['M2']['trigger_pattern'] == MC_TRIGGER_PATTERN
-
-            # shortcut if only single file is loaded:
-            if n_m1_events == 0:
-                if n_m2_events > 0:
-                    mono_ids['M2'] = np.arange(0, n_m2_events)[
-                        m2_data_condition]
-                return mono_ids
-            if n_m2_events == 0:
-                if n_m1_events > 0:
-                    mono_ids['M1'] = np.arange(0, n_m1_events)[
-                        m1_data_condition]
-                return mono_ids
-
-            for m1_id in range(0, n_m1_events):
-                if m1_data_condition[m1_id]:
-                    if self.event_data['M1']['stereo_event_number'][m1_id] == 0:
-                        mono_ids['M1'].append(m1_id)
-            for m2_id in range(0, n_m2_events):
-                if m2_data_condition[m2_id]:
-                    if self.event_data['M2']['stereo_event_number'][m2_id] == 0:
-                        mono_ids['M2'].append(m2_id)
+            m1_data = self.event_data['M1']['stereo_event_number'][np.where(self.event_data['M1']['trigger_pattern'] == MC_TRIGGER_PATTERN)]
+            m2_data = self.event_data['M2']['stereo_event_number'][np.where(self.event_data['M2']['trigger_pattern'] == MC_TRIGGER_PATTERN)]
+            m1_ids = np.where(m1_data == 0)[0]
+            m2_ids = np.where(m2_data == 0)[0]
+
+            mono_ids['M1'] = m1_ids
+            mono_ids['M2'] = m2_ids
 
         return mono_ids