Fix resolving TH2D from numpy.histogram2d

scikit-hep · Aug 6, 2019 · b17ec73 · b17ec73
1 parent b64848e
commit b17ec73
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/uproot_methods/convert.py b/uproot_methods/convert.py
@@ -31,11 +31,11 @@ def types(cls, obj):
             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:
+        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[2].shape[0] - 1:
             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[0] - 1 and obj[0].shape[0] == obj[1][1].shape[0] - 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[1] == 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)):