feat: use uproot4 for file writing (#256)

* use uproot>=4.1 for file writing instead of uproot3 * drop uproot3 dependency from contrib setup extra
scikit-hep · Aug 27, 2021 · 5ed199a · 5ed199a
1 parent 3959aee
commit 5ed199a
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 47 deletions.
diff --git a/setup.py b/setup.py
@@ -1,6 +1,6 @@
 from setuptools import setup
 
-extras_require = {"contrib": ["uproot3", "uproot>=4.0"]}
+extras_require = {"contrib": ["uproot>=4.1"]}  # uproot 4.1 for file writing
 extras_require["test"] = sorted(
     set(
         extras_require["contrib"]

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,13 +1,12 @@
 import pytest
-import uproot3 as uproot
+import uproot
 
 
 class Utils:
     @staticmethod
     def create_ntuple(fname, treename, varname, var_array, weightname, weight_array):
         with uproot.recreate(fname) as f:
-            f[treename] = uproot.newtree({varname: "float64", weightname: "float64"})
-            f[treename].extend({varname: var_array, weightname: weight_array})
+            f[treename] = {varname: var_array, weightname: weight_array}
 
 
 @pytest.fixture

diff --git a/tests/contrib/test_histogram_creation.py b/tests/contrib/test_histogram_creation.py
@@ -7,9 +7,9 @@ def test_from_uproot(tmp_path, utils):
     fnames = [tmp_path / "test.root"]
     treename = "tree"
     varname = "var"
-    var_array = [1.1, 2.3, 3.0, 3.2]
+    var_array = np.asarray([1.1, 2.3, 3.0, 3.2])
     weightname_write = "weight"
-    weight_array = [1.0, 1.0, 2.0, 1.0]
+    weight_array = np.asarray([1.0, 1.0, 2.0, 1.0])
     bins = np.asarray([1, 2, 3, 4])
     # create something to read
     utils.create_ntuple(

diff --git a/util/create_ntuples.py b/util/create_ntuples.py
@@ -3,7 +3,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-import uproot3 as uproot
+import uproot
 
 
 def toy_distribution(noncentral, multiplier, offset, num_events):
@@ -60,73 +60,51 @@ def create_file(file_name, distributions, weights, labels, extra_weights=None):
         for i, label in enumerate(labels):
             lep_charge = create_lepton_charge(n_events)
             if label == "background":
-                f[label] = uproot.newtree(
-                    {
-                        "jet_pt": "float64",
-                        "weight": "float64",
-                        "lep_charge": "int",
-                        "weight_up": "float64",
-                        "weight_down": "float64",
-                    }
-                )
-                f[label].extend(
-                    {
-                        "jet_pt": distributions[i],
-                        "weight": weights[i],
-                        "lep_charge": lep_charge,
-                        "weight_up": extra_weights[0],
-                        "weight_down": extra_weights[1],
-                    }
-                )
+                f[label] = {
+                    "jet_pt": distributions[i],
+                    "weight": weights[i],
+                    "lep_charge": lep_charge,
+                    "weight_up": extra_weights[0],
+                    "weight_down": extra_weights[1],
+                }
             else:
-                f[label] = uproot.newtree(
-                    {"jet_pt": "float64", "weight": "float64", "lep_charge": "int"}
-                )
-                f[label].extend(
-                    {
-                        "jet_pt": distributions[i],
-                        "weight": weights[i],
-                        "lep_charge": lep_charge,
-                    }
-                )
+                f[label] = {
+                    "jet_pt": distributions[i],
+                    "weight": weights[i],
+                    "lep_charge": lep_charge,
+                }
 
 
 def create_file_pseudodata(file_name, pseudodata):
     n_events = len(pseudodata)
     with uproot.recreate(file_name) as f:
         # write pseudodata
         lep_charge = create_lepton_charge(n_events)
-        f["pseudodata"] = uproot.newtree({"jet_pt": "float64", "lep_charge": "int"})
-        f["pseudodata"].extend({"jet_pt": pseudodata, "lep_charge": lep_charge})
+        f["pseudodata"] = {"jet_pt": pseudodata, "lep_charge": lep_charge}
 
 
 def read_file(file_name):
     distributions = []
     weights = []
     labels = []
     with uproot.open(file_name) as f:
-        all_trees = f.allkeys(
-            filterclass=lambda cls: issubclass(cls, uproot.tree.TTreeMethods)
-        )
+        all_trees = f.keys(filter_classname="TTree", recursive=True)
         for tree in all_trees:
-            distributions.append(f[tree].array("jet_pt"))
-            weights.append(f[tree].array("weight"))
-            labels.append(tree)
+            distributions.append(f[tree]["jet_pt"].array(library="np"))
+            weights.append(f[tree]["weight"].array(library="np"))
+            labels.append(f[tree].name)
     return distributions, weights, labels
 
 
 def read_file_pseudodata(file_name):
     with uproot.open(file_name) as f:
-        distribution = f["pseudodata"].array("jet_pt")
+        distribution = f["pseudodata"]["jet_pt"].array(library="np")
     return distribution
 
 
 def plot_distributions(data, weights, labels, pseudodata, bins):
     bin_width_str = str(int(bins[1] - bins[0]))
-
-    # labels = [l.split('\'')[1] for l in labels]
     yield_each = [str(round(np.sum(w), 1)) for w in weights]
-    labels = [label.decode().split(";")[0] for label in labels]
 
     # plot normalized distributions
     for i in reversed(range(len(data))):