Belle II Software  release-05-01-25
makeMCParticlesFromGCRData.py
1 #!/usr/bin/env python
2 # -*- coding: utf-8 -*-
3 
4 # --------------------------------------------------------------------------------
5 # Make MCParticles from reconstructed cosmic tracks and write to basf2 root file
6 # This file can be used as input to simulation, see examples/simGCRData.py
7 # --------------------------------------------------------------------------------
8 
9 from basf2 import *
10 from rawdata import add_unpackers
11 from reconstruction import add_cosmics_reconstruction
12 from ROOT import Belle2
13 import math
14 
15 
16 class MakeMCParticles(Module):
17  ''' make MCParticles from reconstructed cosmic tracks '''
18 
19  def initialize(self):
20  ''' initialization '''
21 
22  mcParticles = Belle2.PyStoreArray('MCParticles')
23  mcParticles.registerInDataStore()
24 
25  def event(self):
26  ''' event processing '''
27 
28  mcParticles = Belle2.PyStoreArray('MCParticles')
29  tracks = Belle2.PyStoreArray('Tracks')
30  absPDG = abs(Belle2.Const.muon.getPDGCode())
31  mass = Belle2.Const.muon.getMass()
32  for track in tracks:
33  extHits = track.getRelationsWith('ExtHits')
34  selExtHits = []
35  for extHit in extHits:
36  if extHit.getDetectorID() != Belle2.Const.TOP:
37  continue
38  if abs(extHit.getPdgCode()) != absPDG:
39  continue
40  if extHit.getPosition().Perp() < 123.5:
41  continue
42  if extHit.getPosition() * extHit.getMomentum() > 0:
43  continue
44  selExtHits.append(extHit)
45  if len(selExtHits) == 0:
46  continue
47  sortedExtHits = sorted(selExtHits, key=lambda x: (x.getTOF()))
48  extHit = sortedExtHits[0]
49  pos = extHit.getPosition()
50  mom = extHit.getMomentum()
51  tof = extHit.getTOF()
52  part = mcParticles.appendNew()
53  part.setPDG(extHit.getPdgCode())
54  part.setMass(mass)
55  part.setProductionVertex(pos)
56  part.setProductionTime(tof)
57  part.setMomentum(mom)
58  part.setEnergy(math.sqrt(mom.Mag2() + mass**2))
59  part.setValidVertex(True)
60  part.setStatus(Belle2.MCParticle.c_PrimaryParticle)
61  part.addStatus(Belle2.MCParticle.c_StableInGenerator)
62 
63  if mcParticles.getEntries() > 0:
64  self.return_value(1)
65  else:
66  self.return_value(0)
67 
68 
69 # Create path
70 main = create_path()
71 emptypath = create_path()
72 
73 # input
74 roinput = register_module('RootInput')
75 # roinput = register_module('SeqRootInput')
76 main.add_module(roinput)
77 
78 # geometry parameters
79 gearbox = register_module('Gearbox')
80 main.add_module(gearbox)
81 
82 # Geometry
83 geometry = register_module('Geometry')
84 main.add_module(geometry)
85 
86 # unpackers
87 add_unpackers(main, components=['CDC'])
88 
89 # reconstruction
90 add_cosmics_reconstruction(main, components=['CDC'], merge_tracks=True)
91 
92 # make MCParticles from reconstructed tracks
93 maker = MakeMCParticles()
94 main.add_module(maker)
95 maker.if_false(emptypath)
96 
97 # output
98 output = register_module('RootOutput')
99 output.param('branchNames', ['MCParticles'])
100 main.add_module(output)
101 
102 # Print progress
103 progress = register_module('Progress')
104 main.add_module(progress)
105 
106 # Process events
107 process(main)
108 
109 # Print statistics
110 print(statistics)
makeMCParticlesFromGCRData.MakeMCParticles.event
def event(self)
Definition: makeMCParticlesFromGCRData.py:25
Belle2::PyStoreArray
a (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:58
makeMCParticlesFromGCRData.MakeMCParticles
Definition: makeMCParticlesFromGCRData.py:16
makeMCParticlesFromGCRData.MakeMCParticles.initialize
def initialize(self)
Definition: makeMCParticlesFromGCRData.py:19