|
| 1 | +#!/usr/bin/env python3 |
| 2 | +# |
| 3 | +# Copyright (c) 2019-2024 Key4hep-Project. |
| 4 | +# |
| 5 | +# This file is part of Key4hep. |
| 6 | +# See https://key4hep.github.io/key4hep-doc/ for further info. |
| 7 | +# |
| 8 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 9 | +# you may not use this file except in compliance with the License. |
| 10 | +# You may obtain a copy of the License at |
| 11 | +# |
| 12 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 13 | +# |
| 14 | +# Unless required by applicable law or agreed to in writing, software |
| 15 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 16 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 17 | +# See the License for the specific language governing permissions and |
| 18 | +# limitations under the License. |
| 19 | +# |
| 20 | + |
| 21 | +import sys |
| 22 | +import logging |
| 23 | + |
| 24 | +from Configurables import ( |
| 25 | + LcioEvent, |
| 26 | + MarlinProcessorWrapper, |
| 27 | + EDM4hep2LcioTool, |
| 28 | + Lcio2EDM4hepTool, |
| 29 | +) |
| 30 | + |
| 31 | +logger = logging.getLogger() |
| 32 | + |
| 33 | + |
| 34 | +def _is_wrapped_proc_without_conv(alg, from_edm, to_edm): |
| 35 | + """Check if this algorithm has a configured (i.e. named) converter attached |
| 36 | + for the direction described by from_edm and to_edm""" |
| 37 | + if isinstance(alg, MarlinProcessorWrapper): |
| 38 | + if not getattr(alg, f"{from_edm}2{to_edm}Tool").getName(): |
| 39 | + return True |
| 40 | + |
| 41 | + return False |
| 42 | + |
| 43 | + |
| 44 | +class IOHandlerHelper: |
| 45 | + """Helper class to facilitate the transparent handling of LCIO or EDM4hep |
| 46 | + inputs and outputs. |
| 47 | +
|
| 48 | + This class allows to |
| 49 | + - add the correct reader as determined from the input file name(s) |
| 50 | + - add (multiple) LCIO writers |
| 51 | + - add an EDM4hep writer |
| 52 | + - make sure that the necessary converters are introduced at the right places |
| 53 | + """ |
| 54 | + |
| 55 | + def __init__(self, alg_list, io_svc): |
| 56 | + """Create a IOHandlerHelper |
| 57 | +
|
| 58 | + Args: |
| 59 | + alg_list (list): The algorithm list which is being populated |
| 60 | + io_svc (IOSvc): The IOSvc that is being used for this run |
| 61 | + """ |
| 62 | + self._alg_list = alg_list |
| 63 | + self._io_svc = io_svc |
| 64 | + self._lcio_input = False |
| 65 | + self._edm4hep_output = False |
| 66 | + |
| 67 | + def add_reader(self, input_files): |
| 68 | + """Add a reader that is equipped to read the passed files |
| 69 | +
|
| 70 | + If the input is LCIO the necessary algorithm will be configured and |
| 71 | + added to the list of algorithms at the current spot. If the input is |
| 72 | + EDM4hep the file names will be passed to the IOSvc.Input |
| 73 | +
|
| 74 | + Args: |
| 75 | + input_files (list): The input files that should be read |
| 76 | + """ |
| 77 | + if input_files[0].endswith(".slcio"): |
| 78 | + if any(not f.endswith(".slcio") for f in input_files): |
| 79 | + logger.error("All input files need to have the same format (LCIO)") |
| 80 | + sys.exit(1) |
| 81 | + |
| 82 | + read = LcioEvent() |
| 83 | + read.Files = input_files |
| 84 | + self._alg_list.append(read) |
| 85 | + self._lcio_input = True |
| 86 | + else: |
| 87 | + if any(not f.endswith(".root") for f in input_files): |
| 88 | + logger.error("All input files need to have the same format (EDM4hep)") |
| 89 | + sys.exit(1) |
| 90 | + |
| 91 | + self._io_svc.Input = input_files |
| 92 | + |
| 93 | + def add_lcio_writer(self, alg_name): |
| 94 | + """Add a writer for LCIO output at the current spot in the algorithm list |
| 95 | +
|
| 96 | + Note: |
| 97 | + This doesn't configure anything yet, that is still left to do outside |
| 98 | +
|
| 99 | + Args: |
| 100 | + alg_name (str): The name this writer should have |
| 101 | +
|
| 102 | + Returns: |
| 103 | + MarlinProcessorWrapper: The wrapped processor that has just been |
| 104 | + inserted into the algorithm list and that now needs to be |
| 105 | + further configured |
| 106 | + """ |
| 107 | + writer = MarlinProcessorWrapper(alg_name, ProcessorType="LCIOOutputProcessor") |
| 108 | + self._alg_list.append(writer) |
| 109 | + return writer |
| 110 | + |
| 111 | + def add_edm4hep_writer(self, output_file, output_cmds=["keep *"]): |
| 112 | + """Add an EDM4hep writer at the very end of the algorithm execution |
| 113 | +
|
| 114 | + This will pass the output file name as well as the output commands to |
| 115 | + the IOSvc.Ouptut and IOSvc.outputCommands respectively |
| 116 | +
|
| 117 | + Args: |
| 118 | + output_file (str): The name of the output file |
| 119 | +
|
| 120 | + output_cmds (list, optional): The list of output commands that |
| 121 | + should be applied. Defaults to ["keep *"] |
| 122 | + """ |
| 123 | + self._io_svc.Output = output_file |
| 124 | + self._io_svc.outputCommands = output_cmds |
| 125 | + self._edm4hep_output = True |
| 126 | + |
| 127 | + def finalize_converters(self): |
| 128 | + """Attach the appropriate converters in all places they are necessary |
| 129 | +
|
| 130 | + Go through the algorithm list and determine where appropriate converters |
| 131 | + need to be inserted such that the algorithms or wrapped processors always |
| 132 | + see a consistent picture of the event in both formats. Basically what |
| 133 | + this does is to go through the complete list of algorithms and introduce |
| 134 | + an LCIO to EDM4hep converter at the start of every run of |
| 135 | + MarlinProcessorWrappers and in the other direction at the end of every |
| 136 | + such run |
| 137 | +
|
| 138 | + Note: |
| 139 | + Call this just before you pass the algorithm list to the ApplicationMgr |
| 140 | +
|
| 141 | + Note: |
| 142 | + This will not change existing converters on wrapped processors |
| 143 | + """ |
| 144 | + for alg, next_alg in zip(self._alg_list, self._alg_list[1:]): |
| 145 | + if not isinstance(next_alg, MarlinProcessorWrapper) and _is_wrapped_proc_without_conv( |
| 146 | + alg, "Lcio", "EDM4hep" |
| 147 | + ): |
| 148 | + # We change from a run of wrapped processors to algorithms and |
| 149 | + # we don't have a converter yet |
| 150 | + output_conv = Lcio2EDM4hepTool(f"{alg.getName()}_OutputConverter") |
| 151 | + output_conv.convertAll = True |
| 152 | + output_conv.collNameMapping = {"MCParticle": "MCParticles"} |
| 153 | + alg.Lcio2EDM4hepTool = output_conv |
| 154 | + logger.info( |
| 155 | + f"Added an output converter (LCIO to EDM4hep) to the {alg.getName()} algorithm" |
| 156 | + ) |
| 157 | + |
| 158 | + if not isinstance( |
| 159 | + alg, (MarlinProcessorWrapper, LcioEvent) |
| 160 | + ) and _is_wrapped_proc_without_conv(next_alg, "EDM4hep", "Lcio"): |
| 161 | + # We change from a run of algorithms to wrapped processors and |
| 162 | + # we do not have a converter yet |
| 163 | + input_conv = EDM4hep2LcioTool(f"{next_alg.getName()}_InputConverter") |
| 164 | + input_conv.convertAll = True |
| 165 | + input_conv.collNameMapping = {"MCParticles": "MCParticle"} |
| 166 | + next_alg.EDM4hep2LcioTool = input_conv |
| 167 | + logger.info( |
| 168 | + f"Added an input converter (EDM4hep to LCIO) to the {next_alg.getName()} algorithm" |
| 169 | + ) |
| 170 | + |
| 171 | + if not self._lcio_input: |
| 172 | + # We need to convert the input to LCIO from EDM4hep. We attach this |
| 173 | + # to the first wrapped processor that does NOT have another converter |
| 174 | + # configured |
| 175 | + for alg in self._alg_list: |
| 176 | + if _is_wrapped_proc_without_conv(alg, "EDM4hep", "Lcio"): |
| 177 | + input_conv = EDM4hep2LcioTool("InputConversion") |
| 178 | + input_conv.convertAll = True |
| 179 | + input_conv.collNameMapping = {"MCParticles": "MCParticle"} |
| 180 | + alg.EDM4hep2LcioTool = input_conv |
| 181 | + logger.info( |
| 182 | + f"Added an input converter (EDM4hep to LCIO) to the {alg.getName()} algorithm" |
| 183 | + ) |
| 184 | + break |
| 185 | + |
| 186 | + if self._edm4hep_output: |
| 187 | + # We need to convert to EDM4hep. We attach the converter to the last |
| 188 | + # wrapped processor that does not have another converter attached |
| 189 | + for alg in reversed(self._alg_list): |
| 190 | + if _is_wrapped_proc_without_conv(alg, "Lcio", "EDM4hep"): |
| 191 | + output_conv = Lcio2EDM4hepTool("OutputConverter") |
| 192 | + output_conv.convertAll = True |
| 193 | + output_conv.collNameMapping = {"MCParticle": "MCParticles"} |
| 194 | + alg.Lcio2EDM4hepTool = output_conv |
| 195 | + logger.info( |
| 196 | + f"Added an output converter (LCIO to EDM4hep) to the {alg.getName()} algorithm" |
| 197 | + ) |
| 198 | + break |
0 commit comments