Stack.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
 
#ifndef ALICEO2_DATA_STACK_H_
#define ALICEO2_DATA_STACK_H_
 
#include "DetectorsCommonDataFormats/DetID.h"
#include "FairGenericStack.h"
#include "SimulationDataFormat/MCTrack.h"
#include "SimulationDataFormat/TrackReference.h"
#include "SimulationDataFormat/MCEventStats.h"
#include "SimulationDataFormat/ParticleStatus.h"
#include "Rtypes.h"
#include "TParticle.h"
 
#include <map>
#include <memory>
#include <stack>
#include <utility>
#include <functional>
 
class TClonesArray;
class TRefArray;
 
namespace o2
{
namespace base
{
class Detector;
}
 
namespace data
{
class Stack : public FairGenericStack
{
 public:
  Stack(Int_t size = 100);
 
  ~Stack() override;
 
  void PushTrack(Int_t toBeDone, Int_t parentID, Int_t pdgCode, Double_t px, Double_t py, Double_t pz, Double_t e,
                 Double_t vx, Double_t vy, Double_t vz, Double_t time, Double_t polx, Double_t poly, Double_t polz,
                 TMCProcess proc, Int_t& ntr, Double_t weight, Int_t is) override;
 
  void PushTrack(Int_t toBeDone, Int_t parentID, Int_t pdgCode, Double_t px, Double_t py, Double_t pz, Double_t e,
                 Double_t vx, Double_t vy, Double_t vz, Double_t time, Double_t polx, Double_t poly, Double_t polz,
                 TMCProcess proc, Int_t& ntr, Double_t weight, Int_t is, Int_t secondParentId) override;
 
  void PushTrack(Int_t toBeDone, Int_t parentID, Int_t pdgCode, Double_t px, Double_t py, Double_t pz, Double_t e,
                 Double_t vx, Double_t vy, Double_t vz, Double_t time, Double_t polx, Double_t poly, Double_t polz,
                 TMCProcess proc, Int_t& ntr, Double_t weight, Int_t is, Int_t secondParentId, Int_t daughter1Id, Int_t daughter2Id,
                 TMCProcess proc2);
 
  // similar function taking a particle
  void PushTrack(Int_t toBeDone, TParticle&);
 
  TParticle* PopNextTrack(Int_t& iTrack) override;
 
  TParticle* PopPrimaryForTracking(Int_t iPrim) override;
 
  void SetCurrentTrack(Int_t iTrack) override;
 
  Int_t GetNtrack() const override { return mNumberOfEntriesInParticles; }
  Int_t GetNprimary() const override { return mNumberOfPrimaryParticles; }
  TParticle* GetCurrentTrack() const override
  {
    // the const cast is necessary ... the interface should have been `const TParticle* GetCurrentParticle() const`
    return const_cast<TParticle*>(&mCurrentParticle);
  }
 
  Int_t GetCurrentTrackNumber() const override { return mIndexOfCurrentTrack; }
  Int_t GetCurrentParentTrackNumber() const override;
 
  TMCProcess GetProdProcessOfCurrentTrack() const;
 
  void FillTrackArray() override;
 
  void UpdateTrackIndex(TRefArray* detArray = nullptr) override;
 
  void FinishPrimary() override;
 
  void Reset() override;
 
  void Register() override;
 
 
  virtual void Print(Int_t iVerbose = 0) const;
 
  void Print(Option_t* option = nullptr) const override;
 
  void ReorderKine(std::vector<MCTrack>& particles, std::vector<int>& reOrderedIndices);
  void StoreSecondaries(Bool_t choice = kTRUE) { mStoreSecondaries = choice; }
  void pruneKinematics(bool choice = true) { mPruneKinematics = choice; }
  void setMinHits(Int_t min) { mMinHits = min; }
  void SetEnergyCut(Double_t eMin) { mEnergyCut = eMin; }
  void StoreMothers(Bool_t choice = kTRUE) { mStoreMothers = choice; }
  void addHit(int iDet);
 
  TClonesArray* GetListOfParticles() override;
 
  std::vector<MCTrack> const* const getMCTracks() const { return mTracks; }
 
  FairGenericStack* CloneStack() const override;
 
  // methods concerning track references
  void addTrackReference(const o2::TrackReference& p);
 
  // get primaries
  const std::vector<TParticle>& getPrimaries() const { return mPrimaryParticles; }
 
  // initialize Stack from external vector containing primaries
 
  void initFromPrimaries(std::vector<TParticle>& primaries)
  {
    Reset();
    for (auto p : primaries) {
      Int_t doTrack = 0;
      if (p.TestBit(ParticleStatus::kToBeDone)) {
        doTrack = 1;
      }
      PushTrack(doTrack, p);
    }
    mNumberOfPrimaryParticles = primaries.size();
    mNumberOfEntriesInParticles = mNumberOfPrimaryParticles;
  }
 
  void setExternalMode(bool m) { mIsExternalMode = m; }
 
  int getMotherTrackId(int /*trackid*/) const;
 
  bool isTrackDaughterOf(int /*trackid*/, int /*parentid*/) const;
 
