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