Belle II Software  release-05-01-25
create_dedx_PDFs.cc
1 /**************************************************************************
2  * BASF2 (Belle Analysis Framework 2) *
3  * Copyright(C) 2012 - Belle II Collaboration *
4  * *
5  * Author: The Belle II Collaboration *
6  * Contributors: Jake Bennett, Christian Pulvermacher
7  * *
8  * This software is provided "as is" without any warranty. *
9  **************************************************************************/
10 
11 #include <framework/gearbox/Const.h>
12 #include <reconstruction/dataobjects/DedxConstants.h>
13 
14 #include <TFile.h>
15 #include <TTree.h>
16 #include <TCanvas.h>
17 #include <TH2F.h>
18 #include <TApplication.h>
19 #include <TROOT.h>
20 #include <TStyle.h>
21 #include <TSystem.h>
22 
23 #include <iostream>
24 
25 using namespace Belle2;
26 
28 const int num_dedx_bins = 100;
29 
30 
36 int main(int argc, char* argv[])
37 {
38  const int num_pdg_codes = Const::ChargedStable::c_SetSize;
39 
40  if (argc < 2 or argc > 3) {
41  std::cerr << "Usage: " << argv[0] << " INPUT_FILE [OUTPUT_FILE]\n\n";
42  std::cerr <<
43  "Generates PDFs for the DedxPID module, requires an input .root file created by running DedxPID with enableDebug=True.\n";
44  std::cerr << "If OUTPUT_FILE is not given, histograms will be drawn instead of being saved.\n";
45  return 1;
46  }
47 
48  const TString input_filename(argv[1]);
49  TString output_filename;
50  if (argc == 3)
51  output_filename = TString(argv[2]);
52 
53  //TApplication eats command line arguments, don't use argc/argv after that
54  TApplication app("noname", &argc, argv);
55  gROOT->SetStyle("Plain");
56  gStyle->SetPalette(1);
57 
58  TFile* file = new TFile(input_filename, "READ");
59  TTree* tree = dynamic_cast<TTree*>(file->Get("tree"));
60  if (!tree) {
61  std::cerr << "'tree' not found, aborting.\n";
62  return 1;
63  }
64 
65  TCanvas* c1 = 0;
66 
67 
68  const bool save_graphs = !output_filename.IsNull();
69  TFile* pdffile = 0;
70  if (save_graphs) {
71  pdffile = new TFile(output_filename, "RECREATE");
72  } else {
73  c1 = new TCanvas("c1", input_filename, 600, 900);
74  c1->Divide(num_pdg_codes, Dedx::c_num_detectors);
75  }
76 
77  //create momentum binning (larger bins for high momentum)
78  const int num_p_bins = 69;
79  double pbins[num_p_bins + 1];
80  pbins[0] = 0.0; pbins[1] = 0.05;
81  std::cout << ": " << pbins[0] << "\n";
82  std::cout << ": " << pbins[1] << "\n";
83  for (int bin = 2; bin <= num_p_bins; bin++) {
84  if (bin <= 19)
85  pbins[bin] = 0.025 + 0.025 * bin;
86  else if (bin <= 59)
87  pbins[bin] = pbins[19] + 0.05 * (bin - 19);
88  else
89  pbins[bin] = pbins[59] + 0.3 * (bin - 59);
90  std::cout << ": " << pbins[bin] << "\n";
91  }
92 
93  for (int i = 0; i < 2; i++) { //normal/truncated
94  const bool use_truncated_mean = (i == 1);
95  const char* suffix = use_truncated_mean ? "_trunc" : "";
96 
97  for (int iPart = 0; iPart < num_pdg_codes; iPart++) {
98  TH2F* hists[Dedx::c_num_detectors]; //one for each Detector
99  const int pdg_code = Const::chargedStableSet.at(iPart).getPDGCode();
100  //const int num_dedx_bins = use_truncated_mean ? (num_dedx_bins_hits) : num_dedx_bins_hits;
101 
102  // first prepare the histograms for the VXD
103  for (int detector = 0; detector <= Dedx::c_SVD; detector++) {
104  std::cout << "i " << i << ", d" << detector << "\n";
105  const char* flayer_var = "VXDDedxTracks.dedxLayer";
106  TString flayer_selection;
107  double dedx_cutoff = 0;
108  switch (Dedx::Detector(detector)) {
109  case Dedx::c_PXD:
110  flayer_selection = TString::Format("%s < 0 && %s >= -2", flayer_var, flayer_var);
111  dedx_cutoff = 10e3;
112  break;
113  case Dedx::c_SVD:
114  flayer_selection = TString::Format("%s < -2", flayer_var);
115  dedx_cutoff = 5e6;
116  break;
117  case Dedx::c_CDC:
118  break;
119  }
120 
121  const TString histname = TString::Format("hist_d%i_%i%s", detector, pdg_code, suffix);
122  const char* varname = use_truncated_mean ? "VXDDedxTracks.m_dedxAvgTruncated" : "VXDDedxTracks.dedx";
123 
124  hists[detector] = new TH2F(histname.Data(), histname.Data(),
125  num_p_bins, pbins,
126  num_dedx_bins, 0, dedx_cutoff);
127  hists[detector]->Sumw2(); //save weights (important for fitting)
128  if (use_truncated_mean)
129  tree->Project(histname.Data(),
130  TString::Format("%s[][%i]:VXDDedxTracks.m_p", varname, detector),
131  TString::Format("%s < %g && abs(VXDDedxTracks.m_pdg) == %i && %s", varname, dedx_cutoff, pdg_code, flayer_selection.Data()));
132  else
133  tree->Project(histname.Data(),
134  TString::Format("%s:VXDDedxTracks.m_p", varname),
135  TString::Format("%s < %g && abs(VXDDedxTracks.m_pdg) == %i && %s", varname, dedx_cutoff, pdg_code, flayer_selection.Data()));
136  } //detector type
137 
138  // now add the CDC histograms
139  int detector = (int)Dedx::c_CDC;
140  std::cout << "i " << i << ", d" << detector << "\n";
141  double dedx_cutoff = 0;
142  dedx_cutoff = 3.0;
143 
144  const TString histname = TString::Format("hist_d%i_%i%s", detector, pdg_code, suffix);
145  const char* varname = use_truncated_mean ? "CDCDedxTracks.m_dedxAvgTruncated" : "CDCDedxTracks.m_lDedx";
146 
147  hists[detector] = new TH2F(histname.Data(), histname.Data(),
148  num_p_bins, pbins,
149  num_dedx_bins, 0, dedx_cutoff);
150  hists[detector]->Sumw2(); //save weights (important for fitting)
151  if (use_truncated_mean)
152  tree->Project(histname.Data(),
153  TString::Format("%s[][%i]:CDCDedxTracks.m_pCDC", varname, detector),
154  TString::Format("CDCDedxTracks.m_lNHitsUsed > 15 && %s < %g && abs(CDCDedxTracks.m_pdg) == %i", varname, dedx_cutoff, pdg_code));
155  else
156  tree->Project(histname.Data(),
157  TString::Format("%s:CDCDedxTracks.m_pCDC", varname),
158  TString::Format("CDCDedxTracks.m_lNHitsUsed > 15 && %s < %g && abs(CDCDedxTracks.m_pdg) == %i", varname, dedx_cutoff, pdg_code));
159 
160  //{{{ for each momentum bin, normalize pdf (disable if you want to keep the orginals, e.g. for fitting)
161  if (true) {
162  for (int d = 0; d < Dedx::c_num_detectors; d++) {
163  for (int pbin = 0; pbin <= num_p_bins + 1; pbin++) {
164  // get number of entries in this pbin
165  double integral = 0;
166  for (int dedxbin = 0; dedxbin <= num_dedx_bins + 1; dedxbin++) {
167  integral += hists[d]->GetBinContent(pbin, dedxbin);
168  }
169 
170  if (integral == 0)
171  continue; //nothing to do
172 
173  // normalize this pbin to 1
174  const double normal_width = -(hists[d]->GetYaxis()->GetBinLowEdge(1) - hists[d]->GetYaxis()->GetBinUpEdge(
175  num_dedx_bins)) / num_dedx_bins;
176  for (int dedxbin = 0; dedxbin <= num_dedx_bins + 1; dedxbin++) {
177  hists[d]->SetBinContent(pbin, dedxbin, hists[d]->GetBinContent(pbin,
178  dedxbin) / integral * hists[d]->GetYaxis()->GetBinWidth(dedxbin) / normal_width);
179  //std::cout << d << ", " << dedxbin << ": " << hists[d]->GetYaxis()->GetBinWidth(dedxbin)/normal_width << "\n";
180  }
181  }
182  }
183  }
184  //}}}
185 
186  for (int d = 0; d < Dedx::c_num_detectors; d++) {
187  if (save_graphs) {
188  hists[d]->Write();
189  } else if (i == 0) {
190  TPad* pad = static_cast<TPad*>(c1->cd(num_pdg_codes * d + iPart + 1));
191  pad->SetLogz();
192  hists[d]->Draw("colz");
193  c1->Update();
194  }
195  }
196  } //particle type
197  } //normal / trunc
198 
199  if (save_graphs) {
200  pdffile->Close();
201  delete pdffile;
202 
203  std::cout << "Saved histograms into '" << output_filename.Data() << "'.\n";
204 
205  //no output, so we'll just exit when we're done
206  gSystem->Exit(0);
207  }
208 
209  app.Run(true);
210  return 0;
211 }
Belle2::Const::ChargedStable::c_SetSize
static const unsigned int c_SetSize
Number of elements (for use in array bounds etc.)
Definition: Const.h:491
Belle2::Const::chargedStableSet
static const ParticleSet chargedStableSet
set of charged stable particles
Definition: Const.h:494
Belle2::Const::ParticleType::getPDGCode
int getPDGCode() const
PDG code.
Definition: Const.h:349
main
int main(int argc, char **argv)
Run all tests.
Definition: test_main.cc:77
Belle2
Abstract base class for different kinds of events.
Definition: MillepedeAlgorithm.h:19
Belle2::Const::ParticleSet::at
const ParticleType & at(unsigned int index) const
Return particle at given index, or end() if out of range.
Definition: Const.h:425