Belle II Software  release-08-01-10
EclCrystalInfoPrinter.py
1 #!/usr/bin/env python3
2 
3 
10 
11 # Doxygen should skip this script
12 # @cond
13 
14 import basf2
15 from ROOT import Belle2
16 import numpy as np
17 import pandas as pd
18 
19 # --------------------------------
20 
21 
22 class 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
173 main_path = basf2.create_path()
174 
175 main_path.add_module('EventInfoSetter', expList=[1003])
176 main_path.add_module('Gearbox')
177 main_path.add_module('Geometry')
178 printCrystals = printCrystals()
179 main_path.add_module(printCrystals)
180 
181 basf2.process(main_path)
182 print(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