Belle II Software  release-05-02-19
MasterClassModule.cc
1 /**************************************************************************
2 * BASF2 (Belle Analysis Framework 2) *
3 * Copyright(C) 2020 - Belle II Collaboration *
4 * *
5 * Author: The Belle II Collaboration *
6 * Contributors: Thomas Kuhr *
7 * *
8 * This software is provided "as is" without any warranty. *
9 **************************************************************************/
10 
11 #include <masterclass/modules/MasterClassModule.h>
12 #include <mdst/dataobjects/PIDLikelihood.h>
13 
14 #include "TDirectory.h"
15 
16 using namespace Belle2;
17 
18 //-----------------------------------------------------------------
19 // Register the Module
20 //-----------------------------------------------------------------
21 REG_MODULE(MasterClass)
22 
23 //-----------------------------------------------------------------
24 // Implementation
25 //-----------------------------------------------------------------
26 
28 {
29  // Set module properties
30  setDescription(R"DOC("Module to write out data in a format for Belle II masterclasses)DOC");
31  addParam("outputFileName", m_filename, "File name of root ntuple output.", std::string("masterclass.root"));
32 }
33 
35 {
36  m_tracks.isRequired();
37  m_clusters.isRequired();
38 
39  m_event = new BEvent;
40  m_file = TFile::Open(m_filename.c_str(), "RECREATE");
41  m_tree = new TTree("T", "Event Tree");
42  m_tree->Branch("BEvent", &m_event);
43 }
44 
46 {
47  m_event->Clear();
48  m_event->EventNo(m_index++);
49 
50  for (auto& track : m_tracks) {
51  auto pid = track.getRelated<PIDLikelihood>();
52  const double priors[] = {0.05, 0.05, 0.65, 0.24, 0.01, 0};
53  auto type = pid->getMostLikely(priors);
54  auto trackFit = track.getTrackFitResultWithClosestMass(type);
55  auto p = trackFit->getMomentum();
56  double m = type.getMass();
57  SIMPLEPID id = ALL;
58  switch (type.getPDGCode()) {
59  case 11: id = ELECTRON; break;
60  case 13: id = MUON; break;
61  case 211: id = PION; break;
62  case 321: id = KAON; break;
63  case 2212: id = PROTON; break;
64  default: id = ALL;
65  }
66  m_event->AddTrack(p.x(), p.y(), p.z(), sqrt(m * m + p.Mag2()), trackFit->getChargeSign(), id);
67  }
68 
69  for (auto& cluster : m_clusters) {
70  if (!cluster.hasHypothesis(ECLCluster::EHypothesisBit::c_nPhotons)) continue;
71  double E = cluster.getEnergy(ECLCluster::EHypothesisBit::c_nPhotons);
72  if (E < 0.1) continue;
73  TVector3 p = cluster.getClusterPosition();
74  p.SetMag(E);
75  m_event->AddTrack(p.x(), p.y(), p.z(), E, 0, PHOTON);
76  }
77 
78  m_tree->Fill();
79 }
80 
82 {
83  m_tree->Write();
84  m_file->Close();
85  delete m_file;
86 }
Belle2::MasterClassModule::m_event
BEvent * m_event
output event object
Definition: MasterClassModule.h:67
Belle2::MasterClassModule::event
virtual void event() override
Write out particles.
Belle2::MasterClassModule::initialize
virtual void initialize() override
Register input and output data.
Belle2::MasterClassModule
Module to write out data in a format for Belle II masterclasses.
Definition: MasterClassModule.h:41
REG_MODULE
#define REG_MODULE(moduleName)
Register the given module (without 'Module' suffix) with the framework.
Definition: Module.h:652
Belle2::MasterClassModule::m_index
int m_index
event number
Definition: MasterClassModule.h:66
Belle2::MasterClassModule::m_file
TFile * m_file
root ntuple file
Definition: MasterClassModule.h:64
Belle2::MasterClassModule::m_tracks
StoreArray< Track > m_tracks
Track objects.
Definition: MasterClassModule.h:61
Belle2::ECLCluster::EHypothesisBit::c_nPhotons
@ c_nPhotons
CR is split into n photons (N1)
Belle2::PIDLikelihood
Class to collect log likelihoods from TOP, ARICH, dEdx, ECL and KLM aimed for output to mdst includes...
Definition: PIDLikelihood.h:37
BEvent
Definition: BEvent.h:16
Belle2::MasterClassModule::terminate
virtual void terminate() override
Close ntuple.
Belle2::Module
Base class for Modules.
Definition: Module.h:74
Belle2
Abstract base class for different kinds of events.
Definition: MillepedeAlgorithm.h:19
Belle2::MasterClassModule::m_tree
TTree * m_tree
output tree
Definition: MasterClassModule.h:65
Belle2::MasterClassModule::m_filename
std::string m_filename
output file name
Definition: MasterClassModule.h:63
Belle2::MasterClassModule::m_clusters
StoreArray< ECLCluster > m_clusters
Cluster objects.
Definition: MasterClassModule.h:62