Refactor III

CLDConfig/CLDReconstruction.py

@@ -22,16 +22,21 @@
 from Configurables import k4DataSvc, MarlinProcessorWrapper
 from k4MarlinWrapper.inputReader import create_reader, attach_edm4hep2lcio_conversion
 from k4FWCore.parseArgs import parser
-from py_utils import SequenceLoader
+from py_utils import SequenceLoader, attach_lcio2edm4hep_conversion, create_writer
 
 import ROOT
 ROOT.gROOT.SetBatch(True)
 
 
-parser.add_argument("--inputFiles", action="extend", nargs="+", metavar=("file1", "file2"), help="One or multiple input files")
-parser.add_argument("--outputBasename", help="Basename of the output file(s)", default="output")
-parser.add_argument("--trackingOnly", action="store_true", help="Run only track reconstruction", default=False)
-parser.add_argument("--enableLCFIJet", action="store_true", help="Enable LCFIPlus jet clustering parts", default=False)
+parser_group = parser.add_argument_group("CLDReconstruction.py custom options")
+parser_group.add_argument("--inputFiles", action="extend", nargs="+", metavar=("file1", "file2"), help="One or multiple input files")
+parser_group.add_argument("--outputBasename", help="Basename of the output file(s)", default="output")
+parser_group.add_argument("--trackingOnly", action="store_true", help="Run only track reconstruction", default=False)
+parser_group.add_argument("--enableLCFIJet", action="store_true", help="Enable LCFIPlus jet clustering parts", default=False)
+parser_group.add_argument("--compactFile", help="Compact detector file to use", type=str, default=os.environ["K4GEO"] + "/FCCee/CLD/compact/CLD_o2_v07/CLD_o2_v07.xml")
+tracking_group = parser_group.add_mutually_exclusive_group()
+tracking_group.add_argument("--conformalTracking", action="store_true", default=True, help="Use conformal tracking pattern recognition")
+tracking_group.add_argument("--truthTracking", action="store_true", default=False, help="Cheat tracking pattern recognition")
 reco_args = parser.parse_known_args()[0]
 
 algList = []
@@ -45,8 +50,6 @@
              "CalorimeterIntegrationTimeWindowChoices": ["10ns", "400ns"],
              "Overlay": "False",
              "OverlayChoices": ["False", "91GeV", "365GeV"],
-             "Tracking": "Conformal",
-             "TrackingChoices": ["Truth", "Conformal"],
              "VertexUnconstrained": "OFF",
              "VertexUnconstrainedChoices": ["ON", "OFF"],
              "OutputMode": "EDM4Hep",
@@ -55,7 +58,7 @@
 
 from Configurables import GeoSvc, TrackingCellIDEncodingSvc
 geoservice = GeoSvc("GeoSvc")
-geoservice.detectors = [os.environ["K4GEO"]+"/FCCee/CLD/compact/CLD_o2_v06/CLD_o2_v06.xml"]
+geoservice.detectors = [reco_args.compactFile]
 geoservice.OutputLevel = INFO
 geoservice.EnableGeant4Geo = False
 svcList.append(geoservice)
@@ -99,17 +102,16 @@
                           "HowOften": ["1"]
                           }
 
-# TODO: put this somewhere else, needs to be in front of the output for now :(
 # setup AIDA histogramming and add eventual background overlay
 algList.append(MyAIDAProcessor)
 sequenceLoader.load("Overlay/Overlay")
 # tracker hit digitisation
 sequenceLoader.load("Tracking/TrackingDigi")
 
 # tracking
-if CONFIG["Tracking"] == "Truth":
+if reco_args.truthTracking:
     sequenceLoader.load("Tracking/TruthTracking")
-elif CONFIG["Tracking"] == "Conformal":
+elif reco_args.conformalTracking:
     sequenceLoader.load("Tracking/ConformalTracking")
 
 sequenceLoader.load("Tracking/Refit")
@@ -131,52 +133,34 @@
 # event number processor, down here to attach the conversion back to edm4hep to it
 algList.append(EventNumber)
 
