Merge pull request #236 from Mu2e/pyutils-dev

Pyutils: first round of pyutils for v6

Author: AndrewEdmonds11

.muse

@@ -3,4 +3,5 @@ ROOT_LIBRARY_PATH
 
 PYTHONPATH utils
 PYTHONPATH utils/helper
+PYTHONPATH utils/pyutils
 PATH bin

example-analysis-scripts/PythonScript.py

@@ -0,0 +1,86 @@
+#! /usr/bin/env python
+import pyimport as evn
+import pyvector as vec
+import pyplot as plot
+import pyprint as prnt
+from pyselect import Select as slct
+
+import awkward as ak
+def main():
+  """ simple test function to run some of the utils """
+
+  # import the files
+  myevn = evn.Import("/exp/mu2e/data/users/sophie/ensembles/MDS1/MDS1av0.root", "EventNtuple", "ntuple")
+
+  # import code and extract branch
+  ntuple = myevn.ImportTree()
+
+  # make a pyselect object
+  mysel = slct()
+  
+  # check if it is an electron
+  is_elec = mysel.isElectron(nutple)
+
+  # import branches associated with trk fits
+  trksegs = myevn.ImportBranches(ntuple,['trksegs','trksegpars_lh'])
+  
+  # check if fit is going down stream
+  is_down = mysel.isDownstream(trksegs)
+
+  # is trk entrance
+  trkent_mask = (trksegs['trksegs']['sid']==0)
+  
+  # is in time
+  trksegs_mask = (trksegs['trksegs']['time'] > 640) & (trksegs['trksegs']['time'] < 1650)
+  
+  # check trk pars
+  trkpars_mask = (trksegs['trksegpars_lh']['t0err'] < 0.9) & (trksegs['trksegpars_lh']['maxr'] < 680) & (trksegs['trksegpars_lh']['maxr'] > 450) 
+  
+  
+  # these are deprecated cuts
+  oldtrkpar_mask = (trksegs['trksegpars_lh']['tanDip'] > 0.5) & (trksegs['trksegpars_lh']['tanDip'] < 1.0) & (trksegs['trksegpars_lh']['d0'] > -100) & (trksegs['trksegpars_lh']['d0'] < 100)
+  
+  # check trk quality
+  trkqual = ntuple["trkqual"] 
+  mytrkqual = trkqual.arrays(library='ak')
+  trkqual_mask = (mytrkqual['trkqual.result']> 0.2)
+  
+  # check for crv coincidences 
+  crvs = ntuple["crvcoincs"]
+  mycrvs = crvs.arrays(library='ak')
+  crv_mask = mysel.hasTrkCrvCoincs( trksegs, mycrvs, 150)
+  
+  # apply joint mask
+  mytrksegs = trksegs['trksegs'].mask[(is_elec) & (is_down) & (trkent_mask)  & (trksegs_mask)   & (crv_mask) &(oldtrkpar_mask) & (active_mask) & (trkqual_mask) & (trkpars_mask)  ]
+
+  # plot time before and after cuts:
+  myhist = plot.Plot()
+  flatarraytime = ak.flatten(trksegs['trksegs']['time'], axis=None)
+  flatarraycut = ak.flatten(mytrksegs['time'], axis=None)
+  dictarrays = { "no cut" : flatarraytime, "with cut" : flatarraycut }
+  myhist.Plot1DOverlay(dictarrays, 100, 450, 1695, "Mu2e Example", "fit time at Trk Ent [ns]", "#events per bin", 'timecut.pdf', 'best', 300,False, True, True)
+  
+  # plot the momentum before and after the cuts
+  myvect = vec.Vector()
+  electrksegs = trkentall['trksegs'].mask[(is_elec) & (is_down)]
+  vector_all = myvect.GetVectorXYZFromLeaf(electrksegs, 'mom')
+  magnitude_all = myvect.Mag(vector_all)
+  
+  vector_cut = myvect.GetVectorXYZFromLeaf(mytrksegs, 'mom')
+  magnitude_cut = myvect.Mag(vector_cut)
+  
+  flatarraymom_all = ak.flatten(magnitude_all, axis=None)
+  flatarraymom_cut = ak.flatten(magnitude_cut, axis=None)
+  
+  myhist.Plot1D(flatarraymom_cut  , None, 100, 95, 115, "Mu2e Example", "fit mom at Trk Ent [MeV/c]", "#events per bin", 'black', 'best', 'momcut.pdf', 300, True, False, True, False, True, True, True)
+  
+  dictarrays = { "all dem" : flatarraymom_all, "dem + trkcuts" : flatarraymom_cut }
+  myhist.Plot1DOverlay(dictarrays, 100, 95,115, "Mu2e Example", "fit mom at Trk Ent [ns]", "#events per bin", 'momcutcompare.pdf', 'best', 300,False, True, True)
+  
+  
+  
+
+
+  
+if __name__ == "__main__":
+    main()

