g4simulation/g4histos/G4SnglTree.cc

0001 #include "G4SnglTree.h"
0002
0003 #include <g4main/PHG4Hit.h>
0004 #include <g4main/PHG4HitContainer.h>  // for PHG4HitContainer, PHG4Hit...
0005 #include <g4main/PHG4Particle.h>
0006 #include <g4main/PHG4TruthInfoContainer.h>
0007
0008 #include <fun4all/SubsysReco.h>  // for SubsysReco
0009
0010 #include <phool/getClass.h>
0011
0012 #include <TFile.h>
0013 #include <TTree.h>
0014
0015 #include <cmath>     // for atan2, sqrt
0016 #include <cstring>   // for strcmp
0017 #include <iostream>  // for ostringstream, operator<<
0018 #include <sstream>
0019 #include <utility>  // for pair
0020
0021 G4SnglTree::G4SnglTree(const std::string &name, const std::string &filename)
0022   : SubsysReco(name)
0023   , _filename(filename)
0024 {
0025 }
0026
0027 int G4SnglTree::Init(PHCompositeNode * /*unused*/)
0028 {
0029   delete outfile; // make cppcheck happy
0030   outfile = new TFile(_filename.c_str(), "RECREATE");
0031   g4tree = new TTree("mG4EvtTree", "g4tree");
0032   g4tree->SetAutoSave(1000000);
0033
0034   std::cout << "Initialize Geant 4 Tree ... << " << std::endl;
0035
0036   /// Event level
0037   g4tree->Branch("energy", &mG4EvtTree.energy, "energy/F");
0038   g4tree->Branch("theta", &mG4EvtTree.theta, "theta/F");
0039   g4tree->Branch("phi", &mG4EvtTree.phi, "phi/F");
0040   g4tree->Branch("px", &mG4EvtTree.px, "px/F");
0041   g4tree->Branch("py", &mG4EvtTree.py, "py/F");
0042   g4tree->Branch("pz", &mG4EvtTree.pz, "pz/F");
0043   g4tree->Branch("cemcactLayers", &mG4EvtTree.cemcactLayers, "cemcactLayers/I");
0044   g4tree->Branch("cemcabsLayers", &mG4EvtTree.cemcabsLayers, "cemcabsLayers/I");
0045   g4tree->Branch("hcalactLayers", &mG4EvtTree.hcalactLayers, "hcalactLayers/I");
0046   g4tree->Branch("hcalabsLayers", &mG4EvtTree.hcalabsLayers, "hcalabsLayers/I");
0047   g4tree->Branch("cemcactESum", mG4EvtTree.cemcactESum, "cemcactESum[cemcactLayers]/F");
0048   g4tree->Branch("cemcabsESum", mG4EvtTree.cemcabsESum, "cemcabsESum[cemcabsLayers]/F");
0049   g4tree->Branch("hcalactESum", mG4EvtTree.hcalactESum, "hcalactESum[hcalactLayers]/F");
0050   g4tree->Branch("hcalabsESum", mG4EvtTree.hcalabsESum, "hcalabsESum[hcalabsLayers]/F");
0051
0052   /// Hit level
0053   g4tree->Branch("nhits", &mG4EvtTree.nhits, "nhits/I");
0054   g4tree->Branch("detid", mG4EvtTree.detid, "detid[nhits]/I");
0055   g4tree->Branch("layer", mG4EvtTree.layer, "layer[nhits]/I");
0056   g4tree->Branch("hitid", mG4EvtTree.hitid, "hitid[nhits]/I");
0057   g4tree->Branch("trkid", mG4EvtTree.trkid, "trkid[nhits]/I");
0058   g4tree->Branch("scintid", mG4EvtTree.scintid, "scintid[nhits]/I");
0059   g4tree->Branch("x0", mG4EvtTree.x0, "x0[nhits]/F");
0060   g4tree->Branch("y0", mG4EvtTree.y0, "y0[nhits]/F");
0061   g4tree->Branch("z0", mG4EvtTree.z0, "z0[nhits]/F");
0062   g4tree->Branch("x1", mG4EvtTree.x1, "x1[nhits]/F");
0063   g4tree->Branch("y1", mG4EvtTree.y1, "y1[nhits]/F");
0064   g4tree->Branch("z1", mG4EvtTree.z1, "z1[nhits]/F");
0065   g4tree->Branch("edep", mG4EvtTree.edep, "edep[nhits]/F");
0066
0067   return 0;
0068 }
0069
0070 int G4SnglTree::process_event(PHCompositeNode *topNode)
0071 {
0072   // get the primary particle which did this to us....
0073   PHG4TruthInfoContainer *truthInfoList = findNode::getClass<PHG4TruthInfoContainer>(topNode, "G4TruthInfo");
0074
0075   const PHG4TruthInfoContainer::Range primRange = truthInfoList->GetPrimaryParticleRange();
0076   double px = primRange.first->second->get_px();
0077   double py = primRange.first->second->get_py();
0078   double pz = primRange.first->second->get_pz();
0079   double e = primRange.first->second->get_e();
0080   double pt = sqrt(px * px + py * py);
0081   double phi = atan2(py, px);
0082   double theta = atan2(pt, pz);
0083
0084   mG4EvtTree.energy = e;
0085   mG4EvtTree.theta = theta;
0086   mG4EvtTree.phi = phi;
0087   mG4EvtTree.px = px;
0088   mG4EvtTree.py = py;
0089   mG4EvtTree.pz = pz;
0090
0091   int nhits = 0;
0092
0093   std::string nodename;
0094
0095   for (const auto &iter : _node_postfix)
0096   {
0097     int detid = (_detid.find(iter))->second;
0098     nodename = "G4HIT_" + iter;
0099     PHG4HitContainer *hits = findNode::getClass<PHG4HitContainer>(topNode, nodename);
0100
0101     if (nodename == "G4HIT_CEMC")  // CEMC scintillator
0102     {
0103       mG4EvtTree.cemcactLayers = process_hit(hits, "G4HIT_CEMC", detid, nhits);
0104     }
0105     else if (nodename == "G4HIT_ABSORBER_CEMC")  // CEMC Aabsorber G4_W
0106     {
0107       mG4EvtTree.cemcabsLayers = process_hit(hits, "G4HIT_ABSORBER_CEMC", detid, nhits);
0108     }
0109     else if (nodename == "G4HIT_HCAL")  // HCAL Active scintilltor
0110     {
0111       mG4EvtTree.hcalactLayers = process_hit(hits, "G4HIT_HCAL", detid, nhits);
0112     }
0113     else if (nodename == "G4HIT_ABSORBER_HCAL")  // HCAL Aabsorber steel
0114     {
0115       mG4EvtTree.hcalabsLayers = process_hit(hits, "G4HIT_ABSORBER_HCAL", detid, nhits);
0116     }
0117   }
0118
0119   mG4EvtTree.nhits = nhits;
0120
0121   if (g4tree)
0122   {
0123     g4tree->Fill();
0124   }
0125
0126   return 0;
0127 }
0128
0129 int G4SnglTree::End(PHCompositeNode * /*topNode*/)
0130 {
0131   outfile->cd();
0132   g4tree->Write();
0133   outfile->Write();
0134   outfile->Close();
0135   delete outfile;
0136   return 0;
0137 }
0138
0139 void G4SnglTree::AddNode(const std::string &name, const int detid)
0140 {
0141   _node_postfix.insert(name);
0142   _detid[name] = detid;
0143   return;
0144 }
0145
0146 int G4SnglTree::process_hit(PHG4HitContainer *hits, const std::string &dName, int detid, int &nhits)
0147 {
0148   std::map<int, double> layer_edep_map;
0149
0150   int nLayers = 0;
0151   if (hits)
0152   {
0153     // double esum = 0;
0154     PHG4HitContainer::ConstRange hit_range = hits->getHits();
0155     for (PHG4HitContainer::ConstIterator hit_iter = hit_range.first; hit_iter != hit_range.second; hit_iter++)
0156     {
0157       layer_edep_map[hit_iter->second->get_layer()] += hit_iter->second->get_edep();
0158       mG4EvtTree.detid[nhits] = detid;
0159       mG4EvtTree.layer[nhits] = hit_iter->second->get_layer();
0160       mG4EvtTree.hitid[nhits] = hit_iter->second->get_hit_id();
0161       mG4EvtTree.trkid[nhits] = hit_iter->second->get_trkid();
0162       mG4EvtTree.scintid[nhits] = hit_iter->second->get_scint_id();
0163       mG4EvtTree.x0[nhits] = hit_iter->second->get_x(0);
0164       mG4EvtTree.y0[nhits] = hit_iter->second->get_y(0);
0165       mG4EvtTree.z0[nhits] = hit_iter->second->get_z(0);
0166       mG4EvtTree.x1[nhits] = hit_iter->second->get_x(1);
0167       mG4EvtTree.y1[nhits] = hit_iter->second->get_y(1);
0168       mG4EvtTree.z1[nhits] = hit_iter->second->get_z(1);
0169       mG4EvtTree.edep[nhits] = hit_iter->second->get_edep();
0170       nhits++;
0171       // esum += hit_iter->second->get_edep();
0172     }
0173     for (auto edepiter : layer_edep_map)
0174     {
0175       nLayers = edepiter.first - 1;
0176       if (dName == "G4HIT_CEMC")
0177       {
0178         mG4EvtTree.cemcactESum[nLayers] = edepiter.second;
0179       }
0180       else if (dName == "G4HIT_ABSORBER_CEMC")
0181       {
0182         mG4EvtTree.cemcabsESum[nLayers] = edepiter.second;
0183       }
0184       else if (dName == "G4HIT_HCAL")
0185       {
0186         mG4EvtTree.hcalactESum[nLayers] = edepiter.second;
0187       }
0188       else if (dName == "G4HIT_ABSORBER_HCAL")
0189       {
0190         mG4EvtTree.hcalabsESum[nLayers] = edepiter.second;
0191       }
0192     }
0193   }
0194
0195   return nLayers + 1;
0196 }