Belle II Software light-2406-ragdoll
cdcplotmodule.py
1#!/usr/bin/env python3
2
3
10
11
17
18import os
19from ROOT import Belle2
20import matplotlib.cm as colormap
21from matplotlib.patches import Circle
22import matplotlib.pyplot as plt
23from basf2 import Module, Path, process, B2INFO, B2WARNING
24from simulation import add_simulation
25
26
27def plot(x, y, col, show=0):
28 """
29 Plot a list of x/y values, plus CDC superlayer boundaries.
30
31 Returns a pyplot.figure that can be saved.
32 """
33 fig = plt.figure(figsize=(8, 8))
34 ax = fig.add_subplot(111)
35
36 # draw the x/y arrays. note that looping over the hits and
37 # drawing them individually would be much slower
38 for i in range(len(col)):
39 ax.plot(x[i], y[i], marker='.', color=col[i], linestyle='None', markersize=1)
40
41 ax.set_title('CDCSimHits')
42 ax.set_xlabel('x [cm]')
43 ax.set_ylabel('y [cm]')
44 ax.axis('scaled')
45
46 # draw CDC superlayer boundaries
47 layers = [16.8, 25.7, 36.5, 47.6, 58.4, 69.5, 80.2, 91.3, 102.0, 111.1]
48 Circs = [Circle((0, 0), a, facecolor='none', edgecolor='lightgrey')
49 for a in layers]
50 for e in Circs:
51 ax.add_artist(e)
52
53 ax.set_xlim(-130, 130)
54 ax.set_ylim(-130, 130)
55 fig.tight_layout()
56 if show:
57 plt.show()
58 return fig
59
60
61class CDCPlotModule(Module):
62 """An example python module.
63
64 It gathers the x/y position off all CDCSimHits and draws them using
65 matplotlib. The result is saved as a PNG.
66 """
67
68
69 num_events = 0
70
71 def event(self):
72 """reimplementation of Module::event().
73
74 loops over the CDCSimHits in the current event.
75 """
76 simhits = Belle2.PyStoreArray('CDCSimHits')
77
78 # list of lists of simhit positions, one list per mcpart
79 trackhits_x = []
80 trackhits_y = []
81
82 mcparts = []
83 for hit in simhits:
84 mcpart = hit.getRelatedFrom("MCParticles")
85 if mcpart not in mcparts:
86 mcparts.append(mcpart)
87 trackhits_x.append([])
88 trackhits_y.append([])
89 # add simhit to the list corresponding to this particle
90 idx = mcparts.index(mcpart)
91 hitpos = hit.getPosWire() # TVector3
92 trackhits_x[idx].append(hitpos.X())
93 trackhits_y[idx].append(hitpos.Y())
94
95 npart = len(mcparts)
96 if npart > 0:
97 # plot the (x,y) list on a matplotlib figure
98 col = [colormap.jet(1.0 * c / (npart - 1)) for c in range(npart)]
99 fig = plot(trackhits_x, trackhits_y, col)
100
101 filename = f'cdchits_{self.num_events}.png'
102 if os.path.lexists(filename):
103 B2WARNING(filename + ' exists, overwriting ...')
104 else:
105 B2INFO('creating ' + filename + ' ...')
106 fig.savefig(filename)
107
108 self.num_events += 1
109
110 def terminate(self):
111 """reimplementation of Module::terminate()."""
112 B2INFO('terminating CDCPlotModule')
113
114
115# Normal steering file part begins here
116
117# choose the particles you want to simulate
118param_pGun = {
119 'pdgCodes': [211, -211],
120 'nTracks': 4,
121 'varyNTracks': 0,
122 'momentumGeneration': 'uniform',
123 'momentumParams': [0.4, 1.6],
124 'thetaGeneration': 'uniform',
125 'thetaParams': [60., 120.],
126 'phiGeneration': 'uniform',
127 'phiParams': [0, 360],
128 'vertexGeneration': 'uniform',
129 'xVertexParams': [0.0, 0.0],
130 'yVertexParams': [0.0, 0.0],
131 'zVertexParams': [0.0, 0.0],
132}
133
134# Create main path
135main = Path()
136main.add_module('EventInfoSetter', evtNumList=[5])
137main.add_module('ParticleGun', **param_pGun)
138add_simulation(main)
139main.add_module(CDCPlotModule())
140
141process(main)
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72
int num_events
event counter
Definition: plot.py:1