+
+DST_KEEPLIST = ["MCParticlesSkimmed", "MCPhysicsParticles", "RecoMCTruthLink", "SiTracks", "SiTracks_Refitted", "PandoraClusters", "PandoraPFOs", "SelectedPandoraPFOs", "LooseSelectedPandoraPFOs", "TightSelectedPandoraPFOs", "RefinedVertexJets", "RefinedVertexJets_rel", "RefinedVertexJets_vtx", "RefinedVertexJets_vtx_RP", "BuildUpVertices", "BuildUpVertices_res", "BuildUpVertices_RP", "BuildUpVertices_res_RP", "BuildUpVertices_V0", "BuildUpVertices_V0_res", "BuildUpVertices_V0_RP", "BuildUpVertices_V0_res_RP", "PrimaryVertices", "PrimaryVertices_res", "PrimaryVertices_RP", "PrimaryVertices_res_RP", "RefinedVertices", "RefinedVertices_RP"]
+
+DST_SUBSETLIST = ["EfficientMCParticles", "InefficientMCParticles", "MCPhysicsParticles"]
+
+# TODO: replace all the ugly strings by something sensible like Enum
 if CONFIG["OutputMode"] == "LCIO":
-    Output_REC = MarlinProcessorWrapper("Output_REC")
-    Output_REC.OutputLevel = WARNING
-    Output_REC.ProcessorType = "LCIOOutputProcessor"
-    Output_REC.Parameters = {
-                             "DropCollectionNames": [],
-                             "DropCollectionTypes": [],
-                             "FullSubsetCollections": ["EfficientMCParticles", "InefficientMCParticles"],
-                             "KeepCollectionNames": [],
-                             "LCIOOutputFile": [f"{reco_args.outputBasename}_REC.slcio"],
-                             "LCIOWriteMode": ["WRITE_NEW"]
-                             }
-
-    Output_DST = MarlinProcessorWrapper("Output_DST")
-    Output_DST.OutputLevel = WARNING
-    Output_DST.ProcessorType = "LCIOOutputProcessor"
-    Output_DST.Parameters = {
-                             "DropCollectionNames": [],
-                             "DropCollectionTypes": ["MCParticle", "LCRelation", "SimCalorimeterHit", "CalorimeterHit", "SimTrackerHit", "TrackerHit", "TrackerHitPlane", "Track", "ReconstructedParticle", "LCFloatVec"],
-                             "FullSubsetCollections": ["EfficientMCParticles", "InefficientMCParticles", "MCPhysicsParticles"],
-                             "KeepCollectionNames": ["MCParticlesSkimmed", "MCPhysicsParticles", "RecoMCTruthLink", "SiTracks", "SiTracks_Refitted", "PandoraClusters", "PandoraPFOs", "SelectedPandoraPFOs", "LooseSelectedPandoraPFOs", "TightSelectedPandoraPFOs", "RefinedVertexJets", "RefinedVertexJets_rel", "RefinedVertexJets_vtx", "RefinedVertexJets_vtx_RP", "BuildUpVertices", "BuildUpVertices_res", "BuildUpVertices_RP", "BuildUpVertices_res_RP", "BuildUpVertices_V0", "BuildUpVertices_V0_res", "BuildUpVertices_V0_RP", "BuildUpVertices_V0_res_RP", "PrimaryVertices", "PrimaryVertices_res", "PrimaryVertices_RP", "PrimaryVertices_res_RP", "RefinedVertices", "RefinedVertices_RP"],
-                             "LCIOOutputFile": [f"{reco_args.outputBasename}_DST.slcio"],
-                             "LCIOWriteMode": ["WRITE_NEW"]
-                             }
+    Output_REC = create_writer("lcio", "Output_REC", f"{reco_args.outputBasename}_REC")
     algList.append(Output_REC)
