Belle II Software development
EclCrystalInfoPrinter.py
1#!/usr/bin/env python3
2
3
10
11# Doxygen should skip this script
12# @cond
13
14import basf2
15from ROOT import Belle2
16import numpy as np
17import pandas as pd
18
19# --------------------------------
20
21
22class printCrystals(basf2.Module):
23 """Module to save crystal information to dataframes"""
24
25 def __init__(self):
26 super().__init__() # don't forget to call parent constructor
27 self.obj_eclgeometrypar = Belle2.ECL.ECLGeometryPar.Instance()
28 self.obj_neighbours1 = Belle2.ECL.ECLNeighbours("N", 1)
29 self.obj_neighbours1C = Belle2.ECL.ECLNeighbours("NC", 1)
30
31 self.obj_neighbours2 = Belle2.ECL.ECLNeighbours("N", 2)
32 self.obj_neighbours2C = Belle2.ECL.ECLNeighbours("NC", 2)
33
34 self.obj_neighbours1legacy = Belle2.ECL.ECLNeighbours("NLegacy", 1)
35 self.obj_neighbours1Clegacy = Belle2.ECL.ECLNeighbours("NCLegacy", 1)
36
37 self.obj_neighbours2legacy = Belle2.ECL.ECLNeighbours("NLegacy", 2)
38 self.obj_neighbours2Clegacy = Belle2.ECL.ECLNeighbours("NCLegacy", 2)
39
40 self.list_pos = []
41 self.list_theta = []
42 self.list_phi = []
43
44 self.list_neighbours1 = []
45 self.list_nneighbours1 = []
46
47 self.list_neighbours1c = []
48 self.list_nneighbours1c = []
49
50 self.list_neighbours2 = []
51 self.list_nneighbours2 = []
52
53 self.list_neighbours2c = []
54 self.list_nneighbours2c = []
55
56 self.list_neighbours1legacy = []
57 self.list_nneighbours1legacy = []
58
59 self.list_neighbours1clegacy = []
60 self.list_nneighbours1clegacy = []
61
62 self.list_neighbours2legacy = []
63 self.list_nneighbours2legacy = []
64
65 self.list_neighbours2clegacy = []
66 self.list_nneighbours2clegacy = []
67
68 def initialize(self):
69 for idx in range(8736):
70 pos = self.obj_eclgeometrypar.getCrystalPos(idx)
71 theta = pos.theta()
72 phi = pos.phi()
73 x = pos.X()
74 y = pos.Y()
75 z = pos.Z()
76
77 self.obj_eclgeometrypar.Mapping(idx)
78 thetaid = int(self.obj_eclgeometrypar.GetThetaID())
79 phiid = int(self.obj_eclgeometrypar.GetPhiID())
80
81 self.list_pos.append(np.array([x, y, z, phi, theta, phiid, thetaid]))
82
83 nn1 = np.array(self.obj_neighbours1.getNeighbours(idx + 1))
84 nn1c = np.array(self.obj_neighbours1C.getNeighbours(idx + 1))
85 nn1legacy = np.array(self.obj_neighbours1legacy.getNeighbours(idx + 1))
86 nn1clegacy = np.array(self.obj_neighbours1Clegacy.getNeighbours(idx + 1))
87 self.list_neighbours1.append(nn1)
88 self.list_nneighbours1.append(len(nn1))
89
90 self.list_neighbours1c.append(nn1c)
91 self.list_nneighbours1c.append(len(nn1c))
92
93 self.list_neighbours1legacy.append(nn1legacy)
94 self.list_nneighbours1legacy.append(len(nn1legacy))
95
96 self.list_neighbours1clegacy.append(nn1clegacy)
97 self.list_nneighbours1clegacy.append(len(nn1clegacy))
98
99 nn2 = np.array(self.obj_neighbours2.getNeighbours(idx + 1))
100 nn2c = np.array(self.obj_neighbours2C.getNeighbours(idx + 1))
101 nn2legacy = np.array(self.obj_neighbours2legacy.getNeighbours(idx + 1))
102 nn2clegacy = np.array(self.obj_neighbours2Clegacy.getNeighbours(idx + 1))
103
104 self.list_neighbours2.append(nn2)
105 self.list_nneighbours2.append(len(nn2))
106
107 self.list_neighbours2legacy.append(nn2legacy)
108 self.list_nneighbours2legacy.append(len(nn2legacy))
109
110 self.list_neighbours2c.append(nn2c)
111 self.list_nneighbours2c.append(len(nn2c))
112
113 self.list_neighbours2clegacy.append(nn2clegacy)
114 self.list_nneighbours2clegacy.append(len(nn2clegacy))
115
116 # crystal coordinates
117 dfcoord = pd.DataFrame(data=self.list_pos)
118 dfcoord.to_csv('crystals_coordinates.csv', index=True, header=False)
119
120 # neighbour maps: N1 and NC1
121 dfneighbours1 = pd.DataFrame(data=self.list_neighbours1)
122 dfneighbours1.to_csv('crystals_neighbours1.csv', index=True, header=False)
123 dfneighbours1c = pd.DataFrame(data=self.list_neighbours1c)
124 dfneighbours1c.to_csv('crystals_neighbours1c.csv', index=True, header=False)
125
126 # number of neighbours
127 dfnneighbours1 = pd.DataFrame(data=self.list_nneighbours1)
128 dfnneighbours1.to_csv('crystals_nneighbours1.csv', index=True, header=False)
129 dfnneighbours1c = pd.DataFrame(data=self.list_nneighbours1c)
130 dfnneighbours1c.to_csv('crystals_nneighbours1.csv', index=True, header=False)
131
132 # neighbour maps: N2 and NC2
133 dfneighbours2 = pd.DataFrame(data=self.list_neighbours2)
134 dfneighbours2.to_csv('crystals_neighbours2.csv', index=True, header=False)
135 dfneighbours2c = pd.DataFrame(data=self.list_neighbours2c)
136 dfneighbours2c.to_csv('crystals_neighbours2c.csv', index=True, header=False)
137
138 # number of neighbours
139 dfnneighbours2 = pd.DataFrame(data=self.list_nneighbours2)
140 dfnneighbours2.to_csv('crystals_nneighbours2.csv', index=True, header=False)
141 dfnneighbours2c = pd.DataFrame(data=self.list_nneighbours2c)
142 dfnneighbours2c.to_csv('crystals_nneighbours2c.csv', index=True, header=False)
143
144 # neighbour maps: N1 and NC1 LEGACY VERSION
145 dfneighbours1legacy = pd.DataFrame(data=self.list_neighbours1legacy)
146 dfneighbours1legacy.to_csv('crystals_neighbours1legacy.csv', index=True, header=False)
147 dfneighbours1clegacy = pd.DataFrame(data=self.list_neighbours1clegacy)
148 dfneighbours1clegacy.to_csv('crystals_neighbours1clegacy.csv', index=True, header=False)
149
150 # number of neighbours
151 dfnneighbours1legacy = pd.DataFrame(data=self.list_nneighbours1legacy)
152 dfnneighbours1legacy.to_csv('crystals_nneighbours1legacy.csv', index=True, header=False)
153 dfnneighbours1clegacy = pd.DataFrame(data=self.list_nneighbours1clegacy)
154 dfnneighbours1clegacy.to_csv('crystals_nneighbours1clegacy.csv', index=True, header=False)
155
156 # neighbour maps: N2 and NC2 LEGACY VERSION
157 dfneighbours2legacy = pd.DataFrame(data=self.list_neighbours2legacy)
158 dfneighbours2legacy.to_csv('crystals_neighbours2legacy.csv', index=True, header=False)
159 dfneighbours2clegacy = pd.DataFrame(data=self.list_neighbours2clegacy)
160 dfneighbours2clegacy.to_csv('crystals_neighbours2clegacy.csv', index=True, header=False)
161
162 # number of neighbours
163 dfnneighbours2legacy = pd.DataFrame(data=self.list_nneighbours2legacy)
164 dfnneighbours2legacy.to_csv('crystals_nneighbours2legacy.csv', index=True, header=False)
165 dfnneighbours2clegacy = pd.DataFrame(data=self.list_nneighbours2clegacy)
166 dfnneighbours2clegacy.to_csv('crystals_nneighbours2clegacy.csv', index=True, header=False)
167
168 def event(self):
169 """Event loop"""
170
171
172# create path
173main_path = basf2.create_path()
174
175main_path.add_module('EventInfoSetter', expList=[1003])
176main_path.add_module('Gearbox')
177main_path.add_module('Geometry')
178printCrystals = printCrystals()
179main_path.add_module(printCrystals)
180
181basf2.process(main_path)
182print(basf2.statistics)
183
184# @endcond
Belle2::ECL::ECLGeometryPar::Instance
static ECLGeometryPar * Instance()
Static method to get a reference to the ECLGeometryPar instance.
Definition: ECLGeometryPar.cc:168
Belle2::ECL::ECLNeighbours
Class to get the neighbours for a given cell id.
Definition: ECLNeighbours.h:25