Skip to content

Commit 781fe69

Browse files
sophiemiddletonsophiemiddleton
authored
Merge pull request #2 from sam-grant/main
Transfer pyutils from EventNtuple to standalone
2 parents 8f1a5dc + c9dda7b commit 781fe69

24 files changed

Lines changed: 5572 additions & 2 deletions

‎README.md‎

Lines changed: 778 additions & 2 deletions
Large diffs are not rendered by default.

‎examples/attic/example_cuts.py‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
#! /usr/bin/env python
2+
import pyimport as evn
3+
import pyvector as vec
4+
import pyplot as plot
5+
import pyprint as prnt
6+
from pyselect import Select as slct
7+
8+
import awkward as ak
9+
def main():
10+
""" simple test function to run some of the utils """
11+
12+
# import the files
13+
myevn = evn.Import("/exp/mu2e/data/users/sophie/ensembles/MDS1/MDS1av0.root", "EventNtuple", "ntuple")
14+
15+
# import code and extract branch
16+
ntuple = myevn.ImportTree()
17+
18+
# make a pyselect object
19+
mysel = slct()
20+
21+
# check if it is an electron
22+
is_elec = mysel.isElectron(nutple)
23+
24+
# import branches associated with trk fits
25+
trksegs = myevn.ImportBranches(ntuple,['trksegs','trksegpars_lh'])
26+
27+
# check if fit is going down stream
28+
is_down = mysel.isDownstream(trksegs)
29+
30+
# is trk middle
31+
trkent_mask = (trksegs['trksegs']['sid']==1)
32+
33+
# is in time
34+
trksegs_mask = (trksegs['trksegs']['time'] > 640) & (trksegs['trksegs']['time'] < 1650)
35+
36+
# check trk pars
37+
trkpars_mask = (trksegs['trksegpars_lh']['t0err'] < 0.9) & (trksegs['trksegpars_lh']['maxr'] < 680) & (trksegs['trksegpars_lh']['maxr'] > 450)
38+
39+
40+
# these are deprecated cuts
41+
oldtrkpar_mask = (trksegs['trksegpars_lh']['tanDip'] > 0.5) & (trksegs['trksegpars_lh']['tanDip'] < 1.0) & (trksegs['trksegpars_lh']['d0'] > -100) & (trksegs['trksegpars_lh']['d0'] < 100)
42+
43+
# check trk quality
44+
trkqual = ntuple["trkqual"]
45+
mytrkqual = trkqual.arrays(library='ak')
46+
trkqual_mask = (mytrkqual['trkqual.result']> 0.2)
47+
48+
# check for crv coincidences
49+
crvs = ntuple["crvcoincs"]
50+
mycrvs = crvs.arrays(library='ak')
51+
crv_mask = mysel.hasTrkCrvCoincs( trksegs, mycrvs, 150)
52+
53+
# apply joint mask
54+
mytrksegs = trksegs['trksegs'].mask[(is_elec) & (is_down) & (trkent_mask) & (trksegs_mask) & (crv_mask) &(oldtrkpar_mask) & (active_mask) & (trkqual_mask) & (trkpars_mask) ]
55+
56+
# plot time before and after cuts:
57+
myhist = plot.Plot()
58+
flatarraytime = ak.flatten(trksegs['trksegs']['time'], axis=None)
59+
flatarraycut = ak.flatten(mytrksegs['time'], axis=None)
60+
dictarrays = { "no cut" : flatarraytime, "with cut" : flatarraycut }
61+
myhist.Plot1DOverlay(dictarrays, 100, 450, 1695, "Mu2e Example", "fit time at Trk Ent [ns]", "#events per bin", 'timecut.pdf', 'best', 300,False, True, True)
62+
63+
# plot the momentum before and after the cuts
64+
myvect = vec.Vector()
65+
electrksegs = trkentall['trksegs'].mask[(is_elec) & (is_down)]
66+
vector_all = myvect.GetVectorXYZFromLeaf(electrksegs, 'mom')
67+
magnitude_all = myvect.Mag(vector_all)
68+
69+
vector_cut = myvect.GetVectorXYZFromLeaf(mytrksegs, 'mom')
70+
magnitude_cut = myvect.Mag(vector_cut)
71+
72+
flatarraymom_all = ak.flatten(magnitude_all, axis=None)
73+
flatarraymom_cut = ak.flatten(magnitude_cut, axis=None)
74+
75+
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)
76+
77+
dictarrays = { "all dem" : flatarraymom_all, "dem + trkcuts" : flatarraymom_cut }
78+
myhist.Plot1DOverlay(dictarrays, 100, 95,115, "Mu2e Example", "fit mom at Trk Ent [ns]", "#events per bin", 'momcutcompare.pdf', 'best', 300,False, True, True)
79+
80+
81+
82+
83+
84+
85+
if __name__ == "__main__":
86+
main()