+
+    Output_DST = create_writer("lcio", "Output_DST", f"{reco_args.outputBasename}_DST", DST_KEEPLIST, DST_SUBSETLIST)
     algList.append(Output_DST)
 
 if CONFIG["OutputMode"] == "EDM4Hep":
-    from Configurables import Lcio2EDM4hepTool
-    lcioConvTool = Lcio2EDM4hepTool("lcio2EDM4hep")
-    lcioConvTool.convertAll = True
-    lcioConvTool.collNameMapping = {
-        "MCParticle": "MCParticles"
-    }
-    lcioConvTool.OutputLevel = DEBUG
-# attach to the last non output processor
-    EventNumber.Lcio2EDM4hepTool = lcioConvTool
-
-    from Configurables import PodioOutput
-    out = PodioOutput("PodioOutput", filename = f"{reco_args.outputBasename}_edm4hep.root")
-    out.outputCommands = ["keep *"]
-    algList.append(out)
+    Output_REC = create_writer("edm4hep", "Output_REC", f"{reco_args.outputBasename}_REC")
+    algList.append(Output_REC)
+
+    # FIXME: needs https://github.com/key4hep/k4FWCore/issues/226
+    # Output_DST = create_writer("edm4hep", "Output_DST", f"{reco_args.outputBasename}_DST", DST_KEEPLIST)
+    # algList.append(Output_DST)
+
 
 # We need to convert the inputs in case we have EDM4hep input
 attach_edm4hep2lcio_conversion(algList, read)
 
+# We need to convert the outputs in case we have EDM4hep output
+attach_lcio2edm4hep_conversion(algList)
+
 from Configurables import ApplicationMgr
 ApplicationMgr( TopAlg = algList,
                 EvtSel = 'NONE',

CLDConfig/HighLevelReco/JetClusteringOrRenaming.py

@@ -25,7 +25,7 @@
 RenameCollection.ProcessorType = "MergeCollections"
 RenameCollection.Parameters = {
                                "CollectionParameterIndex": ["0"],
-                               "InputCollectionIDs": [],
+                               "InputCollectionIDs": ["0"],
                                "InputCollections": ["PandoraPFOs"],
                                "OutputCollection": ["PFOsFromJets"]
                                }

CLDConfig/py_utils.py

@@ -21,6 +21,9 @@
 import importlib.util
 import importlib.abc
 from importlib.machinery import SourceFileLoader
+from Configurables import PodioOutput, MarlinProcessorWrapper
+from typing import Iterable
+from Gaudi.Configuration import WARNING
 
 
 def import_from(
@@ -128,3 +131,78 @@ def load(self, sequence: str) -> None:
 
         seq = getattr(seq_module, seq_name)
         self.alg_list.extend(seq)
+
+
+
+def attach_lcio2edm4hep_conversion(algList: list) -> None:
+    """Attaches a conversion from lcio to edm4hep at the last MarlinWrapper in algList if necessary
+    """
+    # need to attach a conversion if there are edm4hep outputs and marlin wrappers
+    if not any((isinstance(alg, PodioOutput) for alg in algList)):
+        # no edm4hep output -> no conversion :)
+        return
+    # find last marlin wrapper
+    for alg in reversed(algList):
+        if isinstance(alg, MarlinProcessorWrapper):
+            break
+
+    from Configurables import Lcio2EDM4hepTool
+    lcioConvTool = Lcio2EDM4hepTool("lcio2EDM4hep")
+    lcioConvTool.convertAll = True
+    lcioConvTool.collNameMapping = {
+        "MCParticle": "MCParticles",
+    }
+
+    alg.Lcio2EDM4hepTool = lcioConvTool
+
+
+
+def _create_writer_lcio(writer_name: str, output_name: str, keep_list: Iterable = (), full_subset_list: Iterable = ()):
+    writer = MarlinProcessorWrapper(writer_name)
+    writer.OutputLevel = WARNING
+    writer.ProcessorType = "LCIOOutputProcessor"
+
+    # convert iterables to "real" lists
+    _full_subset_list = list(full_subset_list)
+    _keep_list = list(keep_list)
+
+    dropped_types = []
+    if _keep_list:
+        # drop collections of all types
+        dropped_types = ["MCParticle", "LCRelation", "SimCalorimeterHit", "CalorimeterHit", "SimTrackerHit", "TrackerHit", "TrackerHitPlane", "Track", "ReconstructedParticle", "LCFloatVec"]
+
+    writer.Parameters = {
+        "DropCollectionNames": [],
+        "DropCollectionTypes": dropped_types,
+        "FullSubsetCollections": _full_subset_list,
+        "KeepCollectionNames": _keep_list,
+        "LCIOOutputFile": [f"{output_name}.slcio"],
+        "LCIOWriteMode": ["WRITE_NEW"],
+    }
+
+    return writer
+
+
+def _create_writer_edm4hep(writer_name: str, output_name: str, keep_list: Iterable = ()):
+    writer = PodioOutput(writer_name, filename = f"{output_name}.edm4hep.root")
+
+    if keep_list:
+        writer.outputCommands = ["drop *"] + [f"keep {col}" for col in keep_list]
+    else:
+        writer.outputCommands = ["keep *"]
+
+    return writer
+
+
+def create_writer(format: str, writer_name: str, output_name: str, keep_list: Iterable = (), full_subset_list: Iterable = ()):
+    """
+    Creates writer depending on the requested format
+    In contrast to its name an empty keep_list means keep everything
+    """
+    if format == "lcio":
+        return _create_writer_lcio(writer_name, output_name, keep_list, full_subset_list)
+    elif format == "edm4hep":
+        return _create_writer_edm4hep(writer_name, output_name, keep_list) # FIXME: handle edm4hep subset collections!
+    else:
+        return None
+        # TODO: warn about format being unsupported but without killing --help

test/CMakeLists.txt

@@ -18,7 +18,7 @@ limitations under the License.
 ]]
 include(CTest)
 
