diff --git a/uproot_methods/base.py b/uproot_methods/base.py index a952872..6db4480 100644 --- a/uproot_methods/base.py +++ b/uproot_methods/base.py @@ -28,13 +28,12 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -try: - from collections.abc import Iterable -except ImportError: - from collections import Iterable import awkward +from uproot_methods.wrapjagged import normalize_arrays, unwrap_jagged + + class ROOTMethods(awkward.Methods): _arraymethods = None @@ -45,60 +44,15 @@ def __ne__(self, other): @classmethod def _normalize_arrays(cls, arrays): - length = None - for i in range(len(arrays)): - if isinstance(arrays[i], Iterable): - if length is None: - length = len(arrays[i]) - break - if length is None: - raise TypeError("cannot construct an array if all arguments are scalar") - - arrays = list(arrays) - jaggedtype = [cls.awkward.JaggedArray] * len(arrays) - starts, stops = None, None - for i in range(len(arrays)): - if starts is None and isinstance(arrays[i], cls.awkward.JaggedArray): - starts, stops = arrays[i].starts, arrays[i].stops - - if isinstance(arrays[i], cls.awkward.JaggedArray): - jaggedtype[i] = type(arrays[i]) - - if not isinstance(arrays[i], Iterable): - arrays[i] = cls.awkward.numpy.full(length, arrays[i]) - - arrays[i] = cls.awkward.util.toarray(arrays[i], cls.awkward.numpy.float64) - - if starts is None: - return arrays - - for i in range(len(arrays)): - if not isinstance(arrays[i], cls.awkward.JaggedArray) or not (cls.awkward.numpy.array_equal(starts, arrays[i].starts) and cls.awkward.numpy.array_equal(stops, arrays[i].stops)): - content = cls.awkward.numpy.zeros(stops.max(), dtype=cls.awkward.numpy.float64) - arrays[i] = jaggedtype[i](starts, stops, content) + arrays[i] # invoke jagged broadcasting to align arrays - - return arrays + return normalize_arrays(cls, arrays) @classmethod - def _unwrap_jagged(cls, ArrayMethods, arrays): - if not isinstance(arrays[0], cls.awkward.JaggedArray): - return lambda x: x, arrays - else: - if ArrayMethods is None: - awkcls = arrays[0].JaggedArray - else: - awkcls = ArrayMethods.mixin(ArrayMethods, arrays[0].JaggedArray) - counts = arrays[0].counts.reshape(-1) - offsets = awkcls.counts2offsets(counts) - starts, stops = offsets[:-1], offsets[1:] - starts = starts.reshape(arrays[0].starts.shape[:-1] + (-1,)) - stops = stops.reshape(arrays[0].stops.shape[:-1] + (-1,)) - wrap, arrays = cls._unwrap_jagged(ArrayMethods, [x.flatten() for x in arrays]) - return lambda x: awkcls(starts, stops, wrap(x)), arrays + def _unwrap_jagged(cls, awkcls, arrays): + return unwrap_jagged(cls, awkcls, arrays) def _trymemo(self, name, function): memoname = "_memo_" + name - wrap, (array,) = type(self)._unwrap_jagged(None, (self,)) + wrap, (array,) = type(self)._unwrap_jagged(self.JaggedArray, (self,)) if not hasattr(array, memoname): setattr(array, memoname, function(array)) return wrap(getattr(array, memoname)) diff --git a/uproot_methods/classes/TLorentzVector.py b/uproot_methods/classes/TLorentzVector.py index aa265c6..e8df092 100644 --- a/uproot_methods/classes/TLorentzVector.py +++ b/uproot_methods/classes/TLorentzVector.py @@ -39,6 +39,8 @@ import uproot_methods.common.TVector import uproot_methods.classes.TVector3 +from uproot_methods.wrapjagged import wrapjaggedmethod + class Common(object): @property def E(self): @@ -753,7 +755,14 @@ def mass(self): def mass(self, value): self["fMass"] = value +awkward = uproot_methods.base.ROOTMethods.awkward +PtEtaPhiMassJaggedArrayMethods = PtEtaPhiMassArrayMethods.mixin(PtEtaPhiMassArrayMethods, awkward.JaggedArray) + class TLorentzVectorArray(ArrayMethods, uproot_methods.base.ROOTMethods.awkward.ObjectArray): + + jaggedtype = uproot_methods.base.ROOTMethods.awkward.JaggedArray + awkcls = ArrayMethods.mixin(ArrayMethods, jaggedtype) + def __init__(self, x, y, z, t): if isinstance(x, awkward.array.jagged.JaggedArray) or isinstance(y, awkward.array.jagged.JaggedArray) or isinstance(z, awkward.array.jagged.JaggedArray) or isinstance(t, awkward.array.jagged.JaggedArray): raise TypeError("TLorentzVectorArray constructor arguments must not be jagged; use TLorentzVectorArray.from_cartesian for jaggedness-handling") @@ -778,32 +787,32 @@ def origin_like(cls, array): @classmethod def from_p3(cls, p3, t): - wrap, (x, y, z, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((p3.x, p3.y, p3.z, t))) + wrap, (x, y, z, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((p3.x, p3.y, p3.z, t))) return wrap(cls(x, y, z, t)) @classmethod def from_cartesian(cls, x, y, z, t): - wrap, (x, y, z, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y, z, t))) + wrap, (x, y, z, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y, z, t))) return wrap(cls(x, y, z, t)) @classmethod def from_spherical(cls, r, theta, phi, t): - wrap, (r, theta, phi, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((r, theta, phi, t))) + wrap, (r, theta, phi, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((r, theta, phi, t))) return wrap(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_spherical(r, theta, phi), t)) @classmethod def from_cylindrical(cls, rho, phi, z, t): - wrap, (rho, phi, z, t) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((rho, phi, z, t))) + wrap, (rho, phi, z, t) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((rho, phi, z, t))) return wrap(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_cylindrical(rho, phi, z), t)) @classmethod def from_xyzm(cls, x, y, z, m): - wrap, (x, y, z, m) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y, z, m))) + wrap, (x, y, z, m) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y, z, m))) return wrap(cls(x, y, z, cls.awkward.numpy.sqrt(x*x + y*y + z*z + m*m*cls.awkward.numpy.sign(m)))) @classmethod def from_ptetaphi(cls, pt, eta, phi, energy): - wrap, (pt, eta, phi, energy) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((pt, eta, phi, energy))) + wrap, (pt, eta, phi, energy) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((pt, eta, phi, energy))) out = cls(pt * cls.awkward.numpy.cos(phi), pt * cls.awkward.numpy.sin(phi), pt * cls.awkward.numpy.sinh(eta), @@ -814,10 +823,9 @@ def from_ptetaphi(cls, pt, eta, phi, energy): return wrap(out) @classmethod + @wrapjaggedmethod(PtEtaPhiMassJaggedArrayMethods) def from_ptetaphim(cls, pt, eta, phi, mass): - wrap, (pt, eta, phi, mass) = cls._unwrap_jagged(PtEtaPhiMassArrayMethods, cls._normalize_arrays((pt, eta, phi, mass))) - out = PtEtaPhiMassLorentzVectorArray(pt,eta,phi,mass) - return wrap(out) + return PtEtaPhiMassLorentzVectorArray(pt,eta,phi,mass) @property def x(self): diff --git a/uproot_methods/classes/TVector2.py b/uproot_methods/classes/TVector2.py index db299a7..510cbab 100644 --- a/uproot_methods/classes/TVector2.py +++ b/uproot_methods/classes/TVector2.py @@ -145,6 +145,9 @@ def _unary(self, operator): return TVector2(operator(self.x), operator(self.y)) class TVector2Array(ArrayMethods, uproot_methods.base.ROOTMethods.awkward.ObjectArray): + jaggedtype = uproot_methods.base.ROOTMethods.awkward.JaggedArray + awkcls = ArrayMethods.mixin(ArrayMethods, jaggedtype) + def __init__(self, x, y): if isinstance(x, awkward.array.jagged.JaggedArray) or isinstance(y, awkward.array.jagged.JaggedArray): raise TypeError("TVector2Array constructor arguments must not be jagged; use TVector2.from_cartesian for jaggedness-handling") @@ -164,12 +167,12 @@ def origin_like(cls, array): @classmethod def from_cartesian(cls, x, y): - wrap, (x, y) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y))) + wrap, (x, y) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y))) return wrap(cls(x, y)) @classmethod def from_polar(cls, rho, phi): - wrap, (rho, phi) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((rho, phi))) + wrap, (rho, phi) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((rho, phi))) return wrap(cls(rho * cls.awkward.numpy.cos(phi), rho * cls.awkward.numpy.sin(phi))) diff --git a/uproot_methods/classes/TVector3.py b/uproot_methods/classes/TVector3.py index 6689ff1..e119307 100644 --- a/uproot_methods/classes/TVector3.py +++ b/uproot_methods/classes/TVector3.py @@ -257,6 +257,9 @@ def rotate_euler(self, phi=0, theta=0, psi=0): return TVector3(x, y, z) class TVector3Array(ArrayMethods, uproot_methods.base.ROOTMethods.awkward.ObjectArray): + jaggedtype = uproot_methods.base.ROOTMethods.awkward.JaggedArray + awkcls = ArrayMethods.mixin(ArrayMethods, jaggedtype) + def __init__(self, x, y, z): if isinstance(x, awkward.array.jagged.JaggedArray) or isinstance(y, awkward.array.jagged.JaggedArray) or isinstance(z, awkward.array.jagged.JaggedArray): raise TypeError("TVector3Array constructor arguments must not be jagged; use TVector3.from_cartesian for jaggedness-handling") @@ -279,19 +282,19 @@ def origin_like(cls, array): @classmethod def from_cartesian(cls, x, y, z): - wrap, (x, y, z) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((x, y, z))) + wrap, (x, y, z) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((x, y, z))) return wrap(cls(x, y, z)) @classmethod def from_spherical(cls, r, theta, phi): - wrap, (r, theta, phi) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((r, theta, phi))) + wrap, (r, theta, phi) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((r, theta, phi))) return wrap(cls(r * cls.awkward.numpy.sin(theta) * cls.awkward.numpy.cos(phi), r * cls.awkward.numpy.sin(theta) * cls.awkward.numpy.sin(phi), r * cls.awkward.numpy.cos(theta))) @classmethod def from_cylindrical(cls, rho, phi, z): - wrap, (rho, phi, z) = cls._unwrap_jagged(ArrayMethods, cls._normalize_arrays((rho, phi, z))) + wrap, (rho, phi, z) = cls._unwrap_jagged(cls.awkcls, cls._normalize_arrays((rho, phi, z))) return wrap(cls(rho * cls.awkward.numpy.cos(phi), rho * cls.awkward.numpy.sin(phi), z)) diff --git a/uproot_methods/version.py b/uproot_methods/version.py index 7740391..3517c8f 100644 --- a/uproot_methods/version.py +++ b/uproot_methods/version.py @@ -30,7 +30,7 @@ import re -__version__ = "0.4.6" +__version__ = "0.4.7" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__)) diff --git a/uproot_methods/wrapjagged.py b/uproot_methods/wrapjagged.py new file mode 100644 index 0000000..e69fdfa --- /dev/null +++ b/uproot_methods/wrapjagged.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python + +# Copyright (c) 2019, IRIS-HEP +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +try: + from collections.abc import Iterable +except ImportError: + from collections import Iterable + +from functools import wraps + + +def normalize_arrays(cls, arrays): + length = None + for i in range(len(arrays)): + if isinstance(arrays[i], Iterable): + if length is None: + length = len(arrays[i]) + break + if length is None: + raise TypeError("cannot construct an array if all arguments are scalar") + + arrays = list(arrays) + jaggedtype = [cls.awkward.JaggedArray] * len(arrays) + starts, stops = None, None + for i in range(len(arrays)): + if starts is None and isinstance(arrays[i], cls.awkward.JaggedArray): + starts, stops = arrays[i].starts, arrays[i].stops + + if isinstance(arrays[i], cls.awkward.JaggedArray): + jaggedtype[i] = type(arrays[i]) + + if not isinstance(arrays[i], Iterable): + arrays[i] = cls.awkward.numpy.full(length, arrays[i]) + + arrays[i] = cls.awkward.util.toarray(arrays[i], cls.awkward.numpy.float64) + + if starts is None: + return arrays + + for i in range(len(arrays)): + if not isinstance(arrays[i], cls.awkward.JaggedArray) or not (cls.awkward.numpy.array_equal(starts, arrays[i].starts) and cls.awkward.numpy.array_equal(stops, arrays[i].stops)): + content = cls.awkward.numpy.zeros(stops.max(), dtype=cls.awkward.numpy.float64) + arrays[i] = jaggedtype[i](starts, stops, content) + arrays[i] # invoke jagged broadcasting to align arrays + + return arrays + +def unwrap_jagged(cls, awkcls, arrays): + if not isinstance(arrays[0], cls.awkward.JaggedArray): + return lambda x: x, arrays + + counts = arrays[0].counts.reshape(-1) + offsets = awkcls.counts2offsets(counts) + starts, stops = offsets[:-1], offsets[1:] + starts = starts.reshape(arrays[0].starts.shape[:-1] + (-1,)) + stops = stops.reshape(arrays[0].stops.shape[:-1] + (-1,)) + wrap, arrays = unwrap_jagged(cls, awkcls, [x.flatten() for x in arrays]) + return lambda x: awkcls(starts, stops, wrap(x)), arrays + +def wrapjaggedmethod(awkcls): + def wrapjagged_decorator(func): + @wraps(func) + def func_wrapper(cls, *arrays): + wrap, arrays = unwrap_jagged(cls, awkcls, normalize_arrays(cls, arrays)) + return wrap(func(cls, *arrays)) + return func_wrapper + return wrapjagged_decorator