Merge pull request #14 from scikit-hep/sync-awkward-070

This is a pre-release of 0.3.0

uproot_methods/base.py

@@ -28,15 +28,33 @@
 # 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
 import awkward.util
 
 def _normalize_arrays(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)
     starts, stops = None, None
     for i in range(len(arrays)):
         if starts is None and isinstance(arrays[i], awkward.JaggedArray):
             starts, stops = arrays[i].starts, arrays[i].stops
+
+        if not isinstance(arrays[i], Iterable):
+            arrays[i] = awkward.util.numpy.full(length, arrays[i])
+
         arrays[i] = awkward.util.toarray(arrays[i], awkward.util.numpy.float64)
 
     if starts is None:
@@ -51,19 +69,23 @@ def _normalize_arrays(arrays):
 
 def _unwrap_jagged(ArrayMethods, arrays):
     if not isinstance(arrays[0], awkward.JaggedArray):
-        return lambda x: x, lambda x: x, arrays
+        return lambda x: x, arrays
     else:
-        JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray)
+        if ArrayMethods is None:
+            cls = awkward.JaggedArray
+        else:
+            cls = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray)
         starts, stops = arrays[0].starts, arrays[0].stops
-        wrapmethods, wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays])
-        return lambda x: JaggedArrayMethods(starts, stops, wrapmethods(x)), lambda x: awkward.JaggedArray(starts, stops, wrap(x)), arrays
+        wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays])
+        return lambda x: cls(starts, stops, wrap(x)), arrays
 
 def memo(function):
     memoname = "_memo_" + function.__name__
-    def memofunction(self):
-        if not hasattr(self, memoname):
-            setattr(self, memoname, function(self))
-        return getattr(self, memoname)
+    def memofunction(array):
+        wrap, (array,) = _unwrap_jagged(None, (array,))
+        if not hasattr(array, memoname):
+            setattr(array, memoname, function(array))
+        return wrap(getattr(array, memoname))
     return memofunction
 
 class ROOTMethods(awkward.Methods):

uproot_methods/classes/TLorentzVector.py

@@ -75,11 +75,6 @@ def perp(self):
     def pt2(self):
         return self.p3.rho2
 
-    @property
-    @uproot_methods.base.memo
-    def pt(self):
-        return self.p3.rho
-
     @property
     def Et(self):
         return self.energy * self.pt / self.p
@@ -92,19 +87,9 @@ def mag2(self):
     def mass2(self):
         return self.mag2
 
-    @property
-    @uproot_methods.base.memo
-    def mass(self):
-        return self.mag
-
     @property
     def mt2(self):
         return self.energy**2 - self.z**2
-
-    @property
-    @uproot_methods.base.memo
-    def phi(self):
-        return self.p3.phi
 
     @property
     def theta(self):
@@ -193,6 +178,26 @@ def z(self):
     def t(self):
         return self["fE"]
 
+    @property
+    @uproot_methods.base.memo
+    def pt(self):
+        return awkward.util.numpy.sqrt(self.pt2)
+
+    @property
+    @uproot_methods.base.memo
+    def eta(self):
+        return awkward.util.numpy.arcsinh(self.z / awkward.util.numpy.sqrt(self.x**2 + self.y**2))
+
+    @property
+    @uproot_methods.base.memo
+    def phi(self):
+        return self.p3.phi
+
+    @property
+    @uproot_methods.base.memo
+    def mass(self):
+        return awkward.util.numpy.sqrt(self.mag2)
+
     @property
     def mag(self):
         return awkward.util.numpy.sqrt(self.mag2)
@@ -203,11 +208,6 @@ def mt(self):
         sign = awkward.util.numpy.sign(mt2)
         return awkward.util.numpy.sqrt(awkward.util.numpy.absolute(mt2)) * sign
 
-    @property
-    @uproot_methods.base.memo
-    def eta(self):
-        return awkward.util.numpy.arcsinh(self.z / awkward.util.numpy.sqrt(self.x**2 + self.y**2))
-
     @property
     def rapidity(self):
         return awkward.util.numpy.log((self.t + self.z) / (self.t - self.z)) / 2.0
@@ -370,6 +370,22 @@ def _vector(self, operator, vector, reverse=False):
     def _unary(self, operator):
         return TLorentzVector(operator(self.x), operator(self.y), operator(self.z), operator(self.t))
 
+    @property
+    def pt(self):
+        return math.sqrt(self.pt2)
+
+    @property
+    def eta(self):
+        return math.asinh(self.z / math.sqrt(self.x**2 + self.y**2))
+
+    @property
+    def phi(self):
+        return self.p3.phi
+
+    @property
+    def mass(self):
+        return math.sqrt(self.mag2)
+
     @property
     def mag(self):
         return math.sqrt(self.mag2)
@@ -382,10 +398,6 @@ def mt(self):
         else:
             return -math.sqrt(out)
 
-    @property
-    def eta(self):
-        return math.asinh(self.z / math.sqrt(self.x**2 + self.y**2))
-
     @property
     def rapidity(self):
         return math.log((self.t + self.z) / (self.t - self.z)) / 2.0
@@ -559,54 +571,54 @@ def origin_like(cls, array):
 
     @classmethod
     def from_p3(cls, p3, t):
-        wrapmethods, wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((p3.x, p3.y, p3.z, t)))
-        return wrapmethods(cls(x, y, z, t))
+        wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._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):
-        wrapmethods, wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, t)))
-        return wrapmethods(cls(x, y, z, t))
+        wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, t)))
+        return wrap(cls(x, y, z, t))
 
     @classmethod
     def from_spherical(cls, r, theta, phi, t):
