drich-dev/scripts/momentum_scan.rb

0001 #!/usr/bin/env ruby
0002 # run momentum scan test for various particles
0003
0004 require 'open3'
0005 require 'fileutils'
0006 require 'pycall/import'
0007
0008 ## SETTINGS ########################################
0009 NumEvents         = 70    # number of events per fixed momentum (full test: 50)
0010 NumPoints         = 10    # number of momenta to sample (full test: 100)
0011 PoolSize          = 6     # number of parallel threads to run
0012 RunSimulation     = true  # if true, run the simulation step
0013 RunReconstruction = true  # if true, run the reconstruction step
0014 UseRINDEXrange    = false # if true, use range of RINDEX values rather than a single reference value
0015 MaxNphot          = 500   # maximum number of incident photons expected (for plot drawing range)
0016 ####################################################
0017
0018
0019 ## args
0020 if ARGV.length<2
0021   $stderr.puts """
0022   USAGE: #{$0} [d/p] [j/e]
0023      [d/p]: detector
0024        d: dRICH
0025        p: pfRICH
0026      [j/e]: reconstruction
0027        j: juggler [DEPRECATED, do not use]
0028        e: eicrecon
0029   """
0030   exit 2
0031 end
0032
0033 ## detector-specific settings
0034 ## FIXME: :rIndexRef values are averages?
0035 case ARGV[0]
0036 when "d"
0037   zDirection = 1
0038   xRICH      = "dRICH"
0039   xrich      = "drich"
0040   radiator_h = {
0041     :agl => { :id=>0, :testNum=>7, :rIndexRef=>1.0190,  :rIndexRange=>[1.01852,1.02381], :maxMomentum=>20.0, :maxNPE=>20, },
0042     :gas => { :id=>1, :testNum=>8, :rIndexRef=>1.00076, :rIndexRange=>[1.00075,1.00084], :maxMomentum=>60.0, :maxNPE=>40, },
0043   }
0044 when "p"
0045   zDirection = -1
0046   xRICH      = "pfRICH"
0047   xrich      = "pfrich"
0048   radiator_h = {
0049     :agl => { :id=>0, :testNum=>7, :rIndexRef=>1.0190, :rIndexRange=>[1.01852,1.02381], :maxMomentum=>20.0, :maxNPE=>40, },
0050     :gas => { :id=>1, :testNum=>8, :rIndexRef=>1.0013, :rIndexRange=>[1.0013,1.0015],   :maxMomentum=>60.0, :maxNPE=>40, },
0051   }
0052 else
0053   $stderr.puts "ERROR: unknown detector '#{ARGV[0]}'"
0054   exit 1
0055 end
0056
0057 ## reconstruction specific settings
0058 case ARGV[1]
0059 when "j"
0060   reconMethod      = :juggler
0061 when "e"
0062   reconMethod      = :eicrecon
0063 else
0064   $stderr.puts "ERROR: unknown reconstruction '#{ARGV[1]}'"
0065   exit 1
0066 end
0067
0068 ## list of particles to test
0069 particle_h = {
0070   'e-'          => { :mass=>0.00051, },
0071   'pi+'         => { :mass=>0.13957, },
0072   'kaon+'       => { :mass=>0.49368, },
0073   'proton'      => { :mass=>0.93827, },
0074   # 'e+'          => { :mass=>0.00051, },
0075   # 'pi-'         => { :mass=>0.13957, },
0076   # 'kaon-'       => { :mass=>0.49368, },
0077   # 'anti_proton' => { :mass=>0.93827, },
0078 }
0079
0080 ## produce output file dir and names
0081 OutputDir = "out/momentum_scan.#{xrich}" # output directory ( ! will be overwritten ! )
0082 if RunSimulation
0083   FileUtils.rm_r OutputDir, secure: true, verbose: true, force: true
0084   FileUtils.mkdir_p OutputDir
0085 end
0086 def out_file(prefix,ext)
0087   "#{OutputDir}/#{prefix}.#{ext}"
0088 end
0089
0090 ## run momentum scan simulation for each particle in `particle_h`
0091 particle_h.keys.product(radiator_h.keys).each_slice(PoolSize) do |slice|
0092   pool = slice.map do |particle,rad_name|
0093     rad = radiator_h[rad_name]
0094     # preparation
0095     sim_file = out_file particle, "sim.#{rad_name}.root"
0096     rec_file = out_file particle, "rec.#{rad_name}.root"
0097     ana_file = rec_file.sub /\.root/, '.ana.root'
0098     cmds = []
0099     # simulation
0100     if RunSimulation
0101       cmds << [
0102         './simulate.py',
0103         "-t#{rad[:testNum]}",
0104         "-d#{zDirection}",
0105         "-p#{particle}",
0106         "-n#{NumEvents}",
0107         "-k#{NumPoints}",
0108         "-o#{sim_file}",
0109       ]
0110     end
0111     # reconstruction
0112     if RunReconstruction
0113       case reconMethod
0114       when :juggler
0115         cmds << [
0116           './juggler.sh',
0117           "-#{xrich[0]}",
0118           "-i#{sim_file}",
0119           "-o#{rec_file}",
0120         ]
0121       when :eicrecon
0122         cmds << [
0123           './benchmark.rb',
0124           '-r',
0125           "--sim-file #{sim_file}",
0126           "--rec-file #{rec_file}",
0127           "--ana-file #{ana_file}",
0128         ]
0129       end
0130     end
0131     # analysis
0132     plot_file = out_file particle, "rec_plots.#{rad_name}.root"
0133     cmds << [ 'root', '-b', '-q' ]
0134     case reconMethod
0135     when :juggler
0136       cmds.last << "scripts/src/momentum_scan_juggler_draw.C'(\"#{rec_file}\",\"#{plot_file}\",\"#{xrich.upcase}\",#{rad[:id]})'"
0137     when :eicrecon
0138       cmds.last << "scripts/src/momentum_scan_eicrecon_draw.C'(\"#{ana_file}\",\"#{plot_file}\",#{rad[:id]})'"
0139     end
0140     # spawn thread
0141     Thread.new do
0142       cmds.each_with_index do |cmd,i|
0143         cmd_shell = cmd.join ' '
0144         puts cmd_shell
0145         mode = i==0 ? 'w' : 'a'
0146         Open3.pipeline(
0147           cmd_shell,
0148           out: [out_file(particle,"#{rad_name}.log.out"),mode],
0149           err: [out_file(particle,"#{rad_name}.log.err"),mode],
0150         )
0151       end
0152     end
0153   end
0154   trap 'INT' do
0155     pool.each &:kill
0156     exit 1
0157   end
0158   # wait for pool to finish
0159   pool.each &:join
0160 end
0161
0162 # draw 2D hadd plots
0163 if reconMethod == :eicrecon
0164   radiator_h.each do |rad_name,rad|
0165     # calculate n_group_rebin, for rebinning the momentum bins:
0166     # n_group_rebin => ceiling[ num momentum bins * (maxMomentum here) / (maxMomentum in histogram) * (1/NumPoints) ]
0167     # FIXME: automate the hard-coded numbers
0168     n_group_rebin = ( 500.0 * rad[:maxMomentum]/70.0 * 1.0/NumPoints ).to_i + 1
0169     drawArgs = [
0170       "\"#{OutputDir}/*.rec.#{rad_name}.ana.root\"",
0171       "\"#{OutputDir}/_theta_scan_2D.#{rad_name}\"",
0172       rad[:id],
0173       n_group_rebin
0174     ].join ','
0175     system "root -b -q scripts/src/momentum_scan_2D_draw.C'(#{drawArgs})'"
0176   end
0177 end
0178
0179 # print errors for one of the particles
0180 puts "ERRORS (for one job) ======================="
0181 system "cat #{Dir.glob(OutputDir+"/*.err").first}"
0182 puts "END ERRORS ======================================"
0183
0184
0185 #############################################################
0186
0187
0188 ## start ROOT analysis
0189 # - must be done after simulations, otherwise this script hangs
0190 r = PyCall.import_module 'ROOT'
0191 r.gROOT.SetBatch true
0192 r.gStyle.SetOptStat 0
0193 r.gStyle.SetLegendTextSize 0.1
0194
0195 ## loop over radiators
0196 radiator_h.each do |rad_name,rad|
0197
0198   ## draw settings
0199   default_color  = r.kBlack
0200   default_marker = r.kFullCircle
0201   draw_h = particle_h.keys.map do |particle|
0202     [
0203       particle,
0204       {
0205         :root_file => r.TFile.new(out_file(particle, "rec_plots.#{rad_name}.root")),
0206         :color     => default_color,
0207         :marker    => default_marker,
0208       }.to_h
0209     ]
0210   end.to_h
0211   draw_h['e-'][:color]      = r.kBlack            if particle_h.has_key? 'e-'
0212   draw_h['pi+'][:color]     = r.kBlue             if particle_h.has_key? 'pi+'
0213   draw_h['kaon+'][:color]   = r.kGreen+1          if particle_h.has_key? 'kaon+'
0214   draw_h['proton'][:color]  = r.kMagenta          if particle_h.has_key? 'proton'
0215   draw_h['e-'][:marker]     = r.kFullCircle       if particle_h.has_key? 'e-'
0216   draw_h['pi+'][:marker]    = r.kFullSquare       if particle_h.has_key? 'pi+'
0217   draw_h['kaon+'][:marker]  = r.kFullTriangleUp   if particle_h.has_key? 'kaon+'
0218   draw_h['proton'][:marker] = r.kFullTriangleDown if particle_h.has_key? 'proton'
0219
0220   ## loop over plots (loop through the first particle's file, assume the rest have the same)
0221   first_root_file = draw_h[particle_h.keys.first][:root_file]
0222   PyCall.iterable(first_root_file.GetListOfKeys).each do |tkey|
0223     plot_name = tkey.GetName
0224     next unless plot_name.match? /scan_pfx$/
0225
0226     ## canvas
0227     canv = r.TCanvas.new "#{plot_name}_#{rad_name}_canv", "#{plot_name}_#{rad_name}_canv", 1000, 800
0228     leg_hits = r.TLegend.new 0, 0.7, 1, 1
0229     canv.Divide 2,1
0230     pad_plot = r.TPad.new 'pad_plot', 'pad_plot', 0,    0, 0.75, 1
0231     pad_leg  = r.TPad.new 'pad_leg',  'pad_leg',  0.75, 0, 1,    1
0232     pad_plot.SetGrid 1,1
0233     pad_plot.SetLeftMargin 0.2
0234     pad_plot.Draw
0235     pad_leg.Draw
0236
0237     ## loop over particles
0238     rad[:maxTheta] = 0
0239     draw_h.each do |particle,h|
0240
0241       ## get the plot
0242       plot = h[:root_file].Get plot_name
0243
0244       ## variables for this particle and radiator
0245       mass   = particle_h[particle][:mass]
0246       rindex = rad[:rIndexRef]
0247
0248       ## minimum momentum for Cherenkov emission
0249       def calculate_mom_min(m,n)
0250         m / Math.sqrt(n**2-1)
0251       end
0252       mom_min = calculate_mom_min mass, rindex
0253
0254       ## drop points with momentum < mom_min
0255       (1..plot.GetNbinsX).each do |bn|
0256         mom = plot.GetBinCenter bn
0257         if mom < mom_min
0258           plot.SetBinContent bn, 0
0259           plot.SetBinError   bn, 0
0260           plot.SetBinEntries bn, 0
0261         end
0262       end
0263
0264       ## saturation Cherenkov angle
0265       def calculate_max_theta(n)
0266         1000 * Math.acos(1/n)
0267       end
0268       max_theta = calculate_max_theta rindex
0269       rad[:maxTheta] = [rad[:maxTheta],max_theta].max
0270
0271       ## add plot to canvas
0272       pad_plot.cd
0273       plot.SetMarkerStyle h[:marker]
0274       plot.SetMarkerColor h[:color]
0275       plot.SetLineColor   h[:color]
0276       plot.SetMarkerSize  1.5
0277       plot.Draw 'SAME E X0'
0278       plot.GetYaxis.SetTitle 'Average ' + plot.GetTitle.sub(/ vs\..*/,'')
0279       plot.SetTitle          'Average ' + plot.GetTitle
0280       leg_hits.AddEntry plot, particle, 'PE'
0281
0282       ## expected Cherenkov angle curve
0283       if plot_name.match?(/^theta_/)
0284         rindex_list = UseRINDEXrange ? rad[:rIndexRange] : [rindex]
0285         rindex_list.each_with_index do |n,i|
0286           ftn_theta = r.TF1.new(
0287             "ftn_theta_#{particle}_#{rad_name}_#{i}",
0288             "1000*TMath::ACos(TMath::Sqrt(x^2+#{mass}^2)/(#{n}*x))",
0289             calculate_mom_min(mass,n),
0290             plot.GetXaxis.GetXmax,
0291           )
0292           ftn_theta.SetLineColor h[:color]
0293           ftn_theta.SetLineWidth 2
0294           ftn_theta.Draw 'SAME'
0295           # set plot maximum
0296           rad[:maxTheta] = [rad[:maxTheta],calculate_max_theta(n)].max # max expected
0297           rad[:maxTheta] = [rad[:maxTheta],plot.GetMaximum].max        # max point
0298         end
0299       end
0300
0301       # set plot ranges
0302       plot.GetXaxis.SetRangeUser 0, 1.1*rad[:maxMomentum]
0303       case plot_name
0304       when /^thetaResid_/
0305         plot.GetYaxis.SetRangeUser -40, 40
0306       when /^theta_/
0307         plot.GetYaxis.SetRangeUser 0, 1.1*rad[:maxTheta]
0308       when /^npe_/
0309         plot.GetYaxis.SetRangeUser 0, rad[:maxNPE]
0310       when /^nphot_/
0311         plot.GetYaxis.SetRangeUser 0, MaxNphot
0312       end
0313
0314     end
0315
0316     ## save canvas to file
0317     pad_leg.cd
0318     leg_hits.Draw
0319     if plot_name.match?(/^nphot_/) # :gas and :agl plots are the same, just save :gas
0320       canv.SaveAs out_file("_#{plot_name}",'png') if rad_name==:gas
0321     else
0322       canv.SaveAs out_file("_#{plot_name}.#{rad_name}",'png')
0323     end
0324
0325   end # loop over plots
0326
0327   draw_h.values.each{ |h| h[:root_file].Close }
0328
0329 end # loop over radiators