‎examples/attic/example_multicuts.py‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
#! /usr/bin/env python
2+
import sys
3+
sys.path.append("../../utils/pyutils")
4+
5+
import pyimport as evn
6+
import pyvector as vec
7+
import pyplot as plot
8+
import pyprint as prnt
9+
from pyselect import Select as slct
10+
11+
import awkward as ak
12+
import argparse
13+
def example_multicuts(filename):
14+
""" simple test function to run some of the utils """
15+
16+
# import the files
17+
test_evn = evn.Import(str(filename), "EventNtuple", "ntuple")
18+
19+
# import code and extract branch
20+
ntuple = test_evn.ImportTree()
21+
22+
# make a pyselect object
23+
mysel = slct()
24+
25+
# import branches associated with trk fits
26+
trksegs = test_evn.ImportBranches(ntuple,['trksegs','trksegpars_lh'])
27+
28+
# check if it is an electron going downstream
29+
is_elec = mysel.isElectron(ntuple)
30+
is_down = mysel.isDownstream(trksegs)
31+
32+
# active hits
33+
active_mask = mysel.SelectHits(ntuple, 20)
34+
35+
# check trk quality
36+
trkqual_mask = mysel.SelectTrkQual(ntuple, 0.2)
37+
38+
# check for crv coincidences
39+
crv_mask = mysel.hasTrkCrvCoincs( trksegs, ntuple, 150)
40+
41+
# set of trkseg cuts
42+
treenames = [ 'trksegs', 'trksegs', 'trksegpars_lh', 'trksegpars_lh', 'trksegpars_lh', 'trksegpars_lh']
43+
leaves = [ 'sid', 'time', 't0err','maxr','tanDip','d0']
44+
equals = [True, False, False, False, False, False]
45+
v1s = [1, 640, 0, 450, 0.5, -100]
46+
v2s = [None, 1650, 0.9, 680, 1.0, 100]
47+
48+
# make a list of masks
49+
trkseg_mask_list = mysel.MakeMaskList(trksegs, treenames, leaves, equals, v1s, v2s)
50+
51+
# FIXME apply joint mask --> can we make this tidier? like a loop or somethin?
52+
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]) ]
53+
54+
# make some plots to compare before/after cuts
55+
myhist = plot.Plot()
56+
57+
# plot the momentum before and after the cuts
58+
myvect = vec.Vector()
59+
electrksegs = trksegs['trksegs'].mask[(is_elec) & (is_down)]
60+
vector_all = myvect.GetVectorXYZFromLeaf(electrksegs, 'mom')
61+
magnitude_all = myvect.Mag(vector_all)
62+
63+
vector_cut = myvect.GetVectorXYZFromLeaf(mytrksegs, 'mom')
64+
magnitude_cut = myvect.Mag(vector_cut)
65+
66+
flatarraymom_all = ak.flatten(magnitude_all, axis=None)
67+
flatarraymom_cut = ak.flatten(magnitude_cut, axis=None)
68+
69+
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)
70+
71+
dictarrays = { "all dem" : flatarraymom_all, "dem + trkcuts" : flatarraymom_cut }
72+
myhist.Plot1DOverlay(dictarrays, 100, 95,115, "Mu2e Example", "fit mom at Trk Ent [ns]", "#events per bin", 'momcutcompare.pdf', 'best', 300,False, True, True)
73+
74+
def main(args):
75+
example_multicuts(args.filename)
76+
77+
if __name__ == "__main__":
78+
parser = argparse.ArgumentParser(description='command arguments', formatter_class=argparse.RawTextHelpFormatter)
79+
parser.add_argument("--filename", type=str, default="/exp/mu2e/data/users/sophie/ensembles/MDS1/MDS1av0.root", help="filename")
80+
args = parser.parse_args()
81+
(args) = parser.parse_args()
82+
main(args)

