From 33d12eccdeba5e1e668d60b24fc1a1f49c891e54 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sun, 26 May 2019 17:03:49 -0500 Subject: [PATCH 1/4] handle ChunkedArrays in ArrayMethods.p3 --- uproot_methods/classes/TLorentzVector.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/uproot_methods/classes/TLorentzVector.py b/uproot_methods/classes/TLorentzVector.py index 678a4c6..27fa5d2 100644 --- a/uproot_methods/classes/TLorentzVector.py +++ b/uproot_methods/classes/TLorentzVector.py @@ -6,6 +6,7 @@ import numbers import operator +import awkward.array.chunked import awkward.array.jagged import awkward.array.objects import awkward.util @@ -134,15 +135,22 @@ def __awkward_persist__(self, ident, fill, prefix, suffix, schemasuffix, storage fill(z, "TLorentzVectorArray.z", prefix, suffix, schemasuffix, storage, compression, **kwargs), fill(t, "TLorentzVectorArray.t", prefix, suffix, schemasuffix, storage, compression, **kwargs)]} + @staticmethod + def _wrapmethods(node): + if isinstance(node, awkward.array.chunked.ChunkedArray): + node.__class__ = type("ChunkedArrayMethods", (self.awkward.ChunkedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) + for chunk in node.chunks: + ArrayMethods._wrapmethods(chunk) + elif isinstance(node, awkward.array.jagged.JaggedArray): + node.__class__ = type("JaggedArrayMethods", (self.awkward.JaggedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) + ArrayMethods._wrapmethods(node.content) + elif isinstance(node, awkward.array.objects.ObjectArray): + node.__class__ = type("ObjectArrayMethods", (self.awkward.ObjectArray, uproot_methods.classes.TVector3.ArrayMethods), {}) + @property def p3(self): out = self.empty_like(generator=lambda row: uproot_methods.classes.TVector3.TVector3(row["fX"], row["fY"], row["fZ"])) - node = out - while isinstance(node, awkward.array.jagged.JaggedArray): - node.__class__ = type("JaggedArrayMethods", (self.awkward.JaggedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) - node = node.content - if isinstance(node, awkward.array.objects.ObjectArray): - node.__class__ = type("ObjectArrayMethods", (self.awkward.ObjectArray, uproot_methods.classes.TVector3.ArrayMethods), {}) + ArrayMethods._wrapmethods(out) out["fX"] = self.x out["fY"] = self.y out["fZ"] = self.z From b8a437b1da73cc57d381be98e65c3660760dbffc Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sun, 26 May 2019 17:08:13 -0500 Subject: [PATCH 2/4] no self --- uproot_methods/classes/TLorentzVector.py | 14 +++++++------- uproot_methods/version.py | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/uproot_methods/classes/TLorentzVector.py b/uproot_methods/classes/TLorentzVector.py index 27fa5d2..b49dd84 100644 --- a/uproot_methods/classes/TLorentzVector.py +++ b/uproot_methods/classes/TLorentzVector.py @@ -136,21 +136,21 @@ def __awkward_persist__(self, ident, fill, prefix, suffix, schemasuffix, storage fill(t, "TLorentzVectorArray.t", prefix, suffix, schemasuffix, storage, compression, **kwargs)]} @staticmethod - def _wrapmethods(node): + def _wrapmethods(node, awkwardlib): if isinstance(node, awkward.array.chunked.ChunkedArray): - node.__class__ = type("ChunkedArrayMethods", (self.awkward.ChunkedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) + node.__class__ = type("ChunkedArrayMethods", (awkwardlib.ChunkedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) for chunk in node.chunks: - ArrayMethods._wrapmethods(chunk) + ArrayMethods._wrapmethods(chunk, awkwardlib) elif isinstance(node, awkward.array.jagged.JaggedArray): - node.__class__ = type("JaggedArrayMethods", (self.awkward.JaggedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) - ArrayMethods._wrapmethods(node.content) + node.__class__ = type("JaggedArrayMethods", (awkwardlib.JaggedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) + ArrayMethods._wrapmethods(node.content, awkwardlib) elif isinstance(node, awkward.array.objects.ObjectArray): - node.__class__ = type("ObjectArrayMethods", (self.awkward.ObjectArray, uproot_methods.classes.TVector3.ArrayMethods), {}) + node.__class__ = type("ObjectArrayMethods", (awkwardlib.ObjectArray, uproot_methods.classes.TVector3.ArrayMethods), {}) @property def p3(self): out = self.empty_like(generator=lambda row: uproot_methods.classes.TVector3.TVector3(row["fX"], row["fY"], row["fZ"])) - ArrayMethods._wrapmethods(out) + ArrayMethods._wrapmethods(out, self.awkward) out["fX"] = self.x out["fY"] = self.y out["fZ"] = self.z diff --git a/uproot_methods/version.py b/uproot_methods/version.py index a09d4fa..7af924e 100644 --- a/uproot_methods/version.py +++ b/uproot_methods/version.py @@ -4,7 +4,7 @@ import re -__version__ = "0.6.0" +__version__ = "0.6.1" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__)) From 3e1d5c412ffa7f86eb1caca2412bb0d9d03e5e02 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sun, 26 May 2019 18:33:37 -0500 Subject: [PATCH 3/4] corrected signatures for Numpy histograms --- uproot_methods/convert.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/uproot_methods/convert.py b/uproot_methods/convert.py index 684d998..6a344c9 100644 --- a/uproot_methods/convert.py +++ b/uproot_methods/convert.py @@ -26,11 +26,17 @@ def types(cls, obj): if any(x == ("builtins", "bytes") or x == ("builtins", "str") or x == ("__builtin__", "str") or x == ("__builtin__", "unicode") for x in types(obj.__class__, obj)): return (None, None, "uproot.write.objects.TObjString", "TObjString") - elif isinstance(obj, tuple) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[0].__class__, obj[0])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1].__class__, obj[1])) and obj[0].shape == (len(obj[1])-1, len(obj[2])-1): + # made with numpy.histogram + elif isinstance(obj, tuple) and len(obj) == 2 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[0].__class__, obj[0])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1].__class__, obj[1])) and len(obj[0].shape) == 1 and len(obj[1].shape) == 1 and obj[0].shape[0] == obj[1].shape[0] - 1: + return ("uproot_methods.classes.TH1", "from_numpy", "uproot.write.objects.TH", "TH") + + # made with numpy.histogram2d + elif isinstance(obj, tuple) and len(obj) == 3 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[0].__class__, obj[0])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1].__class__, obj[1])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[2].__class__, obj[2])) and len(obj[0].shape) == 2 and len(obj[1].shape) == 1 and len(obj[2].shape) == 1 and obj[0].shape[0] == obj[1].shape[0] - 1 and obj[0].shape[1] == obj[1].shape[0] - 1: return ("uproot_methods.classes.TH2", "from_numpy", "uproot.write.objects.TH", "TH") - elif isinstance(obj, tuple) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[0].__class__, obj[0])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1].__class__, obj[1])) and len(obj[0]) + 1 == len(obj[1]): - return ("uproot_methods.classes.TH1", "from_numpy", "uproot.write.objects.TH", "TH") + # made with numpy.histogramdd (2-dimensional) + elif isinstance(obj, tuple) and len(obj) == 2 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[0].__class__, obj[0])) and isinstance(obj[1], list) and len(obj[1]) == 2 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1][0].__class__, obj[1][0])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1][1].__class__, obj[1][1])) and len(obj[0].shape) == 2 and len(obj[1][0].shape) == 1 and len(obj[1][1].shape) == 1 and obj[0].shape[0] == obj[1][0].shape - 1 and obj[0].shape[0] == obj[1][1].shape - 1: + return ("uproot_methods.classes.TH2", "from_numpy", "uproot.write.objects.TH", "TH") elif any(x[:3] == ("pandas.core.frame", "DataFrame", "IntervalIndex") and "count" in x[3] for x in types(obj.__class__, obj)): return ("uproot_methods.classes.TH1", "from_pandas", "uproot.write.objects.TH", "TH") From 63579126ead07e6442e38f452dcdf316689ce721 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Sun, 26 May 2019 18:41:40 -0500 Subject: [PATCH 4/4] all histogram cases recognized --- uproot_methods/classes/TH2.py | 5 ++++- uproot_methods/convert.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/uproot_methods/classes/TH2.py b/uproot_methods/classes/TH2.py index 6d415fd..3a5b328 100644 --- a/uproot_methods/classes/TH2.py +++ b/uproot_methods/classes/TH2.py @@ -224,7 +224,10 @@ def _histtype(content): return b"TH2D", content.astype(">f8") def from_numpy(histogram): - content, xedges, yedges = histogram[:3] + if isinstance(histogram[1], list) and len(histogram[1]) == 2: + content, (xedges, yedges) = histogram[:2] + else: + content, xedges, yedges = histogram[:3] class TH2(Methods, list): pass diff --git a/uproot_methods/convert.py b/uproot_methods/convert.py index 6a344c9..e661267 100644 --- a/uproot_methods/convert.py +++ b/uproot_methods/convert.py @@ -35,7 +35,7 @@ def types(cls, obj): return ("uproot_methods.classes.TH2", "from_numpy", "uproot.write.objects.TH", "TH") # made with numpy.histogramdd (2-dimensional) - elif isinstance(obj, tuple) and len(obj) == 2 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[0].__class__, obj[0])) and isinstance(obj[1], list) and len(obj[1]) == 2 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1][0].__class__, obj[1][0])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1][1].__class__, obj[1][1])) and len(obj[0].shape) == 2 and len(obj[1][0].shape) == 1 and len(obj[1][1].shape) == 1 and obj[0].shape[0] == obj[1][0].shape - 1 and obj[0].shape[0] == obj[1][1].shape - 1: + elif isinstance(obj, tuple) and len(obj) == 2 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[0].__class__, obj[0])) and isinstance(obj[1], list) and len(obj[1]) == 2 and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1][0].__class__, obj[1][0])) and any(x[:2] == ("numpy", "ndarray") for x in types(obj[1][1].__class__, obj[1][1])) and len(obj[0].shape) == 2 and len(obj[1][0].shape) == 1 and len(obj[1][1].shape) == 1 and obj[0].shape[0] == obj[1][0].shape[0] - 1 and obj[0].shape[0] == obj[1][1].shape[0] - 1: return ("uproot_methods.classes.TH2", "from_numpy", "uproot.write.objects.TH", "TH") elif any(x[:3] == ("pandas.core.frame", "DataFrame", "IntervalIndex") and "count" in x[3] for x in types(obj.__class__, obj)):