drich-dev/scripts/makeTrainingData.rb

0001 #!/usr/bin/env ruby
0002 # make set of training data for GNNs
0003 # - execute from top-level directory
0004
0005 require 'fileutils'
0006 require 'awesome_print'
0007 require 'thread/pool'
0008
0009 #### SETTINGS ################
0010 nThreads = `nproc`.to_i-2 # number of parallel threads (if nproc fails, just hardcode a number)
0011 outputDir = 'out/trainingSample'
0012 hepmcFile = 'hepmc/pythia8NCDIS_10x100_minQ2=1_beamEffects_xAngle=-0.025_hiDiv_vtxfix_1_000.hepmc'
0013 gunTests = {
0014   1 => 'single',
0015   5 => 'spray',
0016 }
0017 numEvents = {
0018   :forHepmc => 1000,
0019   :forGun   => 100,
0020 }
0021 particles = [
0022   'e-',
0023   'pi+',
0024   # 'pi-',
0025   'kaon+',
0026   # 'kaon-',
0027   'proton',
0028   # 'anti_proton',
0029 ]
0030 energies = [ # FIXME: simulate.py now reads these as momenta
0031   1,
0032   2,
0033   4,
0034   8,
0035   16,
0036   30,
0037   40,
0038   50,
0039 ]
0040 ###########################
0041
0042 ### setup
0043 FileUtils.mkdir_p outputDir
0044 cmds = []
0045
0046 ### gun tests
0047 gunTests.keys.product(particles,energies).each do |testNum,particle,energy|
0048   outputFile = "#{outputDir}/" + [
0049     gunTests[testNum],
0050     particle,
0051     "#{energy}GeV",
0052     'root',
0053   ].join('.')
0054   cmd = [
0055     './simulate.py',
0056     "-t#{testNum}",
0057     '-d1',
0058     '-s',
0059     "-p#{particle}",
0060     "-n#{numEvents[:forGun]}",
0061     "-m#{energy}",
0062     '-r',
0063     "-o#{outputFile}",
0064     "&& bin/draw_hits d #{outputFile}"
0065   ]
0066   cmds << cmd.join(' ')
0067 end
0068
0069 ### use hepmc
0070 outputFile = "#{outputDir}/" + [
0071   File.basename(hepmcFile,'.hepmc'),
0072   'root',
0073 ].join('.')
0074 cmd = [
0075   './simulate.py',
0076   "-i'#{hepmcFile}'",
0077   '-d1',
0078   '-s',
0079   "-n#{numEvents[:forHepmc]}",
0080   '-r',
0081   "-o#{outputFile}",
0082   "&& bin/draw_hits d #{outputFile}"
0083 ]
0084 cmds << cmd.join(' ')
0085
0086 ### execution
0087 ap cmds
0088 # exit
0089 pool = Thread.pool(nThreads)
0090 cmds.each{|cmd|pool.process{system cmd}} # multi-threaded
0091 # cmds.each{|cmd|system cmd} # single-threaded
0092 pool.shutdown
0093 puts "DONE: FILES PRODUCED IN #{outputDir}"