‎examples/attic/example_multifiles.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
#! /usr/bin/env python
2+
import sys
3+
sys.path.append("../../utils/pyutils")
4+
5+
import pyimport as evn
6+
import pyvector as vec
7+
import pyplot as plot
8+
import pyprint as prnt
9+
10+
import awkward as ak
11+
import argparse
12+
13+
def example_multifiles(filelist):
14+
""" simple test function to run some of the utils """
15+
16+
# import the files
17+
myevn = evn.Import(None, "EventNtuple", "ntuple")
18+
19+
20+
# import code and extract branch
21+
treename = 'trksegs'
22+
branchname = 'time'
23+
surface_id = 1 # tracker middle
24+
branch = myevn.ImportTreeFromFileList(filepath, treename)
25+
26+
# find fit at chosen ID
27+
trkent = myevn.SelectSurfaceID(branch, treename, surface_id)
28+
29+
# make 1D plot
30+
plotter = plot.Plot()
31+
plotter.Plot1D(
32+
array=flatarraytime,
33+
nbins=100,
34+
xmin=450,
35+
xmax=1695,
36+
title="Example 1D histogram",
37+
xlabel="Fit time at Trk Ent [ns]",
38+
ylabel="Events per bin",
39+
out_path='h1_time_merged.png',
40+
stat_box=True,
41+
error_bars=True
42+
)
43+
44+
def main(args):
45+
example_multifiles(args.filelist)
46+
47+
if __name__ == "__main__":
48+
parser = argparse.ArgumentParser(description='command arguments', formatter_class=argparse.RawTextHelpFormatter)
49+
parser.add_argument("--filelist", type=str, default="/exp/mu2e/data/users/sophie/ensembles/MDS1/file.list", help="filelist")
50+
args = parser.parse_args()
51+
(args) = parser.parse_args()
52+
main(args)

