Merge pull request #48 from guitargeek/concatenate_mixin_bug

Avoid dynamic creation of ArrayMethods classes and new wrapjaggedmethods decorator

uproot_methods/base.py

@@ -28,13 +28,12 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-try:
-    from collections.abc import Iterable
-except ImportError:
-    from collections import Iterable
 
 import awkward
 
+from uproot_methods.wrapjagged import normalize_arrays, unwrap_jagged
+
+
 class ROOTMethods(awkward.Methods):
     _arraymethods = None
 
@@ -45,60 +44,15 @@ def __ne__(self, other):
 
     @classmethod
     def _normalize_arrays(cls, arrays):
-        length = None
-        for i in range(len(arrays)):
-            if isinstance(arrays[i], Iterable):
-                if length is None:
-                    length = len(arrays[i])
-                    break
-        if length is None:
-            raise TypeError("cannot construct an array if all arguments are scalar")
-
-        arrays = list(arrays)
-        jaggedtype = [cls.awkward.JaggedArray] * len(arrays)
-        starts, stops = None, None
-        for i in range(len(arrays)):
-            if starts is None and isinstance(arrays[i], cls.awkward.JaggedArray):
-                starts, stops = arrays[i].starts, arrays[i].stops
-
-            if isinstance(arrays[i], cls.awkward.JaggedArray):
-                jaggedtype[i] = type(arrays[i])
-
-            if not isinstance(arrays[i], Iterable):
-                arrays[i] = cls.awkward.numpy.full(length, arrays[i])
-
-            arrays[i] = cls.awkward.util.toarray(arrays[i], cls.awkward.numpy.float64)
-
-        if starts is None:
-            return arrays
-
-        for i in range(len(arrays)):
-            if not isinstance(arrays[i], cls.awkward.JaggedArray) or not (cls.awkward.numpy.array_equal(starts, arrays[i].starts) and cls.awkward.numpy.array_equal(stops, arrays[i].stops)):
-                content = cls.awkward.numpy.zeros(stops.max(), dtype=cls.awkward.numpy.float64)
-                arrays[i] = jaggedtype[i](starts, stops, content) + arrays[i]    # invoke jagged broadcasting to align arrays
-
-        return arrays
+        return normalize_arrays(cls, arrays)
 
     @classmethod
-    def _unwrap_jagged(cls, ArrayMethods, arrays):
-        if not isinstance(arrays[0], cls.awkward.JaggedArray):
-            return lambda x: x, arrays
-        else:
-            if ArrayMethods is None:
-                awkcls = arrays[0].JaggedArray
-            else:
-                awkcls = ArrayMethods.mixin(ArrayMethods, arrays[0].JaggedArray)
-            counts = arrays[0].counts.reshape(-1)
-            offsets = awkcls.counts2offsets(counts)
-            starts, stops = offsets[:-1], offsets[1:]
-            starts = starts.reshape(arrays[0].starts.shape[:-1] + (-1,))
-            stops = stops.reshape(arrays[0].stops.shape[:-1] + (-1,))
-            wrap, arrays = cls._unwrap_jagged(ArrayMethods, [x.flatten() for x in arrays])
-            return lambda x: awkcls(starts, stops, wrap(x)), arrays
+    def _unwrap_jagged(cls, awkcls, arrays):
+        return unwrap_jagged(cls, awkcls, arrays)
 
     def _trymemo(self, name, function):
         memoname = "_memo_" + name
-        wrap, (array,) = type(self)._unwrap_jagged(None, (self,))
+        wrap, (array,) = type(self)._unwrap_jagged(self.JaggedArray, (self,))
         if not hasattr(array, memoname):
             setattr(array, memoname, function(array))
         return wrap(getattr(array, memoname))

uproot_methods/classes/TLorentzVector.py

@@ -39,6 +39,8 @@
 import uproot_methods.common.TVector
 import uproot_methods.classes.TVector3
 
+from uproot_methods.wrapjagged import wrapjaggedmethod
+
 class Common(object):
     @property
     def E(self):
@@ -753,7 +755,14 @@ def mass(self):
     def mass(self, value):
         self["fMass"] = value
 
+awkward = uproot_methods.base.ROOTMethods.awkward
+PtEtaPhiMassJaggedArrayMethods = PtEtaPhiMassArrayMethods.mixin(PtEtaPhiMassArrayMethods, awkward.JaggedArray)
+
 class TLorentzVectorArray(ArrayMethods, uproot_methods.base.ROOTMethods.awkward.ObjectArray):
