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