├── .gitignore ├── AUTHORS ├── LICENSE ├── README.rst ├── requirements.txt ├── setup.py ├── tests ├── __init__.py ├── test_histogram.py ├── test_issues.py └── test_vector.py └── uproot3_methods ├── __init__.py ├── base.py ├── classes ├── TGraph.py ├── TGraphAsymmErrors.py ├── TGraphErrors.py ├── TH1.py ├── TH2.py ├── TH3.py ├── THnSparse.py ├── TLorentzVector.py ├── TParameter.py ├── TVector2.py ├── TVector3.py └── __init__.py ├── common ├── TVector.py └── __init__.py ├── convert.py ├── profiles ├── __init__.py └── cms │ ├── __init__.py │ └── nanoaod.py └── version.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- 1 | jpivarski Jim Pivarski base package, infrastructure 2 | demarley Daniel Marley TH2 3 | kratsg Giordon Stark access to sumw2 in TH1 4 | marinang Matthieu Marinangeli TGraph* 5 | henryiii Henry Schreiner added test and fixed bug in TVector3.cross 6 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2018-2019, Jim Pivarski 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | * Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | * Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | * Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 30 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | uproot3-methods 2 | =============== 3 | 4 | .. inclusion-marker-1-do-not-remove 5 | 6 | Pythonic mix-ins for ROOT classes. 7 | 8 | .. inclusion-marker-1-5-do-not-remove 9 | 10 | This package is typically used as a dependency for `uproot 3.x `__, to define methods on the classes that are automatically generated from ROOT files. This includes histograms (TH*) and physics objects like TLorentzVectors. The reason it's a separate library is so that we can add physics-specific functionality on a shorter timescale than we can update Uproot 3 itself, which is purely an I/O package. 11 | 12 | Occasionally, this library is used without Uproot 3, as a way to make arrays of TLorentzVectors. 13 | 14 | **Note:** this package is incompatible with ``awkward>=1.0`` and ``uproot>=4.0``! For Lorentz vectors, use `vector `__. Since the versions of Awkward Array and Uproot that this is compatible with are deprecated, **this library is deprecated** as well. 15 | 16 | .. inclusion-marker-2-do-not-remove 17 | 18 | Installation 19 | ============ 20 | 21 | Install uproot3-methods like any other Python package: 22 | 23 | .. code-block:: bash 24 | 25 | pip install uproot3-methods # maybe with sudo or --user, or in virtualenv 26 | 27 | Dependencies: 28 | ------------- 29 | 30 | - `numpy `__ (1.13.1+) 31 | - `Awkward Array 0.x `__ 32 | 33 | .. inclusion-marker-3-do-not-remove 34 | 35 | Reference documentation 36 | ======================= 37 | 38 | TBD. 39 | 40 | Acknowledgements 41 | ================ 42 | 43 | Support for this work was provided by NSF cooperative agreement OAC-1836650 (IRIS-HEP), grant OAC-1450377 (DIANA/HEP) and PHY-1520942 (US-CMS LHC Ops). 44 | 45 | Thanks especially to the gracious help of `uproot3-methods contributors `__! 46 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | numpy>=1.13.1 2 | awkward0 3 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import os.path 6 | 7 | from setuptools import find_packages 8 | from setuptools import setup 9 | 10 | def get_version(): 11 | g = {} 12 | exec(open(os.path.join("uproot3_methods", "version.py")).read(), g) 13 | return g["__version__"] 14 | 15 | def get_description(): 16 | description = open("README.rst", "rb").read().decode("utf8", "ignore") 17 | start = description.index(".. inclusion-marker-1-5-do-not-remove") 18 | stop = description.index(".. inclusion-marker-3-do-not-remove") 19 | 20 | before = "" 21 | after = """ 22 | 23 | Reference documentation 24 | ======================= 25 | 26 | """ 27 | 28 | return description[start:stop].strip() # before + + after 29 | 30 | setup(name = "uproot3-methods", 31 | version = get_version(), 32 | packages = find_packages(exclude = ["tests"]), 33 | scripts = [], 34 | description = "Pythonic mix-ins for ROOT classes.", 35 | long_description = get_description(), 36 | author = "Jim Pivarski (IRIS-HEP)", 37 | author_email = "pivarski@princeton.edu", 38 | maintainer = "Jim Pivarski (IRIS-HEP)", 39 | maintainer_email = "pivarski@princeton.edu", 40 | url = "https://github.com/scikit-hep/uproot3-methods", 41 | download_url = "https://github.com/scikit-hep/uproot3-methods/releases", 42 | license = "BSD 3-clause", 43 | test_suite = "tests", 44 | install_requires = ["numpy>=1.13.1", "awkward0"], 45 | tests_require = [], 46 | classifiers = [ 47 | "Development Status :: 7 - Inactive", 48 | "Intended Audience :: Developers", 49 | "Intended Audience :: Information Technology", 50 | "Intended Audience :: Science/Research", 51 | "License :: OSI Approved :: BSD License", 52 | "Operating System :: MacOS", 53 | "Operating System :: POSIX", 54 | "Operating System :: Unix", 55 | "Programming Language :: Python", 56 | "Programming Language :: Python :: 2.7", 57 | "Programming Language :: Python :: 3.4", 58 | "Programming Language :: Python :: 3.5", 59 | "Programming Language :: Python :: 3.6", 60 | "Programming Language :: Python :: 3.7", 61 | "Programming Language :: Python :: 3.8", 62 | "Programming Language :: Python :: 3.9", 63 | "Topic :: Scientific/Engineering", 64 | "Topic :: Scientific/Engineering :: Information Analysis", 65 | "Topic :: Scientific/Engineering :: Mathematics", 66 | "Topic :: Scientific/Engineering :: Physics", 67 | "Topic :: Software Development", 68 | "Topic :: Utilities", 69 | ], 70 | platforms = "Any", 71 | ) 72 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | -------------------------------------------------------------------------------- /tests/test_histogram.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import unittest 6 | import numpy as np 7 | 8 | class Test(unittest.TestCase): 9 | def runTest(self): 10 | pass 11 | 12 | def test_th1(self): 13 | from uproot3_methods.classes.TH1 import Methods, _histtype, from_numpy 14 | 15 | edges = np.array((0., 1., 2.)) 16 | values = np.array([2, 3]) 17 | 18 | h = from_numpy((values, edges)) 19 | 20 | assert h.name is None 21 | assert h.numbins == 2 22 | assert h.title == b"" 23 | assert h.low == 0 24 | assert h.high == 2 25 | assert h.underflows == 0 26 | assert h.overflows == 0 27 | 28 | np.testing.assert_equal(h.edges, edges) 29 | np.testing.assert_equal(h.values, values) 30 | # Sumw2: variances equal the values 31 | np.testing.assert_equal(h.variances, values) 32 | 33 | np.testing.assert_equal(h.alledges, [-np.inf] + list(edges) + [np.inf]) 34 | np.testing.assert_equal(h.allvalues, [0] + list(values) + [0]) 35 | np.testing.assert_equal(h.allvariances, [0] + list(values) + [0]) 36 | 37 | np.testing.assert_equal(h.bins, ((0, 1), (1, 2))) 38 | np.testing.assert_equal(h.allbins, ((-np.inf, 0), (0, 1), (1, 2), (2, np.inf))) 39 | 40 | assert h.interval(0) == (-np.inf, 0) 41 | assert h.interval(1) == (0, 1) 42 | assert h.interval(2) == (1, 2) 43 | assert h.interval(3) == (2, np.inf) 44 | assert h.interval(-1) == h.interval(3) 45 | -------------------------------------------------------------------------------- /tests/test_issues.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import unittest 6 | 7 | import numpy 8 | 9 | import awkward0 10 | import uproot3_methods 11 | from uproot3_methods import * 12 | 13 | import inspect 14 | 15 | class Test(unittest.TestCase): 16 | def runTest(self): 17 | pass 18 | 19 | def test_issue10(self): 20 | p4 = TLorentzVectorArray.from_ptetaphim(awkward0.JaggedArray.fromiter([[1.0]]), awkward0.JaggedArray.fromiter([[1.0]]), awkward0.JaggedArray.fromiter([[1.0]]), awkward0.JaggedArray.fromiter([[1.0]])) 21 | assert p4.mass.tolist() == [[1.0]] 22 | assert p4[0].mass.tolist() == [1.0] 23 | assert p4[0][0].mass == 1.0 24 | assert p4[0][0]._to_cartesian().mass == 0.9999999999999999 25 | assert type(p4.mass) is awkward0.JaggedArray 26 | assert type(p4.x) is awkward0.JaggedArray 27 | 28 | p3 = TVector3Array.from_cylindrical(awkward0.JaggedArray.fromiter([[1.0]]), awkward0.JaggedArray.fromiter([[1.0]]), awkward0.JaggedArray.fromiter([[1.0]])) 29 | assert p3.rho.tolist() == [[1.0]] 30 | assert p3[0].rho.tolist() == [1.0] 31 | assert p3[0][0].rho == 1.0 32 | assert type(p3.rho) is awkward0.JaggedArray 33 | assert type(p3.x) is awkward0.JaggedArray 34 | 35 | p2 = TVector2Array.from_polar(awkward0.JaggedArray.fromiter([[1.0]]), awkward0.JaggedArray.fromiter([[1.0]])) 36 | assert p2.rho.tolist() == [[1.0]] 37 | assert p2[0].rho.tolist() == [1.0] 38 | assert p2[0][0].rho == 1.0 39 | assert type(p2.rho) is awkward0.JaggedArray 40 | assert type(p2.x) is awkward0.JaggedArray 41 | 42 | def test_issue39(self): 43 | counts = [2,2,2] 44 | mask = [True, False, True] 45 | 46 | pt = awkward0.JaggedArray.fromcounts(counts, [42.71, 31.46, 58.72, 30.19, 47.75, 10.83]) 47 | eta = awkward0.JaggedArray.fromcounts(counts, [0.54, 1.57, -2.33, -1.22, -2.03, -0.37]) 48 | phi = awkward0.JaggedArray.fromcounts(counts, [-2.13, 0.65, 2.74, 0.36, 2.87, -0.47]) 49 | 50 | pt = pt[mask] 51 | eta = eta[mask] 52 | phi = phi[mask] 53 | 54 | electrons = uproot3_methods.TLorentzVectorArray.from_ptetaphim(pt, eta, phi, 0.000511) 55 | 56 | def test_issue61(self): 57 | assert TVector2(2, 0).rotate(numpy.pi/6).rotate(-numpy.pi/6) == TVector2(2, 0) 58 | 59 | _xs = numpy.array([2, 0, 1]) 60 | _ys = numpy.array([0, 2, 1]) 61 | arr = TVector2Array.from_cartesian(_xs, _ys).rotate(numpy.pi/4).rotate(-numpy.pi/4) 62 | 63 | _jxs = awkward0.JaggedArray.fromiter([[2,], [], [0, 1]]) 64 | _jys = awkward0.JaggedArray.fromiter([[0,], [], [2, 1]]) 65 | jarr = TVector2Array.from_cartesian(_jxs, _jys).rotate(numpy.pi/3).rotate(-numpy.pi/3) 66 | -------------------------------------------------------------------------------- /tests/test_vector.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import unittest 6 | 7 | import numpy 8 | 9 | import awkward0 10 | import uproot3_methods 11 | from uproot3_methods import * 12 | 13 | class Test(unittest.TestCase): 14 | def runTest(self): 15 | pass 16 | 17 | def test_vector2(self): 18 | a = TVector2(4.4, 5.5) 19 | assert a.dot(a) == 49.61 20 | assert a + TVector2(1000, 2000) == TVector2(1004.4, 2005.5) 21 | assert a - TVector2(1000, 2000) == TVector2(-995.6, -1994.5) 22 | assert TVector2(1000, 2000) - a == TVector2(995.6, 1994.5) 23 | assert a * 1000 == TVector2(4400, 5500) 24 | assert 1000 * a == TVector2(4400, 5500) 25 | assert a / 1000 == TVector2(0.0044, 0.0055) 26 | assert 1000 / a == TVector2(227.27272727272725, 181.8181818181818) 27 | assert a**2 == 49.61 28 | assert a**1 == 7.043436661176133 29 | assert abs(a) == 7.043436661176133 30 | assert -a == TVector2(-4.4, -5.5) 31 | assert +a == TVector2(4.4, 5.5) 32 | 33 | a += TVector2(100, 200) 34 | assert a == TVector2(104.4, 205.5) 35 | a *= 10 36 | assert a == TVector2(1044, 2055) 37 | 38 | def test_vector2_array(self): 39 | a = TVector2Array(numpy.zeros(10), numpy.arange(10)) 40 | 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)] 41 | assert a[5] == TVector2(0, 5) 42 | assert a.y.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 43 | assert a.mag2.tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] 44 | assert (a + TVector2(1000, 2000))[5] == TVector2(1000, 2005) 45 | 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] 46 | assert (a**2).tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] 47 | assert (a**3).tolist() == [0.0, 1.0, 8.0, 27.0, 64.0, 125.0, 216.0, 343.0, 512.0, 729.0] 48 | 49 | def test_vector2_jagged(self): 50 | TVector2Jagged = type("TVector2Jagged", (awkward0.JaggedArray, uproot3_methods.classes.TVector2.ArrayMethods), {}) 51 | a = TVector2Jagged.fromoffsets([0, 3, 3, 5, 10], TVector2Array(numpy.zeros(10), numpy.arange(10))) 52 | a._generator = uproot3_methods.classes.TVector2.TVector2 53 | a._args = () 54 | a._kwargs = {} 55 | 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)]] 56 | assert a.x.tolist() == [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]] 57 | assert a.y.tolist() == [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]] 58 | 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)]] 59 | 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)]] 60 | 61 | def test_vector3(self): 62 | a = TVector3(4.4, 5.5, 0) 63 | assert a.dot(a) == 49.61 64 | assert a.cross(a) == TVector3(0,0,0) 65 | assert a + TVector3(1000, 2000, 0) == TVector3(1004.4, 2005.5, 0) 66 | assert a - TVector3(1000, 2000, 0) == TVector3(-995.6, -1994.5, 0) 67 | assert TVector3(1000, 2000, 0) - a == TVector3(995.6, 1994.5, 0) 68 | assert a * 1000 == TVector3(4400, 5500, 0) 69 | assert 1000 * a == TVector3(4400, 5500, 0) 70 | assert a / 1000 == TVector3(0.0044, 0.0055, 0) 71 | assert 1000 / (a + TVector3(0, 0, 1)) == TVector3(227.27272727272725, 181.8181818181818, 1000) 72 | assert a**2 == 49.61 73 | assert a**1 == 7.043436661176133 74 | assert abs(a) == 7.043436661176133 75 | assert -a == TVector3(-4.4, -5.5, 0) 76 | assert +a == TVector3(4.4, 5.5, 0) 77 | arot = a.rotatez(numpy.pi) 78 | self.assertAlmostEqual(arot.x,-a.x) 79 | self.assertAlmostEqual(arot.y,-a.y) 80 | self.assertAlmostEqual(arot.z,a.z) 81 | 82 | a += TVector3(100, 200, 0) 83 | assert a == TVector3(104.4, 205.5, 0) 84 | a *= 10 85 | assert a == TVector3(1044, 2055, 0) 86 | assert a.angle(a) == 0 87 | assert a.angle(a * 3) == 0 88 | self.assertAlmostEqual(a.angle(-2 * a), numpy.pi) 89 | self.assertAlmostEqual(a.angle(a.rotatez(0.01)), 0.01) 90 | self.assertAlmostEqual(a.angle(a.rotatez(-0.01)), 0.01) 91 | assert a.angle(a * 0) == 0 92 | 93 | 94 | def test_vector3_array(self): 95 | a = TVector3Array(numpy.zeros(10), numpy.arange(10), numpy.zeros(10)) 96 | 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)] 97 | assert a[5] == TVector3(0, 5, 0) 98 | assert a.y.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 99 | assert a.mag2.tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] 100 | assert (a + TVector3(1000, 2000, 0))[5] == TVector3(1000, 2005, 0) 101 | 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] 102 | assert (a**2).tolist() == [0.0, 1.0, 4.0, 9.0, 16.0, 25.0, 36.0, 49.0, 64.0, 81.0] 103 | assert (a**3).tolist() == [0.0, 1.0, 8.0, 27.0, 64.0, 125.0, 216.0, 343.0, 512.0, 729.0] 104 | arot = a.rotatez(numpy.pi) 105 | for aroti, ai in zip(arot.tolist(),a.tolist()): 106 | self.assertAlmostEqual(aroti.x,-ai.x) 107 | self.assertAlmostEqual(aroti.y,-ai.y) 108 | self.assertAlmostEqual(aroti.z,ai.z) 109 | 110 | numpy.testing.assert_almost_equal(a.angle(a), numpy.zeros_like(a)) 111 | numpy.testing.assert_almost_equal(a.angle(3 * a), numpy.zeros_like(a)) 112 | # first element is null vector, skip it, should return 0 113 | numpy.testing.assert_almost_equal(a[1:].angle(-2 * a[1:]), numpy.ones_like(a[1:]) * numpy.pi) 114 | 115 | def test_vector3_jagged(self): 116 | TVector3Jagged = type("TVector3Jagged", (awkward0.JaggedArray, uproot3_methods.classes.TVector3.ArrayMethods), {}) 117 | a = TVector3Jagged.fromoffsets([0, 3, 3, 5, 10], TVector3Array(numpy.zeros(10), numpy.arange(10), numpy.zeros(10))) 118 | a._generator = uproot3_methods.classes.TVector3.TVector3 119 | a._args = () 120 | a._kwargs = {} 121 | 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)]] 122 | assert a.x.tolist() == [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]] 123 | assert a.y.tolist() == [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]] 124 | 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)]] 125 | 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)]] 126 | arot = a.rotatez(numpy.pi) 127 | for aroti, ai in zip(arot.flatten().tolist(),a.flatten().tolist()): 128 | self.assertAlmostEqual(aroti.x,-ai.x) 129 | self.assertAlmostEqual(aroti.y,-ai.y) 130 | self.assertAlmostEqual(aroti.z,ai.z) 131 | 132 | def test_lorentzvector(self): 133 | a = TLorentzVector(4.4, 5.5, 0, 0) 134 | assert a.dot(a) == -49.61 135 | assert a + TLorentzVector(1000, 2000, 0, 0) == TLorentzVector(1004.4, 2005.5, 0, 0) 136 | assert a - TLorentzVector(1000, 2000, 0, 0) == TLorentzVector(-995.6, -1994.5, 0, 0) 137 | assert TLorentzVector(1000, 2000, 0, 0) - a == TLorentzVector(995.6, 1994.5, 0, 0) 138 | assert a * 1000 == TLorentzVector(4400, 5500, 0, 0) 139 | assert 1000 * a == TLorentzVector(4400, 5500, 0, 0) 140 | assert a / 1000 == TLorentzVector(0.0044, 0.0055, 0, 0) 141 | assert 1000 / (a + TLorentzVector(0, 0, 1, 1)) == TLorentzVector(227.27272727272725, 181.8181818181818, 1000, 1000) 142 | assert a**2, -49.61 143 | assert (a + TLorentzVector(0, 0, 0, 10))**1 == 7.098591409568521 144 | assert abs(a + TLorentzVector(0, 0, 0, 10)) == 7.098591409568521 145 | assert -a == TLorentzVector(-4.4, -5.5, 0, 0) 146 | assert +a == TLorentzVector(4.4, 5.5, 0, 0) 147 | arot = a.rotatez(numpy.pi) 148 | self.assertAlmostEqual(arot.x,-a.x) 149 | self.assertAlmostEqual(arot.y,-a.y) 150 | self.assertAlmostEqual(arot.z,a.z) 151 | assert arot.t == a.t 152 | 153 | a += TLorentzVector(100, 200, 0, 0) 154 | assert a == TLorentzVector(104.4, 205.5, 0, 0) 155 | a *= 10 156 | assert a == TLorentzVector(1044, 2055, 0, 0) 157 | 158 | def test_lorentzvector_array(self): 159 | a = TLorentzVectorArray(numpy.zeros(10), numpy.arange(10), numpy.zeros(10), numpy.zeros(10)) 160 | 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)] 161 | assert a[5] == TLorentzVector(0, 5, 0, 0) 162 | assert a.y.tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] 163 | assert a.mag2.tolist() == [0.0, -1.0, -4.0, -9.0, -16.0, -25.0, -36.0, -49.0, -64.0, -81.0] 164 | assert (a + TLorentzVector(1000, 2000, 0, 0))[5] == TLorentzVector(1000, 2005, 0, 0) 165 | 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] 166 | assert (a**2).tolist() == [0.0, -1.0, -4.0, -9.0, -16.0, -25.0, -36.0, -49.0, -64.0, -81.0] 167 | arot = a.rotatez(numpy.pi) 168 | for aroti, ai in zip(arot.tolist(),a.tolist()): 169 | self.assertAlmostEqual(aroti.x,-ai.x) 170 | self.assertAlmostEqual(aroti.y,-ai.y) 171 | self.assertAlmostEqual(aroti.z,ai.z) 172 | assert aroti.t==ai.t 173 | 174 | def test_ptetaphim_array(self): 175 | a = TLorentzVectorArray.from_ptetaphim( 176 | numpy.full(5, 20.), 177 | numpy.linspace(-5, 5, 5), 178 | numpy.linspace(-numpy.pi, numpy.pi, 6)[:-1], 179 | numpy.linspace(0, 20., 5), 180 | ) 181 | assert (a * 5).tolist() == [ 182 | PtEtaPhiMassLorentzVector(pt=100, eta=-5, phi=-numpy.pi + 0*numpy.pi/5, mass=0), 183 | PtEtaPhiMassLorentzVector(pt=100, eta=-2.5, phi=-numpy.pi + 2*numpy.pi/5, mass=25), 184 | PtEtaPhiMassLorentzVector(pt=100, eta=0, phi=-numpy.pi + 4*numpy.pi/5, mass=50), 185 | PtEtaPhiMassLorentzVector(pt=100, eta=2.5, phi=-numpy.pi + 6*numpy.pi/5, mass=75), 186 | PtEtaPhiMassLorentzVector(pt=100, eta=5, phi=-numpy.pi + 8*numpy.pi/5, mass=100) 187 | ] 188 | assert a.sum().p < 1e-10 189 | assert a.sum().mass == numpy.hypot(20 * numpy.cosh(a.eta), a.mass).sum() 190 | 191 | def test_lorentzvector_jagged(self): 192 | TLorentzVectorJagged = type("TLorentzVectorJagged", (awkward0.JaggedArray, uproot3_methods.classes.TLorentzVector.ArrayMethods), {}) 193 | a = TLorentzVectorJagged.fromoffsets([0, 3, 3, 5, 10], TLorentzVectorArray(numpy.zeros(10), numpy.arange(10), numpy.zeros(10), numpy.zeros(10))) 194 | a._generator = uproot3_methods.classes.TLorentzVector.TLorentzVector 195 | a._args = () 196 | a._kwargs = {} 197 | 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)]] 198 | assert a.x.tolist() == [[0.0, 0.0, 0.0], [], [0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0]] 199 | assert a.y.tolist() == [[0, 1, 2], [], [3, 4], [5, 6, 7, 8, 9]] 200 | 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)]] 201 | 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)]] 202 | arot = a.rotatez(numpy.pi) 203 | for aroti, ai in zip(arot.flatten().tolist(),a.flatten().tolist()): 204 | self.assertAlmostEqual(aroti.x,-ai.x) 205 | self.assertAlmostEqual(aroti.y,-ai.y) 206 | self.assertAlmostEqual(aroti.z,ai.z) 207 | assert aroti.t==ai.t 208 | -------------------------------------------------------------------------------- /uproot3_methods/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | from uproot3_methods.classes.TVector2 import TVector2, TVector2Array 6 | from uproot3_methods.classes.TVector3 import TVector3, TVector3Array 7 | from uproot3_methods.classes.TLorentzVector import TLorentzVector, TLorentzVectorArray, PtEtaPhiMassLorentzVector 8 | 9 | # convenient access to the version number 10 | from uproot3_methods.version import __version__ 11 | -------------------------------------------------------------------------------- /uproot3_methods/base.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import awkward0 6 | import awkward0.util 7 | 8 | class ROOTMethods(awkward0.Methods): 9 | _arraymethods = None 10 | 11 | awkward = awkward0 12 | awkward0 = awkward0 13 | 14 | def __ne__(self, other): 15 | return not self.__eq__(other) 16 | 17 | def _trymemo(self, name, function): 18 | memoname = "_memo_" + name 19 | wrap, (array,) = awkward0.util.unwrap_jagged(type(self), self.JaggedArray, (self,)) 20 | if not hasattr(array, memoname): 21 | setattr(array, memoname, function(array)) 22 | return wrap(getattr(array, memoname)) 23 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TGraph.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import uproot3_methods.base 6 | 7 | class Methods(uproot3_methods.base.ROOTMethods): 8 | def __repr__(self): 9 | if self._fName is None: 10 | return "<{0} at 0x{1:012x}>".format(self._classname, id(self)) 11 | else: 12 | return "<{0} {1} 0x{2:012x}>".format(self._classname, repr(self._fName), id(self)) 13 | 14 | @property 15 | def name(self): 16 | return self._fName 17 | 18 | @property 19 | def title(self): 20 | return self._fTitle 21 | 22 | @property 23 | def maximum(self): 24 | return self._fMaximum 25 | 26 | @property 27 | def minimum(self): 28 | return self._fMinimum 29 | 30 | @property 31 | def npoints(self): 32 | return self._fNpoints 33 | 34 | @property 35 | def xvalues(self): 36 | return self._fX 37 | 38 | @property 39 | def yvalues(self): 40 | return self._fY 41 | 42 | @property 43 | def xlabel(self): 44 | if self._fHistogram is None: 45 | return None 46 | elif getattr(self._fHistogram, "_fXaxis", None) is None: 47 | return None 48 | else: 49 | return getattr(self._fHistogram._fXaxis, "_fTitle", None) 50 | 51 | @property 52 | def ylabel(self): 53 | if self._fHistogram is None: 54 | return None 55 | elif getattr(self._fHistogram, "_fYaxis", None) is None: 56 | return None 57 | else: 58 | return getattr(self._fHistogram._fYaxis, "_fTitle", None) 59 | 60 | def matplotlib(self, showtitle=True, show=False, fmt="", **kwargs): 61 | import matplotlib.pyplot as pyplot 62 | 63 | _xlabel = _decode(self.xlabel if self.xlabel is not None else "") 64 | _ylabel = _decode(self.ylabel if self.ylabel is not None else "") 65 | 66 | pyplot.plot(self.xvalues, self.yvalues, fmt, **kwargs) 67 | pyplot.xlabel(_xlabel) 68 | pyplot.ylabel(_ylabel) 69 | if showtitle: 70 | _title = _decode(self.title) 71 | pyplot.title(_title) 72 | 73 | if show: 74 | pyplot.show() 75 | 76 | def _decode(sequence): 77 | return sequence.decode() if isinstance(sequence, bytes) else sequence 78 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TGraphAsymmErrors.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import uproot3_methods.base 6 | 7 | class Methods(uproot3_methods.base.ROOTMethods): 8 | 9 | @property 10 | def xerrorshigh(self): 11 | return self._fEXhigh 12 | 13 | @property 14 | def xerrorslow(self): 15 | return self._fEXlow 16 | 17 | @property 18 | def yerrorshigh(self): 19 | return self._fEYhigh 20 | 21 | @property 22 | def yerrorslow(self): 23 | return self._fEYlow 24 | 25 | def matplotlib(self, showtitle=True, show=False, **kwargs): 26 | import matplotlib.pyplot as pyplot 27 | 28 | _xerrs = [self.xerrorslow, self.xerrorshigh] 29 | _yerrs = [self.yerrorslow, self.yerrorshigh] 30 | 31 | _xlabel = _decode(self.xlabel if self.xlabel is not None else "") 32 | _ylabel = _decode(self.ylabel if self.ylabel is not None else "") 33 | 34 | pyplot.errorbar(self.xvalues, self.yvalues, xerr=_xerrs, yerr=_yerrs, **kwargs) 35 | pyplot.xlabel(_xlabel) 36 | pyplot.ylabel(_ylabel) 37 | if showtitle: 38 | _title = _decode(self.title) 39 | pyplot.title(_title) 40 | 41 | if show: 42 | pyplot.show() 43 | 44 | def _decode(sequence): 45 | return sequence.decode() if isinstance(sequence, bytes) else sequence 46 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TGraphErrors.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import uproot3_methods.base 6 | 7 | class Methods(uproot3_methods.base.ROOTMethods): 8 | 9 | @property 10 | def xerrors(self): 11 | return self._fEX 12 | 13 | @property 14 | def yerrors(self): 15 | return self._fEY 16 | 17 | def matplotlib(self, showtitle=True, show=False, **kwargs): 18 | import matplotlib.pyplot as pyplot 19 | 20 | _xlabel = _decode(self.xlabel if self.xlabel is not None else "") 21 | _ylabel = _decode(self.ylabel if self.ylabel is not None else "") 22 | 23 | pyplot.errorbar(self.xvalues, self.yvalues, xerr=self.xerrors, yerr=self.yerrors, **kwargs) 24 | pyplot.xlabel(_xlabel) 25 | pyplot.ylabel(_ylabel) 26 | if showtitle: 27 | _title = _decode(self.title) 28 | pyplot.title(_title) 29 | 30 | if show: 31 | pyplot.show() 32 | 33 | def _decode(sequence): 34 | return sequence.decode() if isinstance(sequence, bytes) else sequence 35 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TH1.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import numbers 6 | import math 7 | import sys 8 | 9 | import numpy 10 | 11 | import uproot3_methods.base 12 | 13 | class Methods(uproot3_methods.base.ROOTMethods): 14 | def __repr__(self): 15 | if self.name is None: 16 | return "<{0} at 0x{1:012x}>".format(self._classname, id(self)) 17 | else: 18 | return "<{0} {1} 0x{2:012x}>".format(self._classname, repr(self._fName), id(self)) 19 | 20 | @property 21 | def name(self): 22 | return getattr(self, "_fName", None) 23 | 24 | @property 25 | def title(self): 26 | return self._fTitle 27 | 28 | @property 29 | def numbins(self): 30 | return self._fXaxis._fNbins 31 | 32 | @property 33 | def low(self): 34 | return self._fXaxis._fXmin 35 | 36 | @property 37 | def high(self): 38 | return self._fXaxis._fXmax 39 | 40 | @property 41 | def underflows(self): 42 | return self[0] 43 | 44 | @property 45 | def overflows(self): 46 | return self[-1] 47 | 48 | @property 49 | def edges(self): 50 | axis = self._fXaxis 51 | if len(getattr(axis, "_fXbins", [])) > 0: 52 | return numpy.array(axis._fXbins) 53 | else: 54 | return numpy.linspace(axis._fXmin, axis._fXmax, axis._fNbins + 1) 55 | 56 | @property 57 | def alledges(self): 58 | axis = self._fXaxis 59 | v = numpy.empty(axis._fNbins + 3) 60 | v[0] = -numpy.inf 61 | v[-1] = numpy.inf 62 | v[1:-1] = self.edges 63 | return v 64 | 65 | @property 66 | def bins(self): 67 | edges = self.edges 68 | out = numpy.empty((len(edges) - 1, 2)) 69 | out[:, 0] = edges[:-1] 70 | out[:, 1] = edges[1:] 71 | return out 72 | 73 | @property 74 | def allbins(self): 75 | edges = self.alledges 76 | out = numpy.empty((len(edges) - 1, 2)) 77 | out[:, 0] = edges[:-1] 78 | out[:, 1] = edges[1:] 79 | return out 80 | 81 | @property 82 | def values(self): 83 | return self.allvalues[1:-1] 84 | 85 | @property 86 | def allvalues(self): 87 | return numpy.array(self, dtype=getattr(self, "_dtype", numpy.dtype(numpy.float64)).newbyteorder("=")) 88 | 89 | @property 90 | def variances(self): 91 | return self.allvariances[1:-1] 92 | 93 | @property 94 | def allvariances(self): 95 | if len(getattr(self, "_fSumw2", [])) != len(self): 96 | return numpy.array(self, dtype=numpy.float64) 97 | else: 98 | return numpy.array(self._fSumw2, dtype=numpy.float64) 99 | 100 | def numpy(self): 101 | return self.values, self.edges 102 | 103 | def allnumpy(self): 104 | return self.allvalues, self.alledges 105 | 106 | def interval(self, index): 107 | if index < 0: 108 | index += len(self) 109 | 110 | low = self._fXaxis._fXmin 111 | high = self._fXaxis._fXmax 112 | if index == 0: 113 | return (float("-inf"), low) 114 | elif index == len(self) - 1: 115 | return (high, float("inf")) 116 | elif len(getattr(self._fXaxis, "_fXbins", [])) == self._fXaxis._fNbins + 1: 117 | return (self._fXaxis._fXbins[index - 1], self._fXaxis._fXbins[index]) 118 | else: 119 | norm = (high - low) / self._fXaxis._fNbins 120 | return (index - 1)*norm + low, index*norm + low 121 | 122 | @property 123 | def xlabels(self): 124 | if getattr(self._fXaxis, "_fLabels", None) is None: 125 | return None 126 | else: 127 | return [str(x) for x in self._fXaxis._fLabels] 128 | 129 | def show(self, width=80, minimum=None, maximum=None, stream=sys.stdout): 130 | if minimum is None: 131 | minimum = min(self) 132 | if minimum < 0: 133 | minimum *= 1.05 134 | else: 135 | minimum = 0 136 | 137 | if maximum is None: 138 | maximum = max(self) * 1.05 139 | 140 | if maximum <= minimum: 141 | average = (minimum + maximum) / 2.0 142 | minimum = average - 0.5 143 | maximum = average + 0.5 144 | 145 | if self.xlabels is None: 146 | intervals = ["[{0:<.5g}, {1:<.5g})".format(l, h) for l, h in [self.interval(i) for i in range(len(self))]] 147 | intervals[-1] = intervals[-1][:-1] + "]" # last interval is closed on top edge 148 | else: 149 | intervals = ["(underflow)"] + [self.xlabels[i] if i < len(self.xlabels) else self.interval(i+1) for i in range(self.numbins)] + ["(overflow)"] 150 | 151 | intervalswidth = max(len(x) for x in intervals) 152 | 153 | values = ["{0:<.5g}".format(float(x)) for x in self] 154 | valueswidth = max(len(x) for x in values) 155 | 156 | minimumtext = "{0:<.5g}".format(minimum) 157 | maximumtext = "{0:<.5g}".format(maximum) 158 | 159 | plotwidth = max(len(minimumtext) + len(maximumtext), width - (intervalswidth + 1 + valueswidth + 1 + 2)) 160 | scale = minimumtext + " "*(plotwidth + 2 - len(minimumtext) - len(maximumtext)) + maximumtext 161 | 162 | norm = float(plotwidth) / float(maximum - minimum) 163 | zero = int(round((0.0 - minimum)*norm)) 164 | line = numpy.empty(plotwidth, dtype=numpy.uint8) 165 | 166 | formatter = "{0:<%s} {1:<%s} |{2}|" % (intervalswidth, valueswidth) 167 | line[:] = ord("-") 168 | if minimum != 0 and 0 <= zero < plotwidth: 169 | line[zero] = ord("+") 170 | capstone = " " * (intervalswidth + 1 + valueswidth + 1) + "+" + str(line.tostring().decode("ascii")) + "+" 171 | 172 | out = [" "*(intervalswidth + valueswidth + 2) + scale] 173 | out.append(capstone) 174 | for interval, value, x in zip(intervals, values, self): 175 | line[:] = ord(" ") 176 | 177 | pos = int(round((x - minimum)*norm)) 178 | if x < 0: 179 | line[pos:zero] = ord("*") 180 | else: 181 | line[zero:pos] = ord("*") 182 | 183 | if minimum != 0 and 0 <= zero < plotwidth: 184 | line[zero] = ord("|") 185 | 186 | out.append(formatter.format(interval, value, str(line.tostring().decode("ascii")))) 187 | 188 | out.append(capstone) 189 | out = "\n".join(out) 190 | if stream is None: 191 | return out 192 | else: 193 | stream.write(out) 194 | stream.write("\n") 195 | 196 | def pandas(self, underflow=True, overflow=True, variance=True): 197 | 198 | if not underflow and not overflow: 199 | s = slice(1,-1) 200 | elif not underflow: 201 | s = slice(1, None) 202 | elif not overflow: 203 | s = slice(None,-1) 204 | else: 205 | s = slice(None) 206 | 207 | def get_name(obj): 208 | if getattr(obj, "_fTitle", b"") == b"": 209 | return None 210 | else: 211 | return obj._fTitle.decode("utf-8", "ignore") 212 | 213 | allbins = self.allbins 214 | if type(allbins) != tuple: 215 | allbins = (allbins,) 216 | 217 | dim = len(allbins) 218 | 219 | if dim == 1: 220 | index_name_objects = (self,) 221 | elif dim == 2: 222 | index_name_objects = (self._fXaxis, self._fYaxis) 223 | elif dim == 3: 224 | index_name_objects = (self._fXaxis, self._fYaxis, self._fZaxis) 225 | else: 226 | raise NotImplementedError 227 | 228 | data = {"count" : self.allvalues[(s,)*dim].flatten()} 229 | columns = ["count"] 230 | 231 | if variance: 232 | data["variance"] = self.allvariances[(s,)*dim].flatten() 233 | columns += ["variance"] 234 | 235 | import pandas 236 | 237 | intervals = ((pandas.Interval(a,b, closed="left") for a, b in bins[s]) for bins in allbins) 238 | names = (get_name(obj) for obj in index_name_objects) 239 | index = pandas.MultiIndex.from_product(intervals, names=names) 240 | 241 | return pandas.DataFrame(index=index, data=data, columns=columns) 242 | 243 | def physt(self): 244 | import physt.binnings 245 | import physt.histogram1d 246 | freq = numpy.array(self.allvalues, dtype=getattr(self, "_dtype", numpy.dtype(numpy.float64)).newbyteorder("=")) 247 | if getattr(self._fXaxis, "_fXbins", None): 248 | binning = physt.binnings.NumpyBinning(numpy.array(self._fXaxis._fXbins)) 249 | else: 250 | low = self._fXaxis._fXmin 251 | high = self._fXaxis._fXmax 252 | binwidth = (high - low) / self._fXaxis._fNbins 253 | binning = physt.binnings.FixedWidthBinning(binwidth, bin_count=self._fXaxis._fNbins, min=low) 254 | return physt.histogram1d.Histogram1D( 255 | binning, 256 | frequencies=freq[1:-1], 257 | underflow=freq[0], 258 | overflow=freq[-1], 259 | name=getattr(self, "_fTitle", b"").decode("utf-8", "ignore")) 260 | 261 | def hepdata(self, independent={"name": None, "units": None}, dependent={"name": "counts", "units": None}, qualifiers=[], yamloptions={}): 262 | if independent["name"] is None and getattr(self, "_fTitle", b""): 263 | independent = dict(independent) 264 | if isinstance(self._fTitle, bytes): 265 | independent["name"] = self._fTitle.decode("utf-8", "ignore") 266 | else: 267 | independent["name"] = self._fTitle 268 | 269 | if getattr(self._fXaxis, "_fXbins", None): 270 | independent_values = [{"low": float(low), "high": float(high)} for low, high in zip(self._fXaxis._fXbins[:-1], self._fXaxis._fXbins[1:])] 271 | else: 272 | low = self._fXaxis._fXmin 273 | high = self._fXaxis._fXmax 274 | norm = (high - low) / self._fXaxis._fNbins 275 | independent_values = [{"low": float(i*norm + low), "high": float((i + 1)*norm + low)} for i in range(self.numbins)] 276 | 277 | if getattr(self, "_fSumw2", None): 278 | dependent_values = [{"value": float(value), "errors": [{"symerror": math.sqrt(variance), "label": "stat"}]} for value, variance in zip(self.values, self.variances)] 279 | else: 280 | dependent_values = [{"value": float(value), "errors": [{"symerror": math.sqrt(value), "label": "stat"}]} for value in self.values] 281 | 282 | out = {"independent_variables": [{"header": independent, "values": independent_values}], "dependent_variables": [{"header": dependent, "qualifiers": qualifiers, "values": dependent_values}]} 283 | 284 | if yamloptions is None: 285 | return out 286 | else: 287 | import yaml 288 | return yaml.dump(out, **yamloptions) 289 | 290 | def _histtype(content): 291 | if issubclass(content.dtype.type, (numpy.bool_, numpy.bool)): 292 | return b"TH1C", content.astype(">i1") 293 | elif issubclass(content.dtype.type, numpy.int8): 294 | return b"TH1C", content.astype(">i1") 295 | elif issubclass(content.dtype.type, numpy.uint8) and content.max() <= numpy.iinfo(numpy.int8).max: 296 | return b"TH1C", content.astype(">i1") 297 | elif issubclass(content.dtype.type, numpy.uint8): 298 | return b"TH1S", content.astype(">i2") 299 | elif issubclass(content.dtype.type, numpy.int16): 300 | return b"TH1S", content.astype(">i2") 301 | elif issubclass(content.dtype.type, numpy.uint16) and content.max() <= numpy.iinfo(numpy.int16).max: 302 | return b"TH1S", content.astype(">i2") 303 | elif issubclass(content.dtype.type, numpy.uint16): 304 | return b"TH1I", content.astype(">i4") 305 | elif issubclass(content.dtype.type, numpy.int32): 306 | return b"TH1I", content.astype(">i4") 307 | elif issubclass(content.dtype.type, numpy.uint32) and content.max() <= numpy.iinfo(numpy.int32).max: 308 | return b"TH1I", content.astype(">i4") 309 | elif issubclass(content.dtype.type, numpy.integer) and numpy.iinfo(numpy.int32).min <= content.min() and content.max() <= numpy.iinfo(numpy.int32).max: 310 | return b"TH1I", content.astype(">i4") 311 | elif issubclass(content.dtype.type, numpy.float32): 312 | return b"TH1F", content.astype(">f4") 313 | else: 314 | return b"TH1D", content.astype(">f8") 315 | 316 | def from_numpy(histogram): 317 | content, edges = histogram[:2] 318 | 319 | class TH1(Methods, list): 320 | pass 321 | 322 | class TAxis(object): 323 | def __init__(self, edges): 324 | self._fNbins = len(edges) - 1 325 | self._fXmin = edges[0] 326 | self._fXmax = edges[-1] 327 | if numpy.array_equal(edges, numpy.linspace(self._fXmin, self._fXmax, len(edges), dtype=edges.dtype)): 328 | self._fXbins = numpy.array([], dtype=">f8") 329 | else: 330 | self._fXbins = edges.astype(">f8") 331 | 332 | out = TH1.__new__(TH1) 333 | out._fXaxis = TAxis(edges) 334 | 335 | centers = (edges[:-1] + edges[1:]) / 2.0 336 | out._fEntries = out._fTsumw = out._fTsumw2 = content.sum() 337 | out._fTsumwx = (content * centers).sum() 338 | out._fTsumwx2 = (content * centers**2).sum() 339 | 340 | if len(histogram) >= 3: 341 | out._fTitle = histogram[2] 342 | else: 343 | out._fTitle = b"" 344 | 345 | out._classname, content = _histtype(content) 346 | 347 | valuesarray = numpy.empty(len(content) + 2, dtype=content.dtype) 348 | valuesarray[1:-1] = content 349 | valuesarray[0] = 0 350 | valuesarray[-1] = 0 351 | 352 | out.extend(valuesarray) 353 | # variances equal the entries 354 | out._fSumw2 = valuesarray 355 | 356 | return out 357 | 358 | def from_pandas(histogram): 359 | import pandas 360 | 361 | histogram = histogram.sort_index(ascending=True, inplace=False) 362 | if not histogram.index.is_non_overlapping_monotonic: 363 | raise ValueError("intervals overlap; cannot form a histogram") 364 | 365 | sparse = histogram.index[numpy.isfinite(histogram.index.left) & numpy.isfinite(histogram.index.right)] 366 | if (sparse.right[:-1] == sparse.left[1:]).all(): 367 | dense = sparse 368 | else: 369 | pairs = numpy.empty(len(sparse) * 2, dtype=numpy.float64) 370 | pairs[::2] = sparse.left 371 | pairs[1::2] = sparse.right 372 | nonempty = numpy.empty(len(pairs), dtype=numpy.bool_) 373 | nonempty[:-1] = (pairs[1:] != pairs[:-1]) 374 | nonempty[-1] = True 375 | dense = pandas.IntervalIndex.from_breaks(pairs[nonempty], closed="left") 376 | 377 | densehist = pandas.DataFrame(index=dense.left).join(histogram.reindex(histogram.index.left)) 378 | densehist.fillna(0, inplace=True) 379 | 380 | underflowhist = histogram[numpy.isinf(histogram.index.left)] 381 | overflowhist = histogram[numpy.isinf(histogram.index.right)] 382 | 383 | content = numpy.array(densehist["count"]) 384 | 385 | sumw2 = numpy.empty(len(content) + 2, dtype=numpy.float64) 386 | if "variance" in densehist.columns: 387 | sumw2source = "variance" 388 | else: 389 | sumw2source = "count" 390 | sumw2[1:-1] = densehist[sumw2source] 391 | if len(underflowhist) == 0: 392 | sumw2[0] = 0 393 | else: 394 | sumw2[0] = underflowhist[sumw2source] 395 | if len(overflowhist) == 0: 396 | sumw2[-1] = 0 397 | else: 398 | sumw2[-1] = overflowhist[sumw2source] 399 | 400 | edges = numpy.empty(len(densehist) + 1, dtype=numpy.float64) 401 | edges[:-1] = dense.left 402 | edges[-1] = dense.right[-1] 403 | 404 | class TH1(Methods, list): 405 | pass 406 | 407 | class TAxis(object): 408 | def __init__(self, fNbins, fXmin, fXmax): 409 | self._fNbins = fNbins 410 | self._fXmin = fXmin 411 | self._fXmax = fXmax 412 | 413 | out = TH1.__new__(TH1) 414 | out._fXaxis = TAxis(len(edges) - 1, edges[0], edges[-1]) 415 | out._fXaxis._fXbins = edges 416 | 417 | centers = (edges[:-1] + edges[1:]) / 2.0 418 | out._fEntries = content.sum() 419 | out._fTsumw = content.sum() 420 | out._fTsumw2 = sumw2.sum() 421 | out._fTsumwx = (content * centers).sum() 422 | out._fTsumwx2 = (content * centers**2).sum() 423 | 424 | if histogram.index.name is None: 425 | out._fTitle = b"" 426 | elif isinstance(histogram.index.name, bytes): 427 | out._fTitle = histogram.index.name 428 | else: 429 | out._fTitle = histogram.index.name.encode("utf-8", "ignore") 430 | 431 | out._classname, content = _histtype(content) 432 | 433 | valuesarray = numpy.empty(len(content) + 2, dtype=content.dtype) 434 | valuesarray[1:-1] = content 435 | if len(underflowhist) == 0: 436 | valuesarray[0] = 0 437 | else: 438 | valuesarray[0] = underflowhist["count"] 439 | if len(overflowhist) == 0: 440 | valuesarray[-1] = 0 441 | else: 442 | valuesarray[-1] = overflowhist["count"] 443 | 444 | out.extend(valuesarray) 445 | 446 | return out 447 | 448 | def from_physt(histogram): 449 | import physt.binnings 450 | import physt.histogram1d 451 | 452 | class TH1(Methods, list): 453 | pass 454 | 455 | class TAxis(object): 456 | def __init__(self, fNbins, fXmin, fXmax): 457 | self._fNbins = fNbins 458 | self._fXmin = fXmin 459 | self._fXmax = fXmax 460 | 461 | out = TH1.__new__(TH1) 462 | 463 | if isinstance(histogram.binning, physt.binnings.FixedWidthBinning): 464 | out._fXaxis = TAxis(histogram.binning.bin_count, 465 | histogram.binning.first_edge, 466 | histogram.binning.last_edge) 467 | elif isinstance(histogram.binning, physt.binnings.NumpyBinning): 468 | out._fXaxis = TAxis(histogram.binning.bin_count, 469 | histogram.binning.first_edge, 470 | histogram.binning.last_edge) 471 | if not histogram.binning.is_regular(): 472 | out._fXaxis._fXbins = histogram.binning.numpy_bins.astype(">f8") 473 | else: 474 | raise NotImplementedError(histogram.binning) 475 | 476 | centers = histogram.bin_centers 477 | content = histogram.frequencies 478 | 479 | out._fSumw2 = [0] + list(histogram.errors2) + [0] 480 | 481 | mean = histogram.mean() 482 | variance = histogram.variance() 483 | out._fEntries = content.sum() # is there a #entries independent of weights? 484 | out._fTsumw = content.sum() 485 | out._fTsumw2 = histogram.errors2.sum() 486 | if mean is None: 487 | out._fTsumwx = (content * centers).sum() 488 | else: 489 | out._fTsumwx = mean * out._fTsumw 490 | if mean is None or variance is None: 491 | out._fTsumwx2 = (content * centers**2).sum() 492 | else: 493 | out._fTsumwx2 = (mean**2 + variance) * out._fTsumw2 494 | 495 | if histogram.name is not None: 496 | out._fTitle = histogram.name 497 | else: 498 | out._fTitle = b"" 499 | 500 | out._classname, content = _histtype(content) 501 | 502 | valuesarray = numpy.empty(len(content) + 2, dtype=content.dtype) 503 | valuesarray[1:-1] = content 504 | valuesarray[0] = histogram.underflow 505 | valuesarray[-1] = histogram.overflow 506 | 507 | out.extend(valuesarray) 508 | 509 | return out 510 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TH2.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import numpy 6 | 7 | import uproot3_methods.base 8 | 9 | 10 | class Methods(uproot3_methods.base.ROOTMethods): 11 | @property 12 | def numbins(self): 13 | return self.xnumbins * self.ynumbins 14 | 15 | @property 16 | def xnumbins(self): 17 | return self._fXaxis._fNbins 18 | 19 | @property 20 | def ynumbins(self): 21 | return self._fYaxis._fNbins 22 | 23 | @property 24 | def low(self): 25 | return self.xlow, self.ylow 26 | 27 | @property 28 | def high(self): 29 | return self.xhigh, self.yhigh 30 | 31 | @property 32 | def xlow(self): 33 | return self._fXaxis._fXmin 34 | 35 | @property 36 | def xhigh(self): 37 | return self._fXaxis._fXmax 38 | 39 | @property 40 | def ylow(self): 41 | return self._fYaxis._fXmin 42 | 43 | @property 44 | def yhigh(self): 45 | return self._fYaxis._fXmax 46 | 47 | @property 48 | def underflows(self): 49 | uf = numpy.array(self.allvalues) 50 | xuf = uf[0] 51 | yuf = uf[:, 0] 52 | return xuf, yuf 53 | 54 | @property 55 | def xunderflows(self): 56 | return self.underflows[0] 57 | 58 | @property 59 | def yunderflows(self): 60 | return self.underflows[1] 61 | 62 | @property 63 | def overflows(self): 64 | of = numpy.array(self.allvalues) 65 | xof = of[-1] 66 | yof = of[:, -1] 67 | return xof, yof 68 | 69 | @property 70 | def xoverflows(self): 71 | return self.overflows[0] 72 | 73 | @property 74 | def yoverflows(self): 75 | return self.overflows[1] 76 | 77 | @property 78 | def edges(self): 79 | xaxis = self._fXaxis 80 | yaxis = self._fYaxis 81 | if len(getattr(xaxis, "_fXbins", [])) > 0: 82 | xedges = numpy.array(xaxis._fXbins) 83 | else: 84 | xedges = numpy.linspace(xaxis._fXmin, xaxis._fXmax, xaxis._fNbins + 1) 85 | if len(getattr(yaxis, "_fXbins", [])) > 0: 86 | yedges = numpy.array(yaxis._fXbins) 87 | else: 88 | yedges = numpy.linspace(yaxis._fXmin, yaxis._fXmax, yaxis._fNbins + 1) 89 | return xedges, yedges 90 | 91 | @property 92 | def alledges(self): 93 | xedges, yedges = self.edges 94 | vx = numpy.empty(len(xedges) + 2) 95 | vy = numpy.empty(len(yedges) + 2) 96 | vx[0] = -numpy.inf 97 | vx[-1] = numpy.inf 98 | vx[1:-1] = xedges 99 | vy[0] = -numpy.inf 100 | vy[-1] = numpy.inf 101 | vy[1:-1] = yedges 102 | return vx, vy 103 | 104 | @property 105 | def bins(self): 106 | xedges, yedges = self.edges 107 | vx = numpy.empty((len(xedges) - 1, 2)) 108 | vy = numpy.empty((len(yedges) - 1, 2)) 109 | vx[:, 0] = xedges[:-1] 110 | vx[:, 1] = xedges[1:] 111 | vy[:, 0] = yedges[:-1] 112 | vy[:, 1] = yedges[1:] 113 | return vx, vy 114 | 115 | @property 116 | def allbins(self): 117 | xedges, yedges = self.alledges 118 | vx = numpy.empty((len(xedges) - 1, 2)) 119 | vy = numpy.empty((len(yedges) - 1, 2)) 120 | vx[:, 0] = xedges[:-1] 121 | vx[:, 1] = xedges[1:] 122 | vy[:, 0] = yedges[:-1] 123 | vy[:, 1] = yedges[1:] 124 | return vx, vy 125 | 126 | @property 127 | def values(self): 128 | va = self.allvalues 129 | return va[1:self.xnumbins+1, 1:self.ynumbins+1] 130 | 131 | @property 132 | def allvalues(self): 133 | dtype = getattr(self, "_dtype", numpy.dtype(numpy.float64)) 134 | v = numpy.array(self[:], dtype=dtype.newbyteorder("=")) 135 | v = v.reshape(self.ynumbins + 2, self.xnumbins + 2) 136 | return v.T 137 | 138 | @property 139 | def variances(self): 140 | va = self.allvariances 141 | return va[1:self.xnumbins+1, 1:self.ynumbins+1] 142 | 143 | @property 144 | def allvariances(self): 145 | if len(getattr(self, "_fSumw2", [])) != len(self): 146 | v = numpy.array(self, dtype=numpy.float64) 147 | else: 148 | v = numpy.array(self._fSumw2, dtype=numpy.float64) 149 | v = v.reshape(self.ynumbins + 2, self.xnumbins + 2) 150 | return v.T 151 | 152 | def numpy(self): 153 | return (self.values, [self.edges]) 154 | 155 | def allnumpy(self): 156 | return (self.allvalues, [self.alledges]) 157 | 158 | def interval(self, index, axis): 159 | if axis == "x": 160 | low = self.xlow 161 | high = self.xhigh 162 | nbins = self.xnumbins 163 | bins = self._fXaxis._fXbins 164 | elif axis == "y": 165 | low = self.ylow 166 | high = self.yhigh 167 | nbins = self.ynumbins 168 | bins = self._fYaxis._fXbins 169 | else: 170 | raise ValueError("axis must be 'x' or 'y'") 171 | 172 | if index < 0: 173 | index += nbins 174 | 175 | if index == 0: 176 | return float("-inf"), low 177 | elif index == nbins + 1: 178 | return high, float("inf") 179 | else: 180 | if not bins: 181 | norm = float(high-low) / nbins 182 | xedges = (index-1)*norm + low, index*norm + low 183 | else: 184 | xedges = bins[index-1], bins[index] 185 | return xedges 186 | 187 | def xinterval(self, index): 188 | return self.interval(index, "x") 189 | 190 | def yinterval(self, index): 191 | return self.interval(index, "y") 192 | 193 | @property 194 | def ylabels(self): 195 | if self._fYaxis._fLabels is None: 196 | return None 197 | else: 198 | return [str(x) for x in self._fYaxis._fLabels] 199 | 200 | def _histtype(content): 201 | if issubclass(content.dtype.type, (numpy.bool_, numpy.bool)): 202 | return b"TH2C", content.astype(">i1") 203 | elif issubclass(content.dtype.type, numpy.int8): 204 | return b"TH2C", content.astype(">i1") 205 | elif issubclass(content.dtype.type, numpy.uint8) and content.max() <= numpy.iinfo(numpy.int8).max: 206 | return b"TH2C", content.astype(">i1") 207 | elif issubclass(content.dtype.type, numpy.uint8): 208 | return b"TH2S", content.astype(">i2") 209 | elif issubclass(content.dtype.type, numpy.int16): 210 | return b"TH2S", content.astype(">i2") 211 | elif issubclass(content.dtype.type, numpy.uint16) and content.max() <= numpy.iinfo(numpy.int16).max: 212 | return b"TH2S", content.astype(">i2") 213 | elif issubclass(content.dtype.type, numpy.uint16): 214 | return b"TH2I", content.astype(">i4") 215 | elif issubclass(content.dtype.type, numpy.int32): 216 | return b"TH2I", content.astype(">i4") 217 | elif issubclass(content.dtype.type, numpy.uint32) and content.max() <= numpy.iinfo(numpy.int32).max: 218 | return b"TH2I", content.astype(">i4") 219 | elif issubclass(content.dtype.type, numpy.integer) and numpy.iinfo(numpy.int32).min <= content.min() and content.max() <= numpy.iinfo(numpy.int32).max: 220 | return b"TH2I", content.astype(">i4") 221 | elif issubclass(content.dtype.type, numpy.float32): 222 | return b"TH2F", content.astype(">f4") 223 | else: 224 | return b"TH2D", content.astype(">f8") 225 | 226 | def from_numpy(histogram): 227 | if isinstance(histogram[1], list) and len(histogram[1]) == 2: 228 | content, (xedges, yedges) = histogram[:2] 229 | else: 230 | content, xedges, yedges = histogram[:3] 231 | 232 | class TH2(Methods, list): 233 | pass 234 | 235 | class TAxis(object): 236 | def __init__(self, edges): 237 | self._fNbins = len(edges) - 1 238 | self._fXmin = edges[0] 239 | self._fXmax = edges[-1] 240 | if numpy.array_equal(edges, numpy.linspace(self._fXmin, self._fXmax, len(edges), dtype=edges.dtype)): 241 | self._fXbins = numpy.array([], dtype=">f8") 242 | else: 243 | self._fXbins = edges.astype(">f8") 244 | 245 | out = TH2.__new__(TH2) 246 | out._fXaxis = TAxis(xedges) 247 | out._fYaxis = TAxis(yedges) 248 | out._fEntries = out._fTsumw = out._fTsumw2 = content.sum() 249 | 250 | xcenters = (xedges[:-1] + xedges[1:]) / 2. 251 | out._fTsumwx = xcenters.dot(content.sum(1)) 252 | out._fTsumwx2 = (xcenters**2).dot(content.sum(1)) 253 | 254 | ycenters = (yedges[:-1] + yedges[1:]) / 2. 255 | out._fTsumwy = ycenters.dot(content.sum(0)) 256 | out._fTsumwy2 = (ycenters**2).dot(content.sum(0)) 257 | 258 | out._fTsumwxy = xcenters.dot(content).dot(ycenters) 259 | 260 | if len(histogram) >= 4: 261 | out._fTitle = histogram[3] 262 | else: 263 | out._fTitle = b"" 264 | 265 | out._classname, content = _histtype(content) 266 | 267 | valuesarray = numpy.pad(content.T, (1, 1), mode='constant').flatten() 268 | 269 | out.extend(valuesarray) 270 | out._fSumw2 = valuesarray**2 271 | 272 | return out 273 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TH3.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import numpy 6 | 7 | import uproot3_methods.base 8 | 9 | 10 | class Methods(uproot3_methods.base.ROOTMethods): 11 | @property 12 | def numbins(self): 13 | return self.xnumbins * self.ynumbins * self.znumbins 14 | 15 | @property 16 | def xnumbins(self): 17 | return self._fXaxis._fNbins 18 | 19 | @property 20 | def ynumbins(self): 21 | return self._fYaxis._fNbins 22 | 23 | @property 24 | def znumbins(self): 25 | return self._fZaxis._fNbins 26 | 27 | @property 28 | def low(self): 29 | return self.xlow, self.ylow, self.zlow 30 | 31 | @property 32 | def high(self): 33 | return self.xhigh, self.yhigh, self.zlow 34 | 35 | @property 36 | def xlow(self): 37 | return self._fXaxis._fXmin 38 | 39 | @property 40 | def xhigh(self): 41 | return self._fXaxis._fXmax 42 | 43 | @property 44 | def ylow(self): 45 | return self._fYaxis._fXmin 46 | 47 | @property 48 | def yhigh(self): 49 | return self._fYaxis._fXmax 50 | 51 | @property 52 | def zlow(self): 53 | return self._fZaxis._fXmin 54 | 55 | @property 56 | def zhigh(self): 57 | return self._fZaxis._fXmax 58 | 59 | @property 60 | def underflows(self): 61 | uf = numpy.array(self.allvalues) 62 | xuf = uf[0] 63 | yuf = uf[:, 0, :] 64 | zuf = uf[:, :, 0] 65 | return xuf, yuf, zuf 66 | 67 | @property 68 | def xunderflows(self): 69 | return self.underflows[0] 70 | 71 | @property 72 | def yunderflows(self): 73 | return self.underflows[1] 74 | 75 | @property 76 | def zunderflows(self): 77 | return self.underflows[2] 78 | 79 | @property 80 | def overflows(self): 81 | of = numpy.array(self.allvalues) 82 | xof = of[-1] 83 | yof = of[:, -1, :] 84 | zof = of[:, :, -1] 85 | return xof, yof, zof 86 | 87 | @property 88 | def xoverflows(self): 89 | return self.overflows[0] 90 | 91 | @property 92 | def yoverflows(self): 93 | return self.overflows[1] 94 | 95 | @property 96 | def zoverflows(self): 97 | return self.overflows[2] 98 | 99 | @property 100 | def edges(self): 101 | xaxis = self._fXaxis 102 | yaxis = self._fYaxis 103 | zaxis = self._fZaxis 104 | if len(getattr(xaxis, "_fXbins", [])) > 0: 105 | xedges = numpy.array(xaxis._fXbins) 106 | else: 107 | xedges = numpy.linspace(xaxis._fXmin, xaxis._fXmax, 108 | xaxis._fNbins + 1) 109 | 110 | if len(getattr(yaxis, "_fXbins", [])) > 0: 111 | yedges = numpy.array(yaxis._fXbins) 112 | else: 113 | yedges = numpy.linspace(yaxis._fXmin, yaxis._fXmax, 114 | yaxis._fNbins + 1) 115 | 116 | if len(getattr(zaxis, "_fXbins", [])) > 0: 117 | zedges = numpy.array(zaxis._fXbins) 118 | else: 119 | zedges = numpy.linspace(zaxis._fXmin, zaxis._fXmax, 120 | zaxis._fNbins + 1) 121 | 122 | return xedges, yedges, zedges 123 | 124 | @property 125 | def alledges(self): 126 | xedges, yedges, zedges = self.edges 127 | vx = numpy.empty(len(xedges) + 2) 128 | vy = numpy.empty(len(yedges) + 2) 129 | vz = numpy.empty(len(zedges) + 2) 130 | vx[0] = -numpy.inf 131 | vx[-1] = numpy.inf 132 | vx[1:-1] = xedges 133 | vy[0] = -numpy.inf 134 | vy[-1] = numpy.inf 135 | vy[1:-1] = yedges 136 | vz[0] = -numpy.inf 137 | vz[-1] = numpy.inf 138 | vz[1:-1] = zedges 139 | return vx, vy, vz 140 | 141 | @property 142 | def bins(self): 143 | xedges, yedges, zedges = self.edges 144 | vx = numpy.empty((len(xedges) - 1, 2)) 145 | vy = numpy.empty((len(yedges) - 1, 2)) 146 | vz = numpy.empty((len(zedges) - 1, 2)) 147 | vx[:, 0] = xedges[:-1] 148 | vx[:, 1] = xedges[1:] 149 | vy[:, 0] = yedges[:-1] 150 | vy[:, 1] = yedges[1:] 151 | vz[:, 0] = zedges[:-1] 152 | vz[:, 1] = zedges[1:] 153 | return vx, vy, vz 154 | 155 | @property 156 | def allbins(self): 157 | xedges, yedges, zedges = self.alledges 158 | vx = numpy.empty((len(xedges) - 1, 2)) 159 | vy = numpy.empty((len(yedges) - 1, 2)) 160 | vz = numpy.empty((len(zedges) - 1, 2)) 161 | vx[:, 0] = xedges[:-1] 162 | vx[:, 1] = xedges[1:] 163 | vy[:, 0] = yedges[:-1] 164 | vy[:, 1] = yedges[1:] 165 | vz[:, 0] = zedges[:-1] 166 | vz[:, 1] = zedges[1:] 167 | return vx, vy, vz 168 | 169 | @property 170 | def values(self): 171 | va = self.allvalues 172 | return va[1:self.xnumbins+1, 1:self.ynumbins+1, 1:self.znumbins+1] 173 | 174 | @property 175 | def allvalues(self): 176 | dtype = getattr(self, "_dtype", numpy.dtype(numpy.float64)) 177 | v = numpy.array(self[:], dtype=dtype.newbyteorder("=")) 178 | v = v.reshape(self.znumbins + 2, self.ynumbins + 2, self.xnumbins + 2) 179 | return v.T 180 | 181 | @property 182 | def variances(self): 183 | va = self.allvariances 184 | return va[1:self.xnumbins+1, 1:self.ynumbins+1, 1:self.znumbins+1] 185 | 186 | @property 187 | def allvariances(self): 188 | if len(getattr(self, "_fSumw2", [])) != len(self): 189 | v = numpy.array(self, dtype=numpy.float64) 190 | else: 191 | v = numpy.array(self._fSumw2, dtype=numpy.float64) 192 | v = v.reshape(self.znumbins + 2, self.ynumbins + 2, self.xnumbins + 2) 193 | return v.T 194 | 195 | def numpy(self): 196 | return (self.values, [self.edges]) 197 | 198 | def allnumpy(self): 199 | return (self.allvalues, [self.alledges]) 200 | 201 | def interval(self, index, axis): 202 | if axis == "x": 203 | low = self.xlow 204 | high = self.xhigh 205 | nbins = self.xnumbins 206 | bins = self._fXaxis._fXbins 207 | elif axis == "y": 208 | low = self.ylow 209 | high = self.yhigh 210 | nbins = self.ynumbins 211 | bins = self._fYaxis._fXbins 212 | elif axis == "z": 213 | low = self.zlow 214 | high = self.zhigh 215 | nbins = self.znumbins 216 | bins = self._fZaxis._fXbins 217 | else: 218 | raise ValueError("axis must be 'x','y' or 'z'") 219 | 220 | if index < 0: 221 | index += nbins 222 | 223 | if index == 0: 224 | return float("-inf"), low 225 | elif index == nbins + 1: 226 | return high, float("inf") 227 | else: 228 | if not bins: 229 | norm = float(high-low) / nbins 230 | xedges = (index-1)*norm + low, index*norm + low 231 | else: 232 | xedges = bins[index-1], bins[index] 233 | return xedges 234 | 235 | def xinterval(self, index): 236 | return self.interval(index, "x") 237 | 238 | def yinterval(self, index): 239 | return self.interval(index, "y") 240 | 241 | def zinterval(self, index): 242 | return self.interval(index, "z") 243 | 244 | @property 245 | def ylabels(self): 246 | if self._fYaxis._fLabels is None: 247 | return None 248 | else: 249 | return [str(x) for x in self._fYaxis._fLabels] 250 | 251 | @property 252 | def zlabels(self): 253 | if self._fZaxis._fLabels is None: 254 | return None 255 | else: 256 | return [str(x) for x in self._fZaxis._fLabels] 257 | -------------------------------------------------------------------------------- /uproot3_methods/classes/THnSparse.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import numpy 6 | 7 | import uproot3_methods.base 8 | 9 | class Methods(uproot3_methods.base.ROOTMethods): 10 | def hello(self): 11 | return "world", len(dir(self)) 12 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TLorentzVector.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import math 6 | import numbers 7 | import operator 8 | 9 | import awkward0.array.chunked 10 | import awkward0.array.jagged 11 | import awkward0.array.objects 12 | import awkward0.util 13 | 14 | import uproot3_methods.base 15 | import uproot3_methods.common.TVector 16 | import uproot3_methods.classes.TVector3 17 | 18 | class Common(object): 19 | @property 20 | def E(self): 21 | return self.t 22 | 23 | def dot(self, other): 24 | out = self.t*other.t 25 | out = out - self.x*other.x 26 | out = out - self.y*other.y 27 | out = out - self.z*other.z 28 | return out 29 | 30 | @property 31 | def energy(self): 32 | return self.t 33 | 34 | @property 35 | def p(self): 36 | return self.p3.mag 37 | 38 | @property 39 | def p2(self): 40 | return self.p3.mag2 41 | 42 | @property 43 | def perp2(self): 44 | return self.p3.rho2 45 | 46 | @property 47 | def perp(self): 48 | return self.p3.rho 49 | 50 | @property 51 | def pt2(self): 52 | return self.p3.rho2 53 | 54 | @property 55 | def Et(self): 56 | return self.energy * self.pt / self.p 57 | 58 | @property 59 | def mag2(self): 60 | return self.dot(self) 61 | 62 | @property 63 | def mass2(self): 64 | return self.mag2 65 | 66 | @property 67 | def mt2(self): 68 | return self.energy**2 - self.z**2 69 | 70 | @property 71 | def theta(self): 72 | return self.p3.theta 73 | 74 | @property 75 | def cottheta(self): 76 | return self.p3.cottheta 77 | 78 | @property 79 | def beta(self): 80 | return self.p / self.energy 81 | 82 | def delta_phi(self, other): 83 | return (self.phi - other.phi + math.pi) % (2*math.pi) - math.pi 84 | 85 | def delta_r2(self, other): 86 | return (self.eta - other.eta)**2 + self.delta_phi(other)**2 87 | 88 | def _rotate_axis(self, axis, angle): 89 | if not isinstance(axis, uproot3_methods.classes.TVector3.Common): 90 | raise TypeError("axis must be an (array of) TVector3") 91 | p3 = self.p3._rotate_axis(axis, angle) 92 | return p3, self.t 93 | 94 | def _rotate_euler(self, phi, theta, psi): 95 | return self.p3._rotate_euler(phi, theta, psi), self.t 96 | 97 | def rotatex(self, angle): 98 | return self.rotate_axis(uproot3_methods.classes.TVector3.TVector3(1.0, 0.0, 0.0), angle) 99 | 100 | def rotatey(self, angle): 101 | return self.rotate_axis(uproot3_methods.classes.TVector3.TVector3(0.0, 1.0, 0.0), angle) 102 | 103 | def rotatez(self, angle): 104 | return self.rotate_axis(uproot3_methods.classes.TVector3.TVector3(0.0, 0.0, 1.0), angle) 105 | 106 | def isspacelike(self, tolerance=1e-10): 107 | return self.mag2 < -tolerance 108 | 109 | def istimelike(self, tolerance=1e-10): 110 | return self.mag2 > tolerance 111 | 112 | def __lt__(self, other): 113 | raise TypeError("Lorentz vectors have no natural ordering") 114 | 115 | def __gt__(self, other): 116 | raise TypeError("Lorentz vectors have no natural ordering") 117 | 118 | def __le__(self, other): 119 | raise TypeError("Lorentz vectors have no natural ordering") 120 | 121 | def __ge__(self, other): 122 | raise TypeError("Lorentz vectors have no natural ordering") 123 | 124 | class ArrayMethods(Common, uproot3_methods.base.ROOTMethods): 125 | def _initObjectArray(self, table): 126 | self.awkward0.ObjectArray.__init__(self, table, lambda row: TLorentzVector(row["fX"], row["fY"], row["fZ"], row["fE"])) 127 | 128 | def __awkward_serialize__(self, serializer): 129 | self._valid() 130 | x, y, z, t = self.x, self.y, self.z, self.t 131 | return serializer.encode_call( 132 | ["uproot3_methods.classes.TLorentzVector", "TLorentzVectorArray", "from_cartesian"], 133 | serializer(x, "TLorentzVectorArray.x"), 134 | serializer(y, "TLorentzVectorArray.y"), 135 | serializer(z, "TLorentzVectorArray.z"), 136 | serializer(t, "TLorentzVectorArray.t")) 137 | 138 | @staticmethod 139 | def _wrapmethods(node, awkwardlib): 140 | if isinstance(node, awkward0.array.chunked.ChunkedArray): 141 | node.__class__ = type("ChunkedArrayMethods", (awkwardlib.ChunkedArray, uproot3_methods.classes.TVector3.ArrayMethods), {}) 142 | for chunk in node.chunks: 143 | ArrayMethods._wrapmethods(chunk, awkwardlib) 144 | elif isinstance(node, awkward0.array.jagged.JaggedArray): 145 | node.__class__ = type("JaggedArrayMethods", (awkwardlib.JaggedArray, uproot3_methods.classes.TVector3.ArrayMethods), {}) 146 | ArrayMethods._wrapmethods(node.content, awkwardlib) 147 | elif isinstance(node, awkward0.array.objects.ObjectArray): 148 | node.__class__ = type("ObjectArrayMethods", (awkwardlib.ObjectArray, uproot3_methods.classes.TVector3.ArrayMethods), {}) 149 | 150 | @property 151 | def p3(self): 152 | out = self.empty_like(generator=lambda row: uproot3_methods.classes.TVector3.TVector3(row["fX"], row["fY"], row["fZ"])) 153 | ArrayMethods._wrapmethods(out, self.awkward0) 154 | out["fX"] = self.x 155 | out["fY"] = self.y 156 | out["fZ"] = self.z 157 | return out 158 | 159 | @property 160 | def x(self): 161 | return self["fX"] 162 | 163 | @property 164 | def y(self): 165 | return self["fY"] 166 | 167 | @property 168 | def z(self): 169 | return self["fZ"] 170 | 171 | @property 172 | def t(self): 173 | return self["fE"] 174 | 175 | @property 176 | def pt(self): 177 | return self._trymemo("pt", lambda self: self.awkward0.numpy.sqrt(self.pt2)) 178 | 179 | @property 180 | def eta(self): 181 | return self._trymemo("eta", lambda self: self.awkward0.numpy.arcsinh(self.z / self.awkward0.numpy.sqrt(self.x**2 + self.y**2))) 182 | 183 | @property 184 | def phi(self): 185 | return self._trymemo("phi", lambda self: self.p3.phi) 186 | 187 | @property 188 | def mass(self): 189 | return self._trymemo("mass", lambda self: self.awkward0.numpy.sqrt(self.mag2)) 190 | 191 | @property 192 | def mag(self): 193 | return self.awkward0.numpy.sqrt(self.mag2) 194 | 195 | @property 196 | def mt(self): 197 | mt2 = self.mt2 198 | sign = self.awkward0.numpy.sign(mt2) 199 | return self.awkward0.numpy.sqrt(self.awkward0.numpy.absolute(mt2)) * sign 200 | 201 | @property 202 | def rapidity(self): 203 | return 0.5 * self.awkward0.numpy.log((self.t + self.z) / (self.t - self.z)) 204 | 205 | @property 206 | def unit(self): 207 | return self / self.awkward0.numpy.sqrt(self.mag) 208 | 209 | @property 210 | def boostp3(self): 211 | out = self.empty_like(generator=lambda row: uproot3_methods.classes.TVector3.TVector3(row["fX"], row["fY"], row["fZ"])) 212 | if isinstance(self, self.awkward0.JaggedArray): 213 | out.__class__ = type("JaggedArrayMethods", (self.awkward0.JaggedArray, uproot3_methods.classes.TVector3.ArrayMethods), {}) 214 | else: 215 | out.__class__ = type("ObjectArrayMethods", (self.awkward0.ObjectArray, uproot3_methods.classes.TVector3.ArrayMethods), {}) 216 | out["fX"] = self.x / self.t 217 | out["fY"] = self.y / self.t 218 | out["fZ"] = self.z / self.t 219 | return out 220 | 221 | def boost(self, p3): 222 | if not isinstance(p3, (uproot3_methods.classes.TVector3.ArrayMethods, uproot3_methods.classes.TVector3.Methods)): 223 | raise TypeError("boost p3 must be an (array of) TVector3") 224 | 225 | b2 = p3.mag2 226 | gamma = (1 - b2)**(-0.5) 227 | gamma2 = self.awkward0.numpy.zeros(b2.shape, dtype=self.awkward0.numpy.float64) 228 | mask = (b2 != 0) 229 | gamma2[mask] = (gamma[mask] - 1) / b2[mask] 230 | del mask 231 | 232 | bp = self.p3.dot(p3) 233 | 234 | v = self.p3 + gamma2*bp*p3 + self.t*gamma*p3 235 | out = self.empty_like() 236 | out["fX"] = v.x 237 | out["fY"] = v.y 238 | out["fZ"] = v.z 239 | out["fE"] = gamma*(self.t + bp) 240 | return out 241 | 242 | @property 243 | def gamma(self): 244 | out = self.beta 245 | mask = (out < 1) & (out > -1) 246 | out[mask] = (1 - out[mask]**2)**(-0.5) 247 | out[~mask] = self.awkward0.numpy.inf 248 | return out 249 | 250 | def delta_r(self, other): 251 | return self.awkward0.numpy.sqrt(self.delta_r2(other)) 252 | 253 | def rotate_axis(self, axis, angle): 254 | p3, t = self._rotate_axis(axis, angle) 255 | x, y, z = p3 256 | out = self.empty_like() 257 | out["fX"] = x 258 | out["fY"] = y 259 | out["fZ"] = z 260 | out["fE"] = t 261 | return out 262 | 263 | def rotate_euler(self, phi=0, theta=0, psi=0): 264 | p3, t = self._rotate_euler(phi, theta, psi) 265 | x, y, z = p3 266 | out = self.empty_like() 267 | out["fX"] = x 268 | out["fY"] = y 269 | out["fZ"] = z 270 | out["fE"] = t 271 | return out 272 | 273 | def islightlike(self, tolerance=1e-10): 274 | return self.awkward0.numpy.absolute(self.mag2) < tolerance 275 | 276 | def sum(self): 277 | if isinstance(self, awkward0.AwkwardArray) and self._util_hasjagged(self): 278 | return TLorentzVectorArray.from_cartesian(self.x.sum(), self.y.sum(), self.z.sum(), self.t.sum()) 279 | else: 280 | return TLorentzVector(self.x.sum(), self.y.sum(), self.z.sum(), self.t.sum()) 281 | 282 | def _to_cartesian(self): 283 | return TLorentzVectorArray.from_cartesian(self.x,self.y,self.z,self.t) 284 | 285 | def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): 286 | if "out" in kwargs: 287 | raise NotImplementedError("in-place operations not supported") 288 | 289 | if method != "__call__": 290 | return NotImplemented 291 | 292 | inputs = list(inputs) 293 | for i in range(len(inputs)): 294 | if isinstance(inputs[i], self.awkward0.numpy.ndarray) and inputs[i].dtype == self.awkward0.numpy.dtype(object) and len(inputs[i]) > 0: 295 | idarray = self.awkward0.numpy.frombuffer(inputs[i], dtype=self.awkward0.numpy.uintp) 296 | if (idarray == idarray[0]).all(): 297 | inputs[i] = inputs[i][0] 298 | 299 | if ufunc is self.awkward0.numpy.add or ufunc is self.awkward0.numpy.subtract: 300 | if not all(isinstance(x, (ArrayMethods, Methods)) for x in inputs): 301 | raise TypeError("(arrays of) TLorentzVector can only be added to/subtracted from other (arrays of) TLorentzVector") 302 | cart_inputs = [x._to_cartesian() for x in inputs] 303 | out = cart_inputs[0].empty_like() 304 | out["fX"] = getattr(ufunc, method)(*[x.x for x in cart_inputs], **kwargs) 305 | out["fY"] = getattr(ufunc, method)(*[x.y for x in cart_inputs], **kwargs) 306 | out["fZ"] = getattr(ufunc, method)(*[x.z for x in cart_inputs], **kwargs) 307 | out["fE"] = getattr(ufunc, method)(*[x.t for x in cart_inputs], **kwargs) 308 | return out 309 | 310 | elif ufunc is self.awkward0.numpy.power and len(inputs) >= 2 and isinstance(inputs[1], (numbers.Number, self.awkward0.numpy.number)): 311 | if inputs[1] == 2: 312 | return self.mag2 313 | else: 314 | return self.mag2**(0.5*inputs[1]) 315 | 316 | elif ufunc is self.awkward0.numpy.absolute: 317 | return self.mag 318 | 319 | else: 320 | return super(ArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs) 321 | 322 | JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward0.JaggedArray) 323 | 324 | class PtEtaPhiMassArrayMethods(ArrayMethods): 325 | def _initObjectArray(self, table): 326 | self.awkward0.ObjectArray.__init__(self, table, lambda row: PtEtaPhiMassLorentzVector(row["fPt"], row["fEta"], row["fPhi"], row["fMass"])) 327 | 328 | def __awkward_serialize__(self, serializer): 329 | self._valid() 330 | pt, eta, phi, mass = self.pt, self.eta, self.phi, self.mass 331 | return serializer.encode_call( 332 | ["uproot3_methods.classes.TLorentzVector", "TLorentzVectorArray", "from_ptetaphim"], 333 | serializer(pt, "TLorentzVectorArray.pt"), 334 | serializer(eta, "TLorentzVectorArray.eta"), 335 | serializer(phi, "TLorentzVectorArray.phi"), 336 | serializer(mass, "TLorentzVectorArray.mass")) 337 | 338 | @property 339 | def x(self): 340 | return self._trymemo("x",lambda self: self.pt * self.awkward0.numpy.cos(self.phi)) 341 | 342 | @property 343 | def y(self): 344 | return self._trymemo("y",lambda self: self.pt * self.awkward0.numpy.sin(self.phi)) 345 | 346 | @property 347 | def z(self): 348 | return self._trymemo("z",lambda self: self.pt * self.awkward0.numpy.sinh(self.eta)) 349 | 350 | @property 351 | def t(self): 352 | return self._trymemo("t",lambda self: self.awkward0.numpy.hypot(self.mass, self.p)) 353 | 354 | @property 355 | def pt(self): 356 | return self["fPt"] 357 | 358 | @property 359 | def pt2(self): 360 | return self["fPt"]**2 361 | 362 | @property 363 | def perp(self): 364 | return self["fPt"] 365 | 366 | @property 367 | def perp2(self): 368 | return self["fPt"]**2 369 | 370 | @property 371 | def eta(self): 372 | return self["fEta"] 373 | 374 | @property 375 | def phi(self): 376 | return self["fPhi"] 377 | 378 | @property 379 | def mass(self): 380 | return self["fMass"] 381 | 382 | @property 383 | def mass2(self): 384 | return self["fMass"]**2 385 | 386 | @property 387 | def mag(self): 388 | return self["fMass"] 389 | 390 | @property 391 | def mag2(self): 392 | return self["fMass"]**2 393 | 394 | @property 395 | def mt(self): 396 | return self.awkward0.numpy.sqrt(self.mt2) 397 | 398 | @property 399 | def mt2(self): 400 | return self["fMass"]**2 + self["fPt"]**2 401 | 402 | @property 403 | def p(self): 404 | return self._trymemo("p",lambda self: self["fPt"]*self.awkward0.numpy.cosh(self["fEta"])) 405 | 406 | @property 407 | def p2(self): 408 | return self.p**2 409 | 410 | def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): 411 | if "out" in kwargs: 412 | raise NotImplementedError("in-place operations not supported") 413 | 414 | if method != "__call__": 415 | return NotImplemented 416 | 417 | inputs = list(inputs) 418 | for i in range(len(inputs)): 419 | if isinstance(inputs[i], self.awkward0.numpy.ndarray) and inputs[i].dtype == self.awkward0.numpy.dtype(object) and len(inputs[i]) > 0: 420 | idarray = self.awkward0.numpy.frombuffer(inputs[i], dtype=self.awkward0.numpy.uintp) 421 | if (idarray == idarray[0]).all(): 422 | inputs[i] = inputs[i][0] 423 | 424 | if ufunc is self.awkward0.numpy.multiply or ufunc is self.awkward0.numpy.divide: 425 | if sum(isinstance(x, PtEtaPhiMassArrayMethods) for x in inputs) > 1: 426 | raise ValueError("cannot multiply or divide two PtEtaPhiMassArrayMethods") 427 | this_input = None 428 | for i in range(len(inputs)): 429 | if isinstance(inputs[i], PtEtaPhiMassArrayMethods) and not isinstance(inputs[i], self.awkward0.JaggedArray) and this_input is None: 430 | this_input = inputs[i] 431 | inputs[i] = self.awkward0.Table(fPt=inputs[i]['fPt'], fMass=inputs[i]['fMass']) 432 | 433 | out = super(PtEtaPhiMassArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs) 434 | if this_input is not None: 435 | out['fEta'] = this_input['fEta'] 436 | out['fPhi'] = this_input['fPhi'] 437 | out.__class__ = this_input.__class__ 438 | return out 439 | 440 | else: 441 | return super(PtEtaPhiMassArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs) 442 | 443 | PtEtaPhiMassJaggedArrayMethods = PtEtaPhiMassArrayMethods.mixin(PtEtaPhiMassArrayMethods, awkward0.JaggedArray) 444 | 445 | class Methods(Common, uproot3_methods.base.ROOTMethods): 446 | _arraymethods = ArrayMethods 447 | 448 | @property 449 | def p3(self): 450 | return self._fP 451 | 452 | @property 453 | def x(self): 454 | return self._fP._fX 455 | 456 | @property 457 | def y(self): 458 | return self._fP._fY 459 | 460 | @property 461 | def z(self): 462 | return self._fP._fZ 463 | 464 | @property 465 | def t(self): 466 | return self._fE 467 | 468 | def _to_cartesian(self): 469 | return TLorentzVector(self.x,self.y,self.z,self.t) 470 | 471 | def __repr__(self): 472 | return "TLorentzVector(x={0:.5g}, y={1:.5g}, z={2:.5g}, t={3:.5g})".format(self._fP._fX, self._fP._fY, self._fP._fZ, self._fE) 473 | 474 | def __str__(self): 475 | return repr(self) 476 | 477 | def __eq__(self, other): 478 | return isinstance(other, Methods) and self.x == other.x and self.y == other.y and self.z == other.z and self.t == other.t 479 | 480 | def _scalar(self, operator, scalar, reverse=False): 481 | cart = self._to_cartesian() 482 | if not isinstance(scalar, (numbers.Number, self.awkward0.numpy.number)): 483 | raise TypeError("cannot {0} a TLorentzVector with a {1}".format(operator.__name__, type(scalar).__name__)) 484 | if reverse: 485 | return TLorentzVector(operator(scalar, cart.x), operator(scalar, cart.y), operator(scalar, cart.z), operator(scalar, cart.t)) 486 | else: 487 | return TLorentzVector(operator(cart.x, scalar), operator(cart.y, scalar), operator(cart.z, scalar), operator(cart.t, scalar)) 488 | 489 | def _vector(self, operator, vector, reverse=False): 490 | cart = self._to_cartesian() 491 | if not isinstance(vector, Methods): 492 | raise TypeError("cannot {0} a TLorentzVector with a {1}".format(operator.__name__, type(vector).__name__)) 493 | if reverse: 494 | return TLorentzVector(operator(vector.x, cart.x), operator(vector.y, cart.y), operator(vector.z, cart.z), operator(vector.t, cart.t)) 495 | else: 496 | return TLorentzVector(operator(cart.x, vector.x), operator(cart.y, vector.y), operator(cart.z, vector.z), operator(cart.t, vector.t)) 497 | 498 | def _unary(self, operator): 499 | cart = self._to_cartesian() 500 | return TLorentzVector(operator(cart.x), operator(cart.y), operator(cart.z), operator(cart.t)) 501 | 502 | @property 503 | def pt(self): 504 | return math.sqrt(self.pt2) 505 | 506 | @property 507 | def eta(self): 508 | return math.asinh(self.z / math.sqrt(self.x**2 + self.y**2)) 509 | 510 | @property 511 | def phi(self): 512 | return self.p3.phi 513 | 514 | @property 515 | def mass(self): 516 | return math.sqrt(self.mag2) 517 | 518 | @property 519 | def mag(self): 520 | return math.sqrt(self.mag2) 521 | 522 | @property 523 | def mt(self): 524 | out = self.mt2 525 | if out >= 0: 526 | return math.sqrt(out) 527 | else: 528 | return -math.sqrt(out) 529 | 530 | @property 531 | def rapidity(self): 532 | return math.log((self.t + self.z) / (self.t - self.z)) / 2.0 533 | 534 | @property 535 | def unit(self): 536 | return self / math.sqrt(self.mag) 537 | 538 | @property 539 | def boostp3(self): 540 | return uproot3_methods.classes.TVector3.TVector3(self.x/self.t, self.y/self.t, self.z/self.t) 541 | 542 | def boost(self, p3): 543 | if not isinstance(p3, uproot3_methods.classes.TVector3.Methods): 544 | raise TypeError("boost p3 must be a TVector3") 545 | 546 | b2 = p3.mag2 547 | gamma = (1.0 - b2)**(-0.5) 548 | if b2 != 0: 549 | gamma2 = (gamma - 1.0) / b2 550 | else: 551 | gamma2 = 0.0 552 | 553 | bp = self.p3.dot(p3) 554 | v = self.p3 + gamma2*bp*p3 + gamma*p3*self.t 555 | return self.__class__(v.x, v.y, v.z, gamma*(self.t + bp)) 556 | 557 | @property 558 | def gamma(self): 559 | out = self.beta 560 | if -1 < out < 1: 561 | return (1 - out**2)**(-0.5) 562 | else: 563 | return float("inf") 564 | 565 | def delta_r(self, other): 566 | return math.sqrt(self.delta_r2(other)) 567 | 568 | def rotate_axis(self, axis, angle): 569 | p3, t = self._rotate_axis(axis, angle) 570 | x, y, z = p3 571 | return self.__class__(x, y, z, t) 572 | 573 | def rotate_euler(self, phi=0, theta=0, psi=0): 574 | p3, t = self._rotate_euler(phi, theta, psi) 575 | x, y, z = p3 576 | return self.__class__(x, y, z, t) 577 | 578 | def islightlike(self, tolerance=1e-10): 579 | return abs(self.mag2) < tolerance 580 | 581 | def __add__(self, other): 582 | return self._vector(operator.add, other) 583 | 584 | def __radd__(self, other): 585 | return self._vector(operator.add, other, True) 586 | 587 | def __sub__(self, other): 588 | return self._vector(operator.sub, other) 589 | 590 | def __rsub__(self, other): 591 | return self._vector(operator.sub, other, True) 592 | 593 | def __mul__(self, other): 594 | return self._scalar(operator.mul, other) 595 | 596 | def __rmul__(self, other): 597 | return self._scalar(operator.mul, other, True) 598 | 599 | def __div__(self, other): 600 | return self._scalar(operator.div, other) 601 | 602 | def __rdiv__(self, other): 603 | return self._scalar(operator.div, other, True) 604 | 605 | def __truediv__(self, other): 606 | return self._scalar(operator.truediv, other) 607 | 608 | def __rtruediv__(self, other): 609 | return self._scalar(operator.truediv, other, True) 610 | 611 | def __floordiv__(self, other): 612 | return self._scalar(operator.floordiv, other) 613 | 614 | def __rfloordiv__(self, other): 615 | return self._scalar(operator.floordiv, other, True) 616 | 617 | def __mod__(self, other): 618 | return self._scalar(operator.mod, other) 619 | 620 | def __rmod__(self, other): 621 | return self._scalar(operator.mod, other, True) 622 | 623 | def __divmod__(self, other): 624 | return self._scalar(operator.divmod, other) 625 | 626 | def __rdivmod__(self, other): 627 | return self._scalar(operator.divmod, other, True) 628 | 629 | def __pow__(self, other): 630 | if isinstance(other, (numbers.Number, self.awkward0.numpy.number)): 631 | if other == 2: 632 | return self.mag2 633 | else: 634 | return self.mag2**(0.5*other) 635 | else: 636 | self._scalar(operator.pow, other) 637 | 638 | # no __rpow__ 639 | 640 | def __lshift__(self, other): 641 | return self._scalar(operator.lshift, other) 642 | 643 | def __rlshift__(self, other): 644 | return self._scalar(operator.lshift, other, True) 645 | 646 | def __rshift__(self, other): 647 | return self._scalar(operator.rshift, other) 648 | 649 | def __rrshift__(self, other): 650 | return self._scalar(operator.rshift, other, True) 651 | 652 | def __and__(self, other): 653 | return self._scalar(operator.and_, other) 654 | 655 | def __rand__(self, other): 656 | return self._scalar(operator.and_, other, True) 657 | 658 | def __or__(self, other): 659 | return self._scalar(operator.or_, other) 660 | 661 | def __ror__(self, other): 662 | return self._scalar(operator.or_, other, True) 663 | 664 | def __xor__(self, other): 665 | return self._scalar(operator.xor, other) 666 | 667 | def __rxor__(self, other): 668 | return self._scalar(operator.xor, other, True) 669 | 670 | def __neg__(self): 671 | return self._unary(operator.neg) 672 | 673 | def __pos__(self): 674 | return self._unary(operator.pos) 675 | 676 | def __abs__(self): 677 | return self.mag 678 | 679 | def __invert__(self): 680 | return self._unary(operator.invert) 681 | 682 | class PtEtaPhiMassMethods(Methods): 683 | _arraymethods = PtEtaPhiMassArrayMethods 684 | 685 | @property 686 | def pt(self): 687 | return self._fPt 688 | 689 | @property 690 | def eta(self): 691 | return self._fEta 692 | 693 | @property 694 | def phi(self): 695 | return self._fPhi 696 | 697 | @property 698 | def mass(self): 699 | return self._fMass 700 | 701 | @property 702 | def mag(self): 703 | return self._fMass 704 | 705 | @property 706 | def mag2(self): 707 | return self._fMass**2 708 | 709 | @property 710 | def mt(self): 711 | out = self.mt2 712 | if out >= 0: 713 | return math.sqrt(out) 714 | else: 715 | return -math.sqrt(out) 716 | 717 | @property 718 | def mt2(self): 719 | return self._fMass**2 + self._fPt**2 720 | 721 | @property 722 | def p3(self): 723 | return uproot3_methods.classes.TVector3.TVector3(self.x, self.y, self.z) 724 | 725 | @property 726 | def x(self): 727 | return self.pt * self.awkward0.numpy.cos(self.phi) 728 | 729 | @property 730 | def y(self): 731 | return self.pt * self.awkward0.numpy.sin(self.phi) 732 | 733 | @property 734 | def z(self): 735 | return self.pt * self.awkward0.numpy.sinh(self.eta) 736 | 737 | @property 738 | def t(self): 739 | x = self.x 740 | y = self.y 741 | z = self.z 742 | mass = self.mass 743 | return self.awkward0.numpy.sqrt(x*x + y*y + z*z + mass*mass*self.awkward0.numpy.sign(mass)) 744 | 745 | def __repr__(self): 746 | return "PtEtaPhiMassLorentzVector(pt={0:.5g}, eta={1:.5g}, phi={2:.5g}, mass={3:.5g})".format(self._fPt, self._fEta, self._fPhi, self._fMass) 747 | 748 | class PtEtaPhiMassLorentzVectorArray(PtEtaPhiMassArrayMethods, uproot3_methods.base.ROOTMethods.awkward0.ObjectArray): 749 | def __init__(self, pt, eta, phi, mass): 750 | if isinstance(pt, awkward0.array.jagged.JaggedArray) or isinstance(eta, awkward0.array.jagged.JaggedArray) or isinstance(phi, awkward0.array.jagged.JaggedArray) or isinstance(mass, awkward0.array.jagged.JaggedArray): 751 | raise TypeError("PtEtaPhiMassLorentzVectorArray constructor arguments must not be jagged; use TLorentzVectorArray.from_ptetaphim for jaggedness-handling") 752 | self._initObjectArray(self.awkward0.Table()) 753 | self["fPt"] = pt 754 | self["fEta"] = eta 755 | self["fPhi"] = phi 756 | self["fMass"] = mass 757 | 758 | @property 759 | def pt(self): 760 | return self["fPt"] 761 | 762 | @pt.setter 763 | def pt(self, value): 764 | self["fPt"] = value 765 | 766 | @property 767 | def eta(self): 768 | return self["fEta"] 769 | 770 | @eta.setter 771 | def eta(self, value): 772 | self["fEta"] = value 773 | 774 | @property 775 | def phi(self): 776 | return self["fPhi"] 777 | 778 | @phi.setter 779 | def phi(self, value): 780 | self["fPhi"] = value 781 | 782 | @property 783 | def mass(self): 784 | return self["fMass"] 785 | 786 | @mass.setter 787 | def mass(self, value): 788 | self["fMass"] = value 789 | 790 | class TLorentzVectorArray(ArrayMethods, uproot3_methods.base.ROOTMethods.awkward0.ObjectArray): 791 | 792 | def __init__(self, x, y, z, t): 793 | if isinstance(x, awkward0.array.jagged.JaggedArray) or isinstance(y, awkward0.array.jagged.JaggedArray) or isinstance(z, awkward0.array.jagged.JaggedArray) or isinstance(t, awkward0.array.jagged.JaggedArray): 794 | raise TypeError("TLorentzVectorArray constructor arguments must not be jagged; use TLorentzVectorArray.from_cartesian for jaggedness-handling") 795 | self._initObjectArray(self.awkward0.Table()) 796 | self["fX"] = x 797 | self["fY"] = y 798 | self["fZ"] = z 799 | self["fE"] = t 800 | 801 | @classmethod 802 | def origin(cls, shape, dtype=None): 803 | if dtype is None: 804 | dtype = cls.awkward0.numpy.float64 805 | return cls(cls.awkward0.numpy.zeros(shape, dtype=dtype), 806 | cls.awkward0.numpy.zeros(shape, dtype=dtype), 807 | cls.awkward0.numpy.zeros(shape, dtype=dtype), 808 | cls.awkward0.numpy.zeros(shape, dtype=dtype)) 809 | 810 | @classmethod 811 | def origin_like(cls, array): 812 | return array * 0.0 813 | 814 | @classmethod 815 | def from_p3(cls, p3, t): 816 | return cls.from_cartesian(p3.x, p3.y, p3.z, t) 817 | 818 | @classmethod 819 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 820 | def from_cartesian(cls, x, y, z, t): 821 | return cls(x, y, z, t) 822 | 823 | @classmethod 824 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 825 | def from_spherical(cls, r, theta, phi, t): 826 | return cls.from_p3(uproot3_methods.classes.TVector3.TVector3Array.from_spherical(r, theta, phi), t) 827 | 828 | @classmethod 829 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 830 | def from_cylindrical(cls, rho, phi, z, t): 831 | return cls.from_p3(uproot3_methods.classes.TVector3.TVector3Array.from_cylindrical(rho, phi, z), t) 832 | 833 | @classmethod 834 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 835 | def from_xyzm(cls, x, y, z, m): 836 | return cls(x, y, z, cls.awkward0.numpy.sqrt(x*x + y*y + z*z + m*m*cls.awkward0.numpy.sign(m))) 837 | 838 | @classmethod 839 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 840 | def from_ptetaphi(cls, pt, eta, phi, energy): 841 | out = cls(pt * cls.awkward0.numpy.cos(phi), 842 | pt * cls.awkward0.numpy.sin(phi), 843 | pt * cls.awkward0.numpy.sinh(eta), 844 | energy) 845 | out._memo_pt = pt 846 | out._memo_eta = eta 847 | out._memo_phi = phi 848 | return out 849 | 850 | @classmethod 851 | def from_ptetaphie(cls, pt, eta, phi, energy): 852 | return cls.from_ptetaphi(pt, eta, phi, energy) 853 | 854 | @classmethod 855 | @awkward0.util.wrapjaggedmethod(PtEtaPhiMassJaggedArrayMethods) 856 | def from_ptetaphim(cls, pt, eta, phi, mass): 857 | return PtEtaPhiMassLorentzVectorArray(pt,eta,phi,mass) 858 | 859 | @property 860 | def x(self): 861 | return self["fX"] 862 | 863 | @x.setter 864 | def x(self, value): 865 | self["fX"] = value 866 | 867 | @property 868 | def y(self): 869 | return self["fY"] 870 | 871 | @y.setter 872 | def y(self, value): 873 | self["fY"] = value 874 | 875 | @property 876 | def z(self): 877 | return self["fZ"] 878 | 879 | @z.setter 880 | def z(self, value): 881 | self["fZ"] = value 882 | 883 | @property 884 | def t(self): 885 | return self["fE"] 886 | 887 | @t.setter 888 | def t(self, value): 889 | self["fE"] = value 890 | 891 | @property 892 | def E(self): 893 | return self["fE"] 894 | 895 | @E.setter 896 | def E(self, value): 897 | self["fE"] = value 898 | 899 | class PtEtaPhiMassLorentzVector(PtEtaPhiMassMethods): 900 | def __init__(self, pt, eta, phi, mass): 901 | self._fPt = float(pt) 902 | self._fEta = float(eta) 903 | self._fPhi = float(phi) 904 | self._fMass = float(mass) 905 | 906 | @property 907 | def pt(self): 908 | return self._fPt 909 | 910 | @pt.setter 911 | def pt(self,value): 912 | self._fPt = value 913 | 914 | @property 915 | def eta(self): 916 | return self._fEta 917 | 918 | @eta.setter 919 | def eta(self, value): 920 | self._fEta = value 921 | 922 | @property 923 | def phi(self): 924 | return self._fPhi 925 | 926 | @phi.setter 927 | def phi(self, value): 928 | self._fPhi = value 929 | 930 | @property 931 | def mass(self): 932 | return self._fMass 933 | 934 | @mass.setter 935 | def mass(self, value): 936 | self._fMass = value 937 | 938 | class TLorentzVector(Methods): 939 | def __init__(self, x, y, z, t): 940 | self._fP = uproot3_methods.classes.TVector3.TVector3(float(x), float(y), float(z)) 941 | self._fE = float(t) 942 | 943 | @classmethod 944 | def origin(cls): 945 | return cls(0.0, 0.0, 0.0, 0.0) 946 | 947 | @classmethod 948 | def from_p3(cls, p3, t): 949 | out = cls.__new__(cls) 950 | out._fP = p3 951 | out._fE = t 952 | return out 953 | 954 | @classmethod 955 | def from_spherical(cls, r, theta, phi, t): 956 | return cls.from_p3(uproot3_methods.classes.TVector3.Methods.from_spherical(r, theta, phi), t) 957 | 958 | @classmethod 959 | def from_cylindrical(cls, rho, phi, z, t): 960 | return cls.from_p3(uproot3_methods.classes.TVector3.Methods.from_cylindrical(rho, phi, z), t) 961 | 962 | @classmethod 963 | def from_xyzm(cls, x, y, z, m): 964 | return cls(x, y, z, math.sqrt(x*x + y*y + z*z + m*m*(1 if m >= 0 else -1))) 965 | 966 | @classmethod 967 | def from_ptetaphi(cls, pt, eta, phi, energy): 968 | return cls(pt * math.cos(phi), 969 | pt * math.sin(phi), 970 | pt * math.sinh(eta), 971 | energy) 972 | 973 | @classmethod 974 | def from_ptetaphie(cls, pt, eta, phi, energy): 975 | return cls.from_ptetaphi(pt, eta, phi, energy) 976 | 977 | @classmethod 978 | def from_ptetaphim(cls, pt, eta, phi, mass): 979 | return PtEtaPhiMassLorentzVector(pt,eta,phi,mass) 980 | 981 | @property 982 | def x(self): 983 | return self._fP._fX 984 | 985 | @x.setter 986 | def x(self, value): 987 | self._fP._fX = value 988 | 989 | @property 990 | def y(self): 991 | return self._fP._fY 992 | 993 | @y.setter 994 | def y(self, value): 995 | self._fP._fY = value 996 | 997 | @property 998 | def z(self): 999 | return self._fP._fZ 1000 | 1001 | @z.setter 1002 | def z(self, value): 1003 | self._fP._fZ = value 1004 | 1005 | @property 1006 | def t(self): 1007 | return self._fE 1008 | 1009 | @t.setter 1010 | def t(self, value): 1011 | self._fE = value 1012 | 1013 | @property 1014 | def E(self): 1015 | return self._fE 1016 | 1017 | @E.setter 1018 | def E(self, value): 1019 | self._fE = value 1020 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TParameter.py: -------------------------------------------------------------------------------- 1 | import uproot3_methods.base 2 | 3 | 4 | def _decode(seq): 5 | if isinstance(seq, bytes): 6 | return seq.decode("UTF-8") 7 | else: 8 | return seq 9 | 10 | 11 | class Methods(uproot3_methods.base.ROOTMethods): 12 | def __repr__(self): 13 | if self._fName is None: 14 | return "<{0} at 0x{1:012x}>".format(_decode(self._classname), id(self)) 15 | else: 16 | return "<{0} {1} 0x{2:012x}>".format( 17 | _decode(self._classname), _decode(self._fName), id(self) 18 | ) 19 | 20 | def __str__(self): 21 | return str(self._fVal) 22 | 23 | @property 24 | def name(self): 25 | return _decode(self._fName) 26 | 27 | @property 28 | def value(self): 29 | return self._fVal 30 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TVector2.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import math 6 | import numbers 7 | 8 | import awkward0.array.jagged 9 | import awkward0.util 10 | 11 | import uproot3_methods.common.TVector 12 | import uproot3_methods.base 13 | 14 | class Common(object): 15 | def dot(self, other): 16 | out = self.x*other.x 17 | out = out + self.y*other.y 18 | return out 19 | 20 | def _rotate(self, angle): 21 | c = self.awkward0.numpy.cos(angle) 22 | s = self.awkward0.numpy.sin(angle) 23 | x = self.x*c - self.y*s 24 | y = self.x*s + self.y*c 25 | return x, y 26 | 27 | class ArrayMethods(Common, uproot3_methods.common.TVector.ArrayMethods, uproot3_methods.base.ROOTMethods): 28 | def _initObjectArray(self, table): 29 | self.awkward0.ObjectArray.__init__(self, table, lambda row: TVector2(row["fX"], row["fY"])) 30 | 31 | def __awkward_serialize__(self, serializer): 32 | self._valid() 33 | x, y = self.x, self.y 34 | return serializer.encode_call( 35 | ["uproot3_methods.classes.TVector2", "TVector2Array", "from_cartesian"], 36 | serializer(x, "TVector3Array.x"), 37 | serializer(y, "TVector3Array.y")) 38 | 39 | @property 40 | def x(self): 41 | return self["fX"] 42 | 43 | @property 44 | def y(self): 45 | return self["fY"] 46 | 47 | def rotate(self, angle): 48 | x, y = self._rotate(angle) 49 | out = self.empty_like() 50 | out["fX"] = x 51 | out["fY"] = y 52 | return out 53 | 54 | def sum(self): 55 | if isinstance(self, self.awkward0.JaggedArray): 56 | return TVector2Array.from_cartesian(self.x.sum(), self.y.sum()) 57 | else: 58 | return TVector2(self.x.sum(), self.y.sum()) 59 | 60 | def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): 61 | if "out" in kwargs: 62 | raise NotImplementedError("in-place operations not supported") 63 | 64 | if method != "__call__": 65 | return NotImplemented 66 | 67 | inputs = list(inputs) 68 | for i in range(len(inputs)): 69 | if isinstance(inputs[i], self.awkward0.numpy.ndarray) and inputs[i].dtype == self.awkward0.numpy.dtype(object) and len(inputs[i]) > 0: 70 | idarray = self.awkward0.numpy.frombuffer(inputs[i], dtype=self.awkward0.numpy.uintp) 71 | if (idarray == idarray[0]).all(): 72 | inputs[i] = inputs[i][0] 73 | 74 | if ufunc is self.awkward0.numpy.add or ufunc is self.awkward0.numpy.subtract: 75 | if not all(isinstance(x, (ArrayMethods, Methods)) for x in inputs): 76 | raise TypeError("(arrays of) TVector2 can only be added to/subtracted from other (arrays of) TVector2") 77 | out = self.empty_like() 78 | out["fX"] = getattr(ufunc, method)(*[x.x for x in inputs], **kwargs) 79 | out["fY"] = getattr(ufunc, method)(*[x.y for x in inputs], **kwargs) 80 | return out 81 | 82 | elif ufunc is self.awkward0.numpy.power and len(inputs) >= 2 and isinstance(inputs[1], (numbers.Number, self.awkward0.numpy.number)): 83 | if inputs[1] == 2: 84 | return self.mag2 85 | else: 86 | return self.mag2**(0.5*inputs[1]) 87 | 88 | elif ufunc is self.awkward0.numpy.absolute: 89 | return self.mag 90 | 91 | else: 92 | return super(ArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs) 93 | 94 | JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward0.JaggedArray) 95 | 96 | class Methods(Common, uproot3_methods.common.TVector.Methods, uproot3_methods.base.ROOTMethods): 97 | _arraymethods = ArrayMethods 98 | 99 | @property 100 | def x(self): 101 | return self._fX 102 | 103 | @property 104 | def y(self): 105 | return self._fY 106 | 107 | def __repr__(self): 108 | return "TVector2({0:.5g}, {1:.5g})".format(self.x, self.y) 109 | 110 | def __str__(self): 111 | return repr(self) 112 | 113 | def __eq__(self, other): 114 | return isinstance(other, Methods) and self.x == other.x and self.y == other.y 115 | 116 | def _scalar(self, operator, scalar, reverse=False): 117 | if not isinstance(scalar, (numbers.Number, self.awkward0.numpy.number)): 118 | raise TypeError("cannot {0} a TVector2 with a {1}".format(operator.__name__, type(scalar).__name__)) 119 | if reverse: 120 | return TVector2(operator(scalar, self.x), operator(scalar, self.y)) 121 | else: 122 | return TVector2(operator(self.x, scalar), operator(self.y, scalar)) 123 | 124 | def _vector(self, operator, vector, reverse=False): 125 | if not isinstance(vector, Methods): 126 | raise TypeError("cannot {0} a TVector2 with a {1}".format(operator.__name__, type(vector).__name__)) 127 | if reverse: 128 | return TVector2(operator(vector.x, self.x), operator(vector.y, self.y)) 129 | else: 130 | return TVector2(operator(self.x, vector.x), operator(self.y, vector.y)) 131 | 132 | def _unary(self, operator): 133 | return TVector2(operator(self.x), operator(self.y)) 134 | 135 | def rotate(self, angle): 136 | x, y = self._rotate(angle) 137 | return TVector2(x, y) 138 | 139 | class TVector2Array(ArrayMethods, uproot3_methods.base.ROOTMethods.awkward0.ObjectArray): 140 | 141 | def __init__(self, x, y): 142 | if isinstance(x, awkward0.array.jagged.JaggedArray) or isinstance(y, awkward0.array.jagged.JaggedArray): 143 | raise TypeError("TVector2Array constructor arguments must not be jagged; use TVector2.from_cartesian for jaggedness-handling") 144 | self._initObjectArray(self.awkward0.Table()) 145 | self["fX"] = x 146 | self["fY"] = y 147 | 148 | @classmethod 149 | def origin(cls, shape, dtype=None): 150 | if dtype is None: 151 | dtype = cls.awkward0.numpy.float64 152 | return cls(cls.awkward0.numpy.zeros(shape, dtype=dtype), cls.awkward0.numpy.zeros(shape, dtype=dtype)) 153 | 154 | @classmethod 155 | def origin_like(cls, array): 156 | return array * 0.0 157 | 158 | @classmethod 159 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 160 | def from_cartesian(cls, x, y): 161 | return cls(x, y) 162 | 163 | @classmethod 164 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 165 | def from_polar(cls, rho, phi): 166 | return cls(rho * cls.awkward0.numpy.cos(phi), 167 | rho * cls.awkward0.numpy.sin(phi)) 168 | 169 | @property 170 | def x(self): 171 | return self["fX"] 172 | 173 | @x.setter 174 | def x(self, value): 175 | self["fX"] = value 176 | 177 | @property 178 | def y(self): 179 | return self["fY"] 180 | 181 | @y.setter 182 | def y(self, value): 183 | self["fY"] = value 184 | 185 | class TVector2(Methods): 186 | def __init__(self, x, y): 187 | self._fX = float(x) 188 | self._fY = float(y) 189 | 190 | @classmethod 191 | def origin(cls): 192 | return cls(0.0, 0.0) 193 | 194 | @classmethod 195 | def from_polar(cls, rho, phi): 196 | return cls(rho * math.cos(phi), 197 | rho * math.sin(phi)) 198 | 199 | @property 200 | def x(self): 201 | return self._fX 202 | 203 | @x.setter 204 | def x(self, value): 205 | self._fX = value 206 | 207 | @property 208 | def y(self): 209 | return self._fY 210 | 211 | @y.setter 212 | def y(self, value): 213 | self._fY = value 214 | -------------------------------------------------------------------------------- /uproot3_methods/classes/TVector3.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import math 6 | import numbers 7 | 8 | import awkward0.array.jagged 9 | import awkward0.util 10 | 11 | import uproot3_methods.base 12 | import uproot3_methods.common.TVector 13 | 14 | class Common(object): 15 | def dot(self, other): 16 | out = self.x*other.x 17 | out = out + self.y*other.y 18 | out = out + self.z*other.z 19 | return out 20 | 21 | def _cross(self, other): 22 | return (self.y*other.z - self.z*other.y, 23 | self.z*other.x - self.x*other.z, 24 | self.x*other.y - self.y*other.x) 25 | 26 | @property 27 | def cottheta(self): 28 | out = self.rho 29 | out /= self.z 30 | return out 31 | 32 | def _rotate_axis(self, axis, angle): 33 | u = axis.unit 34 | c = self.awkward0.numpy.cos(angle) 35 | s = self.awkward0.numpy.sin(angle) 36 | c1 = 1 - c 37 | 38 | x = (c + u.x**2 * c1) * self.x + (u.x * u.y * c1 - u.z * s) * self.y + (u.x * u.z * c1 + u.y * s) * self.z 39 | y = (u.x * u.y * c1 + u.z * s) * self.x + (c + u.y**2 * c1) * self.y + (u.y * u.z * c1 - u.x * s) * self.z 40 | z = (u.x * u.z * c1 - u.y * s) * self.x + (u.y * u.z * c1 + u.x * s) * self.y + (c + u.z**2 * c1) * self.z 41 | 42 | return x, y, z 43 | 44 | def _rotate_euler(self, phi, theta, psi): 45 | # Rotate Z (phi) 46 | c1 = self.awkward0.numpy.cos(phi) 47 | s1 = self.awkward0.numpy.sin(phi) 48 | c2 = self.awkward0.numpy.cos(theta) 49 | s2 = self.awkward0.numpy.sin(theta) 50 | c3 = self.awkward0.numpy.cos(psi) 51 | s3 = self.awkward0.numpy.sin(psi) 52 | 53 | # Rotate Y (theta) 54 | fzx2 = -s2*c1 55 | fzy2 = s2*s1 56 | fzz2 = c2 57 | 58 | # Rotate Z (psi) 59 | fxx3 = c3*c2*c1 - s3*s1 60 | fxy3 = -c3*c2*s1 - s3*c1 61 | fxz3 = c3*s2 62 | fyx3 = s3*c2*c1 + c3*s1 63 | fyy3 = -s3*c2*s1 + c3*c1 64 | fyz3 = s3*s2 65 | 66 | # Transform v 67 | x = fxx3*self.x + fxy3*self.y + fxz3*self.z 68 | y = fyx3*self.x + fyy3*self.y + fyz3*self.z 69 | z = fzx2*self.x + fzy2*self.y + fzz2*self.z 70 | 71 | return x, y, z 72 | 73 | def rotatex(self, angle): 74 | return self.rotate_axis(TVector3(1.0, 0.0, 0.0), angle) 75 | 76 | def rotatey(self, angle): 77 | return self.rotate_axis(TVector3(0.0, 1.0, 0.0), angle) 78 | 79 | def rotatez(self, angle): 80 | return self.rotate_axis(TVector3(0.0, 0.0, 1.0), angle) 81 | 82 | class ArrayMethods(Common, uproot3_methods.common.TVector.ArrayMethods, uproot3_methods.base.ROOTMethods): 83 | def _initObjectArray(self, table): 84 | self.awkward0.ObjectArray.__init__(self, table, lambda row: TVector3(row["fX"], row["fY"], row["fZ"])) 85 | 86 | def __awkward_serialize__(self, serializer): 87 | self._valid() 88 | x, y, z = self.x, self.y, self.z 89 | return serializer.encode_call( 90 | ["uproot3_methods.classes.TVector3", "TVector3Array", "from_cartesian"], 91 | serializer(x, "TVector3Array.x"), 92 | serializer(y, "TVector3Array.y"), 93 | serializer(z, "TVector3Array.z")) 94 | 95 | @property 96 | def x(self): 97 | return self["fX"] 98 | 99 | @property 100 | def y(self): 101 | return self["fY"] 102 | 103 | @property 104 | def z(self): 105 | return self["fZ"] 106 | 107 | def cross(self, other): 108 | x, y, z = self._cross(other) 109 | out = self.empty_like() 110 | out["fX"] = x 111 | out["fY"] = y 112 | out["fZ"] = z 113 | return out 114 | 115 | @property 116 | def theta(self): 117 | return self.awkward0.numpy.arctan2(self.rho, self.z) 118 | 119 | def rotate_axis(self, axis, angle): 120 | x, y, z = self._rotate_axis(axis, angle) 121 | out = self.empty_like() 122 | out["fX"] = x 123 | out["fY"] = y 124 | out["fZ"] = z 125 | return out 126 | 127 | def rotate_euler(self, phi=0, theta=0, psi=0): 128 | x, y, z = self._rotate_euler(phi, theta, psi) 129 | out = self.empty_like() 130 | out["fX"] = x 131 | out["fY"] = y 132 | out["fZ"] = z 133 | return out 134 | 135 | def sum(self): 136 | if isinstance(self, self.awkward0.JaggedArray): 137 | return TVector3Array.from_cartesian(self.x.sum(), self.y.sum(), self.z.sum()) 138 | else: 139 | return TVector3(self.x.sum(), self.y.sum(), self.z.sum()) 140 | 141 | def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): 142 | if "out" in kwargs: 143 | raise NotImplementedError("in-place operations not supported") 144 | 145 | if method != "__call__": 146 | return NotImplemented 147 | 148 | inputs = list(inputs) 149 | for i in range(len(inputs)): 150 | if isinstance(inputs[i], self.awkward0.numpy.ndarray) and inputs[i].dtype == self.awkward0.numpy.dtype(object) and len(inputs[i]) > 0: 151 | idarray = self.awkward0.numpy.frombuffer(inputs[i], dtype=self.awkward0.numpy.uintp) 152 | if (idarray == idarray[0]).all(): 153 | inputs[i] = inputs[i][0] 154 | 155 | if ufunc is self.awkward0.numpy.add or ufunc is self.awkward0.numpy.subtract: 156 | if not all(isinstance(x, (ArrayMethods, Methods)) for x in inputs): 157 | raise TypeError("(arrays of) TVector3 can only be added to/subtracted from other (arrays of) TVector3") 158 | out = self.empty_like() 159 | out["fX"] = getattr(ufunc, method)(*[x.x for x in inputs], **kwargs) 160 | out["fY"] = getattr(ufunc, method)(*[x.y for x in inputs], **kwargs) 161 | out["fZ"] = getattr(ufunc, method)(*[x.z for x in inputs], **kwargs) 162 | return out 163 | 164 | elif ufunc is self.awkward0.numpy.power and len(inputs) >= 2 and isinstance(inputs[1], (numbers.Number, self.awkward0.numpy.number)): 165 | if inputs[1] == 2: 166 | return self.mag2 167 | else: 168 | return self.mag2**(0.5*inputs[1]) 169 | 170 | elif ufunc is self.awkward0.numpy.absolute: 171 | return self.mag 172 | 173 | else: 174 | return super(ArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs) 175 | 176 | JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward0.JaggedArray) 177 | 178 | class Methods(Common, uproot3_methods.common.TVector.Methods, uproot3_methods.base.ROOTMethods): 179 | _arraymethods = ArrayMethods 180 | 181 | @property 182 | def x(self): 183 | return self._fX 184 | 185 | @property 186 | def y(self): 187 | return self._fY 188 | 189 | @property 190 | def z(self): 191 | return self._fZ 192 | 193 | def __repr__(self): 194 | return "TVector3({0:.5g}, {1:.5g}, {2:.5g})".format(self.x, self.y, self.z) 195 | 196 | def __str__(self): 197 | return repr(self) 198 | 199 | def __eq__(self, other): 200 | return isinstance(other, Methods) and self.x == other.x and self.y == other.y and self.z == other.z 201 | 202 | def _scalar(self, operator, scalar, reverse=False): 203 | if not isinstance(scalar, (numbers.Number, self.awkward0.numpy.number)): 204 | raise TypeError("cannot {0} a TVector3 with a {1}".format(operator.__name__, type(scalar).__name__)) 205 | if reverse: 206 | return TVector3(operator(scalar, self.x), operator(scalar, self.y), operator(scalar, self.z)) 207 | else: 208 | return TVector3(operator(self.x, scalar), operator(self.y, scalar), operator(self.z, scalar)) 209 | 210 | def _vector(self, operator, vector, reverse=False): 211 | if not isinstance(vector, Methods): 212 | raise TypeError("cannot {0} a TVector3 with a {1}".format(operator.__name__, type(vector).__name__)) 213 | if reverse: 214 | return TVector3(operator(vector.x, self.x), operator(vector.y, self.y), operator(vector.z, self.z)) 215 | else: 216 | return TVector3(operator(self.x, vector.x), operator(self.y, vector.y), operator(self.z, vector.z)) 217 | 218 | def _unary(self, operator): 219 | return TVector3(operator(self.x), operator(self.y), operator(self.z)) 220 | 221 | def cross(self, other): 222 | x, y, z = self._cross(other) 223 | return TVector3(x, y, z) 224 | 225 | def angle(self, other): 226 | denominator = math.sqrt(self.dot(self) * other.dot(other)) 227 | if denominator == 0: 228 | # one of the vector is null 229 | return 0. 230 | cos_angle = self.dot(other) / denominator 231 | if cos_angle > 1: 232 | cos_angle = 1 233 | elif cos_angle < -1: 234 | cos_angle = -1 235 | return math.acos(cos_angle) 236 | 237 | @property 238 | def theta(self): 239 | return math.atan2(self.rho, self.z) 240 | 241 | def rotate_axis(self, axis, angle): 242 | x, y, z = self._rotate_axis(axis, angle) 243 | return TVector3(x, y, z) 244 | 245 | def rotate_euler(self, phi=0, theta=0, psi=0): 246 | return TVector3(x, y, z) 247 | 248 | class TVector3Array(ArrayMethods, uproot3_methods.base.ROOTMethods.awkward0.ObjectArray): 249 | 250 | def __init__(self, x, y, z): 251 | if isinstance(x, awkward0.array.jagged.JaggedArray) or isinstance(y, awkward0.array.jagged.JaggedArray) or isinstance(z, awkward0.array.jagged.JaggedArray): 252 | raise TypeError("TVector3Array constructor arguments must not be jagged; use TVector3.from_cartesian for jaggedness-handling") 253 | self._initObjectArray(self.awkward0.Table()) 254 | self["fX"] = x 255 | self["fY"] = y 256 | self["fZ"] = z 257 | 258 | @classmethod 259 | def origin(cls, shape, dtype=None): 260 | if dtype is None: 261 | dtype = cls.awkward0.numpy.float64 262 | return cls(cls.awkward0.numpy.zeros(shape, dtype=dtype), 263 | cls.awkward0.numpy.zeros(shape, dtype=dtype), 264 | cls.awkward0.numpy.zeros(shape, dtype=dtype)) 265 | 266 | @classmethod 267 | def origin_like(cls, array): 268 | return array * 0.0 269 | 270 | @classmethod 271 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 272 | def from_cartesian(cls, x, y, z): 273 | return cls(x, y, z) 274 | 275 | @classmethod 276 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 277 | def from_spherical(cls, r, theta, phi): 278 | return cls(r * cls.awkward0.numpy.sin(theta) * cls.awkward0.numpy.cos(phi), 279 | r * cls.awkward0.numpy.sin(theta) * cls.awkward0.numpy.sin(phi), 280 | r * cls.awkward0.numpy.cos(theta)) 281 | 282 | @classmethod 283 | @awkward0.util.wrapjaggedmethod(JaggedArrayMethods) 284 | def from_cylindrical(cls, rho, phi, z): 285 | return cls(rho * cls.awkward0.numpy.cos(phi), rho * cls.awkward0.numpy.sin(phi),z) 286 | 287 | @property 288 | def x(self): 289 | return self["fX"] 290 | 291 | @x.setter 292 | def x(self, value): 293 | self["fX"] = value 294 | 295 | @property 296 | def y(self): 297 | return self["fY"] 298 | 299 | @y.setter 300 | def y(self, value): 301 | self["fY"] = value 302 | 303 | @property 304 | def z(self): 305 | return self["fZ"] 306 | 307 | @z.setter 308 | def z(self, value): 309 | self["fZ"] = value 310 | 311 | class TVector3(Methods): 312 | def __init__(self, x, y, z): 313 | self._fX = float(x) 314 | self._fY = float(y) 315 | self._fZ = float(z) 316 | 317 | @classmethod 318 | def origin(cls): 319 | return cls(0.0, 0.0, 0.0) 320 | 321 | @classmethod 322 | def from_spherical(cls, r, theta, phi): 323 | return cls(r * math.sin(theta) * math.cos(phi), 324 | r * math.sin(theta) * math.sin(phi), 325 | r * math.cos(theta)) 326 | 327 | @classmethod 328 | def from_cylindrical(cls, rho, phi, z): 329 | return cls(rho * math.cos(phi), 330 | rho * math.sin(phi), 331 | z) 332 | 333 | @property 334 | def x(self): 335 | return self._fX 336 | 337 | @x.setter 338 | def x(self, value): 339 | self._fX = value 340 | 341 | @property 342 | def y(self): 343 | return self._fY 344 | 345 | @y.setter 346 | def y(self, value): 347 | self._fY = value 348 | 349 | @property 350 | def z(self): 351 | return self._fZ 352 | 353 | @z.setter 354 | def z(self, value): 355 | self._fZ = value 356 | -------------------------------------------------------------------------------- /uproot3_methods/classes/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | def hasmethods(name): 6 | if name not in globals(): 7 | if name in hasmethods.loaders: 8 | globals()[name] = hasmethods.loaders[name].load_module(name) 9 | elif '_3c_' in name and '_3e_' in name: 10 | bare_name = name.split('_3c_')[0] 11 | if bare_name in hasmethods.loaders: 12 | globals()[name] = hasmethods.loaders[bare_name].load_module(bare_name) 13 | 14 | return name in globals() and isinstance(getattr(globals()[name], "Methods", None), type) 15 | 16 | import pkgutil 17 | 18 | hasmethods.loaders = dict([(module_name, loader.find_module(module_name)) for loader, module_name, is_pkg in pkgutil.walk_packages(__path__)]) 19 | 20 | del pkgutil 21 | -------------------------------------------------------------------------------- /uproot3_methods/common/TVector.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import math 6 | import numbers 7 | import operator 8 | 9 | class Common(object): 10 | @property 11 | def mag2(self): 12 | return self.dot(self) 13 | 14 | @property 15 | def mag(self): 16 | return self.awkward0.numpy.sqrt(self.mag2) 17 | 18 | @property 19 | def rho2(self): 20 | out = self.x*self.x 21 | out = out + self.y*self.y 22 | return out 23 | 24 | def delta_phi(self, other): 25 | return (self.phi - other.phi + math.pi) % (2*math.pi) - math.pi 26 | 27 | def isparallel(self, other, tolerance=1e-10): 28 | return 1 - self.cosdelta(other) < tolerance 29 | 30 | def isantiparallel(self, other, tolerance=1e-10): 31 | return self.cosdelta(other) - (-1) < tolerance 32 | 33 | def iscollinear(self, other, tolerance=1e-10): 34 | return 1 - self.awkward0.numpy.absolute(self.cosdelta(other)) < tolerance 35 | 36 | def __lt__(self, other): 37 | raise TypeError("spatial vectors have no natural ordering") 38 | 39 | def __gt__(self, other): 40 | raise TypeError("spatial vectors have no natural ordering") 41 | 42 | def __le__(self, other): 43 | raise TypeError("spatial vectors have no natural ordering") 44 | 45 | def __ge__(self, other): 46 | raise TypeError("spatial vectors have no natural ordering") 47 | 48 | class ArrayMethods(Common): 49 | @property 50 | def unit(self): 51 | return self / self.mag 52 | 53 | @property 54 | def rho(self): 55 | out = self.rho2 56 | return self.awkward0.numpy.sqrt(out) 57 | 58 | @property 59 | def phi(self): 60 | return self.awkward0.numpy.arctan2(self.y, self.x) 61 | 62 | def cosdelta(self, other): 63 | denom = self.mag2 * other.mag2 64 | mask = (denom > 0) 65 | denom = denom[mask] 66 | denom[:] = self.awkward0.numpy.sqrt(denom) 67 | 68 | out = self.dot(other) 69 | out[mask] /= denom 70 | 71 | mask = self.awkward0.numpy.logical_not(mask) 72 | out[mask] = 1 73 | 74 | return self.awkward0.numpy.clip(out, -1, 1) 75 | 76 | def angle(self, other, normal=None, degrees=False): 77 | out = self.awkward0.numpy.arccos(self.cosdelta(other)) 78 | if normal is not None: 79 | a = self.unit 80 | b = other.unit 81 | out = out * self.awkward0.numpy.sign(normal.dot(a.cross(b))) 82 | if degrees: 83 | out = self.awkward0.numpy.multiply(out, 180.0/self.awkward0.numpy.pi) 84 | return out 85 | 86 | def isopposite(self, other, tolerance=1e-10): 87 | tmp = self + other 88 | tmp.x = self.awkward0.numpy.absolute(tmp.x) 89 | tmp.y = self.awkward0.numpy.absolute(tmp.y) 90 | tmp.z = self.awkward0.numpy.absolute(tmp.z) 91 | 92 | out = (tmp.x < tolerance) 93 | out = self.awkward0.numpy.bitwise_and(out, tmp.y < tolerance) 94 | out = self.awkward0.numpy.bitwise_and(out, tmp.z < tolerance) 95 | return out 96 | 97 | def isperpendicular(self, other, tolerance=1e-10): 98 | tmp = self.dot(other) 99 | tmp.x = self.awkward0.numpy.absolute(tmp.x) 100 | tmp.y = self.awkward0.numpy.absolute(tmp.y) 101 | tmp.z = self.awkward0.numpy.absolute(tmp.z) 102 | 103 | out = (tmp.x < tolerance) 104 | out = self.awkward0.numpy.bitwise_and(out, tmp.y < tolerance) 105 | out = self.awkward0.numpy.bitwise_and(out, tmp.z < tolerance) 106 | return out 107 | 108 | class Methods(Common): 109 | @property 110 | def unit(self): 111 | return self / self.mag 112 | 113 | @property 114 | def rho(self): 115 | return math.sqrt(self.rho2) 116 | 117 | @property 118 | def phi(self): 119 | return math.atan2(self.y, self.x) 120 | 121 | def cosdelta(self, other): 122 | m1 = self.mag2 123 | m2 = other.mag2 124 | if m1 == 0 or m2 == 0: 125 | return 1.0 126 | r = self.dot(other) / math.sqrt(m1 * m2) 127 | return max(-1.0, min(1.0, r)) 128 | 129 | def angle(self, other, degrees=False): 130 | out = math.acos(self.cosdelta(other)) 131 | if degrees: 132 | out = out * 180.0/math.pi 133 | return out 134 | 135 | def isopposite(self, other, tolerance=1e-10): 136 | tmp = self + other 137 | return abs(tmp.x) < tolerance and abs(tmp.y) < tolerance and abs(tmp.z) < tolerance 138 | 139 | def isperpendicular(self, other, tolerance=1e-10): 140 | tmp = self.dot(other) 141 | return abs(tmp.x) < tolerance and abs(tmp.y) < tolerance and abs(tmp.z) < tolerance 142 | 143 | def __add__(self, other): 144 | return self._vector(operator.add, other) 145 | 146 | def __radd__(self, other): 147 | return self._vector(operator.add, other, True) 148 | 149 | def __sub__(self, other): 150 | return self._vector(operator.sub, other) 151 | 152 | def __rsub__(self, other): 153 | return self._vector(operator.sub, other, True) 154 | 155 | def __mul__(self, other): 156 | return self._scalar(operator.mul, other) 157 | 158 | def __rmul__(self, other): 159 | return self._scalar(operator.mul, other, True) 160 | 161 | def __div__(self, other): 162 | return self._scalar(operator.div, other) 163 | 164 | def __rdiv__(self, other): 165 | return self._scalar(operator.div, other, True) 166 | 167 | def __truediv__(self, other): 168 | return self._scalar(operator.truediv, other) 169 | 170 | def __rtruediv__(self, other): 171 | return self._scalar(operator.truediv, other, True) 172 | 173 | def __floordiv__(self, other): 174 | return self._scalar(operator.floordiv, other) 175 | 176 | def __rfloordiv__(self, other): 177 | return self._scalar(operator.floordiv, other, True) 178 | 179 | def __mod__(self, other): 180 | return self._scalar(operator.mod, other) 181 | 182 | def __rmod__(self, other): 183 | return self._scalar(operator.mod, other, True) 184 | 185 | def __divmod__(self, other): 186 | return self._scalar(operator.divmod, other) 187 | 188 | def __rdivmod__(self, other): 189 | return self._scalar(operator.divmod, other, True) 190 | 191 | def __pow__(self, other): 192 | if isinstance(other, (numbers.Number, self.awkward0.numpy.number)): 193 | if other == 2: 194 | return self.mag2 195 | else: 196 | return self.mag2**(0.5*other) 197 | else: 198 | self._scalar(operator.pow, other) 199 | 200 | # no __rpow__ 201 | 202 | def __lshift__(self, other): 203 | return self._scalar(operator.lshift, other) 204 | 205 | def __rlshift__(self, other): 206 | return self._scalar(operator.lshift, other, True) 207 | 208 | def __rshift__(self, other): 209 | return self._scalar(operator.rshift, other) 210 | 211 | def __rrshift__(self, other): 212 | return self._scalar(operator.rshift, other, True) 213 | 214 | def __and__(self, other): 215 | return self._scalar(operator.and_, other) 216 | 217 | def __rand__(self, other): 218 | return self._scalar(operator.and_, other, True) 219 | 220 | def __or__(self, other): 221 | return self._scalar(operator.or_, other) 222 | 223 | def __ror__(self, other): 224 | return self._scalar(operator.or_, other, True) 225 | 226 | def __xor__(self, other): 227 | return self._scalar(operator.xor, other) 228 | 229 | def __rxor__(self, other): 230 | return self._scalar(operator.xor, other, True) 231 | 232 | def __neg__(self): 233 | return self._unary(operator.neg) 234 | 235 | def __pos__(self): 236 | return self._unary(operator.pos) 237 | 238 | def __abs__(self): 239 | return self.mag 240 | 241 | def __invert__(self): 242 | return self._unary(operator.invert) 243 | -------------------------------------------------------------------------------- /uproot3_methods/common/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | -------------------------------------------------------------------------------- /uproot3_methods/convert.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import importlib 6 | 7 | # uses numpy, not awkward0.numpy, because it operates on TFile data, not TTree data 8 | import numpy 9 | 10 | def towriteable(obj): 11 | def identity(x): 12 | return x 13 | 14 | def resolve(obj): 15 | def types(cls, obj): 16 | if cls is numpy.ndarray: 17 | yield ("numpy", "ndarray", len(obj.shape), str(obj.dtype)) 18 | elif cls.__module__ == "pandas.core.frame" and cls.__name__ == "DataFrame": 19 | yield ("pandas.core.frame", "DataFrame", obj.index.__class__.__name__, set(obj.columns)) 20 | else: 21 | yield (cls.__module__, cls.__name__) 22 | for x in cls.__bases__: 23 | for y in types(x, obj): 24 | yield y 25 | 26 | if any(x == ("builtins", "bytes") or x == ("builtins", "str") or x == ("__builtin__", "str") or x == ("__builtin__", "unicode") for x in types(obj.__class__, obj)): 27 | return (None, None, "uproot3.write.objects.TObjString", "TObjString") 28 | 29 | # made with numpy.histogram 30 | elif isinstance(obj, tuple) and len(obj) == 2 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 len(obj[0].shape) == 1 and len(obj[1].shape) == 1 and obj[0].shape[0] == obj[1].shape[0] - 1: 31 | return ("uproot3_methods.classes.TH1", "from_numpy", "uproot3.write.objects.TH", "TH") 32 | 33 | # made with numpy.histogram2d 34 | 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: 35 | return ("uproot3_methods.classes.TH2", "from_numpy", "uproot3.write.objects.TH", "TH") 36 | 37 | # made with numpy.histogramdd (2-dimensional) 38 | 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: 39 | return ("uproot3_methods.classes.TH2", "from_numpy", "uproot3.write.objects.TH", "TH") 40 | 41 | elif any(x[:3] == ("pandas.core.frame", "DataFrame", "IntervalIndex") and "count" in x[3] for x in types(obj.__class__, obj)): 42 | return ("uproot3_methods.classes.TH1", "from_pandas", "uproot3.write.objects.TH", "TH") 43 | 44 | elif any(x == ("physt.histogram1d", "Histogram1D") for x in types(obj.__class__, obj)): 45 | return ("uproot3_methods.classes.TH1", "from_physt", "uproot3.write.objects.TH", "TH") 46 | 47 | elif any(x == ("uproot3_methods.classes.TH1", "Methods") or x == ("TH1", "Methods") for x in types(obj.__class__, obj)): 48 | return (None, None, "uproot3.write.objects.TH", "TH") 49 | 50 | elif any(x == ("uproot3_methods.classes.TH2", "Methods") or x == ("TH2", "Methods") for x in types(obj.__class__, obj)): 51 | return (None, None, "uproot3.write.objects.TH", "TH") 52 | 53 | elif any(x == ("uproot3_methods.classes.TH3", "Methods") or x == ("TH3", "Methods") for x in types(obj.__class__, obj)): 54 | return (None, None, "uproot3.write.objects.TH", "TH") 55 | 56 | else: 57 | raise TypeError("type {0} from module {1} is not writeable by uproot3".format(obj.__class__.__name__, obj.__class__.__module__)) 58 | 59 | convertmod, convertfcn, mod, cls = resolve(obj) 60 | 61 | if convertfcn is not None: 62 | convert = getattr(importlib.import_module(convertmod), convertfcn) 63 | obj = convert(obj) 64 | 65 | cls = getattr(importlib.import_module(mod), cls) 66 | return cls(obj) 67 | -------------------------------------------------------------------------------- /uproot3_methods/profiles/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import importlib 6 | import re 7 | 8 | def transformer(name): 9 | m = re.match(r"^([a-zA-Z_][a-zA-Z_0-9]*)(\.[a-zA-Z_][a-zA-Z_0-9]*)*$", name) 10 | if m is None: 11 | raise ValueError("profile name must match \"identifier(.identifier)*\"") 12 | return getattr(importlib.import_module("uproot3_methods.profiles." + m.string), "transform") 13 | -------------------------------------------------------------------------------- /uproot3_methods/profiles/cms/__init__.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | -------------------------------------------------------------------------------- /uproot3_methods/profiles/cms/nanoaod.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import awkward0.type 6 | import awkward0.array.chunked 7 | import awkward0.array.objects 8 | 9 | import uproot3_methods.classes.TLorentzVector 10 | 11 | def getcontent(virtual): 12 | return virtual.array.content 13 | 14 | def jaggedtable(rowname, counts, fields): 15 | Table = counts.Table 16 | JaggedArray = counts.JaggedArray 17 | ChunkedArray = counts.ChunkedArray 18 | VirtualArray = counts.VirtualArray 19 | VirtualTLorentzVectorArray = awkward0.array.objects.Methods.mixin(uproot3_methods.classes.TLorentzVector.PtEtaPhiMassArrayMethods, VirtualArray) 20 | 21 | countsarray = counts.array 22 | if isinstance(countsarray, awkward0.array.chunked.ChunkedArray): 23 | return lazyjagged(countsarray, rowname, [(n, x.array) for n, x in fields]) 24 | else: 25 | offsets = JaggedArray.counts2offsets(countsarray) 26 | table = Table.named(rowname) 27 | for n, x in fields: 28 | table[n] = VirtualArray(getcontent, x, type=awkward0.type.ArrayType(offsets[-1], x.type.to.to), cache=counts.cache, persistvirtual=counts.persistvirtual) 29 | columns = table.columns 30 | if "pt" in columns and "eta" in columns and "phi" in columns and "mass" in columns and "p4" not in columns: 31 | table["p4"] = VirtualTLorentzVectorArray(uproot3_methods.classes.TLorentzVector.TLorentzVectorArray.from_ptetaphim, (table["pt"], table["eta"], table["phi"], table["mass"]), type=awkward0.type.ArrayType(offsets[-1], uproot3_methods.classes.TLorentzVector.PtEtaPhiMassLorentzVectorArray), cache=counts.cache, persistvirtual=counts.persistvirtual) 32 | return JaggedArray.fromoffsets(offsets, table) 33 | 34 | def lazyjagged(countsarray, rowname, fields): 35 | ChunkedArray = countsarray.ChunkedArray 36 | VirtualArray = countsarray.VirtualArray 37 | 38 | chunks = [] 39 | for i, countschunk in enumerate(countsarray.chunks): 40 | fieldschunks = [] 41 | tabletype = awkward0.type.TableType() 42 | for fieldname, field in fields: 43 | assert field.chunksizes[i] == countsarray.chunksizes[i] 44 | fieldschunks.append((fieldname, field.chunks[i])) 45 | tabletype[fieldname] = field.type.to.to 46 | columns = tabletype.columns 47 | if "pt" in columns and "eta" in columns and "phi" in columns and "mass" in columns and "p4" not in columns: 48 | tabletype["p4"] = uproot3_methods.classes.TLorentzVector.TLorentzVectorArray.from_ptetaphim 49 | chunks.append(VirtualArray(jaggedtable, (rowname, countschunk, fieldschunks), type=awkward0.type.ArrayType(len(countschunk), float("inf"), tabletype), cache=countschunk.cache, persistvirtual=countschunk.persistvirtual)) 50 | return ChunkedArray(chunks, countsarray.chunksizes) 51 | 52 | def crossref(fromarray, links, subj): 53 | out = fromarray.array 54 | ChunkedArray = out.ChunkedArray 55 | VirtualArray = out.VirtualArray 56 | 57 | if isinstance(out, awkward0.array.chunked.ChunkedArray): 58 | chunks = [] 59 | for j, chunk in enumerate(out.chunks): 60 | newtype = awkward0.type.ArrayType(out.chunksizes[j], float("inf"), awkward0.type.TableType()) 61 | for n in chunk.type.to.to.columns: 62 | newtype.to.to[n] = chunk.type.to.to[n] 63 | for collection, subname, i, localindex, name, totype in links: 64 | newtype.to.to[name] = totype 65 | chunks.append(VirtualArray(crossref, (chunk, links, j), type=newtype, cache=fromarray.cache, persistvirtual=fromarray.persistvirtual)) 66 | 67 | return ChunkedArray(chunks, out.chunksizes) 68 | 69 | else: 70 | for collection, subname, i, localindex, name, totype in links: 71 | toarray = collection[subname].chunks[i] 72 | out.content[name] = VirtualArray(indexedmask, (toarray, localindex, subj), type=awkward0.type.ArrayType(out.offsets[-1], totype), cache=fromarray.cache, persistvirtual=fromarray.persistvirtual) 73 | return out 74 | 75 | def indexedmask(toarray, localindex, subj): 76 | jagged = toarray.array 77 | localindex = localindex.array 78 | if subj is not None: 79 | jagged = jagged.chunks[subj].array 80 | localindex = localindex.chunks[subj].array 81 | 82 | globalindex = localindex + jagged.starts 83 | globalindex.content[localindex.content < 0] = -1 84 | return toarray.IndexedMaskedArray(globalindex.content, jagged.content) 85 | 86 | def transform(array): 87 | array._valid() 88 | array.check_whole_valid = False 89 | 90 | Table = array.Table 91 | VirtualArray = array.VirtualArray 92 | 93 | stuff = [("run", "run", None), 94 | ("luminosityBlock", "lumi", None), 95 | ("event", "event", None), 96 | ("Electron_", "electrons", []), 97 | ("Muon_", "muons", []), 98 | ("Tau_", "taus", []), 99 | ("Photon_", "photons", []), 100 | ("Jet_", "jets", []), 101 | ("FatJet_", "fatjets", []), 102 | ("SubJet_", "subjets", []), 103 | ("IsoTrack_", "isotracks", []), 104 | ("SoftActivityJet_", "softjets", []), 105 | ("SoftActivityJetHT", "softactivity.HT", None), 106 | ("SoftActivityJetHT2", "softactivity.HT2", None), 107 | ("SoftActivityJetHT5", "softactivity.HT5", None), 108 | ("SoftActivityJetHT10", "softactivity.HT10", None), 109 | ("SoftActivityJetNjets2", "softactivity.njets2", None), 110 | ("SoftActivityJetNjets5", "softactivity.njets5", None), 111 | ("SoftActivityJetNjets10", "softactivity.njets10", None), 112 | ("fixedGridRhoFastjetAll", "fixedGridRhoFastjet.everything", None), 113 | ("fixedGridRhoFastjetCentralCalo", "fixedGridRhoFastjet.centralcalo", None), 114 | ("fixedGridRhoFastjetCentralNeutral", "fixedGridRhoFastjet.centralneutral", None), 115 | ("MET_", "MET", Table.named("MET")), 116 | ("RawMET_", "rawMET", Table.named("RawMET")), 117 | ("CaloMET_", "caloMET", Table.named("CaloMET")), 118 | ("PuppiMET_", "puppiMET", Table.named("PuppiMET")), 119 | ("TkMET_", "tkMET", Table.named("TkMET")), 120 | ("PV_", "PV", Table.named("PV")), 121 | ("SV_", "SVs", []), 122 | ("OtherPV_", "otherPVs", []), 123 | ("Pileup_", "pileup", Table.named("Pileup")), 124 | ("Flag_", "flags", Table.named("Flags")), 125 | ("TrigObj_", "trigobjs", []), 126 | ("HLT_", "HLT", Table.named("HLT")), 127 | ("HLTriggerFirstPath", "HLT.firstpath", None), 128 | ("HLTriggerFinalPath", "HLT.finalpath", None), 129 | ("Generator_", "gen", Table.named("Generator")), 130 | ("GenDressedLepton_", "gen.dressedleptons", []), 131 | ("GenPart_", "gen.partons", []), 132 | ("GenJet_", "gen.jets", []), 133 | ("GenJetAK8_", "gen.jetsAK8", []), 134 | ("SubGenJetAK8_", "gen.subjetsAK8", []), 135 | ("GenVisTau_", "gen.vistaus", []), 136 | ("GenMET_", "gen.MET", Table.named("GenMET")), 137 | ("LHE_", "gen.LHE", Table.named("LHE")), 138 | ("LHEPart_", "gen.LHEpartons", []), 139 | ("genWeight", "gen.genweight", None), 140 | ("LHEPdfWeight", "gen.LHEpdfweight", None), 141 | ("LHEScaleWeight", "gen.LHEscaleweight", None), 142 | ("LHEWeight_originalXWGTUP", "gen.LHEweight_originalXWGTUP", None), 143 | ] 144 | 145 | others = [] 146 | for n in array.columns: 147 | for prefix, rename, data in stuff: 148 | if n.startswith(prefix): 149 | if data is None: 150 | pass 151 | elif isinstance(data, list): 152 | data.append((n[len(prefix):], array[n])) 153 | else: 154 | data[n[len(prefix):]] = array[n] 155 | break 156 | elif n == "n" + prefix.rstrip("_"): 157 | break 158 | else: 159 | others.append(n) 160 | 161 | events = Table.named("Event") 162 | 163 | def makecollection(rename): 164 | if "." in rename: 165 | outer, inner = rename.split(".") 166 | if outer not in events.columns: 167 | events[outer] = Table.named(outer.capitalize()) 168 | return events[outer], inner 169 | else: 170 | return events, rename 171 | 172 | for prefix, rename, data in stuff: 173 | if data is None: 174 | if prefix in array.columns: 175 | collection, rename = makecollection(rename) 176 | collection[rename] = array[prefix] 177 | elif isinstance(data, list): 178 | rowname = prefix[:-1] 179 | countname = "n" + rowname 180 | if len(data) > 0 and countname in array.columns: 181 | collection, rename = makecollection(rename) 182 | collection[rename] = lazyjagged(array[countname], rowname, data) 183 | else: 184 | if len(data.columns) > 0: 185 | collection, rename = makecollection(rename) 186 | collection[rename] = data 187 | 188 | eventtype = events.type 189 | 190 | eventtype.to["electrons"].to["photon"] = awkward0.type.OptionType(eventtype.to["photons"].to) 191 | eventtype.to["electrons"].to["photon"].check = False 192 | eventtype.to["electrons"].to["jet"] = awkward0.type.OptionType(eventtype.to["jets"].to) 193 | eventtype.to["electrons"].to["jet"].check = False 194 | for i, chunk in enumerate(events["electrons"].chunks): 195 | assert events["electrons"].chunksizes[i] == events["jets"].chunksizes[i] == events["photons"].chunksizes[i] 196 | events["electrons"].chunks[i] = VirtualArray(crossref, (chunk, [ 197 | (events, "photons", i, array["Electron_photonIdx"].chunks[i], "photon", eventtype.to["electrons"].to["photon"]), 198 | (events, "jets", i, array["Electron_jetIdx"].chunks[i], "jet", eventtype.to["electrons"].to["jet"]), 199 | ], None), type=awkward0.type.ArrayType(events["electrons"].chunksizes[i], eventtype.to["electrons"]), cache=chunk.cache, persistvirtual=chunk.persistvirtual) 200 | 201 | eventtype.to["muons"].to["jet"] = awkward0.type.OptionType(eventtype.to["jets"].to) 202 | eventtype.to["muons"].to["jet"].check = False 203 | for i, chunk in enumerate(events["muons"].chunks): 204 | assert events["muons"].chunksizes[i] == events["jets"].chunksizes[i] 205 | events["muons"].chunks[i] = VirtualArray(crossref, (chunk, [ 206 | (events, "jets", i, array["Muon_jetIdx"].chunks[i], "jet", eventtype.to["muons"].to["jet"]), 207 | ], None), type=awkward0.type.ArrayType(events["muons"].chunksizes[i], eventtype.to["muons"]), cache=chunk.cache, persistvirtual=chunk.persistvirtual) 208 | 209 | eventtype.to["taus"].to["jet"] = awkward0.type.OptionType(eventtype.to["jets"].to) 210 | eventtype.to["taus"].to["jet"].check = False 211 | for i, chunk in enumerate(events["taus"].chunks): 212 | assert events["taus"].chunksizes[i] == events["jets"].chunksizes[i] 213 | events["taus"].chunks[i] = VirtualArray(crossref, (chunk, [ 214 | (events, "jets", i, array["Tau_jetIdx"].chunks[i], "jet", eventtype.to["taus"].to["jet"]), 215 | ], None), type=awkward0.type.ArrayType(events["jets"].chunksizes[i], eventtype.to["taus"]), cache=chunk.cache, persistvirtual=chunk.persistvirtual) 216 | 217 | eventtype.to["taus"].to["jet"] = awkward0.type.OptionType(eventtype.to["jets"].to) 218 | eventtype.to["taus"].to["jet"].check = False 219 | for i, chunk in enumerate(events["taus"].chunks): 220 | assert events["taus"].chunksizes[i] == events["jets"].chunksizes[i] 221 | events["taus"].chunks[i] = VirtualArray(crossref, (chunk, [ 222 | (events, "jets", i, array["Tau_jetIdx"].chunks[i], "jet", eventtype.to["taus"].to["jet"]), 223 | ], None), type=awkward0.type.ArrayType(events["taus"].chunksizes[i], eventtype.to["taus"]), cache=chunk.cache, persistvirtual=chunk.persistvirtual) 224 | 225 | eventtype.to["photons"].to["electron"] = awkward0.type.OptionType(eventtype.to["electrons"].to) 226 | eventtype.to["photons"].to["electron"].check = False 227 | eventtype.to["photons"].to["jet"] = awkward0.type.OptionType(eventtype.to["jets"].to) 228 | eventtype.to["photons"].to["jet"].check = False 229 | for i, chunk in enumerate(events["photons"].chunks): 230 | assert events["photons"].chunksizes[i] == events["jets"].chunksizes[i] == events["electrons"].chunksizes[i] 231 | events["photons"].chunks[i] = VirtualArray(crossref, (chunk, [ 232 | (events, "electrons", i, array["Photon_electronIdx"].chunks[i], "electron", eventtype.to["photons"].to["electron"]), 233 | (events, "jets", i, array["Photon_jetIdx"].chunks[i], "jet", eventtype.to["photons"].to["jet"]), 234 | ], None), type=awkward0.type.ArrayType(events["photons"].chunksizes[i], eventtype.to["photons"]), cache=chunk.cache, persistvirtual=chunk.persistvirtual) 235 | 236 | eventtype.to["jets"].to["electron1"] = awkward0.type.OptionType(eventtype.to["electrons"].to) 237 | eventtype.to["jets"].to["electron1"].check = False 238 | eventtype.to["jets"].to["electron2"] = awkward0.type.OptionType(eventtype.to["electrons"].to) 239 | eventtype.to["jets"].to["electron2"].check = False 240 | eventtype.to["jets"].to["muon1"] = awkward0.type.OptionType(eventtype.to["muons"].to) 241 | eventtype.to["jets"].to["muon1"].check = False 242 | eventtype.to["jets"].to["muon2"] = awkward0.type.OptionType(eventtype.to["muons"].to) 243 | eventtype.to["jets"].to["muon2"].check = False 244 | for i, chunk in enumerate(events["jets"].chunks): 245 | assert events["jets"].chunksizes[i] == events["electrons"].chunksizes[i] == events["muons"].chunksizes[i] 246 | events["jets"].chunks[i] = VirtualArray(crossref, (chunk, [ 247 | (events, "electrons", i, array["Jet_electronIdx1"].chunks[i], "electron1", eventtype.to["jets"].to["electron1"]), 248 | (events, "electrons", i, array["Jet_electronIdx2"].chunks[i], "electron2", eventtype.to["jets"].to["electron2"]), 249 | (events, "muons", i, array["Jet_muonIdx1"].chunks[i], "muon1", eventtype.to["jets"].to["muon1"]), 250 | (events, "muons", i, array["Jet_muonIdx2"].chunks[i], "muon2", eventtype.to["jets"].to["muon2"]), 251 | ], None), type=awkward0.type.ArrayType(events["jets"].chunksizes[i], eventtype.to["jets"]), cache=chunk.cache, persistvirtual=chunk.persistvirtual) 252 | 253 | eventtype.to["fatjets"].to["subjet1"] = awkward0.type.OptionType(eventtype.to["jets"].to) 254 | eventtype.to["fatjets"].to["subjet1"].check = False 255 | eventtype.to["fatjets"].to["subjet2"] = awkward0.type.OptionType(eventtype.to["jets"].to) 256 | eventtype.to["fatjets"].to["subjet2"].check = False 257 | for i, chunk in enumerate(events["fatjets"].chunks): 258 | assert events["fatjets"].chunksizes[i] == events["jets"].chunksizes[i] 259 | events["fatjets"].chunks[i] = VirtualArray(crossref, (chunk, [ 260 | (events, "jets", i, array["FatJet_subJetIdx1"].chunks[i], "subjet1", eventtype.to["fatjets"].to["subjet1"]), 261 | (events, "jets", i, array["FatJet_subJetIdx2"].chunks[i], "subjet2", eventtype.to["fatjets"].to["subjet2"]), 262 | ], None), type=awkward0.type.ArrayType(events["fatjets"].chunksizes[i], eventtype.to["fatjets"]), cache=chunk.cache, persistvirtual=chunk.persistvirtual) 263 | 264 | if len(others) > 0: 265 | etc = events["etc"] = Table.named("OtherFields") 266 | for n in others: 267 | etc[n] = array[n] 268 | events["raw"] = array 269 | 270 | return events 271 | -------------------------------------------------------------------------------- /uproot3_methods/version.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | # BSD 3-Clause License; see https://github.com/scikit-hep/uproot3-methods/blob/master/LICENSE 4 | 5 | import re 6 | 7 | __version__ = "0.10.1" 8 | version = __version__ 9 | version_info = tuple(re.split(r"[-\.]", __version__)) 10 | 11 | del re 12 | --------------------------------------------------------------------------------