example-analysis-scripts/pyutils-examples/example_cuts.py

@@ -0,0 +1,71 @@
+#! /usr/bin/env python
+import pyimport as evn
+import pyvector as vec
+import pyplot as plot
+import pyprint as prnt
+from pyselect import Select as slct
+import awkward as ak
+
+def main():
+  """ simple test function to run some of the utils """
+
+  # import the files
+  test_evn = evn.Import("/exp/mu2e/data/users/sophie/ensembles/MDS1/MDS1av2.root", "EventNtuple", "ntuple")
+
+  # import code and extract branch
+  ntuple = test_evn.ImportTree()
+
+  # make a pyselect object
+  mysel = slct()
+
+  # import branches associated with trk fits
+  trksegs = test_evn.ImportBranches(ntuple,['trksegs','trksegpars_lh'])
+
+  # check if it is an electron going downstream
+  is_elec = mysel.isElectron(ntuple)
+  is_down = mysel.isDownstream(trksegs)
+   
+  # active hits
+  active_mask = mysel.SelectHits(ntuple, 20)
+  
+  # check trk quality
+  trkqual_mask = mysel.SelectTrkQual(ntuple, 0.2)
+
+  # check for crv coincidences 
+  crv_mask = mysel.hasTrkCrvCoincs( trksegs, ntuple, 150)
+  
+  # set of trkseg cuts
+  treenames = [ 'trksegs', 'trksegs', 'trksegpars_lh', 'trksegpars_lh', 'trksegpars_lh', 'trksegpars_lh']
+  leaves = [ 'sid', 'time', 't0err','maxr','tanDip','d0']
+  equals = [True, False, False, False, False, False]
+  v1s = [0, 640, 0, 450, 0.5, -100]
+  v2s = [None, 1650, 0.9, 680, 1.0, 100]
+  
+  # make a list of masks
+  trkseg_mask_list = mysel.MakeMaskList(trksegs, treenames, leaves, equals, v1s, v2s)
+
+  # FIXME apply joint mask --> can we make this tidier? like a loop or somethin?
+  mytrksegs = trksegs['trksegs'].mask[(is_elec) & (is_down) & (trkqual_mask) & (crv_mask) & (active_mask) & (trkseg_mask_list[0]) & (trkseg_mask_list[1]) &  (trkseg_mask_list[2]) & (trkseg_mask_list[3]) & (trkseg_mask_list[4]) & (trkseg_mask_list[5]) ]
+  
+  # make some plots to compare before/after cuts
+  myhist = plot.Plot()
+  
+  # plot the momentum before and after the cuts
+  myvect = vec.Vector()
+  electrksegs = trksegs['trksegs'].mask[(is_elec) & (is_down)]
+  vector_all = myvect.GetVectorXYZFromLeaf(electrksegs, 'mom')
+  magnitude_all = myvect.Mag(vector_all)
+  
+  vector_cut = myvect.GetVectorXYZFromLeaf(mytrksegs, 'mom')
+  magnitude_cut = myvect.Mag(vector_cut)
+  
+  flatarraymom_all = ak.flatten(magnitude_all, axis=None)
+  flatarraymom_cut = ak.flatten(magnitude_cut, axis=None)
+  
+  myhist.Plot1D(flatarraymom_cut  , None, 100, 95, 115, "Mu2e Example", "fit mom at Trk Ent [MeV/c]", "#events per bin", 'black', 'best', 'momcut.pdf', 300, True, False, True, False, True, True, True)
+  
+  dictarrays = { "all dem" : flatarraymom_all, "dem + trkcuts" : flatarraymom_cut }
+  myhist.Plot1DOverlay(dictarrays, 100, 95,115, "Mu2e Example", "fit mom at Trk Ent [ns]", "#events per bin", 'momcutcompare.pdf', 'best', 300,False, True, True)
+  
+if __name__ == "__main__":
+    main()

example-analysis-scripts/pyutils-examples/example_multicuts.py

