LCIOOutputProcessor.cc

Go to the documentation of this file.

// -- marlinmt headers
#include <marlinmt/Processor.h>
#include <marlinmt/ProcessorApi.h>
#include <marlinmt/PluginManager.h>

// -- lcio headers
#include <lcio.h>
#include <MT/LCWriter.h>
#include <IMPL/LCRunHeaderImpl.h>
#include <UTIL/LCTOOLS.h>
#include <EVENT/LCCollection.h>
#include <IMPL/LCCollectionVec.h>

// -- std headers
#include <iostream>
#include <algorithm>
#include <bitset>

namespace marlinmt {

  class LCIOOutputProcessor : public Processor {
  private:
    typedef std::vector< IMPL::LCCollectionVec* > SubSetVec ;
    typedef std::shared_ptr<MT::LCWriter> Writer ;

  public:
    LCIOOutputProcessor() ;
    LCIOOutputProcessor(const LCIOOutputProcessor&) = delete ;
    LCIOOutputProcessor& operator=(const LCIOOutputProcessor&) = delete ;
    ~LCIOOutputProcessor() = default ;

    void init() ;

    void processRunHeader( RunHeader* run) ;

    void processEvent( EventStore * evt ) ;

    void end() ;

  private:
    std::set<std::string> getWriteCollections( EVENT::LCEvent * evt ) const ;

  private:
    StringParameter _lcioOutputFile {*this, "LCIOOutputFile",
             "Name of the LCIO output file", "outputfile.slcio" } ;

    StringParameter _lcioWriteMode {*this, "LCIOWriteMode",
             "Write mode for output file:  WRITE_APPEND, WRITE_NEW or None", "None" } ;

    StringVectorParameter _dropCollectionNames {*this, "DropCollectionNames" ,
             "drops the named collections from the event", {"TPCHits", "HCalHits"} } ;

    StringVectorParameter _dropCollectionTypes {*this, "DropCollectionTypes" ,
             "drops all collections of the given type from the event", {"SimTrackerHit", "SimCalorimeterHit"} } ;

    StringVectorParameter _keepCollectionNames {*this, "KeepCollectionNames" ,
             "force keep of the named collections - overrules DropCollectionTypes (and DropCollectionNames)", {"MyPreciousSimTrackerHits"} } ;

    // OptionalProperty<int> _splitFileSizekB {this, "SplitFileSizekB" ,
    //          "will split output file if size in kB exceeds given value - doesn't work with APPEND and NEW", 1992294 } ;

    // runtime members
    Writer                _writer {nullptr} ;
    std::atomic<int>      _nRuns {0} ;
    std::atomic<int>      _nEvents {0} ;
  };

  //--------------------------------------------------------------------------
  //--------------------------------------------------------------------------

  LCIOOutputProcessor::LCIOOutputProcessor() :
    Processor("LCIOOutputProcessor") {
    setDescription( "Writes the current event to the specified LCIO outputfile." ) ;
    // no lock, the writer implementation is thread safe
    setRuntimeOption( Processor::ERuntimeOption::eCritical, false ) ;
    // don't duplicate opening/writing of output file
    setRuntimeOption( Processor::ERuntimeOption::eClone, false ) ;
  }

  //--------------------------------------------------------------------------

  void LCIOOutputProcessor::init() {
    printParameters() ;
    _writer = std::make_shared<Writer::element_type>() ;
    if ( _lcioWriteMode == "WRITE_APPEND" ) {
      _writer->open( _lcioOutputFile , EVENT::LCIO::WRITE_APPEND ) ;
    }
    else if ( _lcioWriteMode == "WRITE_NEW" ) {
      _writer->open( _lcioOutputFile , EVENT::LCIO::WRITE_NEW ) ;
    }
    else {
      _writer->open( _lcioOutputFile ) ;
    }
  }

  //--------------------------------------------------------------------------

