DDG4/edm4hep/EDM4hepFileReader.cpp

0001 //==========================================================================
0002 //  AIDA Detector description implementation
0003 //--------------------------------------------------------------------------
0004 // Copyright (C) Organisation europeenne pour la Recherche nucleaire (CERN)
0005 // All rights reserved.
0006 //
0007 // For the licensing terms see $DD4hepINSTALL/LICENSE.
0008 // For the list of contributors see $DD4hepINSTALL/doc/CREDITS.
0009 //
0010 // Author     : A.Sailer
0011 //
0012 //==========================================================================
0013
0014 /** \addtogroup Geant4EventReader
0015  *
0016  @{
0017  \package EDM4hepFileReader
0018  * \brief Plugin to read EDM4hep files
0019  *
0020  *
0021  @}
0022 */
0023
0024 #include <CLHEP/Units/SystemOfUnits.h>
0025
0026 #include <DDG4/EventParameters.h>
0027 #include <DDG4/Factories.h>
0028 #include <DDG4/FileParameters.h>
0029 #include <DDG4/Geant4InputAction.h>
0030 #include <DDG4/RunParameters.h>
0031
0032 #include <edm4hep/EventHeaderCollection.h>
0033 #include <edm4hep/MCParticleCollection.h>
0034
0035 #include <podio/Frame.h>
0036 #include <podio/Reader.h>
0037
0038 typedef dd4hep::detail::ReferenceBitMask<int> PropertyMask;
0039
0040
0041 namespace dd4hep::sim {
0042
0043   // we use the index of the objectID to identify particles
0044   // we will not support MCParticles from different collections
0045   using MCPARTICLE_MAP=std::map<int, int>;
0046
0047   /// get the parameters from the GenericParameters of the input EDM4hep frame and store them in the EventParameters
0048   /// extension
0049   template <class T=podio::GenericParameters> void EventParameters::ingestParameters(T const& source) {
0050     for(auto const& key: source.template getKeys<int>()) {
0051       m_intValues[key] = source.template get<std::vector<int>>(key).value();
0052     }
0053     for(auto const& key: source.template getKeys<float>()) {
0054       m_fltValues[key] = source.template get<std::vector<float>>(key).value();
0055     }
0056     for(auto const& key: source.template getKeys<double>()) {
0057       m_dblValues[key] = source.template get<std::vector<double>>(key).value();
0058     }
0059     for(auto const& key: source.template getKeys<std::string>()) {
0060       m_strValues[key] = source.template get<std::vector<std::string>>(key).value();
0061     }
0062   }
0063
0064   /// get the parameters from the GenericParameters of the input EDM4hep run frame and store them in the RunParameters
0065   /// extension
0066   template <class T=podio::GenericParameters> void RunParameters::ingestParameters(T const& source) {
0067     for(auto const& key: source.template getKeys<int>()) {
0068       m_intValues[key] = source.template get<std::vector<int>>(key).value();
0069     }
0070     for(auto const& key: source.template getKeys<float>()) {
0071       m_fltValues[key] = source.template get<std::vector<float>>(key).value();
0072     }
0073     for(auto const& key: source.template getKeys<double>()) {
0074       m_dblValues[key] = source.template get<std::vector<double>>(key).value();
0075     }
0076     for(auto const& key: source.template getKeys<std::string>()) {
0077       m_strValues[key] = source.template get<std::vector<std::string>>(key).value();
0078     }
0079   }
0080
0081   template <class T=podio::GenericParameters> void FileParameters::ingestParameters(T const& source) {
0082     for(auto const& key: source.template getKeys<int>()) {
0083       m_intValues[key] = source.template get<std::vector<int>>(key).value();
0084     }
0085     for(auto const& key: source.template getKeys<float>()) {
0086       m_fltValues[key] = source.template get<std::vector<float>>(key).value();
0087     }
0088     for(auto const& key: source.template getKeys<double>()) {
0089       m_dblValues[key] = source.template get<std::vector<double>>(key).value();
0090     }
0091     for(auto const& key: source.template getKeys<std::string>()) {
0092       m_strValues[key] = source.template get<std::vector<std::string>>(key).value();
0093     }
0094   }
0095
0096   /// Class to read EDM4hep files
0097   /**
0098    *  \version 1.0
0099    *  \ingroup DD4HEP_SIMULATION
0100    */
0101   class EDM4hepFileReader : public Geant4EventReader {
0102   protected:
0103     /// Reference to reader object
0104     podio::Reader m_reader;
0105     /// Name of the MCParticle collection to read
0106     std::string m_collectionName;
0107     /// Name of the EventHeader collection to read
0108     std::string m_eventHeaderCollectionName;
0109
0110   public:
0111     /// Initializing constructor
0112     EDM4hepFileReader(const std::string& nam);
0113
0114     /// Read parameters set for this action
0115     virtual EventReaderStatus setParameters(std::map< std::string, std::string >& parameters);
0116
0117     /// Read an event and fill a vector of MCParticles.
0118     virtual EventReaderStatus readParticles(int event_number, Vertices& vertices, std::vector<Particle*>& particles);
0119     /// Call to register the run parameters
0120     void registerRunParameters();
0121
0122   };
0123
0124   /// Initializing constructor
0125   dd4hep::sim::EDM4hepFileReader::EDM4hepFileReader(const std::string& nam)
0126     : Geant4EventReader(nam)
0127     , m_reader(podio::makeReader(nam))
0128     , m_collectionName("MCParticles")
0129     , m_eventHeaderCollectionName("EventHeader")
0130   {
0131     printout(INFO,"EDM4hepFileReader","Created file reader. Try to open input %s",nam.c_str());
0132     m_directAccess = true;
0133   }
0134
0135   void EDM4hepFileReader::registerRunParameters() {
0136     try {
0137       // get RunParameters or create new if not existent yet
0138       auto* parameters = context()->run().extension<RunParameters>(false);
0139       if (!parameters) {
0140         parameters = new RunParameters();
0141         context()->run().addExtension<RunParameters>(parameters);
0142       }
0143       try {
0144         podio::Frame runFrame = m_reader.readFrame("runs", 0);
0145         parameters->ingestParameters(runFrame.getParameters());
0146       } catch (std::runtime_error& e) {
0147         // we ignore if we do not have runs information
0148       } catch(std::invalid_argument&) {
0149         // we ignore if we do not have runs information
0150       }
0151     } catch(std::exception &e) {
0152       printout(ERROR,"EDM4hepFileReader::registerRunParameters","Failed to register run parameters: %s", e.what());
0153     }
0154
0155     try {
0156       auto *fileParameters = new FileParameters();
0157       try {
0158         podio::Frame metaFrame = m_reader.readFrame("metadata", 0);
0159         fileParameters->ingestParameters(metaFrame.getParameters());
0160       } catch (std::runtime_error& e) {
0161         // we ignore if we do not have metadata information
0162       } catch(std::invalid_argument&) {
0163         // we ignore if we do not have metadata information
0164       }
0165       context()->run().addExtension<FileParameters>(fileParameters);
0166     } catch(std::exception &e) {
0167       printout(ERROR,"EDM4hepFileReader::registerRunParameters","Failed to register file parameters: %s", e.what());
0168     }
0169   }
0170
0171   namespace {
0172     /// Helper function to look up MCParticles from mapping
0173     inline int GET_ENTRY(MCPARTICLE_MAP const& mcparts, int partID)  {
0174       MCPARTICLE_MAP::const_iterator ip = mcparts.find(partID);
0175       if ( ip == mcparts.end() )  {
0176         throw std::runtime_error("Unknown particle identifier look-up!");
0177       }
0178       return (*ip).second;
0179     }
0180   }
0181
0182   /// Read an event and fill a vector of MCParticles
0183   EDM4hepFileReader::EventReaderStatus
0184   EDM4hepFileReader::readParticles(int event_number, Vertices& vertices, std::vector<Particle*>& particles) {
0185     m_currEvent = event_number;
0186     podio::Frame frame = m_reader.readFrame("events", event_number);
0187     const auto& primaries = frame.get<edm4hep::MCParticleCollection>(m_collectionName);
0188     int eventNumber = event_number, runNumber = 0;
0189     if (primaries.isValid()) {
0190       //Read the event header collection and add it to the context as an extension
0191       const auto& eventHeaderCollection = frame.get<edm4hep::EventHeaderCollection>(m_eventHeaderCollectionName);
0192       if(eventHeaderCollection.isValid() && eventHeaderCollection.size() == 1){
0193         const auto& eh = eventHeaderCollection.at(0);
0194         eventNumber = eh.getEventNumber();
0195         runNumber = eh.getRunNumber();
0196         try {
0197           Geant4Context* ctx = context();
0198           ctx->event().addExtension<edm4hep::MutableEventHeader>(new edm4hep::MutableEventHeader(eh.clone()));
0199         } catch(std::exception& ) {}
0200         // Create input event parameters context
0201         try {
0202           Geant4Context* ctx = context();
0203           EventParameters *parameters = new EventParameters();
0204           parameters->setRunNumber(runNumber);
0205           parameters->setEventNumber(eventNumber);
0206           parameters->ingestParameters(frame.getParameters());
0207           ctx->event().addExtension<EventParameters>(parameters);
0208         } catch(std::exception &) {}
0209       } else {
0210         printout(WARNING,"EDM4hepFileReader","No EventHeader collection found or too many event headers!");
0211         try {
0212           Geant4Context* ctx = context();
0213           EventParameters *parameters = new EventParameters();
0214           parameters->setRunNumber(0);
0215           parameters->setEventNumber(event_number);
0216           parameters->ingestParameters(frame.getParameters());
0217           ctx->event().addExtension<EventParameters>(parameters);
0218         } catch(std::exception &) {}
0219       }
0220       printout(INFO,"EDM4hepFileReader","read collection %s from event %d in run %d ",
0221                m_collectionName.c_str(), eventNumber, runNumber);
0222
0223     } else {
0224       return EVENT_READER_EOF;
0225     }
0226
0227     printout(INFO,"EDM4hepFileReader", "We read the particle collection");
0228     int NHEP = primaries.size();
0229     // check if there is at least one particle
0230     if ( NHEP == 0 ) {
0231       printout(WARNING,"EDM4hepFileReader", "We have no particles");
0232       return EVENT_READER_NO_PRIMARIES;
0233     }
0234
0235     MCPARTICLE_MAP mcparts;
0236     std::vector<int>  mcpcoll;
0237     mcpcoll.resize(NHEP);
0238     for(int i=0; i<NHEP; ++i ) {
0239       edm4hep::MCParticle p = primaries.at(i);
0240       mcparts[p.getObjectID().index] = i;
0241       mcpcoll[i] = p.getObjectID().index;
0242     }
0243
0244     // build collection of MCParticles
0245     for(int i=0; i<NHEP; ++i )   {
0246       const auto& mcp = primaries.at(i);
0247       Geant4ParticleHandle p(new Particle(i));
0248       const auto mom   = mcp.getMomentum();
0249       const auto vsx   = mcp.getVertex();
0250       const auto vex   = mcp.getEndpoint();
0251       const auto spin  = mcp.getSpin();
0252       const int  pdg   = mcp.getPDG();
0253       p->pdgID        = pdg;
0254       p->charge       = int(mcp.getCharge()*3.0);
0255       p->psx          = mom[0]*CLHEP::GeV;
0256       p->psy          = mom[1]*CLHEP::GeV;
0257       p->psz          = mom[2]*CLHEP::GeV;
0258       p->time         = mcp.getTime()*CLHEP::ns;
0259       p->properTime   = mcp.getTime()*CLHEP::ns;
0260       p->vsx          = vsx[0]*CLHEP::mm;
0261       p->vsy          = vsx[1]*CLHEP::mm;
0262       p->vsz          = vsx[2]*CLHEP::mm;
0263       p->vex          = vex[0]*CLHEP::mm;
0264       p->vey          = vex[1]*CLHEP::mm;
0265       p->vez          = vex[2]*CLHEP::mm;
0266       p->process      = 0;
0267       p->spin[0]      = spin[0];
0268       p->spin[1]      = spin[1];
0269       p->spin[2]      = spin[2];
0270       p->colorFlow[0] = 0;
0271       p->colorFlow[1] = 0;
0272       p->mass         = mcp.getMass()*CLHEP::GeV;
0273       const auto par = mcp.getParents(), &dau=mcp.getDaughters();
0274       for(int num=dau.size(),k=0; k<num; ++k) {
0275         edm4hep::MCParticle dau_k = dau[k];
0276         p->daughters.insert(GET_ENTRY(mcparts,dau_k.getObjectID().index));
0277       }
0278       for(int num=par.size(),k=0; k<num; ++k) {
0279         auto par_k = par[k];
0280         p->parents.insert(GET_ENTRY(mcparts, par_k.getObjectID().index));
0281       }
0282
0283       PropertyMask status(p->status);
0284       int genStatus = mcp.getGeneratorStatus();
0285       // Copy raw generator status
0286       p->genStatus = genStatus&G4PARTICLE_GEN_STATUS_MASK;
0287       if(m_inputAction) {
0288         // in some tests we do not set up the inputAction
0289         m_inputAction->setGeneratorStatus(genStatus, status);
0290       }
0291
0292       //fg: we simply add all particles without parents as with their own vertex.
0293       //    This might include the incoming beam particles, e.g. in
0294       //    the case of lcio files written with Whizard2, which is slightly odd,
0295       //    however should be treated correctly in Geant4InputHandling.cpp.
0296       //    We no longer make an attempt to identify the incoming particles
0297       //    based on the generator status, as this varies widely with different
0298       //    generators.
0299
0300       if ( p->parents.size() == 0 )  {
0301
0302         Geant4Vertex* vtx = new Geant4Vertex ;
0303         vertices.emplace_back( vtx );
0304         vtx->x = p->vsx;
0305         vtx->y = p->vsy;
0306         vtx->z = p->vsz;
0307         vtx->time = p->time;
0308
0309         vtx->out.insert(p->id) ;
0310       }
0311
0312       if ( mcp.isCreatedInSimulation() )       status.set(G4PARTICLE_SIM_CREATED);
0313       if ( mcp.isBackscatter() )               status.set(G4PARTICLE_SIM_BACKSCATTER);
0314       if ( mcp.vertexIsNotEndpointOfParent() ) status.set(G4PARTICLE_SIM_PARENT_RADIATED);
0315       if ( mcp.isDecayedInTracker() )          status.set(G4PARTICLE_SIM_DECAY_TRACKER);
0316       if ( mcp.isDecayedInCalorimeter() )      status.set(G4PARTICLE_SIM_DECAY_CALO);
0317       if ( mcp.hasLeftDetector() )             status.set(G4PARTICLE_SIM_LEFT_DETECTOR);
0318       if ( mcp.isStopped() )                   status.set(G4PARTICLE_SIM_STOPPED);
0319       if ( mcp.isOverlay() )                   status.set(G4PARTICLE_SIM_OVERLAY);
0320       particles.emplace_back(p);
0321     }
0322     return EVENT_READER_OK;
0323   }
0324
0325   /// Set the parameters for the class, in particular the name of the MCParticle
0326   /// list
0327   Geant4EventReader::EventReaderStatus
0328   dd4hep::sim::EDM4hepFileReader::setParameters( std::map< std::string, std::string > & parameters ) {
0329     _getParameterValue(parameters, "MCParticleCollectionName", m_collectionName, m_collectionName);
0330     _getParameterValue(parameters, "EventHeaderCollectionName", m_eventHeaderCollectionName, m_eventHeaderCollectionName);
0331     return EVENT_READER_OK;
0332   }
0333
0334 } //end dd4hep::sim
0335
0336 DECLARE_GEANT4_EVENT_READER_NS(dd4hep::sim,EDM4hepFileReader)