+
+    jaggedtype = uproot_methods.base.ROOTMethods.awkward.JaggedArray
+    awkcls = ArrayMethods.mixin(ArrayMethods, jaggedtype)
+
     def __init__(self, x, y, z, t):
         if isinstance(x, awkward.array.jagged.JaggedArray) or isinstance(y, awkward.array.jagged.JaggedArray) or isinstance(z, awkward.array.jagged.JaggedArray) or isinstance(t, awkward.array.jagged.JaggedArray):
             raise TypeError("TLorentzVectorArray constructor arguments must not be jagged; use TLorentzVectorArray.from_cartesian for jaggedness-handling")
@@ -778,32 +787,32 @@ def origin_like(cls, array):
 
     @classmethod
     def from_p3(cls, p3, t):
-        wrap, (x, y, z, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((p3.x, p3.y, p3.z, t)))
+        wrap, (x, y, z, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((p3.x, p3.y, p3.z, t)))
         return wrap(cls(x, y, z, t))
 
     @classmethod
     def from_cartesian(cls, x, y, z, t):
-        wrap, (x, y, z, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y, z, t)))
+        wrap, (x, y, z, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y, z, t)))
         return wrap(cls(x, y, z, t))
 
     @classmethod
     def from_spherical(cls, r, theta, phi, t):
-        wrap, (r, theta, phi, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((r, theta, phi, t)))
+        wrap, (r, theta, phi, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((r, theta, phi, t)))
         return wrap(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_spherical(r, theta, phi), t))
 
     @classmethod
     def from_cylindrical(cls, rho, phi, z, t):
-        wrap, (rho, phi, z, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((rho, phi, z, t)))
+        wrap, (rho, phi, z, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((rho, phi, z, t)))
         return wrap(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_cylindrical(rho, phi, z), t))
 
     @classmethod
     def from_xyzm(cls, x, y, z, m):
-        wrap, (x, y, z, m) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y, z, m)))
+        wrap, (x, y, z, m) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y, z, m)))
         return wrap(cls(x, y, z, cls.awkward.numpy.sqrt(x*x + y*y + z*z + m*m*cls.awkward.numpy.sign(m))))
 
     @classmethod
     def from_ptetaphi(cls, pt, eta, phi, energy):
-        wrap, (pt, eta, phi, energy) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((pt, eta, phi, energy)))
+        wrap, (pt, eta, phi, energy) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((pt, eta, phi, energy)))
         out = cls(pt * cls.awkward.numpy.cos(phi),
                   pt * cls.awkward.numpy.sin(phi),
                   pt * cls.awkward.numpy.sinh(eta),
@@ -814,10 +823,9 @@ def from_ptetaphi(cls, pt, eta, phi, energy):
         return wrap(out)
 
     @classmethod
+    @wrapjaggedmethod(PtEtaPhiMassJaggedArrayMethods)
     def from_ptetaphim(cls, pt, eta, phi, mass):
-        wrap, (pt, eta, phi, mass) = cls._unwrap_jagged(PtEtaPhiMassArrayMethods, cls._normalize_arrays((pt, eta, phi, mass)))
-        out = PtEtaPhiMassLorentzVectorArray(pt,eta,phi,mass)
-        return wrap(out)
+        return PtEtaPhiMassLorentzVectorArray(pt,eta,phi,mass)
 
     @property
     def x(self):

uproot_methods/classes/TVector2.py

@@ -145,6 +145,9 @@ def _unary(self, operator):
         return TVector2(operator(self.x), operator(self.y))
 
 class TVector2Array(ArrayMethods, uproot_methods.base.ROOTMethods.awkward.ObjectArray):
+    jaggedtype = uproot_methods.base.ROOTMethods.awkward.JaggedArray
+    awkcls = ArrayMethods.mixin(ArrayMethods, jaggedtype)
+
     def __init__(self, x, y):
         if isinstance(x, awkward.array.jagged.JaggedArray) or isinstance(y, awkward.array.jagged.JaggedArray):
             raise TypeError("TVector2Array constructor arguments must not be jagged; use TVector2.from_cartesian for jaggedness-handling")
@@ -164,12 +167,12 @@ def origin_like(cls, array):
 
     @classmethod
     def from_cartesian(cls, x, y):
-        wrap, (x, y) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y)))
+        wrap, (x, y) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y)))
         return wrap(cls(x, y))
 
     @classmethod
     def from_polar(cls, rho, phi):
-        wrap, (rho, phi) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((rho, phi)))
+        wrap, (rho, phi) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((rho, phi)))
         return wrap(cls(rho * cls.awkward.numpy.cos(phi),
                         rho * cls.awkward.numpy.sin(phi)))
 

uproot_methods/classes/TVector3.py

@@ -257,6 +257,9 @@ def rotate_euler(self, phi=0, theta=0, psi=0):
         return TVector3(x, y, z)
 
 class TVector3Array(ArrayMethods, uproot_methods.base.ROOTMethods.awkward.ObjectArray):
