MC-Calibrations/src/JESMCTreeGen.h

0001 // Tell emacs that this is a C++ source
0002 //  -*- C++ -*-.
0003 #ifndef JESMCTREEGEN_H
0004 #define JESMCTREEGEN_H
0005
0006 #include <string>
0007 #include <vector>
0008 #include <limits>
0009
0010 #include <TFile.h>
0011 #include <TTree.h>
0012
0013 #include <fun4all/SubsysReco.h>
0014
0015 class PHCompositeNode;
0016
0017 class JESMCTreeGen : public SubsysReco
0018 {
0019  public:
0020   JESMCTreeGen(const std::string &name = "JESMCTreeGen", const std::string &outfilename = "output.root");
0021   ~JESMCTreeGen() override = default;
0022   int Init(PHCompositeNode *topNode) override;
0023   int process_event(PHCompositeNode *topNode) override;
0024   int ResetEvent(PHCompositeNode *topNode) override;
0025   int End(PHCompositeNode *topNode) override;
0026
0027   // ********** Setters ********** //
0028   void SetVerbosity(int verbo) {verbosity = verbo;}
0029   void DoZVertexReadout(bool operate) {doZVertexReadout = operate;}
0030   void DoTowerInfoReadout(bool operate) {doTowerInfoReadout = operate;}
0031     void DoTowerInfoReadout_RetowerEMCal(bool operate) {doTowerInfoReadout_RetowerEMCal = operate;}
0032   void DoClusterReadout(bool operate) {doClusterReadout = operate;}
0033     void DoClusterIsoReadout(bool operate) {doClusterIsoReadout = operate;}
0034     void DoClusterCNNReadout(bool operate) {doClusterCNNReadout = operate;}
0035   void DoRecoJetReadout(bool operate) {doRecoJetReadout = operate;}
0036     void DoRecoJet_radius(std::vector<float> radius_vector) {doRecoJet_radius.clear(); doRecoJet_radius = radius_vector;}
0037     void DoSeedJetRawReadout(bool operate) {doSeedJetRawReadout = operate;}
0038     void DoSeedJetSubReadout(bool operate) {doSeedJetSubReadout = operate;}
0039   void DoTruthJetReadout(bool operate) {doTruthJetReadout = operate;}
0040     void DoTruthJet_radius(std::vector<float> radius_vector) {doTruthJet_radius.clear(); doTruthJet_radius = radius_vector;}
0041   void DoTruthParticleReadout(bool operate) {doTruthParticleReadout = operate;}
0042
0043   // ********** Functions ********** //
0044   void Initialize_z_vertex();
0045   void Initialize_calo_tower();
0046   void Initialize_calo_tower_RetowerEMCal();
0047   void Initialize_cluster();
0048   void Initialize_clusterIso();
0049   void Initialize_clusterCNN();
0050   void Initialize_jet(std::string jet_prefix , float radius);
0051   void Initialize_truthparticle();
0052   void Fill_z_vertex(PHCompositeNode *topNode);
0053   void Fill_calo_tower(PHCompositeNode *topNode, std::string calorimeter);
0054   void Fill_emcal_retower(PHCompositeNode *topNode);
0055   void Fill_cluster(PHCompositeNode *topNode);
0056   void Fill_recojet(PHCompositeNode *topNode, std::string jet_prefix, std::string node_prefix, float radius);
0057   void Fill_truthjet(PHCompositeNode *topNode, float radius);
0058   void Fill_truthparticle(PHCompositeNode *topNode);
0059   void Reset_cluster();
0060   void Reset_jet();
0061   void Reset_truthparticle();
0062  private:
0063
0064   // ********** General variables ********** //
0065   TFile *file{nullptr};
0066   TTree *tree{nullptr};
0067   std::string foutname{"output.root"};
0068   int verbosity{0};
0069   int ievent{0};
0070   // Operation flags.
0071   bool doZVertexReadout{true};
0072   bool doTowerInfoReadout{false};
0073     bool doTowerInfoReadout_RetowerEMCal{false};
0074   bool doClusterReadout{false};
0075     bool doClusterIsoReadout{false};
0076     bool doClusterCNNReadout{false};
0077   bool doRecoJetReadout{true};
0078     std::vector<float> doRecoJet_radius{0.4};
0079     bool doSeedJetRawReadout{false};
0080     bool doSeedJetSubReadout{false};
0081   bool doTruthJetReadout{true};
0082     std::vector<float> doTruthJet_radius{0.4};
0083   bool doTruthParticleReadout{false};
0084
0085   // ********** Constants ********** //
0086   // Tower information.
0087   static const int n_emcal_tower = 24576;
0088   static const int n_emcal_tower_etabin = 96;
0089   static const int n_emcal_tower_phibin = 256;
0090   static const int n_hcal_tower = 1536;
0091   static const int n_hcal_tower_etabin = 24;
0092   static const int n_hcal_tower_phibin = 64;
0093
0094   // ********** Tree variables ********** //
0095   // Vertex information.
0096   float z_vertex{std::numeric_limits<float>::quiet_NaN()};
0097   // Calo tower information.
0098   float emcal_tower_e[n_emcal_tower_etabin][n_emcal_tower_phibin]{}, emcal_tower_time[n_emcal_tower_etabin][n_emcal_tower_phibin]{};
0099   float ihcal_tower_e[n_hcal_tower_etabin][n_hcal_tower_phibin]{}, ihcal_tower_time[n_hcal_tower_etabin][n_hcal_tower_phibin]{};
0100   float ohcal_tower_e[n_hcal_tower_etabin][n_hcal_tower_phibin]{}, ohcal_tower_time[n_hcal_tower_etabin][n_hcal_tower_phibin]{};
0101   float emcal_retower_e[n_hcal_tower_etabin][n_hcal_tower_phibin]{};
0102   int emcal_tower_status[n_emcal_tower_etabin][n_emcal_tower_phibin]{};
0103   int ihcal_tower_status[n_hcal_tower_etabin][n_hcal_tower_phibin]{};
0104   int ohcal_tower_status[n_hcal_tower_etabin][n_hcal_tower_phibin]{};
0105   // Cluster information.
0106   std::vector<float> cluster_e={}, cluster_pt={}, cluster_eta={}, cluster_phi={};
0107   std::vector<float> cluster_e_core={}, cluster_pt_core={}, cluster_eta_core={}, cluster_phi_core={};
0108   std::vector<float> cluster_probability={};
0109   std::vector<float> cluster_iso={};
0110   std::vector<float> cluster_e_cnn={}, cluster_probability_cnn={};
0111   // Jet general information.
0112   std::map<std::string, std::vector<float>> jet_e_map={};
0113   std::map<std::string, std::vector<float>> jet_et_map={};
0114   std::map<std::string, std::vector<float>> jet_pt_map={};
0115   std::map<std::string, std::vector<float>> jet_eta_map={};
0116   std::map<std::string, std::vector<float>> jet_phi_map={};
0117   // Jet tower information.
0118   std::map<std::string, std::vector<int>> jet_tower_ncomponent_map={};
0119   std::map<std::string, std::vector<float>> jet_tower_e_map={};
0120   std::map<std::string, std::vector<int>> jet_tower_ieta_map={};
0121   std::map<std::string, std::vector<int>> jet_tower_iphi_map={};
0122   // Jet particle information.
0123   std::map<std::string, std::vector<float>> jet_truthparticle_e_map={};
0124   std::map<std::string, std::vector<float>> jet_truthparticle_pt_map={};
0125   std::map<std::string, std::vector<float>> jet_truthparticle_eta_map={};
0126   std::map<std::string, std::vector<float>> jet_truthparticle_phi_map={};
0127   std::map<std::string, std::vector<int>> jet_truthparticle_pid_map={};
0128   // Truth particle information.
0129   std::vector<float> truthparticle_e={}, truthparticle_pt={}, truthparticle_eta={}, truthparticle_phi={};
0130   std::vector<int> truthparticle_pid={};
0131 };
0132
0133 #endif