DDDigi/io/DigiEdm4hepOutput.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     : M.Frank
0011 //
0012 //==========================================================================
0013
0014 /// Framework include files
0015 #include <DD4hep/InstanceCount.h>
0016 #include <DDDigi/DigiContext.h>
0017 #include <DDDigi/DigiPlugins.h>
0018 #include <DDDigi/DigiKernel.h>
0019 #include "DigiEdm4hepOutput.h"
0020 #include "DigiIO.h"
0021
0022 /// podio include files
0023 #include <podio/podioVersion.h>
0024 #if PODIO_BUILD_VERSION >= PODIO_VERSION(1, 0, 0)
0025 #include <podio/Writer.h>
0026 #elif PODIO_BUILD_VERSION >= PODIO_VERSION(0, 99, 0)
0027 #include <podio/ROOTWriter.h>
0028 #else
0029 #include <podio/ROOTFrameWriter.h>
0030 namespace podio {
0031   using ROOTWriter = podio::ROOTFrameWriter;
0032 }
0033 #endif
0034 #include <podio/Frame.h>
0035 #include <edm4hep/SimTrackerHit.h>
0036 #include <edm4hep/MCParticleCollection.h>
0037 #include <edm4hep/EventHeaderCollection.h>
0038 #include <edm4hep/CalorimeterHitCollection.h>
0039 #include <edm4hep/CaloHitContributionCollection.h>
0040
0041
0042 /// Namespace for the AIDA detector description toolkit
0043 namespace dd4hep {
0044
0045   /// Namespace for the Digitization part of the AIDA detector description toolkit
0046   namespace digi {
0047
0048 #if PODIO_BUILD_VERSION >= PODIO_VERSION(1, 0, 0)
0049     using writer_t = podio::Writer;
0050 #else
0051     using writer_t = podio::ROOTWriter;
0052 #endif
0053
0054     /// Helper class to create output in edm4hep format
0055     /** Helper class to create output in edm4hep format
0056      *
0057      *  \author  M.Frank
0058      *  \version 1.0
0059      *  \ingroup DD4HEP_DIGITIZATION
0060      */
0061     class DigiEdm4hepOutput::internals_t {
0062     public:
0063       using particlecollection_t = std::pair<std::string,std::unique_ptr<edm4hep::MCParticleCollection> >;
0064       using headercollection_t   = std::pair<std::string,std::unique_ptr<edm4hep::EventHeaderCollection> >;
0065       DigiEdm4hepOutput*                      m_parent    { nullptr };
0066       /// Reference to podio writer
0067       std::unique_ptr<writer_t>               m_writer    { };
0068       /// edm4hep event header collection
0069       headercollection_t                      m_header    { };
0070       /// MC particle collection
0071       particlecollection_t                    m_particles { };
0072       /// Collection of all edm4hep tracker object collections
0073       std::map<std::string, std::unique_ptr<edm4hep::TrackerHit3DCollection> > m_tracker_collections;
0074       /// Collection of all edm4hep calorimeter object collections
0075       std::map<std::string, std::unique_ptr<edm4hep::CalorimeterHitCollection> > m_calo_collections;
0076       /// Output section name
0077       std::string                             m_section_name{ "EVENT" };
0078       /// Output mutex
0079       std::mutex                              m_lock;
0080       /// Total numbe rof events to be processed
0081       long num_events  { -1 };
0082       /// Running event counter
0083       long event_count {  0 };
0084
0085     public:
0086       /// Default constructor
0087       internals_t(DigiEdm4hepOutput* parent);
0088       /// Default destructor
0089       ~internals_t();
0090
0091       /// Clear local data content
0092       void clear();
0093       /// Commit data at end of filling procedure
0094       void commit();
0095       /// Open new output stream
0096       void open();
0097       /// Commit data to disk and close output stream
0098       void close();
0099
0100       /// Create all collections according to the parent setup (locked)
0101       void create_collections();
0102       /// Access named collection: throws exception ifd the collection is not present (unlocked!)
0103       template <typename T> podio::CollectionBase* get_collection(const T&);
0104     };
0105
0106     /// Default constructor
0107     DigiEdm4hepOutput::internals_t::internals_t(DigiEdm4hepOutput* parent) : m_parent(parent)
0108     {
0109     }
0110
0111     /// Default destructor
0112     DigiEdm4hepOutput::internals_t::~internals_t()    {
0113       if ( m_writer ) close();
0114       m_tracker_collections.clear();
0115       m_calo_collections.clear();
0116       m_particles.second.reset();
0117       m_header.second.reset();
0118     }
0119
0120     /// Create all collections according to the parent setup
0121     void DigiEdm4hepOutput::internals_t::create_collections()    {
0122       if ( !m_header.second )   {
0123         m_header = std::make_pair("EventHeader", std::make_unique<edm4hep::EventHeaderCollection>());
0124         for( auto& cont : m_parent->m_containers )   {
0125           const std::string& nam = cont.first;
0126           const std::string& typ = cont.second;
0127           if ( typ == "MCParticles" )   {
0128             m_particles = std::make_pair(nam, std::make_unique<edm4hep::MCParticleCollection>());
0129           }
0130           else if ( typ == "TrackerHits" )   {
0131             m_tracker_collections.emplace(nam, std::make_unique<edm4hep::TrackerHit3DCollection>());
0132           }
0133           else if ( typ == "CalorimeterHits" )   {
0134             m_calo_collections.emplace(nam, std::make_unique<edm4hep::CalorimeterHitCollection>());
0135           }
0136         }
0137         m_parent->info("+++ Will save %ld events to %s", num_events, m_parent->m_output.c_str());
0138       }
0139     }
0140     /// Access named collection: throws exception ifd the collection is not present
0141     template <typename T>
0142     podio::CollectionBase* DigiEdm4hepOutput::internals_t::get_collection(const T& cont)  {
0143       switch(cont.data_type)   {
0144       case SegmentEntry::TRACKER_HITS:   {
0145         auto iter = m_tracker_collections.find(cont.name);
0146         if ( iter == m_tracker_collections.end() )
0147           m_parent->except("Error");
0148         return iter->second.get();
0149       }
0150       case SegmentEntry::CALORIMETER_HITS:   {
0151         auto iter = m_calo_collections.find(cont.name);
0152         if ( iter == m_calo_collections.end() )
0153           m_parent->except("Error");
0154         return iter->second.get();
0155       }
0156       default:
0157         return nullptr;
0158       }
0159     };
0160
0161     /// Clear local data content
0162     void DigiEdm4hepOutput::internals_t::clear()   {
0163 #if 0
0164       m_header.second->clear();
0165       m_particles.second->clear();
0166       for( const auto& c : m_tracker_collections )
0167         c.second->clear();
0168       for( const auto& c : m_calo_collections )
0169         c.second->clear();
0170 #endif
0171       *m_header.second = {};
0172       *m_particles.second = {};
0173       for( const auto& c : m_tracker_collections )
0174         *c.second = {};
0175       for( const auto& c : m_calo_collections )
0176         *c.second = {};
0177     }
0178
0179     /// Commit data at end of filling procedure
0180     void DigiEdm4hepOutput::internals_t::commit()   {
0181       if ( m_writer )   {{
0182           std::lock_guard<std::mutex> protection(m_lock);
0183           podio::Frame frame { };
0184           frame.put( std::move(*m_header.second), m_header.first);
0185           frame.put( std::move(*m_particles.second), m_particles.first);
0186           for( const auto& c : m_tracker_collections )
0187             frame.put( std::move(*c.second), c.first);
0188           for( const auto& c : m_calo_collections )
0189             frame.put( std::move(*c.second), c.first);
0190
0191           m_writer->writeFrame(frame, m_section_name);
0192         }
0193         clear();
0194         return;
0195       }
0196       m_parent->except("+++ Failed to write output file. [Stream is not open]");
0197     }
0198
0199     /// Open new output stream
0200     void DigiEdm4hepOutput::internals_t::open()    {
0201       if ( m_writer )   {
0202         close();
0203       }
0204       clear();
0205       m_writer.reset();
0206       std::string fname = m_parent->next_stream_name();
0207 #if PODIO_BUILD_VERSION >= PODIO_VERSION(1, 0, 0)
0208       m_writer = std::make_unique<writer_t>(podio::makeWriter(fname));
0209 #else
0210       m_writer = std::make_unique<writer_t>(fname);
0211 #endif
0212       m_parent->info("+++ Opened EDM4HEP output file %s", fname.c_str());
0213     }
0214
0215     /// Commit data to disk and close output stream
0216     void DigiEdm4hepOutput::internals_t::close()   {
0217       m_parent->info("+++ Closing EDM4HEP output file.");
0218       if ( m_writer )   {
0219         m_writer->finish();
0220       }
0221       m_writer.reset();
0222     }
0223
0224     /// Standard constructor
0225     DigiEdm4hepOutput::DigiEdm4hepOutput(const DigiKernel& krnl, const std::string& nam)
0226       : DigiOutputAction(krnl, nam)
0227     {
0228       internals = std::make_shared<internals_t>(this);
0229       InstanceCount::increment(this);
0230     }
0231
0232     /// Default destructor
0233     DigiEdm4hepOutput::~DigiEdm4hepOutput()   {
0234       internals.reset();
0235       InstanceCount::decrement(this);
0236     }
0237
0238     /// Initialization callback
0239     void DigiEdm4hepOutput::initialize()   {
0240       this->DigiOutputAction::initialize();
0241       for ( auto& c : m_registered_processors )   {
0242         auto* act = dynamic_cast<DigiEdm4hepOutputProcessor*>(c.second);
0243         if ( act )   { // This is not nice! Need to think about something better.
0244           act->internals = this->internals;
0245           continue;
0246         }
0247         except("Error: Invalid processor type for EDM4HEP output: %s", c.second->c_name());
0248       }
0249       m_parallel = false;
0250       internals->create_collections();
0251     }
0252
0253     /// Check for valid output stream
0254     bool DigiEdm4hepOutput::have_output()  const  {
0255       return internals->m_writer.get() != nullptr;
0256     }
0257
0258     /// Open new output stream
0259     void DigiEdm4hepOutput::open_output() const   {
0260       internals->open();
0261     }
0262
0263     /// Close possible open stream
0264     void DigiEdm4hepOutput::close_output()  const  {
0265       internals->close();
0266     }
0267
0268     /// Commit event data to output stream
0269     void DigiEdm4hepOutput::commit_output() const  {
0270       internals->commit();
0271     }
0272
0273     /// Standard constructor
0274     DigiEdm4hepOutputProcessor::DigiEdm4hepOutputProcessor(const DigiKernel& krnl, const std::string& nam)
0275       : DigiContainerProcessor(krnl, nam)
0276     {
0277       declareProperty("point_resolution_RPhi", m_pointResoutionRPhi);
0278       declareProperty("point_resolution_Z",    m_pointResoutionZ);
0279       declareProperty("hit_type",              m_hit_type = 0);
0280     }
0281
0282     void DigiEdm4hepOutputProcessor::convert_particles(DigiContext& ctxt,
0283                                                        const ParticleMapping& cont)  const
0284     {
0285       auto& parts = internals->m_particles.second;
0286       data_io<edm4hep_input>::_to_edm4hep(cont, parts.get());
0287       info("%s+++ %-24s added %6ld entries from mask: %04X to %s",
0288            ctxt.event->id(), cont.name.c_str(), parts->size(), cont.key.mask(),
0289            parts->getTypeName().data());
0290     }
0291
0292     template <typename T> void
0293     DigiEdm4hepOutputProcessor::convert_depos(const T& cont,
0294                                               const predicate_t& predicate,
0295                                               edm4hep::TrackerHit3DCollection* collection)  const
0296     {
0297       std::array<float,6> covMat = {0., 0., m_pointResoutionRPhi*m_pointResoutionRPhi,
0298         0., 0., m_pointResoutionZ*m_pointResoutionZ
0299       };
0300       for ( const auto& depo : cont )   {
0301         if ( predicate(depo) )   {
0302           data_io<edm4hep_input>::_to_edm4hep(depo, covMat, *collection, m_hit_type /* edm4hep::SIMTRACKERHIT */);
0303         }
0304       }
0305     }
0306
0307     template <typename T> void
0308     DigiEdm4hepOutputProcessor::convert_depos(const T& cont,
0309                                               const predicate_t& predicate,
0310                                               edm4hep::CalorimeterHitCollection* collection)  const
0311     {
0312       for ( const auto& depo : cont )   {
0313         if ( predicate(depo) )   {
0314           data_io<edm4hep_input>::_to_edm4hep(depo, *collection, m_hit_type /* edm4hep::SIMCALORIMETERHIT */);
0315         }
0316       }
0317     }
0318
0319     template <typename T> void
0320     DigiEdm4hepOutputProcessor::convert_deposits(DigiContext&       ctxt,
0321                                                  const T&           cont,
0322                                                  const predicate_t& predicate)  const
0323     {
0324       podio::CollectionBase* coll = internals->get_collection(cont);
0325       std::size_t start = coll->size();
0326       if ( !cont.empty() )   {
0327         switch(cont.data_type)    {
0328         case SegmentEntry::TRACKER_HITS:
0329           convert_depos(cont, predicate, static_cast<edm4hep::TrackerHit3DCollection*>(coll));
0330           break;
0331         case SegmentEntry::CALORIMETER_HITS:
0332           convert_depos(cont, predicate, static_cast<edm4hep::CalorimeterHitCollection*>(coll));
0333           break;
0334         default:
0335           except("Error: Unknown energy deposit type: %d", int(cont.data_type));
0336           break;
0337         }
0338       }
0339       std::size_t end = coll->size();
0340       info("%s+++ %-24s added %6ld/%6ld entries from mask: %04X to %s",
0341            ctxt.event->id(), cont.name.c_str(), end-start, end, cont.key.mask(),
0342            coll->getTypeName().data());
0343     }
0344
0345     void DigiEdm4hepOutputProcessor::convert_history(DigiContext&           ctxt,
0346                                                      const DepositsHistory& cont,
0347                                                      work_t&                work,
0348                                                      const predicate_t&     predicate)  const
0349     {
0350       info("%s+++ %-32s Segment: %d Predicate:%s Conversion to edm4hep not implemented!",
0351            ctxt.event->id(), cont.name.c_str(), int(work.input.segment->id),
0352            typeName(typeid(predicate)).c_str());
0353     }
0354
0355     /// Main functional callback
0356     void DigiEdm4hepOutputProcessor::execute(DigiContext& ctxt, work_t& work, const predicate_t& predicate)  const  {
0357       if ( const auto* p = work.get_input<ParticleMapping>() )
0358         convert_particles(ctxt, *p);
0359       else if ( const auto* m = work.get_input<DepositMapping>() )
0360         convert_deposits(ctxt, *m, predicate);
0361       else if ( const auto* v = work.get_input<DepositVector>() )
0362         convert_deposits(ctxt, *v, predicate);
0363       else if ( const auto* h = work.get_input<DepositsHistory>() )
0364         convert_history(ctxt, *h, work, predicate);
0365       else
0366         except("Request to handle unknown data type: %s", work.input_type_name().c_str());
0367     }
0368
0369   }    // End namespace digi
0370 }      // End namespace dd4hep
0371
0372 /// Factory instantiation:
0373 #include <DDDigi/DigiFactories.h>
0374 DECLARE_DIGIACTION_NS(dd4hep::digi,DigiEdm4hepOutput)
0375 DECLARE_DIGIACTION_NS(dd4hep::digi,DigiEdm4hepOutputProcessor)