@@ -0,0 +1,32 @@
+#! /usr/bin/env python
+import pyimport as evn
+import pyvector as vec
+import pyplot as plot
+import pyprint as prnt
+import awkward as ak
+
+def main():
+  """ simple test function to run some of the utils """
+  
+  # import the files
+  myevn = evn.Import(None, "EventNtuple", "ntuple")
+  
+  #import a list of files
+  filepath = "/exp/mu2e/data/users/sophie/ensembles/MDS1/file.list"
+
+  # import code and extract branch
+  treename = 'trksegs'
+  branchname = 'time'
+  surface_id = 0 # tracker entrance FIXME - we need a better way for this
+  branch = myevn.ImportTreeFromFileList(filepath, treename)
+
+  # find fit at chosen ID
+  trkent = myevn.SelectSurfaceID(branch, treename, surface_id)
+
+  # make 1D plot
+  myhist = plot.Plot()
+  flatarraytime = ak.flatten(trkent[str(branchname)], axis=None)
+  myhist.Plot1D(flatarraytime, None, 100, 450, 1695, "Mu2e Example", "fit time at Trk Ent [ns]", "#events per bin", 'black', 'best', 'time_merged.pdf', 300, True, False, False, False, True, True, True)
+  
+if __name__ == "__main__":
+    main()

example-analysis-scripts/pyutils-examples/example_multifiles.py

@@ -0,0 +1,49 @@
+#! /usr/bin/env python
+import pyimport as evn
+import pyvector as vec
+import pyplot as plot
+import pyprint as prnt
+import pyselect as slct
+import awkward as ak
+def main():
+  """ simple test function to run some of the utils """
+  
+  # import the files
+  test_evn = evn.Import("/exp/mu2e/data/users/sophie/ensembles/MDS1/MDS1av0.root", "EventNtuple", "ntuple")
+
+  # import code and extract branch
+  treename = 'trksegs'
+  branchname = 'time'
+  surface_id = 0 # tracker entrance FIXME - we need a better way for this
+  ntuple = test_evn.ImportTree()
+  branch = test_evn.ImportBranches(ntuple,[str(treename)])
+  
+  # find fit at chosen ID
+  mysel = slct.Select()
+  trkent = mysel.SelectSurfaceID(branch, treename, surface_id)
+  
+  # print out the first 10 events:
+  myprnt = prnt.Print()
+  myprnt.PrintNEvents(branch,10)
+  
+  # make 1D plot
+  myhist = plot.Plot()
+  flatarraytime = ak.flatten(trkent[str(branchname)], axis=None)
+  myhist.Plot1D(flatarraytime, None, 100, 450, 1695, "Mu2e Example", "fit time at Trk Ent [ns]", "#events per bin", 'black', 'best', 'time.pdf', 300, True, False, False, False, True, True, True)
+  
+  # access vectors
+  myvect = vec.Vector()
+  vecbranchname = 'mom'
+  trkentall = mysel.SelectSurfaceID(branch, treename, surface_id)
+  vector_test = myvect.GetVectorXYZFromLeaf(trkentall, vecbranchname)
+  magnitude = myvect.Mag(vector_test)
+
+  # make 1D plot of magnitudes
+  flatarraymom = ak.flatten(magnitude, axis=None)
+  
+  # 2D mom time plot
+  myhist.Plot2D( flatarraymom, flatarraytime, None, 100, 95, 115, 100, 450, 1650, "Mu2e Example", "fit mom at Trk Ent [MeV/c]", "fit time at Trk Ent [ns]", None, 'timevmom.pdf', 'inferno',300,False, False, False, True,True)
+        
+  
+if __name__ == "__main__":
+    main()

example-analysis-scripts/pyutils-examples/example_plotting.py

@@ -0,0 +1,45 @@
+# mu2e.mplstyle
+
+# Generic figure tweaks 
+axes.prop_cycle: cycler("color", ['red','blue','green','orange','violet','cyan','pink','gray','yellow'])
+
+mathtext.default: regular
+font.size: 18
+axes.labelsize: medium
+axes.unicode_minus: False
+axes.labelpad: 10
+axes.titlesize: 15
+axes.linewidth: 2
+grid.alpha: 0.8
+grid.linestyle: :
+savefig.transparent: False
+
+xaxis.labellocation: right
+yaxis.labellocation: top
+xtick.labelsize: small
+ytick.labelsize: small
+xtick.direction: in
+xtick.major.size: 12
+xtick.minor.size: 6
+xtick.major.pad: 6
+xtick.top: True
+xtick.major.top: True
+xtick.major.bottom: True
+xtick.minor.top: True
+xtick.minor.bottom: True
+xtick.minor.visible: True
+ytick.direction: in
+ytick.major.size: 12
+ytick.minor.size: 6.0
+ytick.right: True
+ytick.major.left: True
+ytick.major.right: True
+ytick.minor.left: True
+ytick.minor.right: True
+ytick.minor.visible: True
+
+legend.fontsize: small
+legend.handlelength: 1.5
+legend.borderpad: 0.5
+legend.frameon: False
+