diff --git a/.travis.yml b/.travis.yml index 6bb3ad6..445cc71 100644 --- a/.travis.yml +++ b/.travis.yml @@ -33,7 +33,7 @@ addons: - python-setuptools script: - python setup.py test + python setup.py pytest install: - pip install --upgrade setuptools_scm @@ -41,6 +41,7 @@ install: - 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 diff --git a/appveyor.yml b/appveyor.yml index 1c19ada..093ea2a 100644 --- a/appveyor.yml +++ b/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" diff --git a/tests/test_issues.py b/tests/test_issues.py new file mode 100644 index 0000000..58b460e --- /dev/null +++ b/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 diff --git a/tests/test_vector.py b/tests/test_vector.py index ea33306..af57bb7 100644 --- a/tests/test_vector.py +++ b/tests/test_vector.py @@ -42,35 +42,35 @@ def runTest(self): def test_vector2(self): a = TVector2(4.4, 5.5) - self.assertEqual(a.dot(a), 49.61) - self.assertEqual(a + TVector2(1000, 2000), TVector2(1004.4, 2005.5)) - self.assertEqual(a - TVector2(1000, 2000), TVector2(-995.6, -1994.5)) - self.assertEqual(TVector2(1000, 2000) - a, TVector2(995.6, 1994.5)) - self.assertEqual(a * 1000, TVector2(4400, 5500)) - self.assertEqual(1000 * a, TVector2(4400, 5500)) - self.assertEqual(a / 1000, TVector2(0.0044, 0.0055)) - self.assertEqual(1000 / a, TVector2(227.27272727272725, 181.8181818181818)) - self.assertEqual(a**2, 49.61) - self.assertEqual(a**1, 7.043436661176133) - self.assertEqual(abs(a), 7.043436661176133) - self.assertEqual(-a, TVector2(-4.4, -5.5)) - self.assertEqual(+a, TVector2(4.4, 5.5)) + assert a.dot(a) == 49.61 + assert a + TVector2(1000, 2000) == TVector2(1004.4, 2005.5) + assert a - TVector2(1000, 2000) == TVector2(-995.6, -1994.5) + assert TVector2(1000, 2000) - a == TVector2(995.6, 1994.5) + assert a * 1000 == TVector2(4400, 5500) + assert 1000 * a == TVector2(4400, 5500) + assert a / 1000 == TVector2(0.0044, 0.0055) + assert 1000 / a == TVector2(227.27272727272725, 181.8181818181818) + assert a**2 == 49.61 + assert a**1 == 7.043436661176133 + assert abs(a) == 7.043436661176133 + assert -a == TVector2(-4.4, -5.5) + assert +a == TVector2(4.4, 5.5) a += TVector2(100, 200) - self.assertEqual(a, TVector2(104.4, 205.5)) + assert a == TVector2(104.4, 205.5) a *= 10 - self.assertEqual(a, TVector2(1044, 2055)) + assert a == TVector2(1044, 2055) def test_vector2_array(self): a = TVector2Array(numpy.zeros(10), numpy.arange(10)) - self.assertEqual(a.tolist(), [TVector2(0, 0), TVector2(0, 1), TVector2(0, 2), TVector2(0, 3), TVector2(0, 4), TVector2(0, 5), TVector2(0, 6), TVector2(0, 7), TVector2(0, 8), TVector2(0, 9)]) - self.assertEqual(a[5], TVector2(0, 5)) - self.assertEqual(a.y.tolist(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) - self.assertEqual(a.mag2.tolist(), [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0]) - self.assertEqual((a + TVector2(1000, 2000))[5], TVector2(1000, 2005)) - self.assertEqual((a + TVector2(1000, 2000) == TVector2Array(numpy.full(10, 1000), numpy.arange(2000, 2010))).tolist(), [True, True, True, True, True, True, True, True, True, True]) - self.assertEqual((a**2).tolist(), [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0]) - self.assertEqual((a**3).tolist(), [0.0, 1.0, 8.0, 27.0, 64.0, 125.0, 216.0, 343.0, 512.0, 729.0]) + assert a.tolist() == [TVector2(0, 0), TVector2(0, 1), TVector2(0, 2), TVector2(0, 3), TVector2(0, 4), TVector2(0, 5), TVector2(0, 6), TVector2(0, 7), TVector2(0, 8), TVector2(0, 9)] + assert a[5] == TVector2(0, 5) + assert a.y.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + assert a.mag2.tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] + assert (a + TVector2(1000, 2000))[5] == TVector2(1000, 2005) + assert (a + TVector2(1000, 2000) == TVector2Array(numpy.full(10, 1000), numpy.arange(2000, 2010))).tolist() == [True, True, True, True, True, True, True, True, True, True] + assert (a**2).tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] + assert (a**3).tolist() == [0.0, 1.0, 8.0, 27.0, 64.0, 125.0, 216.0, 343.0, 512.0, 729.0] def test_vector2_jagged(self): TVector2Jagged = type("TVector2Jagged", (awkward.JaggedArray, uproot_methods.classes.TVector2.ArrayMethods), {}) @@ -78,44 +78,44 @@ def test_vector2_jagged(self): a._generator = uproot_methods.classes.TVector2.TVector2 a._args = () a._kwargs = {} - self.assertEqual(a.tolist(), [[TVector2(0, 0), TVector2(0, 1), TVector2(0, 2)], [], [TVector2(0, 3), TVector2(0, 4)], [TVector2(0, 5), TVector2(0, 6), TVector2(0, 7), TVector2(0, 8), TVector2(0, 9)]]) - self.assertEqual(a.x.tolist(), [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]]) - self.assertEqual(a.y.tolist(), [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]]) - self.assertEqual((a + TVector2(1000, 2000)).tolist(), [[TVector2(1000, 2000), TVector2(1000, 2001), TVector2(1000, 2002)], [], [TVector2(1000, 2003), TVector2(1000, 2004)], [TVector2(1000, 2005), TVector2(1000, 2006), TVector2(1000, 2007), TVector2(1000, 2008), TVector2(1000, 2009)]]) - self.assertEqual((a + TVector2Array(numpy.full(4, 1000), numpy.arange(1000, 5000, 1000))).tolist(), [[TVector2(1000, 1000), TVector2(1000, 1001), TVector2(1000, 1002)], [], [TVector2(1000, 3003), TVector2(1000, 3004)], [TVector2(1000, 4005), TVector2(1000, 4006), TVector2(1000, 4007), TVector2(1000, 4008), TVector2(1000, 4009)]]) + assert a.tolist() == [[TVector2(0, 0), TVector2(0, 1), TVector2(0, 2)], [], [TVector2(0, 3), TVector2(0, 4)], [TVector2(0, 5), TVector2(0, 6), TVector2(0, 7), TVector2(0, 8), TVector2(0, 9)]] + assert a.x.tolist() == [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]] + assert a.y.tolist() == [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]] + assert (a + TVector2(1000, 2000)).tolist() == [[TVector2(1000, 2000), TVector2(1000, 2001), TVector2(1000, 2002)], [], [TVector2(1000, 2003), TVector2(1000, 2004)], [TVector2(1000, 2005), TVector2(1000, 2006), TVector2(1000, 2007), TVector2(1000, 2008), TVector2(1000, 2009)]] + assert (a + TVector2Array(numpy.full(4, 1000), numpy.arange(1000, 5000, 1000))).tolist() == [[TVector2(1000, 1000), TVector2(1000, 1001), TVector2(1000, 1002)], [], [TVector2(1000, 3003), TVector2(1000, 3004)], [TVector2(1000, 4005), TVector2(1000, 4006), TVector2(1000, 4007), TVector2(1000, 4008), TVector2(1000, 4009)]] def test_vector3(self): a = TVector3(4.4, 5.5, 0) - self.assertEqual(a.dot(a), 49.61) - self.assertEqual(a.cross(a), TVector3(0,0,0)) - self.assertEqual(a + TVector3(1000, 2000, 0), TVector3(1004.4, 2005.5, 0)) - self.assertEqual(a - TVector3(1000, 2000, 0), TVector3(-995.6, -1994.5, 0)) - self.assertEqual(TVector3(1000, 2000, 0) - a, TVector3(995.6, 1994.5, 0)) - self.assertEqual(a * 1000, TVector3(4400, 5500, 0)) - self.assertEqual(1000 * a, TVector3(4400, 5500, 0)) - self.assertEqual(a / 1000, TVector3(0.0044, 0.0055, 0)) - self.assertEqual(1000 / (a + TVector3(0, 0, 1)), TVector3(227.27272727272725, 181.8181818181818, 1000)) - self.assertEqual(a**2, 49.61) - self.assertEqual(a**1, 7.043436661176133) - self.assertEqual(abs(a), 7.043436661176133) - self.assertEqual(-a, TVector3(-4.4, -5.5, 0)) - self.assertEqual(+a, TVector3(4.4, 5.5, 0)) + assert a.dot(a) == 49.61 + assert a.cross(a) == TVector3(0,0,0) + assert a + TVector3(1000, 2000, 0) == TVector3(1004.4, 2005.5, 0) + assert a - TVector3(1000, 2000, 0) == TVector3(-995.6, -1994.5, 0) + assert TVector3(1000, 2000, 0) - a == TVector3(995.6, 1994.5, 0) + assert a * 1000 == TVector3(4400, 5500, 0) + assert 1000 * a == TVector3(4400, 5500, 0) + assert a / 1000 == TVector3(0.0044, 0.0055, 0) + assert 1000 / (a + TVector3(0, 0, 1)) == TVector3(227.27272727272725, 181.8181818181818, 1000) + assert a**2 == 49.61 + assert a**1 == 7.043436661176133 + assert abs(a) == 7.043436661176133 + assert -a == TVector3(-4.4, -5.5, 0) + assert +a == TVector3(4.4, 5.5, 0) a += TVector3(100, 200, 0) - self.assertEqual(a, TVector3(104.4, 205.5, 0)) + assert a == TVector3(104.4, 205.5, 0) a *= 10 - self.assertEqual(a, TVector3(1044, 2055, 0)) + assert a == TVector3(1044, 2055, 0) def test_vector3_array(self): a = TVector3Array(numpy.zeros(10), numpy.arange(10), numpy.zeros(10)) - self.assertEqual(a.tolist(), [TVector3(0, 0, 0), TVector3(0, 1, 0), TVector3(0, 2, 0), TVector3(0, 3, 0), TVector3(0, 4, 0), TVector3(0, 5, 0), TVector3(0, 6, 0), TVector3(0, 7, 0), TVector3(0, 8, 0), TVector3(0, 9, 0)]) - self.assertEqual(a[5], TVector3(0, 5, 0)) - self.assertEqual(a.y.tolist(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) - self.assertEqual(a.mag2.tolist(), [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0]) - self.assertEqual((a + TVector3(1000, 2000, 0))[5], TVector3(1000, 2005, 0)) - self.assertEqual((a + TVector3(1000, 2000, 0) == TVector3Array(numpy.full(10, 1000), numpy.arange(2000, 2010), numpy.zeros(10))).tolist(), [True, True, True, True, True, True, True, True, True, True]) - self.assertEqual((a**2).tolist(), [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0]) - self.assertEqual((a**3).tolist(), [0.0, 1.0, 8.0, 27.0, 64.0, 125.0, 216.0, 343.0, 512.0, 729.0]) + assert a.tolist() == [TVector3(0, 0, 0), TVector3(0, 1, 0), TVector3(0, 2, 0), TVector3(0, 3, 0), TVector3(0, 4, 0), TVector3(0, 5, 0), TVector3(0, 6, 0), TVector3(0, 7, 0), TVector3(0, 8, 0), TVector3(0, 9, 0)] + assert a[5] == TVector3(0, 5, 0) + assert a.y.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + assert a.mag2.tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] + assert (a + TVector3(1000, 2000, 0))[5] == TVector3(1000, 2005, 0) + assert (a + TVector3(1000, 2000, 0) == TVector3Array(numpy.full(10, 1000), numpy.arange(2000, 2010), numpy.zeros(10))).tolist() == [True, True, True, True, True, True, True, True, True, True] + assert (a**2).tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] + assert (a**3).tolist() == [0.0, 1.0, 8.0, 27.0, 64.0, 125.0, 216.0, 343.0, 512.0, 729.0] def test_vector3_jagged(self): TVector3Jagged = type("TVector3Jagged", (awkward.JaggedArray, uproot_methods.classes.TVector3.ArrayMethods), {}) @@ -123,42 +123,42 @@ def test_vector3_jagged(self): a._generator = uproot_methods.classes.TVector3.TVector3 a._args = () a._kwargs = {} - self.assertEqual(a.tolist(), [[TVector3(0, 0, 0), TVector3(0, 1, 0), TVector3(0, 2, 0)], [], [TVector3(0, 3, 0), TVector3(0, 4, 0)], [TVector3(0, 5, 0), TVector3(0, 6, 0), TVector3(0, 7, 0), TVector3(0, 8, 0), TVector3(0, 9, 0)]]) - self.assertEqual(a.x.tolist(), [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]]) - self.assertEqual(a.y.tolist(), [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]]) - self.assertEqual((a + TVector3(1000, 2000, 0)).tolist(), [[TVector3(1000, 2000, 0), TVector3(1000, 2001, 0), TVector3(1000, 2002, 0)], [], [TVector3(1000, 2003, 0), TVector3(1000, 2004, 0)], [TVector3(1000, 2005, 0), TVector3(1000, 2006, 0), TVector3(1000, 2007, 0), TVector3(1000, 2008, 0), TVector3(1000, 2009, 0)]]) - self.assertEqual((a + TVector3Array(numpy.full(4, 1000), numpy.arange(1000, 5000, 1000), numpy.zeros(4))).tolist(), [[TVector3(1000, 1000, 0), TVector3(1000, 1001, 0), TVector3(1000, 1002, 0)], [], [TVector3(1000, 3003, 0), TVector3(1000, 3004, 0)], [TVector3(1000, 4005, 0), TVector3(1000, 4006, 0), TVector3(1000, 4007, 0), TVector3(1000, 4008, 0), TVector3(1000, 4009, 0)]]) + assert a.tolist() == [[TVector3(0, 0, 0), TVector3(0, 1, 0), TVector3(0, 2, 0)], [], [TVector3(0, 3, 0), TVector3(0, 4, 0)], [TVector3(0, 5, 0), TVector3(0, 6, 0), TVector3(0, 7, 0), TVector3(0, 8, 0), TVector3(0, 9, 0)]] + assert a.x.tolist() == [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]] + assert a.y.tolist() == [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]] + assert (a + TVector3(1000, 2000, 0)).tolist() == [[TVector3(1000, 2000, 0), TVector3(1000, 2001, 0), TVector3(1000, 2002, 0)], [], [TVector3(1000, 2003, 0), TVector3(1000, 2004, 0)], [TVector3(1000, 2005, 0), TVector3(1000, 2006, 0), TVector3(1000, 2007, 0), TVector3(1000, 2008, 0), TVector3(1000, 2009, 0)]] + assert (a + TVector3Array(numpy.full(4, 1000), numpy.arange(1000, 5000, 1000), numpy.zeros(4))).tolist() == [[TVector3(1000, 1000, 0), TVector3(1000, 1001, 0), TVector3(1000, 1002, 0)], [], [TVector3(1000, 3003, 0), TVector3(1000, 3004, 0)], [TVector3(1000, 4005, 0), TVector3(1000, 4006, 0), TVector3(1000, 4007, 0), TVector3(1000, 4008, 0), TVector3(1000, 4009, 0)]] def test_lorentzvector(self): a = TLorentzVector(4.4, 5.5, 0, 0) - self.assertEqual(a.dot(a), -49.61) - self.assertEqual(a + TLorentzVector(1000, 2000, 0, 0), TLorentzVector(1004.4, 2005.5, 0, 0)) - self.assertEqual(a - TLorentzVector(1000, 2000, 0, 0), TLorentzVector(-995.6, -1994.5, 0, 0)) - self.assertEqual(TLorentzVector(1000, 2000, 0, 0) - a, TLorentzVector(995.6, 1994.5, 0, 0)) - self.assertEqual(a * 1000, TLorentzVector(4400, 5500, 0, 0)) - self.assertEqual(1000 * a, TLorentzVector(4400, 5500, 0, 0)) - self.assertEqual(a / 1000, TLorentzVector(0.0044, 0.0055, 0, 0)) - self.assertEqual(1000 / (a + TLorentzVector(0, 0, 1, 1)), TLorentzVector(227.27272727272725, 181.8181818181818, 1000, 1000)) - self.assertEqual(a**2, -49.61) - self.assertEqual((a + TLorentzVector(0, 0, 0, 10))**1, 7.098591409568521) - self.assertEqual(abs(a + TLorentzVector(0, 0, 0, 10)), 7.098591409568521) - self.assertEqual(-a, TLorentzVector(-4.4, -5.5, 0, 0)) - self.assertEqual(+a, TLorentzVector(4.4, 5.5, 0, 0)) + assert a.dot(a) == -49.61 + assert a + TLorentzVector(1000, 2000, 0, 0) == TLorentzVector(1004.4, 2005.5, 0, 0) + assert a - TLorentzVector(1000, 2000, 0, 0) == TLorentzVector(-995.6, -1994.5, 0, 0) + assert TLorentzVector(1000, 2000, 0, 0) - a == TLorentzVector(995.6, 1994.5, 0, 0) + assert a * 1000 == TLorentzVector(4400, 5500, 0, 0) + assert 1000 * a == TLorentzVector(4400, 5500, 0, 0) + assert a / 1000 == TLorentzVector(0.0044, 0.0055, 0, 0) + assert 1000 / (a + TLorentzVector(0, 0, 1, 1)) == TLorentzVector(227.27272727272725, 181.8181818181818, 1000, 1000) + assert a**2, -49.61 + assert (a + TLorentzVector(0, 0, 0, 10))**1 == 7.098591409568521 + assert abs(a + TLorentzVector(0, 0, 0, 10)) == 7.098591409568521 + assert -a == TLorentzVector(-4.4, -5.5, 0, 0) + assert +a == TLorentzVector(4.4, 5.5, 0, 0) a += TLorentzVector(100, 200, 0, 0) - self.assertEqual(a, TLorentzVector(104.4, 205.5, 0, 0)) + assert a == TLorentzVector(104.4, 205.5, 0, 0) a *= 10 - self.assertEqual(a, TLorentzVector(1044, 2055, 0, 0)) + assert a == TLorentzVector(1044, 2055, 0, 0) def test_lorentzvector_array(self): a = TLorentzVectorArray(numpy.zeros(10), numpy.arange(10), numpy.zeros(10), numpy.zeros(10)) - self.assertEqual(a.tolist(), [TLorentzVector(0, 0, 0, 0), TLorentzVector(0, 1, 0, 0), TLorentzVector(0, 2, 0, 0), TLorentzVector(0, 3, 0, 0), TLorentzVector(0, 4, 0, 0), TLorentzVector(0, 5, 0, 0), TLorentzVector(0, 6, 0, 0), TLorentzVector(0, 7, 0, 0), TLorentzVector(0, 8, 0, 0), TLorentzVector(0, 9, 0, 0)]) - self.assertEqual(a[5], TLorentzVector(0, 5, 0, 0)) - self.assertEqual(a.y.tolist(), [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) - self.assertEqual(a.mag2.tolist(), [0.0, -1.0, -4.0, -9.0, -16.0, -25.0, -36.0, -49.0, -64.0, -81.0]) - self.assertEqual((a + TLorentzVector(1000, 2000, 0, 0))[5], TLorentzVector(1000, 2005, 0, 0)) - self.assertEqual((a + TLorentzVector(1000, 2000, 0, 0) == TLorentzVectorArray(numpy.full(10, 1000), numpy.arange(2000, 2010), numpy.zeros(10), numpy.zeros(10))).tolist(), [True, True, True, True, True, True, True, True, True, True]) - self.assertEqual((a**2).tolist(), [0.0, -1.0, -4.0, -9.0, -16.0, -25.0, -36.0, -49.0, -64.0, -81.0]) + assert a.tolist() == [TLorentzVector(0, 0, 0, 0), TLorentzVector(0, 1, 0, 0), TLorentzVector(0, 2, 0, 0), TLorentzVector(0, 3, 0, 0), TLorentzVector(0, 4, 0, 0), TLorentzVector(0, 5, 0, 0), TLorentzVector(0, 6, 0, 0), TLorentzVector(0, 7, 0, 0), TLorentzVector(0, 8, 0, 0), TLorentzVector(0, 9, 0, 0)] + assert a[5] == TLorentzVector(0, 5, 0, 0) + assert a.y.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + assert a.mag2.tolist() == [0.0, -1.0, -4.0, -9.0, -16.0, -25.0, -36.0, -49.0, -64.0, -81.0] + assert (a + TLorentzVector(1000, 2000, 0, 0))[5] == TLorentzVector(1000, 2005, 0, 0) + assert (a + TLorentzVector(1000, 2000, 0, 0) == TLorentzVectorArray(numpy.full(10, 1000), numpy.arange(2000, 2010), numpy.zeros(10), numpy.zeros(10))).tolist() == [True, True, True, True, True, True, True, True, True, True] + assert (a**2).tolist() == [0.0, -1.0, -4.0, -9.0, -16.0, -25.0, -36.0, -49.0, -64.0, -81.0] def test_lorentzvector_jagged(self): TLorentzVectorJagged = type("TLorentzVectorJagged", (awkward.JaggedArray, uproot_methods.classes.TLorentzVector.ArrayMethods), {}) @@ -166,8 +166,8 @@ def test_lorentzvector_jagged(self): a._generator = uproot_methods.classes.TLorentzVector.TLorentzVector a._args = () a._kwargs = {} - self.assertEqual(a.tolist(), [[TLorentzVector(0, 0, 0, 0), TLorentzVector(0, 1, 0, 0), TLorentzVector(0, 2, 0, 0)], [], [TLorentzVector(0, 3, 0, 0), TLorentzVector(0, 4, 0, 0)], [TLorentzVector(0, 5, 0, 0), TLorentzVector(0, 6, 0, 0), TLorentzVector(0, 7, 0, 0), TLorentzVector(0, 8, 0, 0), TLorentzVector(0, 9, 0, 0)]]) - self.assertEqual(a.x.tolist(), [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]]) - self.assertEqual(a.y.tolist(), [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]]) - self.assertEqual((a + TLorentzVector(1000, 2000, 0, 0)).tolist(), [[TLorentzVector(1000, 2000, 0, 0), TLorentzVector(1000, 2001, 0, 0), TLorentzVector(1000, 2002, 0, 0)], [], [TLorentzVector(1000, 2003, 0, 0), TLorentzVector(1000, 2004, 0, 0)], [TLorentzVector(1000, 2005, 0, 0), TLorentzVector(1000, 2006, 0, 0), TLorentzVector(1000, 2007, 0, 0), TLorentzVector(1000, 2008, 0, 0), TLorentzVector(1000, 2009, 0, 0)]]) - self.assertEqual((a + TLorentzVectorArray(numpy.full(4, 1000), numpy.arange(1000, 5000, 1000), numpy.zeros(4), numpy.zeros(4))).tolist(), [[TLorentzVector(1000, 1000, 0, 0), TLorentzVector(1000, 1001, 0, 0), TLorentzVector(1000, 1002, 0, 0)], [], [TLorentzVector(1000, 3003, 0, 0), TLorentzVector(1000, 3004, 0, 0)], [TLorentzVector(1000, 4005, 0, 0), TLorentzVector(1000, 4006, 0, 0), TLorentzVector(1000, 4007, 0, 0), TLorentzVector(1000, 4008, 0, 0), TLorentzVector(1000, 4009, 0, 0)]]) + assert a.tolist() == [[TLorentzVector(0, 0, 0, 0), TLorentzVector(0, 1, 0, 0), TLorentzVector(0, 2, 0, 0)], [], [TLorentzVector(0, 3, 0, 0), TLorentzVector(0, 4, 0, 0)], [TLorentzVector(0, 5, 0, 0), TLorentzVector(0, 6, 0, 0), TLorentzVector(0, 7, 0, 0), TLorentzVector(0, 8, 0, 0), TLorentzVector(0, 9, 0, 0)]] + assert a.x.tolist() == [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]] + assert a.y.tolist() == [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]] + assert (a + TLorentzVector(1000, 2000, 0, 0)).tolist() == [[TLorentzVector(1000, 2000, 0, 0), TLorentzVector(1000, 2001, 0, 0), TLorentzVector(1000, 2002, 0, 0)], [], [TLorentzVector(1000, 2003, 0, 0), TLorentzVector(1000, 2004, 0, 0)], [TLorentzVector(1000, 2005, 0, 0), TLorentzVector(1000, 2006, 0, 0), TLorentzVector(1000, 2007, 0, 0), TLorentzVector(1000, 2008, 0, 0), TLorentzVector(1000, 2009, 0, 0)]] + assert (a + TLorentzVectorArray(numpy.full(4, 1000), numpy.arange(1000, 5000, 1000), numpy.zeros(4), numpy.zeros(4))).tolist() == [[TLorentzVector(1000, 1000, 0, 0), TLorentzVector(1000, 1001, 0, 0), TLorentzVector(1000, 1002, 0, 0)], [], [TLorentzVector(1000, 3003, 0, 0), TLorentzVector(1000, 3004, 0, 0)], [TLorentzVector(1000, 4005, 0, 0), TLorentzVector(1000, 4006, 0, 0), TLorentzVector(1000, 4007, 0, 0), TLorentzVector(1000, 4008, 0, 0), TLorentzVector(1000, 4009, 0, 0)]] diff --git a/uproot_methods/base.py b/uproot_methods/base.py index ff46bb8..d81bb3e 100644 --- a/uproot_methods/base.py +++ b/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 diff --git a/uproot_methods/classes/TLorentzVector.py b/uproot_methods/classes/TLorentzVector.py index 5d37a25..f101216 100644 --- a/uproot_methods/classes/TLorentzVector.py +++ b/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): diff --git a/uproot_methods/classes/TVector2.py b/uproot_methods/classes/TVector2.py index 6531c55..91e16d3 100644 --- a/uproot_methods/classes/TVector2.py +++ b/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)) diff --git a/uproot_methods/classes/TVector3.py b/uproot_methods/classes/TVector3.py index 6446c46..fd0fd63 100644 --- a/uproot_methods/classes/TVector3.py +++ b/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): diff --git a/uproot_methods/version.py b/uproot_methods/version.py index 045d700..edd4148 100644 --- a/uproot_methods/version.py +++ b/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__))