ATLAS Offline Software
Loading...
Searching...
No Matches
resample_meson.py
Go to the documentation of this file.
1# Copyright (C) 2002-2022 CERN for the benefit of the ATLAS collaboration
2# a function for resampling meson in the lhe file
3# for specific HeavyN_meson model decay to meson+lepton (only)
4# the meson by default is generated as an on shell particle
5# with this adhoc solution a Breit-Wigner width in introduce to meson particle and
6# the meson+lepton 4 momentum has to be replace with new consistent values
7# for more information mbahmani@cern.ch
8
9import sys
10import math
11import ROOT
12
13
14def calculate_momentum_magnitude(E, Q2):
15 return math.sqrt(E*E - Q2)
16
17# Function to generate random virtuality according to Breit-Wigner distribution
18def generate_virtuality(mass, width):
19 Q = ROOT.gRandom.BreitWigner(mass, width)
20 return Q*Q
21
22def build_four_momentum(Energy, p3_magnitude, theta, phi):
23 # Calculate the components of the new four-momentum vector
24 px = p3_magnitude * math.sin(theta) * math.cos(phi)
25 py = p3_magnitude * math.sin(theta) * math.sin(phi)
26 pz = p3_magnitude * math.cos(theta)
27 # Return the new four-momentum vector
28 return [Energy, px, py, pz]
29
30def build_TLorentz(energy, px,py,pz):
31 # Build the new four-momentum using the energy and momentum components
32 return ROOT.TLorentzVector(px, py, pz, energy)
33
34# Function to resample meson particles in LHE file
35def resample_meson(input_filename, output_filename, meson_id, meson_width ):
36 event = False
37 findW=False
38
39 with open(input_filename, 'r') as input_file:
40 with open(output_filename, 'w') as output_file:
41 for line in input_file:
42 Q2=0
43 if not event and '<event>' not in line:
44 output_file.write(line)
45 continue
46 # Check if we are just starting an event
47 if not event and '<event>' in line:
48 output_file.write(line)
49 event = True
50 findW = False
51 continue
52 # Check if we have finished and are doing something else
53 if '<' in line or '>' in line:
54 event = False
55 if line.startswith('<mgrwt>') or line.startswith('</mgrwt>'):
56 output_file.write(line)
57 continue
58
59 if abs(int(line.split()[0])) == 24:
60 #find W boson
61 findW=True
62
63
64
65 if (abs(int(line.split()[0])) == abs(int(meson_id)) and abs(int(line.split()[2])) == 4) or (not findW and abs(int(line.split()[0])) == abs(int(meson_id)) and abs(int(line.split()[2]))==3 ): # Check for meson
66 #check the charge
67
68 line_parts = line.split()
69 line_parts[0] = str(int(line.split()[0]))
70
71
72 # Resample meson mass using Breit-Wigner distribution
73
74 momentum_x = float(line.split()[6])
75 momentum_y = float(line.split()[7])
76 momentum_z= float(line.split()[8])
77 Energy = float(line.split()[9])
78
79 meson=build_TLorentz(Energy,momentum_x, momentum_y, momentum_z)
80
81 continue
82
83 if ((abs(int(line.split()[0])) == 13 or abs(int(line.split()[0])) ==11) and int(line.split()[2])==4) or ( not findW and (abs(int(line.split()[0])) == 13 or abs(int(line.split()[0])) ==11) and int(line.split()[2])==3):
84 #check for the displace lepton
85 E_l = float(line.split()[9])
86 px_l= float(line.split()[6])
87 py_l= float(line.split()[7])
88 pz_l= float(line.split()[8])
89
90 lep1=build_TLorentz(E_l,px_l, py_l, pz_l)
91
92 # Boost into the rest frame of meson+lep1
93 combined = meson+lep1
94
95 beta=combined.BoostVector()
96 meson.Boost(-beta)
97 lep1.Boost(-beta)
98
99
100
101 combined_B =meson+lep1
102
103 meson_massp =meson.M()
104
105 # Calculate theta
106 theta = math.acos(meson.Pz() / math.sqrt(meson.Px()**2 + meson.Py()**2 + meson.Pz()**2))
107 # Calculate phi
108 phi = math.atan2(meson.Py(), meson.Px())
109 # Generate random virtuality
110 Q2 = generate_virtuality(meson_massp, float(meson_width))
111
112 while (Q2 > meson.E()*meson.E()+meson.Px()**2 + meson.Py()**2 + meson.Pz()**2 or not (0.3 < math.sqrt(Q2) < 1.5)):
113 # Generate another virtuality until it's smaller
114 Q2 = generate_virtuality(meson_massp, float(meson_width))
115 # build the meson energy
116 meson_energy = combined_B.E()/2 +((Q2-lep1.M()**2)/(2*combined_B.E()))
117
118 p3_magnitude=calculate_momentum_magnitude(meson_energy, Q2)
119 four_momentum = build_four_momentum(meson_energy, p3_magnitude, theta, phi)
120 meson.SetPxPyPzE(four_momentum[1],four_momentum[2],four_momentum[3],four_momentum[0])
121
122 px_lep= combined_B.Px()-meson.Px()
123 py_lep= combined_B.Py()-meson.Py()
124 pz_lep= combined_B.Pz()-meson.Pz()
125 # Ensure momentum and energy conservation for mu1
126 lep1_energy = math.sqrt(px_lep**2+py_lep**2+pz_lep**2+lep1.M()**2)
127 lep1.SetPxPyPzE(px_lep,py_lep,pz_lep,lep1_energy)
128
129 # Boost meson and mu1 particles back to the lab frame
130 meson.Boost(beta)
131 lep1.Boost(beta)
132
133 line_parts[9] = str("{:+.10e}".format(meson.E()))
134 line_parts[6] =str("{:+.10e}".format(meson.Px()))
135 line_parts[7] =str("{:+.10e}".format(meson.Py()))
136 line_parts[8] =str("{:+.10e}".format(meson.Pz()))
137 line_parts[10] =str("{:+.10e}".format(meson.M()))
138
139 updated_line = ' '.join(line_parts) + '\n'
140 output_file.write(' ' +''.join(updated_line) )
141
142 if ((abs(int(line.split()[0])) == 13 or abs(int(line.split()[0])) ==11) and int(line.split()[2])==4) or (not findW and (abs(int(line.split()[0])) == 13 or abs(int(line.split()[0])) ==11) and int(line.split()[2])==3):
143 line_parts_l = line.split()
144 line_parts_l[9] = str("{:+.10e}".format(lep1.E()))
145 line_parts_l[6] =str("{:+.10e}".format(lep1.Px()))
146 line_parts_l[7] =str("{:+.10e}".format(lep1.Py()))
147 line_parts_l[8] =str("{:+.10e}".format(lep1.Pz()))
148 line_parts_l[10] =str("{:+.10e}".format(lep1.M()))
149
150 updated_line_l = ' '.join(line_parts_l) + '\n'
151 output_file.write(' ' +''.join(updated_line_l) + '\n')
152
153
154 else:
155 output_file.write(line)
156
157
158
159if __name__ == "__main__":
160 print("len(sys.argv)", len(sys.argv))
161 if len(sys.argv) !=5:
162 print("Usage: python resample_meson.py input_file.lhe output_file.lhe meson_id meson_width")
163 else:
164 input_file = sys.argv[1]
165 output_file = sys.argv[2]
166 meson_id = sys.argv[3]
167 meson_width = sys.argv[4]
168
169 resample_meson(input_file, output_file, meson_id, meson_width)
170
print
void print(char *figname, TCanvas *c1)
Definition TRTCalib_StrawStatusPlots.cxx:26
python.resample_meson.generate_virtuality
generate_virtuality(mass, width)
Definition resample_meson.py:18
python.resample_meson.build_four_momentum
build_four_momentum(Energy, p3_magnitude, theta, phi)
Definition resample_meson.py:22
python.resample_meson.build_TLorentz
build_TLorentz(energy, px, py, pz)
Definition resample_meson.py:30
python.resample_meson.calculate_momentum_magnitude
calculate_momentum_magnitude(E, Q2)
Definition resample_meson.py:14
Generated on for ATLAS Offline Software by doxygen 1.14.0