def test_lz4_leveldown(tmp_path):
filename = join(str(tmp_path), "example.root")
with uproot.recreate(filename, compression=uproot.LZ4(5)) as f:
f["hello"] = "a"*2000
f = ROOT.TFile.Open(filename)
assert (f.GetCompressionAlgorithm()) == uproot.const.kLZ4
assert str(f.Get("hello")) == "a"*2000
f.Close()
def WriteTree(df, cols, treeName, fileName):
'''
Helper method to write a tree with uproot
Arguments
----------
- pandas data frame to be written as tree in a root file
- name of the columns
- name of the output tree
- name of the output file
'''
outBranches = {}
for colName in cols:
outBranches[
colName] = np.float32 #define all branches as float for the moment
with uproot.recreate(fileName, compression=uproot.LZ4(4)) as outFile:
outFile[treeName] = uproot.newtree(outBranches,
compression=uproot.LZ4(4))
outFile[treeName].extend(dict(df[cols]))
def test_flattree_LZ4(tmp_path):
pytest.importorskip("lz4")
newfile = os.path.join(tmp_path, "newfile.root")
branch1 = np.arange(100)
branch2 = 1.1 * np.arange(100)
with uproot.recreate(newfile, compression=uproot.LZ4(5)) as fout:
fout["tree"] = {"branch1": branch1, "branch2": branch2}
fout["tree"].extend({"branch1": branch1, "branch2": branch2})
with uproot.open(newfile) as fin:
assert fin["tree/branch1"].array(library="np").tolist() == branch1.tolist() * 2
assert fin["tree/branch2"].array(library="np").tolist() == branch2.tolist() * 2
f3 = ROOT.TFile(newfile)
t3 = f3.Get("tree")
assert [x.branch1 for x in t3] == branch1.tolist() * 2
assert [x.branch2 for x in t3] == branch2.tolist() * 2
f3.Close()
def test_jaggedtree_LZ4(tmp_path):
pytest.importorskip("lz4")
ak = pytest.importorskip("awkward")
newfile = os.path.join(tmp_path, "newfile.root")
branch1 = ak.Array([[1, 2, 3], [], [4, 5]] * 10)
branch2 = ak.Array([[1.1, 2.2, 3.3], [], [4.4, 5.5]] * 10)
with uproot.recreate(newfile, compression=uproot.LZ4(5)) as fout:
fout["tree"] = {"branch1": branch1, "branch2": branch2}
fout["tree"].extend({"branch1": branch1, "branch2": branch2})
with uproot.open(newfile) as fin:
assert fin["tree/branch1"].array().tolist() == branch1.tolist() * 2
assert fin["tree/branch2"].array().tolist() == branch2.tolist() * 2
f3 = ROOT.TFile(newfile)
t3 = f3.Get("tree")
assert [list(x.branch1) for x in t3] == branch1.tolist() * 2
assert [list(x.branch2) for x in t3] == branch2.tolist() * 2
f3.Close()
def test_histogram_LZ4(tmp_path):
pytest.importorskip("lz4")
newfile = os.path.join(tmp_path, "newfile.root")
SIZE = 2 ** 21
histogram = (np.random.randint(0, 10, SIZE), np.linspace(0, 1, SIZE + 1))
last = histogram[0][-1]
with uproot.recreate(newfile, compression=uproot.LZ4(1)) as fout:
fout["out"] = histogram
with uproot.open(newfile) as fin:
content, edges = fin["out"].to_numpy()
assert len(content) == SIZE
assert len(edges) == SIZE + 1
assert content[-1] == last
f3 = ROOT.TFile(newfile)
h3 = f3.Get("out")
assert h3.GetNbinsX() == SIZE
assert h3.GetBinContent(SIZE) == last
f3.Close()
import uproot
b1 = uproot.newbranch("i4", compression=uproot.ZLIB(5))
b2 = uproot.newbranch("i8", compression=uproot.LZMA(4))
b3 = uproot.newbranch("f4")
branchdict = {"branch1": b1, "branch2": b2, "branch3": b3}
tree = uproot.newtree(branchdict, compression=uproot.LZ4(4))
with uproot.recreate("example.root", compression=uproot.LZMA(5)) as f:
f["t"] = tree
f["t"].extend({
"branch1": [1] * 1000,
"branch2": [2] * 1000,
"branch3": [3] * 1000
})