‎examples/attic/example_plotting.py‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
#! /usr/bin/env python
2+
import sys
3+
sys.path.append("../../utils/pyutils")
4+
5+
import pyimport as evn
6+
import pyvector as vec
7+
import pyplot as plot
8+
import pyprint as prnt
9+
import pyselect as slct
10+
11+
import awkward as ak
12+
import argparse
13+
14+
def example_plotting(filename):
15+
""" Simple demo function to run some of the utils """
16+
17+
# Import the files
18+
test_evn = evn.Import(str(filename), "EventNtuple", "ntuple")
19+
20+
# Import code and extract branch
21+
treename = 'trksegs'
22+
branchname = 'time'
23+
surface_id = 1 # tracker middle
24+
ntuple = test_evn.ImportTree()
25+
branches = test_evn.ImportBranches(ntuple,[str(treename)])
26+
27+
# Find fit at chosen ID
28+
selector = slct.Select()
29+
trkent = selector.SelectSurfaceID(branches, treename, surface_id)
30+
31+
# Print out the first event
32+
printer = prnt.Print()
33+
# Print PrintNEvents help
34+
help(printer.PrintNEvents)
35+
printer.PrintNEvents(branches, n_events=1)
36+
37+
# Make 1D plot from flattened array
38+
flatarraytime = ak.flatten(trkent[str(branchname)], axis=None)
39+
plotter = plot.Plot()
40+
41+
# Print Plot1D help (press "q" to exit)
42+
help(plotter.Plot1D)
43+
44+
# Make plot
45+
plotter.Plot1D(
46+
array=flatarraytime,
47+
nbins=100,
48+
xmin=450,
49+
xmax=1695,
50+
title="Example 1D histogram",
51+
xlabel="Fit time at Trk Ent [ns]",
52+
ylabel="Events per bin",
53+
out_path='h1_time.png',
54+
stat_box=True,
55+
error_bars=True
56+
)
57+
58+
# Access vectors
59+
vector = vec.Vector()
60+
vecbranchname = 'mom'
61+
trkentall = selector.SelectSurfaceID(branches, treename, surface_id)
62+
vector_test = vector.GetVectorXYZFromLeaf(trkentall, vecbranchname)
63+
magnitude = vector.Mag(vector_test)
64+
65+
# make 1D plot of magnitudes
66+
flatarraymom = ak.flatten(magnitude, axis=None)
67+
68+
# Make plot
69+
plotter.Plot1D(
70+
array=flatarraymom,
71+
nbins=100,
72+
xmin=0,
73+
xmax=300,
74+
title="Example 1D histogram",
75+
xlabel="Fit mom at Trk Ent [MeV/c]",
76+
ylabel="Events per bin",
77+
out_path='h1_mom.png',
78+
stat_box=True,
79+
error_bars=True,
80+
log_y=True
81+
)
82+
83+
84+
# Print Plot2D help (press "q" to exit)
85+
help(plotter.Plot2D)
86+
87+
# Make 2D plot
88+
plotter.Plot2D(
89+
x=flatarraymom,
90+
y=flatarraytime,
91+
nbins_x=100,
92+
xmin=95,
93+
xmax=115,
94+
nbins_y=100,
95+
ymin=450,
96+
ymax=1650,
97+
title="Example 2D histogram",
98+
xlabel="Fit mom at Trk Ent [MeV/c]",
99+
ylabel="Fit time at Trk Ent [ns]",
100+
out_path='h2_timevmom.png'
101+
)
102+
103+
def main(args):
104+
example_plotting(args.filename)
105+
106+
if __name__ == "__main__":
107+
parser = argparse.ArgumentParser(description='command arguments', formatter_class=argparse.RawTextHelpFormatter)
108+
parser.add_argument("--filename", type=str, default="/exp/mu2e/data/users/sophie/ensembles/MDS1/MDS1av0.root", help="filename")
109+
args = parser.parse_args()
110+
(args) = parser.parse_args()
111+
main(args)

‎examples/mu2e.mplstyle‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# mu2e.mplstyle
2+
3+
# Generic figure tweaks
4+
axes.prop_cycle: cycler("color", ['red','blue','green','orange','violet','cyan','pink','gray','yellow'])
5+
6+
mathtext.default: regular
7+
font.size: 18
8+
axes.labelsize: medium
9+
axes.unicode_minus: False
10+
axes.labelpad: 10
11+
axes.titlesize: 15
12+
axes.linewidth: 2
13+
grid.alpha: 0.8
14+
grid.linestyle: :
15+
savefig.transparent: False
16+
17+
xaxis.labellocation: right
18+
yaxis.labellocation: top
19+
xtick.labelsize: small
20+
ytick.labelsize: small
21+
xtick.direction: in
22+
xtick.major.size: 12
23+
xtick.minor.size: 6
24+
xtick.major.pad: 6
25+
xtick.top: True
26+
xtick.major.top: True
27+
xtick.major.bottom: True
28+
xtick.minor.top: True
29+
xtick.minor.bottom: True
30+
xtick.minor.visible: True
31+
ytick.direction: in
32+
ytick.major.size: 12
33+
ytick.minor.size: 6.0
34+
ytick.right: True
35+
ytick.major.left: True
36+
ytick.major.right: True
37+
ytick.minor.left: True
38+
ytick.minor.right: True
39+
ytick.minor.visible: True
40+
41+
legend.fontsize: small
42+
legend.handlelength: 1.5
43+
legend.borderpad: 0.5
44+
legend.frameon: False
45+

0 commit comments

Comments
 (0)