Merge pull request #57 from scikit-hep/fixes-for-LPC-HATS

prepare for LPC HATS in all three libraries

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

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, awkwardlib):
+        if isinstance(node, awkward.array.chunked.ChunkedArray):
+            node.__class__ = type("ChunkedArrayMethods", (awkwardlib.ChunkedArray, uproot_methods.classes.TVector3.ArrayMethods), {})
+            for chunk in node.chunks:
+                ArrayMethods._wrapmethods(chunk, awkwardlib)
+        elif isinstance(node, awkward.array.jagged.JaggedArray):
+            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", (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"]))
-        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, self.awkward)
         out["fX"] = self.x
         out["fY"] = self.y
         out["fZ"] = self.z

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[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)):
             return ("uproot_methods.classes.TH1", "from_pandas", "uproot.write.objects.TH", "TH")

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__))