+    jaggedtype = uproot_methods.base.ROOTMethods.awkward.JaggedArray
+    awkcls = ArrayMethods.mixin(ArrayMethods, jaggedtype)
+
     def __init__(self, x, y, z):
         if isinstance(x, awkward.array.jagged.JaggedArray) or isinstance(y, awkward.array.jagged.JaggedArray) or isinstance(z, awkward.array.jagged.JaggedArray):
             raise TypeError("TVector3Array constructor arguments must not be jagged; use TVector3.from_cartesian for jaggedness-handling")
@@ -279,19 +282,19 @@ def origin_like(cls, array):
 
     @classmethod
     def from_cartesian(cls, x, y, z):
-        wrap, (x, y, z) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y, z)))
+        wrap, (x, y, z) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y, z)))
         return wrap(cls(x, y, z))
 
     @classmethod
     def from_spherical(cls, r, theta, phi):
-        wrap, (r, theta, phi) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((r, theta, phi)))
+        wrap, (r, theta, phi) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((r, theta, phi)))
         return wrap(cls(r * cls.awkward.numpy.sin(theta) * cls.awkward.numpy.cos(phi),
                         r * cls.awkward.numpy.sin(theta) * cls.awkward.numpy.sin(phi),
                         r * cls.awkward.numpy.cos(theta)))
 
     @classmethod
     def from_cylindrical(cls, rho, phi, z):
-        wrap, (rho, phi, z) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((rho, phi, z)))
+        wrap, (rho, phi, z) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((rho, phi, z)))
         return wrap(cls(rho * cls.awkward.numpy.cos(phi),
                         rho * cls.awkward.numpy.sin(phi),
                         z))

uproot_methods/version.py

@@ -30,7 +30,7 @@
 
 import re
 
-__version__ = "0.4.6"
+__version__ = "0.4.7"
 version = __version__
 version_info = tuple(re.split(r"[-\.]", __version__))
 

uproot_methods/wrapjagged.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2019, IRIS-HEP
+# All rights reserved.
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+# 
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+# 
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+# 
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+# 
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+try:
+    from collections.abc import Iterable
+except ImportError:
+    from collections import Iterable
+
+from functools import wraps
+
+
+def normalize_arrays(cls, arrays):
+    length = None
+    for i in range(len(arrays)):
+        if isinstance(arrays[i], Iterable):
+            if length is None:
+                length = len(arrays[i])
+                break
+    if length is None:
+        raise TypeError("cannot construct an array if all arguments are scalar")
+
+    arrays = list(arrays)
+    jaggedtype = [cls.awkward.JaggedArray] * len(arrays)
+    starts, stops = None, None
+    for i in range(len(arrays)):
+        if starts is None and isinstance(arrays[i], cls.awkward.JaggedArray):
+            starts, stops = arrays[i].starts, arrays[i].stops
+
+        if isinstance(arrays[i], cls.awkward.JaggedArray):
+            jaggedtype[i] = type(arrays[i])
+
+        if not isinstance(arrays[i], Iterable):
+            arrays[i] = cls.awkward.numpy.full(length, arrays[i])
+
+        arrays[i] = cls.awkward.util.toarray(arrays[i], cls.awkward.numpy.float64)
+
+    if starts is None:
+        return arrays
+
+    for i in range(len(arrays)):
+        if not isinstance(arrays[i], cls.awkward.JaggedArray) or not (cls.awkward.numpy.array_equal(starts, arrays[i].starts) and cls.awkward.numpy.array_equal(stops, arrays[i].stops)):
+            content = cls.awkward.numpy.zeros(stops.max(), dtype=cls.awkward.numpy.float64)
+            arrays[i] = jaggedtype[i](starts, stops, content) + arrays[i]    # invoke jagged broadcasting to align arrays
+
+    return arrays
+
+def unwrap_jagged(cls, awkcls, arrays):
+    if not isinstance(arrays[0], cls.awkward.JaggedArray):
+        return lambda x: x, arrays
+
+    counts = arrays[0].counts.reshape(-1)
+    offsets = awkcls.counts2offsets(counts)
+    starts, stops = offsets[:-1], offsets[1:]
+    starts = starts.reshape(arrays[0].starts.shape[:-1] + (-1,))
+    stops = stops.reshape(arrays[0].stops.shape[:-1] + (-1,))
+    wrap, arrays = unwrap_jagged(cls, awkcls, [x.flatten() for x in arrays])
+    return lambda x: awkcls(starts, stops, wrap(x)), arrays
+
+def wrapjaggedmethod(awkcls):
+    def wrapjagged_decorator(func):
+        @wraps(func)
+        def func_wrapper(cls, *arrays):
+            wrap, arrays = unwrap_jagged(cls, awkcls, normalize_arrays(cls, arrays))
+            return wrap(func(cls, *arrays))
+        return func_wrapper
+    return wrapjagged_decorator