def test_compressed_tprofile3d(tmp_path): filename = join(str(tmp_path), "example.root") testfile = join(str(tmp_path), "test.root") f = ROOT.TFile.Open(testfile, "RECREATE") h = ROOT.TProfile3D("hvar", "title", 5, 1, 10, 6, 1, 20, 8, 2, 8) h.Sumw2() h.Fill(1.0, 5.0, 3, 6) h.Fill(2.0, 10.0, 4, 7) h.Write() f.Close() t = uproot.open(testfile) hist = t["hvar"] with uproot.recreate(filename, compression=uproot.LZMA(6)) as f: f["test"] = hist f = ROOT.TFile.Open(filename) h = f.Get("test") sums = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 36.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 49.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] bincontents = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 6.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 7.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] count = 0 for x in range(1, 6): for y in range(1, 7): for z in range(1, 9): assert h.GetBinContent(x, y, z) == bincontents[count] count += 1 assert list(h.GetSumw2()) == sums assert h.GetMean() == 1.5 assert h.GetRMS() == 0.5 assert h.GetNbinsX() == 5 assert h.GetNbinsY() == 6 assert h.GetNbinsZ() == 8
def test_compressed_tprofile(tmp_path): filename = join(str(tmp_path), "example.root") testfile = join(str(tmp_path), "test.root") f = ROOT.TFile.Open(testfile, "RECREATE") h = ROOT.TProfile("hvar", "title", 5, 1, 10) h.Sumw2() h.Fill(1.0, 3) h.Fill(2.0, 4) h.Write() f.Close() t = uproot.open(testfile) hist = t["hvar"] with uproot.recreate(filename, compression=uproot.LZMA(5)) as f: f["test"] = hist f = ROOT.TFile.Open(filename) h = f.Get("test") sums = [0.0, 25.0, 0.0, 0.0, 0.0, 0.0, 0.0] bincontents = [3.5, 0.0, 0.0, 0.0, 0.0] assert list(h.GetSumw2()) == sums assert h.GetMean() == 1.5 assert h.GetRMS() == 0.5 count = 0 for x in range(1, 6): assert h.GetBinContent(x) == bincontents[count] count += 1
def test_lzma(tmp_path): filename = join(str(tmp_path), "example.root") with uproot.recreate(filename, compression=uproot.LZMA(1)) as f: f["hello"] = "a"*2000 f = ROOT.TFile.Open(filename) assert f.GetCompressionAlgorithm() == uproot.const.kLZMA assert f.GetCompressionLevel() == 1 assert str(f.Get("hello")) == "a"*2000 f.Close()
def test_flattree_LZMA(tmp_path): pytest.importorskip("lzma") newfile = os.path.join(tmp_path, "newfile.root") branch1 = np.arange(100) branch2 = 1.1 * np.arange(100) with uproot.recreate(newfile, compression=uproot.LZMA(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_LZMA(tmp_path): pytest.importorskip("lzma") 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.LZMA(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_LZMA(tmp_path): pytest.importorskip("lzma") newfile = os.path.join(tmp_path, "newfile.root") SIZE = 2 ** 20 histogram = (np.random.randint(0, 10, SIZE), np.linspace(0, 1, SIZE + 1)) last = histogram[0][-1] with uproot.recreate(newfile, compression=uproot.LZMA(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 })