Skip to content

Commit e61d175

Browse files
DanielSamitzDanielSamitz
committed
python script
1 parent d3b93c9 commit e61d175

File tree

1 file changed

+307
-0
lines changed

1 file changed

+307
-0
lines changed

‎Tutorials/PWGEM/Cocktail/run_lmee_lf_cocktail.py‎

Lines changed: 307 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,307 @@
1+
#!/usr/bin/env python3
2+
3+
import argparse
4+
import os
5+
import json
6+
from os.path import join, abspath, exists
7+
import re
8+
import subprocess
9+
import ROOT as root
10+
11+
12+
def clean(dir,name,printtext):
13+
if printtext:
14+
print("==> Delete simulation artefacts in the output directory\n")
15+
for f in os.listdir(dir):
16+
pattern_del=name+"_|o2simtopology|MCStepLogger"
17+
#pattern_keep="Kine.root|serverlog|mergerlog|workerlog"
18+
pattern_keep="Kine.root|serverlog"
19+
if (re.search(pattern_del, f)) and (not re.search(pattern_keep, f) ):
20+
os.remove(abspath(join(dir,f)))
21+
22+
23+
def read_sim_config(args):
24+
IniName = "GeneratorEMCocktail.ini"
25+
generatorName = "${O2DPG_ROOT}/MC/config/PWGEM/external/generator/GeneratorEMCocktailV2.C"
26+
generatorFunc = "GenerateEMCocktail"
27+
28+
# matches GeneratorParamEMlibV2.h
29+
GeneratorParamEMlibV2_CollisionSystem={"kpp900GeV":0x000, "kpp2760GeV":0x64, "kpp7TeV":0xC8, "kpPb":0x12C, "kPbPb":0x190}
30+
GeneratorParamEMlibV2_Centrality={"kpp":0x0, "k0005":0x1, "k0510":0x2, "k1020":0x3, "k2030":0x4, "k3040":0x5, "k4050":0x6, "k5060":0x7, "k0010":0x8, "k2040":0x9, "k4060":0xA, "k6080":0xB, "k0020":0xC, "k0040":0xD, "k2080":0xE, "k4080":0xF, "k2050":0x10, "kCentralities":0x11}
31+
32+
allKeys = ["collisionSystem", "centrality", "decayMode","selectedMothers", "paramFile", "paramFileDir", "numberOfParticles", "minPt", "maxPt", "pythiaErrorTolerance", "externalDecayer", "decayLongLived", "dynamicalPtRange", "useYWeights", "paramV2FileDir", "toFixEP", "yGenRange", "useLMeeDecaytable", "weightingMode"]
33+
if not exists(args.sim_config_file):
34+
print("ERROR: File "+args.sim_config_file+" not found")
35+
return 1
36+
37+
with open(args.sim_config_file, "r") as f:
38+
simConfigs = json.load(f)
39+
40+
print("==> Reading simualtion configuration from "+abspath(args.sim_config_file))
41+
42+
missing_key = False
43+
for key in allKeys:
44+
if not key in simConfigs:
45+
print("ERROR: Missing key \'"+key+"\'")
46+
missing_key = True
47+
if missing_key:
48+
return 1
49+
50+
collisionSystem = GeneratorParamEMlibV2_CollisionSystem[simConfigs["collisionSystem"]]
51+
centrality = GeneratorParamEMlibV2_Centrality[simConfigs["centrality"]]
52+
decayMode = simConfigs["decayMode"]
53+
selectedMothers = simConfigs["selectedMothers"]
54+
paramFile = simConfigs["paramFile"]
55+
paramFileDir = simConfigs["paramFileDir"]
56+
numberOfParticles = simConfigs["numberOfParticles"]
57+
minPt = simConfigs["minPt"]
58+
maxPt = simConfigs["maxPt"]
59+
pythiaErrorTolerance = simConfigs["pythiaErrorTolerance"]
60+
externalDecayer = simConfigs["externalDecayer"]
61+
decayLongLived = simConfigs["decayLongLived"]
62+
dynamicalPtRange = simConfigs["dynamicalPtRange"]
63+
useYWeights = simConfigs["useYWeights"]
64+
paramV2FileDir = simConfigs["paramV2FileDir"]
65+
toFixEP = simConfigs["toFixEP"]
66+
yGenRange = simConfigs["yGenRange"]
67+
useLMeeDecaytable = simConfigs["useLMeeDecaytable"]
68+
weightingMode = simConfigs["weightingMode"]
69+
70+
print(f"collisionSystem\t\t{collisionSystem}\ncentrality\t\t{centrality}\ndecayMode\t\t{decayMode}\nselectedMothers\t\t{selectedMothers}\nparamFile\t\t{paramFile}\nparamFileDir\t\t{paramFileDir}\nnumberOfParticles\t{numberOfParticles}\nminPt\t\t\t{minPt}\nmaxPt\t\t\t{maxPt}\npythiaErrorTolerance\t{pythiaErrorTolerance}")
71+
print(f"externalDecayer\t\t{externalDecayer}\ndecayLongLived\t\t{decayLongLived}\ndynamicalPtRange\t{dynamicalPtRange}\nuseYWeights\t\t{useYWeights}\nparamV2FileDir\t{paramV2FileDir}\ntoFixEP\t\t\t{toFixEP}\nyGenRange\t\t{yGenRange}")
72+
print(f"useLMeeDecaytable\t{useLMeeDecaytable}\nweightingMode\t\t{weightingMode}")
73+
ini_file = abspath(join(args.output,IniName))
74+
print("==> Writing simulation configuration to "+ini_file)
75+
os.makedirs(os.path.dirname(ini_file), exist_ok=True)
76+
with open(f"{ini_file}",'w') as f:
77+
f.write(f"[GeneratorExternal]\nfileName = {generatorName}\nfuncName={generatorFunc}(")
78+
f.write(f"{collisionSystem},{centrality},{decayMode},{selectedMothers},\"{paramFile}\",\"{paramFileDir}\",{numberOfParticles},{minPt},{maxPt},{pythiaErrorTolerance},{externalDecayer},{decayLongLived},{dynamicalPtRange},{useYWeights},\"{paramV2FileDir}\",{toFixEP},{yGenRange},\"{useLMeeDecaytable}\",{weightingMode}")
79+
f.write(")")
80+
return ini_file
81+
82+
def run_only_simulation(args):
83+
sim_log_name = "simlog"
84+
name="o2sim"
85+
86+
output_dir = abspath(args.output)
87+
output_name = abspath(join(args.output,name))
88+
ini_file = abspath(args.sim_ini_file)
89+
log_file = abspath(join(args.output,sim_log_name))
90+
91+
current_dir = os.getcwd()
92+
if not exists(output_dir):
93+
os.makedirs(output_dir)
94+
os.chdir(output_dir)
95+
96+
97+
print("\n==> Running cocktail simulation")
98+
command_to_run = ['o2-sim','-n',args.nEvents,'-g','external','-o',output_name,'--configFile',ini_file,'--noGeant']
99+
print("running with: ",end='')
100+
print(*command_to_run)
101+
print("redirect output to "+log_file)
102+
print("\nrunning...")
103+
f = open(log_file, "w")
104+
simproc = subprocess.Popen(['o2-sim','-n',args.nEvents,'-g','external','-o',output_name,'--configFile',ini_file,'--noGeant'], stdout=f, stderr=subprocess.STDOUT)
105+
simproc.communicate()
106+
f.close()
107+
print("...done\n")
108+
109+
os.chdir(current_dir)
110+
111+
if args.clean:
112+
clean(args.output,name,True)
113+
114+
if not simproc.returncode==0:
115+
print("ERROR: o2sim finished with error. See "+abspath(join(args.output,name+"_serverlog"))+" and "+log_file+"for details\n")
116+
return 1
117+
else:
118+
print("Cocktail simulation finished successful. Result in "+output_name+"_Kine.root\n")
119+
return 0
120+
121+
122+
def run_only_analysis_task(args):
123+
sim_log_name = "readerlog"
124+
ana_log_name = "analog"
125+
name="kineReader"
126+
127+
print("\n==> Running cocktail analysis task")
128+
129+
output_dir = abspath(args.output)
130+
output_name = join(output_dir,name)
131+
current_dir = os.getcwd()
132+
133+
kin_file = abspath(args.input)
134+
json_file = abspath(args.ana_config_file)
135+
sim_log_file = join(output_dir,sim_log_name)
136+
ana_log_file = join(output_dir,ana_log_name)
137+
138+
if not exists(output_dir):
139+
os.makedirs(output_dir)
140+
os.chdir(output_dir)
141+
142+
f = root.TFile(kin_file)
143+
tree = f.Get("o2sim")
144+
nEvents = tree.GetEntries()
145+
print(f"Found {nEvents} events in {kin_file}")
146+
147+
sim_command_to_run = ['o2-sim','-g','extkinO2','-n',str(nEvents),'-o',output_name,'--extKinFile',kin_file,'--noGeant','--forwardKine','--noDiscOutput']
148+
print("running reader with: ",end='')
149+
print(*sim_command_to_run)
150+
print("redirect reader output to "+sim_log_file)
151+
152+
fs = open(sim_log_file, "w")
153+
simproc = subprocess.Popen(sim_command_to_run, stdout=fs, stderr=subprocess.STDOUT)
154+
155+
proxy_command_to_run = ['o2-sim-mctracks-proxy','--nevents',str(nEvents),'--o2sim-pid',str(simproc.pid)]
156+
ana_command_to_run = ['o2-analysis-em-lmee-lf-cocktail', '--configuration','json://'+json_file,'-b']
157+
print("running analysis with: ",end='')
158+
print(*proxy_command_to_run,end='')
159+
print(" | ",end='')
160+
print(*ana_command_to_run)
161+
print("redirect analysis output to "+ana_log_file)
162+
163+
fa = open(ana_log_file,"w")
164+
proxyproc = subprocess.Popen(proxy_command_to_run, stdout=subprocess.PIPE)
165+
anaproc = subprocess.Popen(ana_command_to_run, stdin=proxyproc.stdout, stdout=fa, stderr=subprocess.STDOUT)
166+
167+
print("\nrunning...\n")
168+
simproc.communicate()
169+
if not simproc.returncode==0:
170+
print("ERROR: Reader finished with error. See "+abspath(join(args.output,name+"_serverlog"))+" and "+sim_log_file+" for details\n")
171+
proxyproc.kill();
172+
anaproc.kill();
173+
proxyproc.communicate()
174+
anaproc.communicate()
175+
fs.close()
176+
fa.close()
177+
print("...done\n")
178+
179+
180+
if not anaproc.returncode==0:
181+
print("ERROR: Analysis task finished with error. See "+abspath(join(args.output,"analog"))+" for details\n")
182+
else:
183+
print("Analysis task finished successful. Histograms in "+output_dir+"/AnalysisResults.root\n")
184+
185+
os.chdir(current_dir)
186+
187+
clean(args.output,name,False)
188+
189+
if (not simproc.returncode) and (not anaproc.returncode):
190+
return 0
191+
else:
192+
return 1
193+
194+
def run_full(args):
195+
sim_log_name = "simlog"
196+
ana_log_name = "analog"
197+
name="o2sim"
198+
199+
print("\n==> Running cocktail simulation + analysis task")
200+
201+
output_dir = abspath(args.output)
202+
output_name = join(output_dir,name)
203+
current_dir = os.getcwd()
204+
205+
ini_file = abspath(args.sim_ini_file)
206+
json_file = abspath(args.ana_config_file)
207+
sim_log_file = join(output_dir,sim_log_name)
208+
ana_log_file = join(output_dir,ana_log_name)
209+
210+
if not exists(output_dir):
211+
os.makedirs(output_dir)
212+
os.chdir(output_dir)
213+
214+
sim_command_to_run = ['o2-sim','-g','external','-n',args.nEvents,'-o',output_name,'--configFile',ini_file,'--noGeant','--forwardKine']
215+
if not args.save_kine:
216+
sim_command_to_run.append('--noDiscOutput')
217+
print("running simulation with: ",end='')
218+
print(*sim_command_to_run)
219+
print("redirect sim output to "+sim_log_file)
220+
221+
fs = open(sim_log_file, "w")
222+
simproc = subprocess.Popen(sim_command_to_run, stdout=fs, stderr=subprocess.STDOUT)
223+
224+
proxy_command_to_run = ['o2-sim-mctracks-proxy','--nevents',args.nEvents,'--o2sim-pid',str(simproc.pid)]
225+
ana_command_to_run = ['o2-analysis-em-lmee-lf-cocktail','--configuration','json://'+json_file,'-b']
226+
print("running analysis with: ",end='')
227+
print(*proxy_command_to_run,end='')
228+
print(" | ",end='')
229+
print(*ana_command_to_run)
230+
print("redirect analysis output to "+ana_log_file)
231+
232+
fa = open(ana_log_file,"w")
233+
proxyproc = subprocess.Popen(proxy_command_to_run, stdout=subprocess.PIPE)
234+
anaproc = subprocess.Popen(ana_command_to_run, stdin=proxyproc.stdout, stdout=fa, stderr=subprocess.STDOUT)
235+
236+
print("\nrunning...\n")
237+
simproc.communicate()
238+
if not simproc.returncode==0:
239+
print("ERROR: o2sim finished with error. See "+abspath(join(args.output,name+"_serverlog"))+" and "+sim_log_file+" for details\n")
240+
proxyproc.kill();
241+
anaproc.kill();
242+
else:
243+
print("Cocktail simulation finished successful.",end='')
244+
if args.save_kine:
245+
print(" Output written to "+output_name+"_Kine.root.",end='')
246+
print(" Waiting for analysis task to finish...\n")
247+
proxyproc.communicate()
248+
anaproc.communicate()
249+
fs.close()
250+
fa.close()
251+
print("...done\n")
252+
253+
254+
if not anaproc.returncode==0:
255+
print("ERROR: Analysis task finished with error. See "+abspath(join(args.output,"analog"))+" for details\n")
256+
else:
257+
print("Analysis task finished successful. Histograms in "+output_dir+"/AnalysisResults.root\n")
258+
259+
os.chdir(current_dir)
260+
261+
if args.clean:
262+
clean(args.output,name,True)
263+
264+
if (not simproc.returncode) and (not anaproc.returncode):
265+
return 0
266+
else:
267+
return 1
268+
269+
def main():
270+
parser = argparse.ArgumentParser()
271+
common_sim_parser = argparse.ArgumentParser(add_help=False)
272+
sim_input = common_sim_parser.add_mutually_exclusive_group(required=True)
273+
common_sim_parser.add_argument("--nEvents","-n",help="Number of events to be generated",required=True)
274+
sim_input.add_argument("--sim-ini-file",dest="sim_ini_file",help="INI file for the generator")
275+
sim_input.add_argument("--sim-config-file",dest="sim_config_file",help="JSON config file for the generator (will be converted to INI file)")
276+
common_sim_parser.add_argument("--output","-o",default=os.getcwd(),help="Output directory")
277+
common_sim_parser.add_argument("--clean", dest="clean", action="store_true",help="Delete unwanted files after finishing the job")
278+
common_ana_parser = argparse.ArgumentParser(add_help=False)
279+
common_ana_parser.add_argument("--ana-config-file",dest="ana_config_file",help="Config file for the analysis task",required=True)
280+
sub_parsers = parser.add_subparsers(dest="command",required=True)
281+
sim_parser = sub_parsers.add_parser("sim-only",parents=[common_sim_parser])
282+
ana_parser = sub_parsers.add_parser("ana-only",parents=[common_ana_parser])
283+
ana_parser.add_argument("--input","-i",help="Analysis input file",required=True)
284+
ana_parser.add_argument("--output","-o",default=os.getcwd(),help="Output directory")
285+
full_parser = sub_parsers.add_parser("full",parents=[common_sim_parser,common_ana_parser])
286+
full_parser.add_argument("--save-kine",dest="save_kine",action="store_true",help="Write the generator output to o2sim_Kine.root")
287+
args = parser.parse_args()
288+
289+
print("\n######################")
290+
print("## LMee LF Cocktail ##")
291+
print("######################")
292+
293+
print("\n==> Running in mode \'"+args.command+"\'\n")
294+
295+
if ( ((args.command=="full") or (args.command=="sim-only")) and args.sim_config_file):
296+
args.sim_ini_file=read_sim_config(args)
297+
if args.sim_ini_file == 1:
298+
return 1
299+
300+
if (args.command=="full"):
301+
return run_full(args)
302+
if (args.command=="sim-only"):
303+
return run_only_simulation(args)
304+
if (args.command=="ana-only"):
305+
return run_only_analysis_task(args)
306+
307+
main()

0 commit comments

Comments
 (0)