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