  void LCIOOutputProcessor::processRunHeader( RunHeader* run ) {
    auto rhdr = std::make_unique<IMPL::LCRunHeaderImpl>() ;
    rhdr->setRunNumber( run->runNumber() ) ;
    rhdr->setDetectorName( run->detectorName() ) ;
    rhdr->setDescription( run->description() ) ;
    auto activeSubdets = run->parameter<std::vector<std::string>>( "ActiveSubdetectors" ).value() ;
    for( auto &det : activeSubdets ) {
      rhdr->addActiveSubdetector( det ) ;
    }
    auto intKeys = run->parameter<std::vector<std::string>>( "LCIntKeys" ).value() ;
    auto floatKeys = run->parameter<std::vector<std::string>>( "LCFloatKeys" ).value() ;
    auto strKeys = run->parameter<std::vector<std::string>>( "LCStrKeys" ).value() ;
    auto keys = details::keys( run->parameters() ) ;
    for( auto &key : keys ) {
      if( key == "ActiveSubdetectors" ) {
        continue ;
      }
      if( std::find( intKeys.begin(), intKeys.end(), key ) != intKeys.end() ) {
        auto values = run->parameter<std::vector<int>>( key ).value() ;
        rhdr->parameters().setValues( key, values ) ;
        continue ;
      }
      if( std::find( floatKeys.begin(), floatKeys.end(), key ) != floatKeys.end() ) {
        auto values = run->parameter<std::vector<float>>( key ).value() ;
        rhdr->parameters().setValues( key, values ) ;
        continue ;
      }
      if( std::find( strKeys.begin(), strKeys.end(), key ) != strKeys.end() ) {
        auto values = run->parameter<std::vector<std::string>>( key ).value() ;
        rhdr->parameters().setValues( key, values ) ;
        continue ;
      }
    }

    _writer->writeRunHeader( rhdr.release() ) ;
    _nRuns++ ;
  }

  //--------------------------------------------------------------------------

  std::set<std::string> LCIOOutputProcessor::getWriteCollections( EVENT::LCEvent * evt ) const {
    auto colNames = evt->getCollectionNames() ;
    std::set<std::string> writeCollections {} ;
    // loop over collections and collect the ones to actually write
    for( auto colName : *colNames ) {
      auto col = evt->getCollection( colName ) ;
      std::string collectionType  = col->getTypeName() ;
      auto typeIter = std::find( _dropCollectionTypes.get().begin(), _dropCollectionTypes.get().end(), collectionType ) ;
      auto nameIter = std::find( _dropCollectionNames.get().begin(), _dropCollectionNames.get().end(), colName ) ;
      auto keepIter = std::find( _keepCollectionNames.get().begin(), _keepCollectionNames.get().end(), colName ) ;
      if ( _dropCollectionTypes.isSet() && typeIter != _dropCollectionTypes.get().end() ) {
         continue ;
      }
      if ( _dropCollectionNames.isSet() && nameIter != _dropCollectionNames.get().end() ) {
         continue ;
      }
      if( _keepCollectionNames.isSet() && keepIter != _keepCollectionNames.get().end() ) {
         continue ;
      }
      writeCollections.insert( colName ) ;
    }
    return writeCollections ;
  }

  //--------------------------------------------------------------------------

  void LCIOOutputProcessor::processEvent( EventStore * evt ) {
    auto lcevent = evt->event<EVENT::LCEvent>() ;
    if( nullptr == lcevent ) {
      ProcessorApi::abort( this, "Event is not an LCEvent" ) ;
    }
    auto writeCols = getWriteCollections( lcevent.get() ) ;
    _writer->writeEvent( lcevent.get(), writeCols ) ;
    _nEvents ++ ;
  }

  //--------------------------------------------------------------------------

  void LCIOOutputProcessor::end() {
    log<MESSAGE4>() << std::endl
              << "LCIOOutputProcessor::end()  " << name()
              << ": " << _nEvents.load() << " events in " << _nRuns.load() << " runs written to file  "
              <<  _lcioOutputFile
              << std::endl
              << std::endl ;
    _writer->close() ;
    _writer = nullptr ;
  }

  MARLINMT_DECLARE_PROCESSOR( LCIOOutputProcessor )
}