-set(DETECTOR $ENV{K4GEO}/FCCee/CLD/compact/CLD_o2_v06/CLD_o2_v06.xml)
+set(DETECTOR $ENV{K4GEO}/FCCee/CLD/compact/CLD_o2_v07/CLD_o2_v07.xml)
 set(CLDConfig_DIR ${CMAKE_CURRENT_LIST_DIR}/../CLDConfig)
 
 add_test(NAME ddsim_lcio
@@ -61,6 +61,24 @@ add_test(NAME LCFIJet
 )
 set_property(TEST LCFIJet APPEND PROPERTY DEPENDS ddsim_edm4hep)
 
+add_test(NAME tracking_truth
+         WORKING_DIRECTORY ${CLDConfig_DIR}
+         COMMAND k4run --trackingOnly --truthTracking --inputFiles=test.edm4hep.root --outputBasename=trkOnly_truth_test_edm4hep CLDReconstruction.py --GeoSvc.detectors=${DETECTOR}
+)
+set_property(TEST tracking_truth APPEND PROPERTY DEPENDS ddsim_edm4hep)
+
+add_test(NAME tracking_conformal
+         WORKING_DIRECTORY ${CLDConfig_DIR}
+         COMMAND k4run --trackingOnly --conformalTracking --inputFiles=test.edm4hep.root --outputBasename=trkOnly_conformal_test_edm4hep CLDReconstruction.py --GeoSvc.detectors=${DETECTOR}
+)
+set_property(TEST tracking_conformal APPEND PROPERTY DEPENDS ddsim_edm4hep)
+
+add_test(NAME tracking_both
+         WORKING_DIRECTORY ${CLDConfig_DIR}
+         COMMAND k4run --trackingOnly --conformalTracking --truthTracking CLDReconstruction.py
+)
+set_property(TEST tracking_both APPEND PROPERTY PASS_REGULAR_EXPRESSION "error: argument --truthTracking: not allowed with argument --conformalTracking")
+
 add_test(NAME Help
          WORKING_DIRECTORY ${CLDConfig_DIR}
          COMMAND k4run --help CLDReconstruction.py