  bool isCurrentTrackDaughterOf(int parentid) const;
 
  // returns the index of the currently transported primary
  int getCurrentPrimaryIndex() const;
 
  // Fill container with all parent ids for current track
  // The resulting ids will be in strictly monotonously decreasing order
  void fillParentIDs(std::vector<int>& ids) const;
 
  void setMCEventStats(o2::dataformats::MCEventStats* header);
 
  void updateEventStats();
 
  bool currentTrackLeftHit() const { return mCurrentParticle.TestBit(ParticleStatus::kHasHits); }
 
  bool currentTrackLeftTrackRef() const { return mCurrentParticle.TestBit(ParticleStatus::kHasTrackRefs); }
 
  void setTrackSeedingMode(bool m) { mDoTrackSeeding = m; }
 
  typedef std::function<bool(const TParticle& p, const std::vector<TParticle>& particles)> TransportFcn;
 
 private:
  std::stack<TParticle> mStack; 
 
  std::vector<o2::MCTrack> mParticles;       
  std::vector<int> mTransportedIDs;          
  std::vector<int> mIndexOfPrimaries;        
  std::vector<int> mTrackIDtoParticlesEntry; 
  // the current TParticle object
  TParticle mCurrentParticle;
  TParticle mCurrentParticle0;
 
  // keep primary particles in its original form
  // (mainly for the PopPrimaryParticleInterface
  std::vector<TParticle> mPrimaryParticles;
 
  std::vector<o2::MCTrack>* mTracks;
 
  std::map<Int_t, Int_t> mIndexMap; 
 
  std::vector<o2::base::Detector*> mActiveDetectors; 
 
  Int_t mIndexOfCurrentTrack;          
  Int_t mIndexOfCurrentPrimary;        
  Int_t mNumberOfPrimaryParticles;     
  Int_t mNumberOfEntriesInParticles;   
  Int_t mNumberOfEntriesInTracks;      
  Int_t mNumberOfPrimariesforTracking; 
  Int_t mNumberOfPrimariesPopped;      
  Int_t mIndex;                        
 
  Bool_t mStoreMothers;
  Bool_t mStoreSecondaries;
  bool mPruneKinematics = false; // whether or not we filter the output kinematics
  Int_t mMinHits;
  Int_t mHitCounter = 0; 
  Double32_t mEnergyCut;
 
  // variables for the cleanup / filtering procedure
  Int_t mCleanupCounter = 0;   
  Int_t mCleanupThreshold = 1; 
  Int_t mPrimariesDone = 0;    
  Int_t mTracksDone = 0;       
 
  bool mIsG4Like = false; 
 
  bool mIsExternalMode = false; // is stack an external factory or directly used inside simulation?
  bool mDoTrackSeeding = false; // whether to do track based seeding
 
  TransportFcn mTransportPrimary = [](const TParticle& p, const std::vector<TParticle>& particles) { return false; }; 
 
  // storage for track references
  std::vector<o2::TrackReference>* mTrackRefs = nullptr; 
 
  o2::dataformats::MCEventStats* mMCEventStats = nullptr; 
 
  bool selectTracks();
 
  bool isPrimary(const MCTrack& part);
  bool isFromPrimaryDecayChain(const MCTrack& part);
  bool isFromPrimaryPairProduction(const MCTrack& part);
 
  bool keepPhysics(const MCTrack& part);
 
  Stack(const Stack&);
 
  Stack& operator=(const Stack&);
 
  void addHit(int iDet, Int_t iTrack);
 
  void handleTransportPrimary(TParticle& p);
 
  ClassDefOverride(Stack, 1);
};
 
inline void Stack::addTrackReference(const o2::TrackReference& ref)
{
  if (mIndexOfCurrentTrack >= mNumberOfPrimaryParticles) {
    Int_t iTrack = mTrackIDtoParticlesEntry[mIndexOfCurrentTrack];
    auto& part = mParticles[iTrack];
    part.setStore(true);
  }
  mCurrentParticle.SetBit(ParticleStatus::kHasTrackRefs, 1); // mark that this particle has track refs
  mTrackRefs->push_back(ref);
}
 
inline int Stack::getCurrentPrimaryIndex() const { return mPrimaryParticles.size() - 1 - mPrimariesDone; }
 
inline int Stack::getMotherTrackId(int trackid) const
{
  const auto entryinParticles = mTrackIDtoParticlesEntry[trackid];
  return mParticles[entryinParticles].getMotherTrackId();
}
 
inline bool Stack::isCurrentTrackDaughterOf(int parentid) const
{
  // if parentid is current primary the answer is certainly yes
  if (parentid == getCurrentPrimaryIndex()) {
    return true;
  }
 
  // otherwise ...
  return isTrackDaughterOf(mIndexOfCurrentTrack, parentid);
}
 
inline void Stack::setMCEventStats(o2::dataformats::MCEventStats* header)
{
  mMCEventStats = header;
}
 
inline TMCProcess Stack::GetProdProcessOfCurrentTrack() const
{
  return (TMCProcess)o2::data::Stack::GetCurrentTrack()->GetUniqueID();
}
 
} // namespace data
} // namespace o2
 
#endif