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