def test_get_subdir(tmp_path):
empty_filename = os.path.join(tmp_path, "empty.root")
nonempty_filename = os.path.join(tmp_path, "nonempty.root")
empty = ROOT.TFile(empty_filename, "recreate")
empty.mkdir("subdir")
empty.Close()
nonempty = ROOT.TFile(nonempty_filename, "recreate")
nonempty.mkdir("subdir")
nonempty.cd("subdir")
x = ROOT.TObjString("hello")
x.Write()
nonempty.Close()
with uproot.update(empty_filename) as f1:
f1["subdir"].mkdir("another")
with uproot.update(nonempty_filename) as f2:
f2["subdir"].mkdir("another")
assert uproot.open(empty_filename).keys() == [
"subdir;1", "subdir/another;1"
]
assert uproot.open(nonempty_filename).keys() == [
"subdir;1",
"subdir/hello;1",
"subdir/another;1",
]
f3 = ROOT.TFile(empty_filename, "update")
f3.cd("subdir/another")
y = ROOT.TObjString("there")
y.Write()
f3.Close()
f4 = ROOT.TFile(nonempty_filename, "update")
f4.cd("subdir/another")
z = ROOT.TObjString("you")
z.Write()
f4.Close()
assert uproot.open(empty_filename).keys() == [
"subdir;1",
"subdir/another;1",
"subdir/another/there;1",
]
assert uproot.open(nonempty_filename).keys() == [
"subdir;1",
"subdir/hello;1",
"subdir/another;1",
"subdir/another/you;1",
]
def test_writable_vs_readable_tree(tmp_path):
newfile = os.path.join(tmp_path, "newfile.root")
f1 = ROOT.TFile(newfile, "recreate")
f1.SetCompressionLevel(0)
t1 = ROOT.TTree("t1", "title")
d1 = array.array("i", [0])
t1.Branch("branch1", d1, "branch1/I")
d1[0] = 5
t1.Fill()
d1[0] = 4
t1.Fill()
d1[0] = 3
t1.Fill()
d1[0] = 2
t1.Fill()
d1[0] = 1
t1.Fill()
d1[0] = 5
t1.Fill()
t1.Write()
f1.Close()
b1 = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
b2 = [0.0, 1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9]
with uproot.update(newfile) as fin:
with pytest.raises(TypeError):
oldtree = fin["t1"]
fin.mktree("t2", {"b1": np.int32, "b2": np.float64}, "title")
for _ in range(5):
fin["t2"].extend({"b1": b1, "b2": b2})
f1 = ROOT.TFile(newfile)
t2root = f1.Get("t2")
assert t2root.GetEntries() == len(b1) * 5
assert [x.b1 for x in t2root] == b1 * 5
assert [x.b2 for x in t2root] == b2 * 5
t1root = f1.Get("t1")
assert [x.branch1 for x in t1root] == [5, 4, 3, 2, 1, 5]
with uproot.open(newfile) as finagin:
t2uproot = finagin["t2"]
assert t2uproot.num_entries == len(b1) * 5
assert t2uproot["b1"].array(library="np").tolist() == b1 * 5
assert t2uproot["b2"].array(library="np").tolist() == b2 * 5
assert finagin["t1/branch1"].array(library="np").tolist() == [
5, 4, 3, 2, 1, 5
]
f1.Close()
def test_get_string(tmp_path):
filename = os.path.join(tmp_path, "whatever.root")
f1 = ROOT.TFile(filename, "recreate")
x = ROOT.TObjString("hello")
x.Write()
f1.Close()
with uproot.update(filename) as f2:
assert str(f2["hello"]) == "hello"
def test_get_histogram(tmp_path):
filename = os.path.join(tmp_path, "whatever.root")
f1 = ROOT.TFile(filename, "recreate")
x = ROOT.TH1F("name", "title", 100, -5, 5)
x.Write()
f1.Close()
with uproot.update(filename) as f2:
h = f2["name"]
assert h.name == "name"
assert h.title == "title"
def test_delete(tmp_path):
newfile = os.path.join(tmp_path, "newfile.root")
h1 = uproot.writing.identify.to_TH1x(
fName="h1",
fTitle="title",
data=np.array([1.0, 2.0, 5.0, 4.0], np.float64),
fEntries=5.0,
fTsumw=7.0,
fTsumw2=27.0,
fTsumwx=7.3,
fTsumwx2=55.67,
fSumw2=np.array([1.0, 2.0, 25.0, 16.0], np.float64),
fXaxis=uproot.writing.identify.to_TAxis(
fName="xaxis",
fTitle="",
fNbins=2,
fXmin=-3.14,
fXmax=2.71,
),
)
with uproot.recreate(newfile) as fout:
fout["one"] = h1
fout["two"] = h1
with uproot.update(newfile) as fin:
del fin["one"]
f3 = ROOT.TFile(newfile)
h3 = f3.Get("two")
assert h3.GetEntries() == 5
assert h3.GetSumOfWeights() == 7
assert h3.GetBinLowEdge(1) == pytest.approx(-3.14)
assert h3.GetBinWidth(1) == pytest.approx((2.71 - -3.14) / 2)
assert h3.GetBinContent(0) == pytest.approx(1)
assert h3.GetBinContent(1) == pytest.approx(2)
assert h3.GetBinContent(2) == pytest.approx(5)
assert h3.GetBinContent(3) == pytest.approx(4)
assert h3.GetBinError(0) == pytest.approx(1)
assert h3.GetBinError(1) == pytest.approx(1.4142135623730951)
assert h3.GetBinError(2) == pytest.approx(5)
assert h3.GetBinError(3) == pytest.approx(4)
f3.Close()
def test_get_nested(tmp_path):
filename = os.path.join(tmp_path, "whatever.root")
f1 = ROOT.TFile(filename, "recreate")
one = f1.mkdir("one")
one.cd()
two = one.mkdir("two")
two.cd()
x = ROOT.TObjString("hello")
x.Write()
f1.Close()
with uproot.update(filename) as f2:
assert str(f2["one/two/hello"]) == "hello"
assert f2.keys() == ["one;1", "one/two;1", "one/two/hello;1"]
assert f2.classnames() == {
"one;1": "TDirectory",
"one/two;1": "TDirectory",
"one/two/hello;1": "TObjString",
}
def test_update(tmp_path):
filename = os.path.join(tmp_path, "whatever.root")
f1 = ROOT.TFile(filename, "recreate")
f1.mkdir("subdir")
f1.cd("subdir")
x = ROOT.TObjString("wowie")
x.Write()
f1.Close()
with uproot.open(filename) as f2:
assert f2["subdir/wowie"] == "wowie"
assert list(f2.file.streamers) == ["TObjString"]
with uproot.update(filename) as f3:
f3["hey"] = "you"
f3["subdir/there"] = "you guys"
with uproot.open(filename) as f4:
assert f4["hey"] == "you"
assert f4["subdir/there"] == "you guys"
assert f4["subdir/wowie"] == "wowie"
assert list(f4.file.streamers) == ["TObjString"]
f5 = ROOT.TFile(filename, "update")
assert [x.GetName() for x in f5.GetStreamerInfoList()] == ["TObjString"]
assert f5.Get("hey") == "you"
assert f5.Get("subdir/there") == "you guys"
f5.cd("subdir")
y = ROOT.TObjString("zowie")
y.Write()
f5.Close()
with uproot.open(filename) as f6:
assert f6["hey"] == "you"
assert f6["subdir/there"] == "you guys"
assert f6["subdir/wowie"] == "wowie"
assert f6["subdir/zowie"] == "zowie"
assert list(f6.file.streamers) == ["TObjString"]
def data_mc_comparison_plot(tag,
ymin=0,
ymax=1.1,
distribution='recoil',
jeteta_config=None,
output_format='pdf'):
if 'gamma' in tag:
regions = [
'g_HLT_PFHT1050', 'g_HLT_PFHT590', 'g_HLT_PFHT680',
'g_HLT_PFHT780', 'g_HLT_PFHT890'
]
elif 'recoil' in tag:
regions = ['1m', '2m']
else:
regions = ['1m', '2m', '1e', '2m_hlt']
opts = markers('data')
# opts['markersize'] = 5
# opts['fillstyle'] = 'none'
emarker = opts.pop('emarker', '')
outdir = f"./output/{tag}"
outpath = pjoin(outdir, f'trig_sf.root')
try:
outfile = uproot.recreate(outpath)
except OSError:
outfile = uproot.update(outpath)
for year in [2017, 2018]:
for region in regions:
fig, ax, rax = fig_ratio()
if '1e' in region:
fnum = f'output/{tag}/table_{region}_met_EGamma_{year}.txt'
fden = f'output/{tag}/table_{region}_met_WJetsToLNu_HT_MLM_{year}.txt'
xlabel = "$p_{T}^{miss}$ (GeV)"
elif '1m' in region:
fnum = f'output/{tag}/table_{region}_recoil_SingleMuon_{year}{"_"+jeteta_config if jeteta_config else ""}.txt'
fden = f'output/{tag}/table_{region}_recoil_WJetsToLNu_HT_MLM_{year}{"_"+jeteta_config if jeteta_config else ""}.txt'
xlabel = "Recoil (GeV)"
elif '2m' in region:
fnum = f'output/{tag}/table_{region}_recoil_SingleMuon_{year}{"_"+jeteta_config if jeteta_config else ""}.txt'
fden = f'output/{tag}/table_{region}_recoil_VDYJetsToLL_M-50_HT_MLM_{year}{"_"+jeteta_config if jeteta_config else ""}.txt'
xlabel = "Recoil (GeV)"
elif 'g_' in region:
fnum = f'output/{tag}/table_{region}_photon_pt0_JetHT_{year}.txt'
fden = f'output/{tag}/table_{region}_photon_pt0_GJets_HT_MLM_{year}.txt'
xlabel = "Photon $p_{T}$ (GeV)"
if not os.path.exists(fnum):
print(f"File not found {fnum}")
continue
if not os.path.exists(fden):
print(f"File not found {fden}")
continue
xnum, xedgnum, ynum, yerrnum = get_xy(fnum)
xden, xedgden, yden, yerrden = get_xy(fden)
xsf = xnum
ysf = ynum / yden
ysferr = ratio_unc(ynum, yden, yerrnum, yerrden)
opts['color'] = 'k'
ax.errorbar(xnum,
ynum,
yerr=yerrnum,
label=f'Data, {region} region',
**opts)
opts['color'] = 'r'
ax.errorbar(xden,
yden,
yerr=yerrden,
label=f'MC, {region} region',
**opts)
# rax.plot([0,1000],[0.98,0.98],color='blue')
# rax.plot([0,1000],[0.99,0.99],color='blue',linestyle='--')
if 'g_' in region:
ax.plot([215, 215], [0.9, 1.1], color='blue')
rax.plot([215, 215], [0.95, 1.05], color='blue')
elif distribution == 'recoil':
ax.plot([250, 250], [0.0, 1.1], color='blue')
rax.plot([250, 250], [0.95, 1.05], color='blue')
else:
ax.plot([250, 250], [0.9, 1.1], color='blue')
rax.plot([250, 250], [0.95, 1.05], color='blue')
opts['color'] = 'k'
rax.errorbar(xsf, ysf, ysferr, **opts)
ax.legend()
ax.set_ylabel("Efficiency")
ax.xaxis.set_major_locator(MultipleLocator(200))
ax.xaxis.set_minor_locator(MultipleLocator(50))
ax.set_ylim(ymin, ymax)
ax.grid(1)
if distribution == 'mjj':
ax.yaxis.set_major_locator(MultipleLocator(0.05))
ax.yaxis.set_minor_locator(MultipleLocator(0.01))
elif distribution == 'recoil':
ax.yaxis.set_major_locator(MultipleLocator(0.1))
ax.yaxis.set_minor_locator(MultipleLocator(0.05))
rax.set_xlabel(xlabel)
rax.set_ylabel("Data / MC SF")
rax.set_ylim(0.95, 1.05)
rax.yaxis.set_major_locator(MultipleLocator(0.05))
rax.yaxis.set_minor_locator(MultipleLocator(0.01))
rax.grid(1)
plt.text(1.,
1.,
r"$\approx$ %.1f fb$^{-1}$ (13 TeV)" %
lumi_by_region(region, year),
fontsize=16,
horizontalalignment='right',
verticalalignment='bottom',
transform=ax.transAxes)
plt.text(1.,
0.95,
f'{jeteta_config if jeteta_config else ""}',
fontsize=12,
horizontalalignment='right',
verticalalignment='bottom',
transform=ax.transAxes)
plt.text(0.,
1.,
f'{year}',
fontsize=16,
horizontalalignment='left',
verticalalignment='bottom',
transform=ax.transAxes)
fig.savefig(
pjoin(
outdir,
f'data_mc_comparison_{region}_{year}{"_"+jeteta_config if jeteta_config else ""}.{output_format}'
))
fig.clear()
plt.close(fig)
vals = np.array(sorted(list(set(list(xedgnum.flatten())))))
ysf[np.isnan(ysf) | np.isinf(np.abs(ysf))] = 1
outfile[f'{tag}_{region}_{year}'] = (ysf, vals)