Merge pull request #11 from scikit-hep/issue-10

Issue 10

.travis.yml

@@ -33,14 +33,15 @@ addons:
       - python-setuptools
 
 script:
-  python setup.py test
+  python setup.py pytest
 
 install:
   - pip install --upgrade setuptools_scm
   - pip install $NUMPY
   - python -c 'import numpy; print(numpy.__version__)'
   - pip install awkward
   - python -c 'import awkward; print(awkward.__version__)'
+  - pip install pytest pytest-runner
 
 notifications:
   slack: scikit-hep:b6cgBXwccPoaCNLn5VKFJFVy

appveyor.yml

@@ -45,4 +45,5 @@ install:
 
 build_script:
   - "pip install %NUMPY%"
-  - "python setup.py test"
+  - "pip install pytest pytest-runner"
+  - "python setup.py pytest"

tests/test_issues.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2018, DIANA-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.
+
+import unittest
+
+import numpy
+
+import awkward
+import uproot_methods
+from uproot_methods import *
+
+class Test(unittest.TestCase):
+    def runTest(self):
+        pass
+
+    def test_issue10(self):
+        p4 = TLorentzVectorArray.from_ptetaphim(awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]))
+        assert p4.mass.tolist() == [[0.9999999999999999]]
+        assert p4[0].mass.tolist() == [0.9999999999999999]
+        assert p4[0][0].mass == 0.9999999999999999
+        assert type(p4.mass) is awkward.JaggedArray
+        assert type(p4.x) is awkward.JaggedArray
+
+        p3 = TVector3Array.from_cylindrical(awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]))
+        assert p3.rho.tolist() == [[1.0]]
+        assert p3[0].rho.tolist() == [1.0]
+        assert p3[0][0].rho == 1.0
+        assert type(p3.rho) is awkward.JaggedArray
+        assert type(p3.x) is awkward.JaggedArray
+
+        p2 = TVector2Array.from_polar(awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]))
+        assert p2.rho.tolist() == [[1.0]]
+        assert p2[0].rho.tolist() == [1.0]
+        assert p2[0][0].rho == 1.0
+        assert type(p2.rho) is awkward.JaggedArray
+        assert type(p2.x) is awkward.JaggedArray

uproot_methods/base.py

@@ -29,6 +29,34 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import awkward
+import awkward.util
+
+def _normalize_arrays(arrays):
+    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
+        arrays[i] = awkward.util.toarray(arrays[i], awkward.util.numpy.float64)
+
+    if starts is None:
+        return arrays
+
+    for i in range(len(arrays)):
+        if not isinstance(arrays[i], awkward.JaggedArray) or not (awkward.util.numpy.array_equal(starts, arrays[i].starts) and awkward.util.numpy.array_equal(stops, arrays[i].stops)):
+            content = awkward.util.numpy.zeros(stops.max(), dtype=awkward.util.numpy.float64)
+            arrays[i] = awkward.JaggedArray(starts, stops, content) + arrays[i]    # invoke jagged broadcasting to align arrays
+
+    return arrays
+
+def _unwrap_jagged(ArrayMethods, arrays):
+    if not isinstance(arrays[0], awkward.JaggedArray):
+        return lambda x: x, arrays
+    else:
+        JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray)
+        starts, stops = arrays[0].starts, arrays[0].stops
+        wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays])
+        return lambda x: JaggedArrayMethods(starts, stops, wrap(x)), arrays
 
 class ROOTMethods(awkward.Methods):
     _arraymethods = None

uproot_methods/classes/TLorentzVector.py

@@ -545,44 +545,49 @@ def origin(cls, shape, dtype=None):
 
     @classmethod
     def origin_like(cls, array):
-        return cls.origin(array.shape, array.dtype)
+        return array * 0.0
 
     @classmethod
     def from_p3(cls, p3, t):
-        out = cls.__new__(cls)
-        out._initObjectArray(awkward.Table())
-        out["fX"] = p3.x
-        out["fY"] = p3.y
-        out["fZ"] = p3.z
-        out["fE"] = t
-        return out
+        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):
+        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):
-        return 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):
-        return 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):
-        return 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):
-        return cls(pt * awkward.util.numpy.cos(phi),
-                   pt * awkward.util.numpy.sin(phi),
-                   pt * awkward.util.numpy.sinh(eta),
-                   energy)
+        wrap, (pt, eta, phi, energy) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, energy)))
+        return wrap(cls(pt * awkward.util.numpy.cos(phi),
+                        pt * awkward.util.numpy.sin(phi),
+                        pt * awkward.util.numpy.sinh(eta),
+                        energy))
 
     @classmethod
     def from_ptetaphim(cls, 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)
-        return cls.from_p3(p3, awkward.util.numpy.sqrt(x*x + y*y + z*z + mass*mass*awkward.util.numpy.sign(mass)))
+        return wrap(cls.from_p3(p3, awkward.util.numpy.sqrt(x*x + y*y + z*z + mass*mass*awkward.util.numpy.sign(mass))))
 
     @property
     def x(self):

uproot_methods/classes/TVector2.py

@@ -88,7 +88,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             return self.mag
 
         else:
-            return awkward.ObjectArray.__array_ufunc__(self, ufunc, method, *inputs, **kwargs)
+            return super(ArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs)
 
 class Methods(Common, uproot_methods.common.TVector.Methods, uproot_methods.base.ROOTMethods):
     _arraymethods = ArrayMethods
@@ -143,12 +143,18 @@ def origin(cls, shape, dtype=None):
 
     @classmethod
     def origin_like(cls, array):
-        return cls.origin(array.shape, array.dtype)
+        return array * 0.0
 
     @classmethod
-    def from_circular(cls, rho, phi):
-        return cls(rho * awkward.util.numpy.cos(phi),
-                   rho * awkward.util.numpy.sin(phi))
+    def from_cartesian(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):
+        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):
@@ -176,7 +182,7 @@ def origin(cls):
         return cls(0.0, 0.0)
 
     @classmethod
-    def from_circular(cls, rho, phi):
+    def from_polar(cls, rho, phi):
         return cls(rho * math.cos(phi),
                    rho * math.sin(phi))
 

uproot_methods/classes/TVector3.py

@@ -180,7 +180,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
             return self.mag
 
         else:
-            return awkward.ObjectArray.__array_ufunc__(self, ufunc, method, *inputs, **kwargs)
+            return super(ArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs)
 
 class Methods(Common, uproot_methods.common.TVector.Methods, uproot_methods.base.ROOTMethods):
     _arraymethods = ArrayMethods
@@ -257,19 +257,26 @@ def origin(cls, shape, dtype=None):
 
     @classmethod
     def origin_like(cls, array):
-        return cls.origin(array.shape, array.dtype)
+        return array * 0.0
+
+    @classmethod
+    def from_cartesian(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):
-        return 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, r, rho, phi, z):
-        return cls(rho * awkward.util.numpy.cos(phi),
-                   rho * awkward.util.numpy.sin(phi),
-                   z)
+    def from_cylindrical(cls, rho, 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.7"
+__version__ = "0.2.8"
 version = __version__
 version_info = tuple(re.split(r"[-\.]", __version__))