Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
SCorrelatorJetTree.h
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file SCorrelatorJetTree.h
1 // ----------------------------------------------------------------------------
2 // 'SCorrelatorJetTree.h'
3 // Derek Anderson
4 // 12.04.2022
5 //
6 // A module to produce a tree of jets for the sPHENIX
7 // Cold QCD Energy-Energy Correlator analysis.
8 //
9 // Derived from code by Antonio Silva (thanks!!)
10 // ----------------------------------------------------------------------------
11 
12 #ifndef SCORRELATORJETTREE_H
13 #define SCORRELATORJETTREE_H
14 
15 #pragma GCC diagnostic push
16 #pragma GCC diagnostic ignored "-Wdeprecated-declarations"
17 
18 // standard c include
19 #include <map>
20 #include <array>
21 #include <string>
22 #include <vector>
23 #include <cassert>
24 #include <sstream>
25 #include <cstdlib>
26 #include <utility>
27 // f4a include
28 #include <fun4all/SubsysReco.h>
31 // phool includes
32 #include <phool/phool.h>
33 #include <phool/getClass.h>
34 #include <phool/PHIODataNode.h>
35 #include <phool/PHNodeIterator.h>
36 #include <phool/PHCompositeNode.h>
37 // main geant4 includes
38 #include <g4main/PHG4Hit.h>
39 #include <g4main/PHG4Particle.h>
41 // jet includes
42 #include <jetbase/Jet.h>
43 #include <jetbase/JetMap.h>
44 #include <jetbase/JetMapv1.h>
45 #include <jetbase/FastJetAlgo.h>
46 // track evaluator includes
47 #include <g4eval/SvtxTrackEval.h>
48 #include <g4eval/SvtxEvalStack.h>
49 // vtx includes
52 // tracking includes
56 // calo includes
57 #include <calobase/RawCluster.h>
58 #include <calobase/RawClusterUtility.h>
59 #include <calobase/RawClusterContainer.h>
60 #include <calobase/RawTower.h>
61 #include <calobase/RawTowerGeom.h>
62 #include <calobase/RawTowerContainer.h>
63 #include <calobase/RawTowerGeomContainer.h>
64 #include <calotrigger/CaloTriggerInfo.h>
65 // particle flow includes
66 #include <particleflowreco/ParticleFlowElement.h>
67 #include <particleflowreco/ParticleFlowElementContainer.h>
68 // fastjet includes
69 #include <fastjet/PseudoJet.hh>
70 #include <fastjet/JetDefinition.hh>
71 #include <fastjet/ClusterSequence.hh>
72 #include <fastjet/FunctionOfPseudoJet.hh>
73 // hepmc includes
74 #include <HepMC/GenEvent.h>
75 #include <HepMC/GenVertex.h>
76 #include <HepMC/GenParticle.h>
79 // root includes
80 #include <TH1.h>
81 #include <TH2.h>
82 #include <TF1.h>
83 #include <TFile.h>
84 #include <TTree.h>
85 #include <TMath.h>
86 #include <TNtuple.h>
87 #include <TDirectory.h>
88 
89 #pragma GCC diagnostic pop
90 
91 using namespace std;
92 using namespace fastjet;
93 using namespace findNode;
94 
95 // forward declarations
96 class TH1;
97 class TFile;
98 class TTree;
99 class PHG4Particle;
100 class PHCompositeNode;
101 class PHHepMCGenEvent;
102 class PHHepMCGenEventMap;
104 class Fun4AllHistoManager;
105 class RawClusterContainer;
106 class RawCluster;
107 class GlobalVertex;
108 class SvtxTrackMap;
109 class JetRecoEval;
110 class SvtxTrackEval;
111 class SvtxTrack;
112 class ParticleFlowElement;
113 
114 
115 
116 
117 namespace SColdQcdCorrelatorAnalysis {
118 
119  // SCorrelatorJetTree definition --------------------------------------------
120 
122 
123  public:
124 
125  // public enums
126  enum ALGO {
127  ANTIKT = 0,
128  KT = 1,
129  CAMBRIDGE = 2
130  };
131  enum RECOMB {
132  E_SCHEME = 0,
133  PT_SCHEME = 1,
134  PT2_SCHEME = 2,
135  ET_SCHEME = 3,
136  ET2_SCHEME = 4
137  };
138 
139  // ctor/dtor
140  SCorrelatorJetTree(const string& name = "SCorrelatorJetTree", const string& outFile = "correlator_jet_tree.root", const bool isMC = false, const bool isEmbed = false, const bool debug = false);
141  ~SCorrelatorJetTree() override;
142 
143  // F4A methods
144  int Init(PHCompositeNode*) override;
145  int process_event(PHCompositeNode*) override;
146  int End(PHCompositeNode*) override;
147 
148  // setters (inline)
149  void SetAddTracks(const bool addTracks) {m_addTracks = addTracks;}
150  void SetAddFlow(const bool addFlow) {m_addFlow = addFlow;}
151  void SetAddECal(const bool addECal) {m_addECal = addECal;}
152  void SetAddHCal(const bool addHCal) {m_addHCal = addHCal;}
153  void SetDoVertexCut(const bool doVtx) {m_doVtxCut = doVtx;}
154  void SetDoQualityPlots(const bool doQA) {m_doQualityPlots = doQA;}
155  void SetRequireSiSeeds(const bool require) {m_requireSiSeeds = require;}
156  void SetUseOnlyPrimVtx(const bool primary) {m_useOnlyPrimVtx = primary;}
157  void SetMaskTpcSectors(const bool mask) {m_maskTpcSectors = mask;}
158  void SetCheckWeirdTrks(const bool check) {m_checkWeirdTrks = check;}
159  void SetSaveDST(const bool doSave) {m_saveDST = doSave;}
160  void SetIsMC(const bool isMC) {m_isMC = isMC;}
161  void SetIsEmbed(const bool isEmbed) {m_isEmbed = isEmbed;}
162  void SetJetR(const double jetR) {m_jetR = jetR;}
163  void SetJetType(const uint32_t type) {m_jetType = type;}
164  void SetJetTreeName(const string name) {m_jetTreeName = name;}
165 
166  // setters (*.io.h)
167  // TODO consolidate parameters into less methods
168  void SetEvtVzRange(const pair<double, double> vzRange);
169  void SetEvtVrRange(const pair<double, double> vrRange);
170  void SetParPtRange(const pair<double, double> ptRange);
171  void SetParEtaRange(const pair<double, double> etaRange);
172  void SetTrackPtRange(const pair<double, double> ptRange);
173  void SetTrackEtaRange(const pair<double, double> etaRange);
174  void SetTrackQualityRange(const pair<double, double> qualRange);
175  void SetTrackNMvtxRange(const pair<double, double> nMvtxRange);
176  void SetTrackNInttRange(const pair<double, double> nInttRange);
177  void SetTrackNTpcRange(const pair<double, double> nTpcRange);
178  void SetTrackDcaRangeXY(const pair<double, double> dcaRangeXY);
179  void SetTrackDcaRangeZ(const pair<double, double> dcaRangeZ);
180  void SetTrackDeltaPtRange(const pair<double, double> deltaPtRange);
181  void SetTrackDcaSigmaParameters(const bool doDcaSigmaCut, const pair<double, double> ptFitMax, const pair<double, double> nSigma, const vector<double> paramDcaXY, const vector<double> paramDcaZ);
182  void SetFlowPtRange(const pair<double, double> ptRange);
183  void SetFlowEtaRange(const pair<double, double> etaRange);
184  void SetECalPtRange(const pair<double, double> ptRange);
185  void SetECalEtaRange(const pair<double, double> etaRange);
186  void SetHCalPtRange(const pair<double, double> ptRange);
187  void SetHCalEtaRange(const pair<double, double> etaRange);
188  void SetJetAlgo(const ALGO jetAlgo);
189  void SetRecombScheme(const RECOMB recombScheme);
190  void SetJetParameters(const double rJet, const uint32_t jetType, const ALGO jetAlgo, const RECOMB recombScheme);
191 
192  // system getters
193  bool GetDoVtxCut() {return m_doVtxCut;}
194  bool GetDoQualityPlots() {return m_doQualityPlots;}
195  bool GetRequireSiSeeds() {return m_requireSiSeeds;}
196  bool GetUseOnlyPrimVtx() {return m_useOnlyPrimVtx;}
197  bool GetDoDcaSigmaCut() {return m_doDcaSigmaCut;}
198  bool GetMaskTpcSectors() {return m_maskTpcSectors;}
199  bool GetCheckWeirdTrks() {return m_checkWeirdTrks;}
200  bool GetSaveDST() {return m_saveDST;}
201  bool GetIsMC() {return m_isMC;}
202  bool GetIsEmbed() {return m_isEmbed;}
203  bool GetDoDebug() {return m_doDebug;}
204  bool GetAddTracks() {return m_addTracks;}
205  bool GetAddFlow() {return m_addFlow;}
206  bool GetAddECal() {return m_addECal;}
207  bool GetAddHCal() {return m_addHCal;}
208  string GetJetTreeName() {return m_jetTreeName;}
209 
210  // acceptance getters
211  double GetEvtMinVz() {return m_evtVzRange[0];}
212  double GetEvtMaxVz() {return m_evtVzRange[1];}
213  double GetEvtMinVr() {return m_evtVrRange[0];}
214  double GetEvtMaxVr() {return m_evtVrRange[1];}
215  double GetParMinPt() {return m_parPtRange[0];}
216  double GetParMaxPt() {return m_parPtRange[1];}
217  double GetParMinEta() {return m_parEtaRange[0];}
218  double GetParMaxEta() {return m_parEtaRange[1];}
219  double GetTrackMinPt() {return m_trkPtRange[0];}
220  double GetTrackMaxPt() {return m_trkPtRange[1];}
221  double GetTrackMinEta() {return m_trkEtaRange[0];}
222  double GetTrackMaxEta() {return m_trkEtaRange[1];}
223  double GetTrackMinQual() {return m_trkQualRange[0];}
224  double GetTrackMaxQual() {return m_trkQualRange[1];}
225  double GetTrackMinNMvtx() {return m_trkNMvtxRange[0];}
226  double GetTrackMaxNMvtx() {return m_trkNMvtxRange[1];}
227  double GetTrackMinNIntt() {return m_trkNInttRange[0];}
228  double GetTrackMaxNIntt() {return m_trkNInttRange[1];}
229  double GetTrackMinNTpc() {return m_trkNTpcRange[0];}
230  double GetTrackMaxNTpc() {return m_trkNTpcRange[1];}
231  double GetTrackMinDcaXY() {return m_trkDcaRangeXY[0];}
232  double GetTrackMaxDcaXY() {return m_trkDcaRangeXY[1];}
233  double GetTrackMinDcaZ() {return m_trkDcaRangeZ[0];}
234  double GetTrackMaxDcaZ() {return m_trkDcaRangeZ[1];}
235  double GetTrackMinDeltaPt() {return m_trkDeltaPtRange[0];}
236  double GetTrackMaxDeltaPt() {return m_trkDeltaPtRange[1];}
237  double GetFlowMinPt() {return m_flowPtRange[0];}
238  double GetFlowMaxPt() {return m_flowPtRange[1];}
239  double GetFlowMinEta() {return m_flowEtaRange[0];}
240  double GetFlowMaxEta() {return m_flowEtaRange[1];}
241  double GetECalMinPt() {return m_ecalPtRange[0];}
242  double GetECalMaxPt() {return m_ecalPtRange[1];}
243  double GetECalMinEta() {return m_ecalEtaRange[0];}
244  double GetECalMaxEta() {return m_ecalEtaRange[1];}
245  double GetHCalMinPt() {return m_hcalPtRange[0];}
246  double GetHCalMaxPt() {return m_hcalPtRange[1];}
247  double GetHCalMinEta() {return m_hcalEtaRange[0];}
248  double GetHCalMaxEta() {return m_hcalEtaRange[1];}
249 
250  // jet getters
251  double GetJetR() {return m_jetR;}
252  uint32_t GetJetType() {return m_jetType;}
253  JetAlgorithm GetJetAlgo() {return m_jetAlgo;}
254  RecombinationScheme GetRecombScheme() {return m_recombScheme;}
255 
256  private:
257 
258  // constants
259  enum CONST {
260  NPart = 2,
261  NComp = 3,
262  NParam = 4,
263  NRange = 2,
264  NMoment = 2,
265  NInfoQA = 9,
266  NJetType = 2,
267  NCstType = 5,
268  NObjType = 9,
269  NDirectory = 6,
270  NMvtxLayer = 3,
271  NInttLayer = 8,
272  NTpcLayer = 48,
273  NTpcSector = 12
274  };
275 
276  // qa info & tracking subsystems
277  enum SUBSYS {MVTX, INTT, TPC};
278  enum CST_TYPE {PART_CST, TRACK_CST, FLOW_CST, ECAL_CST, HCAL_CST};
279  enum OBJECT {TRACK, ECLUST, HCLUST, FLOW, PART, TJET, RJET, TCST, RCST};
280  enum INFO {PT, ETA, PHI, ENE, QUAL, DCAXY, DCAZ, DELTAPT, NTPC};
281 
282  // event methods (*.evt.h)
283  bool IsGoodVertex(const CLHEP::Hep3Vector vtx);
284  void GetEventVariables(PHCompositeNode* topNode);
285  void GetPartonInfo(PHCompositeNode* topNode);
286  long GetNumTrks(PHCompositeNode* topNode);
287  long GetNumChrgPars(PHCompositeNode* topNode);
288  double GetSumECalEne(PHCompositeNode* topNode);
289  double GetSumHCalEne(PHCompositeNode* topNode);
290  double GetSumParEne(PHCompositeNode* topNode);
291  CLHEP::Hep3Vector GetRecoVtx(PHCompositeNode* topNode);
292 
293  // jet methods (*.jet.h)
294  void FindTrueJets(PHCompositeNode* topNode);
295  void FindRecoJets(PHCompositeNode* topNode);
296  void AddParticles(PHCompositeNode* topNode, vector<PseudoJet>& particles, map<int, pair<Jet::SRC, int>>& fjMap);
297  void AddTracks(PHCompositeNode* topNode, vector<PseudoJet>& particles, map<int, pair<Jet::SRC, int>>& fjMap);
298  void AddFlow(PHCompositeNode* topNode, vector<PseudoJet>& particles, map<int, pair<Jet::SRC, int>>& fjMap);
299  void AddECal(PHCompositeNode* topNode, vector<PseudoJet>& particles, map<int, pair<Jet::SRC, int>>& fjMap);
300  void AddHCal(PHCompositeNode* topNode, vector<PseudoJet>& particles, map<int, pair<Jet::SRC, int>>& fjMap);
301 
302  // constituent methods (*.cst.h)
303  bool IsGoodParticle(HepMC::GenParticle* par, const bool ignoreCharge = false);
304  bool IsGoodTrack(SvtxTrack* track, PHCompositeNode* topNode);
305  bool IsGoodFlow(ParticleFlowElement* flow);
306  bool IsGoodECal(CLHEP::Hep3Vector& hepVecECal);
307  bool IsGoodHCal(CLHEP::Hep3Vector& hepVecHCal);
308  bool IsGoodTrackSeed(SvtxTrack* track);
309  bool IsGoodTrackPhi(SvtxTrack* track, const float phiMaskSize = 0.01); // FIXME make user configurable
310  bool IsFromPrimaryVtx(SvtxTrack* track, PHCompositeNode* topNode);
311  bool IsOutgoingParton(HepMC::GenParticle* par);
312  pair<double, double> GetTrackDcaPair(SvtxTrack* track, PHCompositeNode* topNode);
313  CLHEP::Hep3Vector GetTrackVertex(SvtxTrack* track, PHCompositeNode* topNode);
314  double GetTrackDeltaPt(SvtxTrack* track);
315  float GetParticleCharge(const int pid);
316  int GetNumLayer(SvtxTrack* track, const uint8_t subsys = 0);
317  int GetNumClust(SvtxTrack* track, const uint8_t subsys = 0);
318  int GetMatchID(SvtxTrack* track);
319 
320  // system methods (*.sys.h)
321  void InitVariables();
322  void InitHists();
323  void InitTuples();
324  void InitTrees();
325  void InitFuncs();
326  void InitEvals(PHCompositeNode* topNode);
327  void FillTrueTree();
328  void FillRecoTree();
329  void SaveOutput();
330  void ResetVariables();
331  void DetermineEvtsToGrab(PHCompositeNode* topNode);
332  int CreateJetNode(PHCompositeNode* topNode);
333  int GetEmbedID(PHCompositeNode* topNode, const int iEvtToGrab = 1);
335  GlobalVertex* GetGlobalVertex(PHCompositeNode* topNode, const int iVtxToGrab = -1);
337  HepMC::GenEvent* GetMcEvent(PHCompositeNode* topNode, const int iEvtToGrab = 1);
338  RawClusterContainer* GetClusterStore(PHCompositeNode* topNode, const TString sNodeName);
340 
341  // F4A/utility members
342  Fun4AllHistoManager* m_histMan = NULL;
343  SvtxEvalStack* m_evalStack = NULL;
344  SvtxTrackEval* m_trackEval = NULL;
345 
346  // io members
347  TFile* m_outFile = NULL;
348  TTree* m_recoTree = NULL;
349  TTree* m_trueTree = NULL;
350  string m_outFileName = "";
351  string m_jetTreeName = "";
352  JetMapv1* m_recoJetMap = NULL;
353  JetMapv1* m_trueJetMap = NULL;
354 
355  // tracking members
356  bool isMvtxLayerHit[CONST::NMvtxLayer] = {false};
357  bool isInttLayerHit[CONST::NInttLayer] = {false};
358  bool isTpcLayerHit[CONST::NTpcLayer] = {false};
359 
360  // QA members
361  // TODO factorize out QA-related operations
362  // TODO replace all QA histograms with tuples
363  TH1D* m_hJetArea[CONST::NJetType];
364  TH1D* m_hJetNumCst[CONST::NJetType];
365  TH1D* m_hNumObject[CONST::NObjType];
366  TH1D* m_hSumCstEne[CONST::NCstType];
367  TH1D* m_hObjectQA[CONST::NObjType][CONST::NInfoQA];
368  TH1D* m_hNumCstAccept[CONST::NCstType][CONST::NMoment];
369  TNtuple* m_ntTrkQA = NULL;
370  TNtuple* m_ntWeirdTracks = NULL;
371 
372  // system members
373  bool m_doVtxCut = false;
374  bool m_doQualityPlots = true;
375  bool m_requireSiSeeds = true;
376  bool m_useOnlyPrimVtx = true;
377  bool m_doDcaSigmaCut = false;
378  bool m_maskTpcSectors = false;
379  bool m_saveDST = false;
380  bool m_isMC = true;
381  bool m_isEmbed = false;
382  bool m_doDebug = false;
383  bool m_checkWeirdTrks = false;
384  bool m_addTracks = true;
385  bool m_addFlow = false;
386  bool m_addECal = false;
387  bool m_addHCal = false;
388  vector<int> m_vecEvtsToGrab;
389  map<int, int> m_mapCstToEmbedID;
390 
391  // event acceptance parameters
392  // TODO convert most acceptances to pairs/pairs of structs
393  double m_evtVzRange[CONST::NRange] = {-10., 10.};
394  double m_evtVrRange[CONST::NRange] = {0.0, 0.418};
395 
396  // particle acceptance parameters
397  double m_parPtRange[CONST::NRange] = {0.1, 9999.};
398  double m_parEtaRange[CONST::NRange] = {-1.1, 1.1};
399 
400  // track acceptance parameters
401  double m_trkPtRange[CONST::NRange] = {0.1, 100.};
402  double m_trkEtaRange[CONST::NRange] = {-1.1, 1.1};
403  double m_trkQualRange[CONST::NRange] = {-1., 10.};
404  double m_trkNMvtxRange[CONST::NRange] = {2., 100.};
405  double m_trkNInttRange[CONST::NRange] = {1., 100.};
406  double m_trkNTpcRange[CONST::NRange] = {25., 100.};
407  double m_trkDcaRangeXY[CONST::NRange] = {-5., 5.};
408  double m_trkDcaRangeZ[CONST::NRange] = {-5., 5.};
409  double m_trkDeltaPtRange[CONST::NRange] = {0., 0.5};
410 
411  // particle flow acceptance parameters
412  double m_flowPtRange[CONST::NRange] = {0., 9999.};
413  double m_flowEtaRange[CONST::NRange] = {-1.1, 1.1};
414 
415  // calorimeter acceptance parameters
416  double m_ecalPtRange[CONST::NRange] = {0., 9999.};
417  double m_ecalEtaRange[CONST::NRange] = {-1.1, 1.1};
418  double m_hcalPtRange[CONST::NRange] = {0., 9999.};
419  double m_hcalEtaRange[CONST::NRange] = {-1.1, 1.1};
420 
421  // for pt-dependent dca cuts
422  TF1* m_fSigDcaXY = NULL;
423  TF1* m_fSigDcaZ = NULL;
424  double m_dcaPtFitMaxXY = 15.;
425  double m_dcaPtFitMaxZ = 15.;
426  double m_nSigCutXY = 1.;
427  double m_nSigCutZ = 1.;
428  array<double, CONST::NParam> m_parSigDcaXY = {1., 1., 1., 1.};
429  array<double, CONST::NParam> m_parSigDcaZ = {1., 1., 1., 1.};
430 
431  // jet parameters
432  double m_jetR = 0.4;
433  uint32_t m_jetType = 0;
435  JetDefinition* m_trueJetDef = NULL;
436  JetDefinition* m_recoJetDef = NULL;
437  ClusterSequence* m_trueClust = NULL;
438  ClusterSequence* m_recoClust = NULL;
439  RecombinationScheme m_recombScheme = pt_scheme;
440 
441  // event, jet members
442  long long m_partonID[CONST::NPart];
443  CLHEP::Hep3Vector m_partonMom[CONST::NPart];
444  CLHEP::Hep3Vector m_trueVtx;
445  CLHEP::Hep3Vector m_recoVtx;
446  vector<PseudoJet> m_trueJets;
447  vector<PseudoJet> m_recoJets;
448 
449  // output truth event variables
450  unsigned long m_trueNumJets;
451  long long m_truePartonID[CONST::NPart];
452  double m_truePartonMomX[CONST::NPart];
453  double m_truePartonMomY[CONST::NPart];
454  double m_truePartonMomZ[CONST::NPart];
455  double m_trueVtxX;
456  double m_trueVtxY;
457  double m_trueVtxZ;
458  double m_trueSumPar;
460  // output truth jet variables
461  vector<unsigned long> m_trueJetNCst;
462  vector<unsigned int> m_trueJetID;
463  vector<double> m_trueJetE;
464  vector<double> m_trueJetPt;
465  vector<double> m_trueJetEta;
466  vector<double> m_trueJetPhi;
467  vector<double> m_trueJetArea;
468  // output truth constituent variables
469  vector<vector<int>> m_trueCstID;
470  vector<vector<int>> m_trueCstEmbedID;
471  vector<vector<double>> m_trueCstZ;
472  vector<vector<double>> m_trueCstDr;
473  vector<vector<double>> m_trueCstE;
474  vector<vector<double>> m_trueCstPt;
475  vector<vector<double>> m_trueCstEta;
476  vector<vector<double>> m_trueCstPhi;
477 
478  // output reco event variables
479  unsigned long m_recoNumJets;
480  double m_recoVtxX;
481  double m_recoVtxY;
482  double m_recoVtxZ;
486  // output reco jet variables
487  vector<unsigned long> m_recoJetNCst;
488  vector<unsigned int> m_recoJetID;
489  vector<double> m_recoJetE;
490  vector<double> m_recoJetPt;
491  vector<double> m_recoJetEta;
492  vector<double> m_recoJetPhi;
493  vector<double> m_recoJetArea;
494  // output reco constituent variables
495  vector<vector<int>> m_recoCstMatchID;
496  vector<vector<double>> m_recoCstZ;
497  vector<vector<double>> m_recoCstDr;
498  vector<vector<double>> m_recoCstE;
499  vector<vector<double>> m_recoCstPt;
500  vector<vector<double>> m_recoCstEta;
501  vector<vector<double>> m_recoCstPhi;
502 
503  };
504 
505 } // end SColdQcdCorrelatorAnalysis namespace
506 
507 #endif
508 
509 // end ------------------------------------------------------------------------