ATLAS Offline Software
Gen_tf.py
Go to the documentation of this file.
1 #! /usr/bin/env python
2 
3 # Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
4 #
5 """
6 # Run event generation and produce an EVNT file.
7 """
8 
9 import os, sys, time, shutil
10 from PyJobTransforms.trfLogger import msg
11 from PyJobTransforms.transform import transform
12 from PyJobTransforms.trfExe import athenaExecutor
13 from PyJobTransforms.trfArgs import addAthenaArguments
14 from PyJobTransforms.trfDecorators import stdTrfExceptionHandler, sigUsrStackTrace
15 from EvgenJobTransforms.evgenTrfArgs import addStdEvgenArgs
16 
17 
19 ListOfDefaultPositionalKeys=['--AMIConfig', '--AMITag', '--argJSON', '--asetup', '--athena', '--athenaMPMergeTargetSize', '--athenaopts', '--attempt', '--checkEventCount', '--command', '--dumpJSON', '--dumpPickle', '--ecmEnergy', '--env', '--eventAcceptanceEfficiency', '--evgenJobOpts', '--execOnly', '--fileValidation', '--firstEvent', '--ignoreErrors', '--ignoreFiles', '--ignorePatterns', '--imf', '--inputEVNT_PreFile', '--inputFileValidation', '--inputGenConfFile', '--inputGeneratorFile', '--jobConfig', '--jobid', '--maxEvents', '--orphanKiller', '--outputEVNTFile', '--outputEVNT_PreFile', '--outputFileValidation', '--outputNTUP_TRUTHFile', '--outputTXTFile', '--parallelFileValidation', '--postExec', '--postInclude', '--preExec', '--preInclude', '--printEvts', '--randomSeed', '--reportName', '--reportType', '--rivetAnas', '--runNumber', '--showGraph', '--showPath', '--showSteps', '--skipEvents', '--skipFileValidation', '--skipInputFileValidation', '--skipOutputFileValidation', '--steering', '--taskid', '--tcmalloc', '--valgrind', '--valgrindbasicopts', '--valgrindextraopts', '--lheOnly', '--localPath', '--cleanOut', '--saveList']
20 
21 class EvgenExecutor(athenaExecutor):
22  "Specialised trf executor class for event generation jobs"
23  def __init__(self, name="generate", skeleton=None, skeletonCA=None, substep=None, inData=set(), outData=set()):
24  athenaExecutor.__init__(self, name=name, skeletonFile=skeleton, skeletonCA=skeletonCA,
25  substep=substep, inputEventTest=False, tryDropAndReload=False,
26  inData=inData, outData=outData)
27 
28  def preExecute(self, input=set(), output=set()):
29  "Get input tarball, unpack and set up env if an evgenJobOpts arg was provided."
30 
31 
32  super(EvgenExecutor, self).preExecute(input, output)
33 
34  def expand_if_archive(filename):
35  "Function to expand a file if it is a zip archive or tarball"
36  if ".tar" in filename:
37  import tarfile
38  with tarfile.open(filename) as tf:
39  tf.extractall()
40  elif filename.endswith(".zip"):
41  import zipfile
42  zf = zipfile.ZipFile(filename)
43  zf.extractall()
44  zf.close()
45 
46  def get_immediate_subdirectories(a_dir):
47  return [name for name in os.listdir(a_dir)
48  if os.path.isdir(os.path.join(a_dir, name))]
49 
50 
52  os.environ['LOCAL_INSTALL_DIR'] = (os.environ['JOBOPTSEARCHPATH']).split(":")[0]
53  CommonCvmfsDir = '/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/common'
54  os.environ["JOBOPTSEARCHPATH"] = CommonCvmfsDir+":"+os.environ["JOBOPTSEARCHPATH"]
55  if os.path.exists('/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/common/MadGraphControl/dat/'):
56  datCvmfsDir = '/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/common/MadGraphControl/dat/'
57  os.environ["DATAPATH"] = datCvmfsDir+":"+os.environ["DATAPATH"]
58 # dsidparam = (self._trf.argdict["jobConfig"].value).values()[0][0]
59  dsidpar = (self._trf.argdict["jobConfig"].value).values()
60  dsidparam = list(dsidpar)[0][0]
61  # Adding cvmfs path to JOBOPTSEARCHPATH
62  BaseCvmfsPath = "/cvmfs/atlas.cern.ch/repo/sw/Generators/MCJobOptions/"
63 
64  if dsidparam.isdigit() and (len(dsidparam)==6 or len(dsidparam)==7):
65  #only dsid is provided, add cvmfs folder like 123xxx to JOBOPTSEARCHPATH
66  Jodir = dsidparam[:3]+'xxx'
67  if len(dsidparam)==7:
68  Jodir = dsidparam[0] + '/' + dsidparam[:4]+'xxx'
69  cwdir = os.getcwd()
70  cwd_ful = os.path.join(cwdir, dsidparam)
71  if (os.path.isdir(cwd_ful)):
72  os.environ["JOBOPTSEARCHPATH"] = cwd_ful+":"+os.environ["JOBOPTSEARCHPATH"]
73  os.environ["DATAPATH"] = cwd_ful+":"+os.environ["DATAPATH"]
74  else:
75  cwd_Jodir = os.path.join(cwdir,Jodir)
76  cwd_Jodir_ful = os.path.join(cwd_Jodir,dsidparam)
77  if (os.path.isdir(cwd_Jodir_ful)):
78  os.environ["JOBOPTSEARCHPATH"] = cwd_Jodir_ful+":"+os.environ["JOBOPTSEARCHPATH"]
79  os.environ["DATAPATH"] = cwd_Jodir_ful+":"+os.environ["DATAPATH"]
80  else:
81  JoCvmfsPath = os.path.join(BaseCvmfsPath, Jodir)
82  JoCvmfsPath_ful = os.path.join(JoCvmfsPath, dsidparam)
83  os.environ["JOBOPTSEARCHPATH"] = JoCvmfsPath_ful+":"+os.environ["JOBOPTSEARCHPATH"]
84  os.environ["DATAPATH"] = JoCvmfsPath_ful+":"+os.environ["DATAPATH"]
85 
86 
87  else: #Suppose full path of dsid folder is provided(/afs/.../123xxx/123456), add cvmfs floder and local path(/afs/.../123xxx) to JOBOPTSEARCHPATH
88  if dsidparam.startswith("Test"): # for testing
89  for token in (os.environ['JOBOPTSEARCHPATH']).split(":"):
90  if "jobOptions" in token:
91  search_token = token
92  break
93  search_token += "/EvgenJobTransforms/EvgenTest/" + dsidparam.split("Test")[-1]
94  os.environ["JOBOPTSEARCHPATH"] = search_token+":"+os.environ["JOBOPTSEARCHPATH"]
95  os.environ["DATAPATH"] = search_token+":"+os.environ["DATAPATH"]
96  elif (os.path.isdir(dsidparam)):
97  os.environ["JOBOPTSEARCHPATH"] = dsidparam+":"+os.environ["JOBOPTSEARCHPATH"]
98  os.environ["DATAPATH"] = dsidparam+":"+os.environ["DATAPATH"]
99  else:
100  msg.error("JOs not found, please check = '%s'" % dsidparam)
101  raise RuntimeError("JOs not found")
102 
103  msg.info("Using JOBOPTSEARCHPATH = '%s'" % os.environ["JOBOPTSEARCHPATH"])
104  msg.info("Using DATAPATH = '%s'" % os.environ["DATAPATH"])
105 
106  if "evgenJobOpts" in self._trf.argdict:
107  tarball = self._trf.argdict["evgenJobOpts"].value
108 
110  if tarball.startswith("http"):
111  url = tarball
112  tarball = os.basename(tarball)
113  else:
114  url = "http://cern.ch/atlas-computing/links/kitsDirectory/EvgenJobOpts/" + tarball
115 
116  if not os.path.exists(tarball):
117  from EvgenJobTransforms.download import downloadUsingProxy
118  status, output = downloadUsingProxy(url)
119  if status != 0:
120  raise EnvironmentError('Error downloading tarball %s. Downloader reports: %s' % (tarball, output))
121  msg.info('Evgen tarball download success: %s' % output)
122 
123  expand_if_archive(tarball)
124 
126 
127  # copy config files to cwd
128  FIRST_DIR = (os.environ['JOBOPTSEARCHPATH']).split(":")[0]
129  configFiles = [f for f in os.listdir(FIRST_DIR) if ( "GRID" in f)]
130  confFile=None
131  if len(configFiles) == 1:
132  confFile = os.path.join(FIRST_DIR, configFiles[0])
133  elif len(configFiles) >1:
134  msg.info("more then one gridpack ! ")
135  if "--ecmEnergy" in str(sys.argv[1:]):
136  split_args=str(sys.argv[1:]).split("ecmEnergy=",1)[1]
137  ener_GeV=split_args.split(",")[0].strip("\'")
138  energy=str(float(ener_GeV)/1000.0).replace('.','p').strip(" =0\p']")
139  msg.info("Should be used gridpack for energy "+energy)
140  else:
141  energy="13"
142  for x in configFiles:
143  gridS="mc_"+energy+"TeV"
144  msg.info("Gridpack should start from "+gridS)
145  if x.startswith(gridS):
146  confFile = os.path.join(FIRST_DIR, x)
147  msg.info("using gridpack = "+confFile)
148  if confFile is None:
149  msg.error("No *GRID* config files, for requested energy = '%s' please check = '%s'" %(energy,dsidparam))
150 
151  if confFile is not None:
152  expand_if_archive(confFile)
153 # os.system("cp %s ." % confFile)
154  msg.info("Configuration input gridpack found " + confFile)
155 
156  #Expand if a tarball is found in local directory
157  loc_files = os.listdir(os.getcwd())
158  for loc_file in loc_files:
159  if "GRID" not in loc_file:
160  expand_if_archive(loc_file)
161 
162 
163  if "inputGeneratorFile" in self._trf.argdict:
164 # expand_if_archive(self._trf.argdict["inputGeneratorFile"].value)
165  myinputfiles = self._trf.argdict["inputGeneratorFile"].value
166  genInputFiles = myinputfiles.split(',')
167  for file in genInputFiles:
168  expand_if_archive(file)
169  if "inputGenConfFile" in self._trf.argdict:
170  expand_if_archive(self._trf.argdict["inputGenConfFile"].value)
171 
172 def move_files(main_dir,tmp_dir,whitelist):
173  files = os.listdir(tmp_dir)
174  files.sort()
175  for f in files:
176  for i in whitelist:
177  if i in f:
178  src = tmp_dir+"/"+f
179  dest = main_dir+"/"+f
180  os.rename(src,dest)
181  break
182 
183 
184 def getTransform():
185  exeSet = set()
186  msg.info("Transform arguments %s" % sys.argv[1:])
187  if "--outputEVNTFile" in str(sys.argv[1:]):
188  exeSet.add(EvgenExecutor(name="generate", skeleton="EvgenJobTransforms/skel.GENtoEVGEN.py", skeletonCA="EvgenJobTransforms.GENtoEVGEN_Skeleton", inData=["inNULL"], outData=["YODA", "EVNT", "EVNT_Pre", "TXT"]))
189  msg.info("Output EVNT file")
190  elif "--outputYODAFile" in str(sys.argv[1:]):
191  exeSet.add(EvgenExecutor(name="generate", skeleton="EvgenJobTransforms/skel.GENtoEVGEN.py", inData=["inNULL"], outData=["YODA", "TXT"]))
192  msg.info("Output EVNT file")
193  elif "--outputTXTFile" in str(sys.argv[1:]):
194  exeSet.add(EvgenExecutor(name="generate", skeleton="EvgenJobTransforms/skel.GENtoTXT.py", inData=["inNULL"], outData=["TXT"]))
195  msg.info("Output TXT file")
196  else:
197  msg.error("Output cannot be recognised")
198 
199  exeSet.add(EvgenExecutor(name="afterburn", skeleton="EvgenJobTransforms/skel.ABtoEVGEN.py", inData=["EVNT_Pre"], outData=["EVNT"]))
200  exeSet.add(athenaExecutor(name = "AODtoDPD", skeletonFile = "PATJobTransforms/skeleton.AODtoDPD_tf.py",
201  substep = "a2d", inData = ["EVNT"], outData = ["NTUP_TRUTH"], perfMonFile = "ntuple_AODtoDPD.pmon.gz"))
202  trf = transform(executor=exeSet)
203  addAthenaArguments(trf.parser, maxEventsDefaultSubstep='all')
204  addStdEvgenArgs(trf.parser)
205  return trf
206 
207 
208 @stdTrfExceptionHandler
209 @sigUsrStackTrace
210 def main():
211  msg.info("This is %s" % sys.argv[0])
212 
213  main_dir = os.getcwd()
214  trf = getTransform()
215  trf.parseCmdLineArgs(sys.argv[1:])
216  if (("cleanOut" in trf.argdict) and (trf.argdict["cleanOut"].value != 0)):
217  name_tmpdir = "tmprun"
218  tmp_dir = os.path.join(main_dir, name_tmpdir)
219  if os.path.isdir(tmp_dir):
220  shutil.rmtree(tmp_dir, ignore_errors=True)
221  os.mkdir("tmprun")
222  os.chdir("tmprun")
223  tmp_dir = os.getcwd()
224  whitelist_in = ['MC','group','TXT']
225  move_files(tmp_dir,main_dir,whitelist_in)
226 
227  trf.execute()
228  trf.generateReport()
229  msg.info("%s stopped at %s, trf exit code %d" % (sys.argv[0], time.asctime(), trf.exitCode))
230 
231 
232 # read files/dirs that should be saved and if present in cwd - remove
233 
234  if (("cleanOut" in trf.argdict) and (trf.argdict["cleanOut"].value!=0)):
235  whitelist_out = ['log.generate','.root']
236  if "outputTXTFile" in trf.argdict:
237  whitelist_out.append('TXT')
238  if "saveList" in trf.argdict:
239  saveList_dic= trf.argdict["saveList"].value
240  saveList_str= str(saveList_dic)
241  saveList_str=saveList_str[10:-3]
242  saveList= saveList_str.split(",")
243  for item in saveList:
244  test_ex = os.path.join(main_dir,str(item))
245  if os.path.isdir(test_ex):
246  shutil.rmtree(test_ex, ignore_errors=True)
247  elif os.path.isfile(test_ex):
248  os.remove(test_ex)
249  if not saveList[0].isdigit():
250  whitelist_out=whitelist_out+saveList
251 
252  move_files(main_dir,tmp_dir,whitelist_out)
253  os.chdir(main_dir)
254  if "saveList" not in trf.argdict:
255  shutil.rmtree(tmp_dir, ignore_errors=True)
256  elif not saveList[0].isdigit():
257  shutil.rmtree(tmp_dir, ignore_errors=True)
258 # if cleanOut is not defined and multipleinput preset, remove the merged file
259  elif ("inputGeneratorFile" in trf.argdict):
260  myinputfiles = trf.argdict["inputGeneratorFile"].value
261  genInputFiles = myinputfiles.split(',')
262  numberOfFiles = len(genInputFiles)
263  merge_file = 'merged_lhef._0.events'
264  if((numberOfFiles>1) and (os.path.exists(merge_file))):
265  os.remove(merge_file)
266 #
267  if (("lheOnly" in trf.argdict ) and (trf.argdict["lheOnly"].value == 1)):
268  outputName = ''.join(trf.argdict["outputEVNTFile"].value)
269  os.remove(outputName)
270  sys.exit(trf.exitCode)
271 
272 
273 # TODO: Open resulting EVNT file to extract cross-section, generator names+versions, etc. from the HepMC::GenRun or whatever... in an executor postExecute?
274 
275 
276 if __name__ == "__main__":
277  main()
replace
std::string replace(std::string s, const std::string &s2, const std::string &s3)
Definition: hcg.cxx:307
Gen_tf.move_files
def move_files(main_dir, tmp_dir, whitelist)
Definition: Gen_tf.py:172
python.jo_proxy.get_immediate_subdirectories
def get_immediate_subdirectories(a_dir)
Definition: jo_proxy.py:6
python.Bindings.values
values
Definition: Control/AthenaPython/python/Bindings.py:797
Gen_tf.getTransform
def getTransform()
Definition: Gen_tf.py:184
python.trfArgs.addAthenaArguments
def addAthenaArguments(parser, maxEventsDefaultSubstep='first', addValgrind=True, addPerfMon=True, addVTune=True)
Options related to running athena in general TODO: Some way to mask certain options (perExec,...
Definition: trfArgs.py:59
Amg::transform
Amg::Vector3D transform(Amg::Vector3D &v, Amg::Transform3D &tr)
Transform a point from a Trasformation3D.
Definition: GeoPrimitivesHelpers.h:156
histSizes.list
def list(name, path='/')
Definition: histSizes.py:38
Gen_tf.EvgenExecutor.preExecute
def preExecute(self, input=set(), output=set())
Definition: Gen_tf.py:28
python.download.downloadUsingProxy
def downloadUsingProxy(url, filename=None)
Definition: download.py:5
CxxUtils::set
constexpr std::enable_if_t< is_bitmask_v< E >, E & > set(E &lhs, E rhs)
Convenience function to set bits in a class enum bitmask.
Definition: bitmask.h:224
Gen_tf.EvgenExecutor.__init__
def __init__(self, name="generate", skeleton=None, skeletonCA=None, substep=None, inData=set(), outData=set())
Definition: Gen_tf.py:23
TCS::join
std::string join(const std::vector< std::string > &v, const char c=',')
Definition: Trigger/TrigT1/L1Topo/L1TopoCommon/Root/StringUtils.cxx:10
PyJobTransforms.trfExe
Transform execution functions.
python.evgenTrfArgs.addStdEvgenArgs
def addStdEvgenArgs(parser)
Definition: evgenTrfArgs.py:5
PyJobTransforms.trfLogger
Logging configuration for ATLAS job transforms.
if
if(febId1==febId2)
Definition: LArRodBlockPhysicsV0.cxx:569
PyJobTransforms.transform
Main package for new style ATLAS job transforms.
str
Definition: BTagTrackIpAccessor.cxx:11
readCCLHist.float
float
Definition: readCCLHist.py:83
Trk::split
@ split
Definition: LayerMaterialProperties.h:38
Gen_tf.EvgenExecutor
Definition: Gen_tf.py:21
Gen_tf.main
def main()
Definition: Gen_tf.py:210
Generated on Tue Jul 2 2024 21:16:03 for ATLAS Offline Software by   doxygen 1.8.18