Modernize event display example steering

k4MarlinWrapper/examples/event_display.py

@@ -18,58 +18,71 @@
 # limitations under the License.
 #
 
-from Gaudi.Configuration import *
+import os
+import sys
+
+from Gaudi.Configuration import INFO
+from Configurables import MarlinProcessorWrapper, k4DataSvc, PodioInput, EDM4hep2LcioTool, LcioEvent, GeoSvc
+from k4FWCore.parseArgs import parser
+
+parser.add_argument(
+    "--inputFiles",
+    action="extend",
+    nargs="+",
+    metavar=["file1", "file2"],
+    help="One or multiple input files",
+)
+parser.add_argument(
+    "--compactFile", help="Compact detector file to use", type=str, default=""
+)
+
+reco_args = parser.parse_known_args()[0]
 
-from Configurables import MarlinProcessorWrapper, k4DataSvc, PodioInput, EDM4hep2LcioTool
 algList = []
+svcList = []
+
+evtsvc = k4DataSvc("EventDataSvc")
+svcList.append(evtsvc)
+
+if reco_args.compactFile:
+    compact_file = reco_args.compactFile
+else:
+    # add the same default as before
+    # TODO: maybe modernize this
+    compact_file = "CLICPerformance/Visualisation/CLIC_o3_v06_CED/CLIC_o3_v06_CED.xml"
+
+geoSvc = GeoSvc("GeoSvc")
+geoSvc.detectors = [compact_file]
+geoSvc.OutputLevel = INFO
+geoSvc.EnableGeant4Geo = False
+svcList.append(geoSvc)
+
+def create_reader(input_files):
+    """Create the appropriate reader for the input files"""
+    if input_files[0].endswith(".slcio"):
+        if any(not f.endswith(".slcio") for f in input_files):
+            print("All input files need to have the same format (LCIO)")
+            sys.exit(1)
 
-evtsvc = k4DataSvc('EventDataSvc')
-evtsvc.input = ''
-
-inp = PodioInput('InputReader')
-inp.collections = [
-  'MCParticles',
-  'VertexBarrelCollection',
-  'VertexEndcapCollection',
-  'InnerTrackerBarrelCollection',
-  'OuterTrackerBarrelCollection',
-  'InnerTrackerEndcapCollection',
-  'OuterTrackerEndcapCollection',
-  'ECalEndcapCollection',
-  'ECalEndcapCollectionContributions',
-  'ECalBarrelCollection',
-  'ECalBarrelCollectionContributions',
-  'ECalPlugCollection',
-  'ECalPlugCollectionContributions',
-  'HCalBarrelCollection',
-  'HCalBarrelCollectionContributions',
-  'HCalEndcapCollection',
-  'HCalEndcapCollectionContributions',
-  'HCalRingCollection',
-  'HCalRingCollectionContributions',
-  'YokeBarrelCollection',
-  'YokeBarrelCollectionContributions',
-  'YokeEndcapCollection',
-  'YokeEndcapCollectionContributions',
-  'LumiCalCollection',
-  'LumiCalCollectionContributions',
-  'BeamCalCollection',
-  'BeamCalCollectionContributions',
-]
-
-MyInitializeDD4hep = MarlinProcessorWrapper("MyInitializeDD4hep")
-MyInitializeDD4hep.OutputLevel = INFO
-MyInitializeDD4hep.ProcessorType = "InitializeDD4hep"
-MyInitializeDD4hep.Parameters = {
-                                 "DD4hepXMLFile": ["CLICPerformance/Visualisation/CLIC_o3_v06_CED/CLIC_o3_v06_CED.xml"]
-                                 }
-
-MyEventSelector = MarlinProcessorWrapper("MyEventSelector")
-MyEventSelector.OutputLevel = INFO
-MyEventSelector.ProcessorType = "EventSelector"
-MyEventSelector.Parameters = {
-                              "EventList": ["28", "0", "33", "0", "52", "0", "63", "0", "73", "0", "78", "0"]
-                              }
+        read = LcioEvent()
+        read.Files = input_files
+    else:
+        if any(not f.endswith(".root") for f in input_files):
+            print("All input files need to have the same format (EDM4hep)")
+            sys.exit(1)
+        read = PodioInput("PodioInput")
+        global evtsvc
+        evtsvc.inputs = input_files
+
+    return read
+
+
+if reco_args.inputFiles:
+    read = create_reader(reco_args.inputFiles)
+    read.OutputLevel = INFO
+    algList.append(read)
+else:
+    read = None
 
 MyCEDViewer = MarlinProcessorWrapper("MyCEDViewer")
 MyCEDViewer.OutputLevel = INFO
@@ -82,7 +95,8 @@
                           "ColorByEnergyMin": ["0.0"],
                           "ColorByEnergySaturation": ["1.0"],
                           "ColorScheme": ["10"],
-                          "DetailledDrawing": ["VXD", "VertexBarrel"],
+                          # "DetailledDrawing": ["VXD", "VertexBarrel"],
+                          "DetailledDrawing": [],
                           "DrawDetector": ["true"],
                           "DrawDetectorID": ["0"],
                           "DrawHelixForPFOs": ["0"],
@@ -197,6 +211,7 @@
                               "PandoraPFOCollection", "0", "3", "7",
                               "JetOut", "0", "0", "3"],
                           "DrawSurfaces": ["true"],
+                          # "DrawSurfaces": ["false"],
                           "HelixMaxR": ["2000"],
                           "HelixMaxZ": ["2500"],
                           "MCParticleEnergyCut": ["0.001"],
@@ -208,21 +223,21 @@
                           "WaitForKeyboard": ["1"]
                           }
 
-# EDM4hep to LCIO converter
-edmConvTool = EDM4hep2LcioTool("EDM4hep2lcio")
-edmConvTool.convertAll = True
-edmConvTool.collNameMapping = {'MCParticles': 'MCParticle'}
-edmConvTool.OutputLevel = DEBUG
-MyCEDViewer.EDM4hep2LcioTool = edmConvTool
-
-algList.append(inp)
-algList.append(MyInitializeDD4hep)
 algList.append(MyCEDViewer)
 
+# We need to convert the inputs in case we have EDM4hep input
+if isinstance(read, PodioInput):
+    EDM4hep2LcioInput = EDM4hep2LcioTool("InputConversion")
+    EDM4hep2LcioInput.convertAll = True
+    # Adjust for the different naming conventions
+    EDM4hep2LcioInput.collNameMapping = {"MCParticles": "MCParticle"}
+    MyCEDViewer.EDM4hep2LcioTool = EDM4hep2LcioInput
+
+
 from Configurables import ApplicationMgr
 ApplicationMgr( TopAlg = algList,
                 EvtSel = 'NONE',
                 EvtMax   = 10,
-                ExtSvc = [evtsvc],
+                ExtSvc = svcList,
                 OutputLevel=INFO
               )