ATLAS Offline Software
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 
9 import sys
10 import math
11 import ROOT
12 
13 
14 def calculate_momentum_magnitude(E, Q2):
15  return math.sqrt(E*E - Q2)
16 
17 # Function to generate random virtuality according to Breit-Wigner distribution
18 def generate_virtuality(mass, width):
19  Q = ROOT.gRandom.BreitWigner(mass, width)
20  return Q*Q
21 
22 def 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 
30 def 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
35 def 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 
159 if __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 
vtune_athena.format
format
Definition: vtune_athena.py:14
CaloCellPos2Ntuple.int
int
Definition: CaloCellPos2Ntuple.py:24
python.resample_meson.calculate_momentum_magnitude
def calculate_momentum_magnitude(E, Q2)
Definition: resample_meson.py:14
python.resample_meson.generate_virtuality
def generate_virtuality(mass, width)
Definition: resample_meson.py:18
python.resample_meson.build_four_momentum
def build_four_momentum(Energy, p3_magnitude, theta, phi)
Definition: resample_meson.py:22
python.resample_meson.resample_meson
def resample_meson(input_filename, output_filename, meson_id, meson_width)
Definition: resample_meson.py:35
TCS::join
std::string join(const std::vector< std::string > &v, const char c=',')
Definition: Trigger/TrigT1/L1Topo/L1TopoCommon/Root/StringUtils.cxx:10
Trk::open
@ open
Definition: BinningType.h:40
python.resample_meson.build_TLorentz
def build_TLorentz(energy, px, py, pz)
Definition: resample_meson.py:30
str
Definition: BTagTrackIpAccessor.cxx:11
dbg::print
void print(std::FILE *stream, std::format_string< Args... > fmt, Args &&... args)
Definition: SGImplSvc.cxx:70
readCCLHist.float
float
Definition: readCCLHist.py:83
Generated on Thu Nov 7 2024 21:24:42 for ATLAS Offline Software by   doxygen 1.8.18