-        wrapmethods, wrap, (r, theta, phi, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((r, theta, phi, t)))
-        return wrapmethods(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_spherical(r, theta, phi), t))
+        wrap, (r, theta, phi, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._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):
-        wrapmethods, wrap, (rho, phi, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z, t)))
-        return wrapmethods(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_cylindrical(rho, phi, z), t))
+        wrap, (rho, phi, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._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):
-        wrapmethods, wrap, (x, y, z, m) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, m)))
-        return wrapmethods(cls(x, y, z, awkward.util.numpy.sqrt(x*x + y*y + z*z + m*m*awkward.util.numpy.sign(m))))
+        wrap, (x, y, z, m) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, m)))
+        return wrap(cls(x, y, z, awkward.util.numpy.sqrt(x*x + y*y + z*z + m*m*awkward.util.numpy.sign(m))))
 
     @classmethod
     def from_ptetaphi(cls, pt, eta, phi, energy):
-        wrapmethods, wrap, (pt, eta, phi, energy) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, energy)))
-        out = wrapmethods(cls(pt * awkward.util.numpy.cos(phi),
-                              pt * awkward.util.numpy.sin(phi),
-                              pt * awkward.util.numpy.sinh(eta),
-                              energy))
-        out._memo_pt = wrap(pt)
-        out._memo_eta = wrap(eta)
-        out._memo_phi = wrap(phi)
-        return out
+        wrap, (pt, eta, phi, energy) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, energy)))
+        out = cls(pt * awkward.util.numpy.cos(phi),
+                  pt * awkward.util.numpy.sin(phi),
+                  pt * awkward.util.numpy.sinh(eta),
+                  energy)
+        out._memo_pt = pt
+        out._memo_eta = eta
+        out._memo_phi = phi
+        return wrap(out)
 
     @classmethod
     def from_ptetaphim(cls, pt, eta, phi, mass):
-        wrapmethods, wrap, (pt, eta, phi, mass) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, mass)))
+        wrap, (pt, eta, phi, mass) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, mass)))
         x = pt * awkward.util.numpy.cos(phi)
         y = pt * awkward.util.numpy.sin(phi)
         z = pt * awkward.util.numpy.sinh(eta)
         p3 = uproot_methods.classes.TVector3.TVector3Array(x, y, z)
-        out = wrapmethods(cls.from_p3(p3, awkward.util.numpy.sqrt(x*x + y*y + z*z + mass*mass*awkward.util.numpy.sign(mass))))
-        out._memo_pt = wrap(pt)
-        out._memo_eta = wrap(eta)
-        out._memo_phi = wrap(phi)
-        out._memo_mass = wrap(mass)
-        return out
+        out = cls.from_p3(p3, awkward.util.numpy.sqrt(x*x + y*y + z*z + mass*mass*awkward.util.numpy.sign(mass)))
+        out._memo_pt = pt
+        out._memo_eta = eta
+        out._memo_phi = phi
+        out._memo_mass = mass
+        return wrap(out)
 
     @property
     def x(self):

uproot_methods/classes/TVector2.py

@@ -153,14 +153,14 @@ def origin_like(cls, array):
 
     @classmethod
     def from_cartesian(cls, x, y):
-        wrapmethods, wrap, (x, y) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y)))
-        return wrapmethods(cls(x, y))
+        wrap, (x, y) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y)))
+        return wrap(cls(x, y))
 
     @classmethod
     def from_polar(cls, rho, phi):
-        wrapmethods, wrap, (rho, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi)))
-        return wrapmethods(cls(rho * awkward.util.numpy.cos(phi),
-                               rho * awkward.util.numpy.sin(phi)))
+        wrap, (rho, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi)))
+        return wrap(cls(rho * awkward.util.numpy.cos(phi),
+                        rho * awkward.util.numpy.sin(phi)))
 
     @property
     def x(self):

uproot_methods/classes/TVector3.py

@@ -267,22 +267,22 @@ def origin_like(cls, array):
 
     @classmethod
     def from_cartesian(cls, x, y, z):
-        wrapmethods, wrap, (x, y, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z)))
-        return wrapmethods(cls(x, y, z))
+        wrap, (x, y, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z)))
+        return wrap(cls(x, y, z))
 
     @classmethod
     def from_spherical(cls, r, theta, phi):
-        wrapmethods, wrap, (r, theta, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((r, theta, phi)))
-        return wrapmethods(cls(r * awkward.util.numpy.sin(theta) * awkward.util.numpy.cos(phi),
-                               r * awkward.util.numpy.sin(theta) * awkward.util.numpy.sin(phi),
-                               r * awkward.util.numpy.cos(theta)))
+        wrap, (r, theta, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((r, theta, phi)))
+        return wrap(cls(r * awkward.util.numpy.sin(theta) * awkward.util.numpy.cos(phi),
+                        r * awkward.util.numpy.sin(theta) * awkward.util.numpy.sin(phi),
+                        r * awkward.util.numpy.cos(theta)))
 
     @classmethod
     def from_cylindrical(cls, rho, phi, z):
-        wrapmethods, wrap, (rho, phi, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z)))
-        return wrapmethods(cls(rho * awkward.util.numpy.cos(phi),
-                               rho * awkward.util.numpy.sin(phi),
-                               z))
+        wrap, (rho, phi, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z)))
+        return wrap(cls(rho * awkward.util.numpy.cos(phi),
+                        rho * awkward.util.numpy.sin(phi),
+                        z))
 
     @property
     def x(self):

uproot_methods/version.py

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