Merge pull request #184 from kosack/fix/n_ssts_incorrect

correctly count SSTs, MSTs, and LSTs.
cta-observatory · Apr 1, 2022 · 9c0932c · 9c0932c
2 parents 1f392ea + 16ab0b5
commit 9c0932c
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 38 deletions.
diff --git a/protopipe/scripts/data_training.py b/protopipe/scripts/data_training.py
@@ -31,9 +31,7 @@ def main():
     parser = make_argparser()
 
     parser.add_argument(
-        "--debug",
-        action="store_true",
-        help="Print debugging information",
+        "--debug", action="store_true", help="Print debugging information",
     )
 
     parser.add_argument(
@@ -43,9 +41,7 @@ def main():
     )
 
     parser.add_argument(
-        "--save_images",
-        action="store_true",
-        help="Save also all images",
+        "--save_images", action="store_true", help="Save also all images",
     )
 
     parser.add_argument(
@@ -247,6 +243,7 @@ def main():
         #    TEMP
         N_reco_LST=tb.Int16Col(dflt=-1, pos=63),
         N_reco_MST=tb.Int16Col(dflt=-1, pos=64),
+        N_reco_SST=tb.Int16Col(dflt=-1, pos=64),
         image_extraction=tb.Int16Col(dflt=-1, pos=65),
     )
 
@@ -507,9 +504,7 @@ def main():
                         )  # not in ctapipe
 
                     outTable[cam_id] = outfile.create_table(
-                        "/",
-                        cam_id,
-                        DataTrainingOutput,
+                        "/", cam_id, DataTrainingOutput,
                     )
                     outData[cam_id] = outTable[cam_id].row
 
@@ -527,19 +522,26 @@ def main():
                 outData[cam_id]["impact_dist"] = impact_dict[tel_id].to("m").value
                 outData[cam_id]["max_signal_cam"] = max_signals[tel_id]
                 outData[cam_id]["hillas_intensity"] = moments.intensity
-                outData[cam_id]["N_LST"] = n_tels["LST_LST_LSTCam"]
-                outData[cam_id]["N_MST"] = (
-                    n_tels["MST_MST_NectarCam"]
-                    + n_tels["MST_MST_FlashCam"]
-                    + n_tels["MST_SCT_SCTCam"]
+
+                outData[cam_id]["N_LST"] = np.sum(
+                    [v for k, v in n_tels.items() if k.startswith("LST")]
+                )
+                outData[cam_id]["N_MST"] = np.sum(
+                    [v for k, v in n_tels.items() if k.startswith("MST")]
                 )
-                outData[cam_id]["N_SST"] = (
-                    n_tels["SST_1M_DigiCam"]
-                    + n_tels["SST_ASTRI_ASTRICam"]
-                    + n_tels["SST_GCT_CHEC"]
+                outData[cam_id]["N_SST"] = np.sum(
+                    [v for k, v in n_tels.items() if k.startswith("SST")]
                 )
-                outData[cam_id]["N_reco_LST"] = n_tels_reco["LST_LST_LSTCam"]
-                outData[cam_id]["N_reco_MST"] = n_tels_reco["MST_MST_NectarCam"]
+                outData[cam_id]["N_reco_LST"] = np.sum(
+                    [v for k, v in n_tels_reco.items() if k.startswith("LST")]
+                )
+                outData[cam_id]["N_reco_MST"] = np.sum(
+                    [v for k, v in n_tels_reco.items() if k.startswith("MST")]
+                )
+                outData[cam_id]["N_reco_SST"] = np.sum(
+                    [v for k, v in n_tels_reco.items() if k.startswith("SST")]
+                )
+
                 outData[cam_id]["hillas_width"] = moments.width.to("deg").value
                 outData[cam_id]["hillas_length"] = moments.length.to("deg").value
                 outData[cam_id]["hillas_psi"] = moments.psi.to("deg").value

diff --git a/protopipe/scripts/write_dl2.py b/protopipe/scripts/write_dl2.py
@@ -734,28 +734,26 @@ class RecoEvent(tb.IsDescription):
             reco_event["event_id"] = event.index.event_id
             reco_event["obs_id"] = event.index.obs_id
             reco_event["NTels_trig"] = n_tels["Triggered"]
-            reco_event["N_LST"] = n_tels["LST_LST_LSTCam"]
-            reco_event["N_MST"] = (
-                n_tels["MST_MST_NectarCam"]
-                + n_tels["MST_MST_FlashCam"]
-                + n_tels["MST_SCT_SCTCam"]
+            reco_event["N_LST"] = np.sum(
+                [v for k, v in n_tels.items() if k.startswith("LST")]
             )
-            reco_event["N_SST"] = (
-                n_tels["SST_1M_DigiCam"]
-                + n_tels["SST_ASTRI_ASTRICam"]
-                + n_tels["SST_GCT_CHEC"]
+            reco_event["N_MST"] = np.sum(
+                [v for k, v in n_tels.items() if k.startswith("MST")]
             )
-            reco_event["N_reco_LST"] = n_tels_reco["LST_LST_LSTCam"]
-            reco_event["N_reco_MST"] = (
-                n_tels_reco["MST_MST_NectarCam"]
-                + n_tels_reco["MST_MST_FlashCam"]
-                + n_tels_reco["MST_SCT_SCTCam"]
+            reco_event["N_SST"] = np.sum(
+                [v for k, v in n_tels.items() if k.startswith("SST")]
             )
-            reco_event["N_reco_SST"] = (
-                n_tels_reco["SST_1M_DigiCam"]
-                + n_tels_reco["SST_ASTRI_ASTRICam"]
-                + n_tels_reco["SST_GCT_CHEC"]
+
+            reco_event["N_reco_LST"] = np.sum(
+                [v for k, v in n_tels_reco.items() if k.startswith("LST")]
+            )
+            reco_event["N_reco_MST"] = np.sum(
+                [v for k, v in n_tels_reco.items() if k.startswith("MST")]
             )
+            reco_event["N_reco_SST"] = np.sum(
+                [v for k, v in n_tels_reco.items() if k.startswith("SST")]
+            )
+
             reco_event["NTels_reco"] = (
                 reco_event["N_reco_LST"]
                 + reco_event["N_reco_MST"]