BeamBkgMixerModule.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
 
#include <framework/core/Module.h>
#include <framework/datastore/StoreArray.h>
#include <simulation/background/BeamBGTypes.h>
#include <framework/dataobjects/BackgroundMetaData.h>
#include <string>
#include <map>
 
#include "TChain.h"
#include "TClonesArray.h"
 
 
namespace Belle2 {
  class BeamBkgMixerModule : public Module {
 
  public:
 
    BeamBkgMixerModule();
 
    virtual ~BeamBkgMixerModule();
 
    virtual void initialize() override;
 
    virtual void beginRun() override;
 
    virtual void event() override;
 
    virtual void endRun() override;
 
    virtual void terminate() override;
 
  private:
 
    struct BkgHits {
      TClonesArray* PXD; 
      TClonesArray* SVD; 
      TClonesArray* CDC; 
      TClonesArray* TOP; 
      TClonesArray* ARICH; 
      TClonesArray* ECL; 
      TClonesArray* BKLM; 
      TClonesArray* EKLM; 
      TClonesArray* BeamBackHits; 
      BkgHits():
        PXD(0), SVD(0), CDC(0), TOP(0), ARICH(0), ECL(0), BKLM(0), EKLM(0),
        BeamBackHits(0)
      {}
    };
 
    struct BkgFiles {
      BackgroundMetaData::BG_TAG tag;  
      std::string type;        
      double realTime;         
      double scaleFactor;      
      std::vector<std::string> fileNames;     
      BackgroundMetaData::EFileType fileType; 
      std::unique_ptr<TChain> tree; 
      unsigned numFiles;       
      unsigned numEvents;      
      unsigned eventCount;     
      double rate;             
      unsigned index;          
      BkgFiles(): tag(BackgroundMetaData::bg_none), realTime(0.0), scaleFactor(1.0),
        fileType(BackgroundMetaData::c_Usual),
        tree(nullptr), numFiles(0), numEvents(0), eventCount(0), rate(0.0), index(0)
      {}
      BkgFiles(BackgroundMetaData::BG_TAG bkgTag,
               const std::string& bkgType,
               const std::string& fileName,
               double time,
               double scaleFac,
               BackgroundMetaData::EFileType fileTyp = BackgroundMetaData::c_Usual,
               unsigned indx = 0):
        tag(bkgTag), type(bkgType), realTime(time), scaleFactor(scaleFac),
        fileType(fileTyp),
        tree(nullptr), numFiles(0), numEvents(0), eventCount(0), rate(0.0), index(indx)
      {
        fileNames.push_back(fileName);
      }
    };
 
 
    template<class SIMHIT>
    void addSimHits(StoreArray<SIMHIT>& simHits,
                    TClonesArray* cloneArray,
                    double timeShift,
                    double minTime,
                    double maxTime)
    {
      if (!cloneArray) return;
      if (!simHits.isValid()) return;
 
      int numEntries = cloneArray->GetEntriesFast();
      for (int i = 0; i < numEntries; i++) {
        SIMHIT* bkgSimHit = static_cast<SIMHIT*>(cloneArray->AddrAt(i));
        SIMHIT* simHit = simHits.appendNew(*bkgSimHit);
        simHit->shiftInTime(timeShift);
        if (simHit->getBackgroundTag() == 0) // should be properly set at bkg simulation
          simHit->setBackgroundTag(BackgroundMetaData::bg_other);
        if (m_wrapAround) {
          double time = simHit->getGlobalTime();
          if (time > maxTime) {
            double windowSize = maxTime - minTime;
            double shift = int((time - minTime) / windowSize) * windowSize;
            simHit->shiftInTime(-shift);
          }
        }
      }
 
    }
 
    template<class HIT>
    void addBeamBackHits(StoreArray<HIT>& hits, TClonesArray* cloneArray,
                         double timeShift, double minTime, double maxTime)
    {
      // Match SubDet id from BeamBackHits to whether we keep it or not
      // KLM is a single component, but it has different hits for BKLM and EKLM.
      bool keep[] = {false, m_PXD, m_SVD, m_CDC, m_ARICH, m_TOP, m_ECL, m_KLM, m_KLM};
      if (!cloneArray) return;
      if (!hits.isValid()) return;
      // this is basically a copy of addSimHits but we only add the
      // BeamBackHits from the specified sub detectors so we have to check
      // each if it is from one of the enabled subdetectors
      int numEntries = cloneArray->GetEntriesFast();
      for (int i = 0; i < numEntries; i++) {
        HIT* bkgHit =  static_cast<HIT*>(cloneArray->AddrAt(i));
        //Only keep selected
        if (!keep[bkgHit->getSubDet()]) continue;
        HIT* hit = hits.appendNew(*bkgHit);
        hit->shiftInTime(timeShift);
        //TODO: BeamBackHits does not have a setBackgroundTag so we do not
        //check or set it
        if (m_wrapAround) {
          double time = hit->getTime();
          if (time > maxTime) {
            double windowSize = maxTime - minTime;
            double shift = int((time - minTime) / windowSize) * windowSize;
            hit->shiftInTime(-shift);
          }
        }
      }
    }
 
    bool isComponentIncluded(std::vector<std::string>& components,
                             const std::string& component);
 
    void appendSample(BackgroundMetaData::BG_TAG bkgTag,
                      const std::string& bkgType,
                      const std::string& fileName,
                      double realTime,
                      BackgroundMetaData::EFileType fileTyp);
 
    bool acceptEvent(TClonesArray* cloneArrayECL);
 
 
    std::vector<std::string> m_backgroundFiles; 
    double m_overallScaleFactor; 
    std::vector<std::tuple<std::string, double> > m_scaleFactors; 
    double m_minTime;  
    double m_maxTime;  
    std::vector<std::string> m_components; 
    bool m_wrapAround; 
    double m_minTimeECL;  
    double m_maxTimeECL;  
    double m_minTimePXD;  
    double m_maxTimePXD;  
    double m_maxEdepECL;  
    int m_cacheSize;  
    std::vector<BkgFiles> m_backgrounds;  
    BkgHits m_simHits;         
    bool m_PXD = false; 
    bool m_SVD = false; 
    bool m_CDC = false; 
    bool m_TOP = false; 
    bool m_ARICH = false; 
    bool m_ECL = false; 
    bool m_KLM = false; 
    bool m_BeamBackHits = false; 
    background::BeamBGTypes m_bgTypes;  
    std::map<std::string, int> m_rejected; 
    std::map<std::string, int> m_reused;  
    int m_rejectedCount = 0;  
  };
 
} // Belle2 namespace