From c453924e5616cf8128b57165493008c9f2d39b49 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 11 Dec 2018 15:08:04 -0600 Subject: [PATCH 1/4] better _unwrap_jagged --- uproot_methods/base.py | 27 +++++++++++++++++++++++++-- uproot_methods/version.py | 2 +- 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/uproot_methods/base.py b/uproot_methods/base.py index 9783b42..ed57a44 100644 --- a/uproot_methods/base.py +++ b/uproot_methods/base.py @@ -28,6 +28,11 @@ # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +try: + from collections.abc import Iterable +except ImportError: + from collections import Iterable + import awkward import awkward.util @@ -50,11 +55,29 @@ def _normalize_arrays(arrays): return arrays def _unwrap_jagged(ArrayMethods, arrays): - if not isinstance(arrays[0], awkward.JaggedArray): + 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") + + jagged = None + arrays = list(arrays) + for i in range(len(arrays)): + if not isinstance(arrays[i], Iterable): + arrays[i] = awkward.util.numpy.full(length, arrays[i]) + elif isinstance(arrays[i], awkward.JaggedArray): + jagged = arrays[i] + + if jagged is None: return lambda x: x, lambda x: x, arrays else: JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray) - starts, stops = arrays[0].starts, arrays[0].stops + starts, stops = jagged.starts, jagged.stops wrapmethods, wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays]) return lambda x: JaggedArrayMethods(starts, stops, wrapmethods(x)), lambda x: awkward.JaggedArray(starts, stops, wrap(x)), arrays diff --git a/uproot_methods/version.py b/uproot_methods/version.py index 12b43c7..3ce216a 100644 --- a/uproot_methods/version.py +++ b/uproot_methods/version.py @@ -30,7 +30,7 @@ import re -__version__ = "0.2.10" +__version__ = "0.3.0" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__)) From 5506ea11e409551aede2e42ff8ebe72f42dc5044 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Tue, 11 Dec 2018 15:13:10 -0600 Subject: [PATCH 2/4] oops; that belonged in _normalize_arrays --- uproot_methods/base.py | 35 +++++++++++++++-------------------- 1 file changed, 15 insertions(+), 20 deletions(-) diff --git a/uproot_methods/base.py b/uproot_methods/base.py index ed57a44..50a0260 100644 --- a/uproot_methods/base.py +++ b/uproot_methods/base.py @@ -37,11 +37,24 @@ import awkward.util def _normalize_arrays(arrays): + length = None + for i in range(len(arrays)): + if isinstance(arrays[i], Iterable): + if length is None: + length = len(arrays[i]) + break + if length is None: + raise TypeError("cannot construct an array if all arguments are scalar") + arrays = list(arrays) starts, stops = None, None for i in range(len(arrays)): if starts is None and isinstance(arrays[i], awkward.JaggedArray): starts, stops = arrays[i].starts, arrays[i].stops + + if not isinstance(arrays[i], Iterable): + arrays[i] = awkward.util.numpy.full(length, arrays[i]) + arrays[i] = awkward.util.toarray(arrays[i], awkward.util.numpy.float64) if starts is None: @@ -55,29 +68,11 @@ def _normalize_arrays(arrays): return arrays def _unwrap_jagged(ArrayMethods, 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") - - jagged = None - arrays = list(arrays) - for i in range(len(arrays)): - if not isinstance(arrays[i], Iterable): - arrays[i] = awkward.util.numpy.full(length, arrays[i]) - elif isinstance(arrays[i], awkward.JaggedArray): - jagged = arrays[i] - - if jagged is None: + if not isinstance(arrays[0], awkward.JaggedArray): return lambda x: x, lambda x: x, arrays else: JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray) - starts, stops = jagged.starts, jagged.stops + starts, stops = arrays[0].starts, arrays[0].stops wrapmethods, wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays]) return lambda x: JaggedArrayMethods(starts, stops, wrapmethods(x)), lambda x: awkward.JaggedArray(starts, stops, wrap(x)), arrays From 62dc7613b5eefff55463b86733fc131d6935a384 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 12 Dec 2018 15:52:43 -0600 Subject: [PATCH 3/4] cache properties only at the deepest level of jaggedness (fixes scikit-hep/awkward-array#56) --- uproot_methods/base.py | 20 ++-- uproot_methods/classes/TLorentzVector.py | 112 +++++++++++++---------- uproot_methods/classes/TVector2.py | 10 +- uproot_methods/classes/TVector3.py | 20 ++-- 4 files changed, 89 insertions(+), 73 deletions(-) diff --git a/uproot_methods/base.py b/uproot_methods/base.py index 50a0260..afed533 100644 --- a/uproot_methods/base.py +++ b/uproot_methods/base.py @@ -69,19 +69,23 @@ def _normalize_arrays(arrays): def _unwrap_jagged(ArrayMethods, arrays): if not isinstance(arrays[0], awkward.JaggedArray): - return lambda x: x, lambda x: x, arrays + return lambda x: x, arrays else: - JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray) + if ArrayMethods is None: + cls = awkward.JaggedArray + else: + cls = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray) starts, stops = arrays[0].starts, arrays[0].stops - wrapmethods, wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays]) - return lambda x: JaggedArrayMethods(starts, stops, wrapmethods(x)), lambda x: awkward.JaggedArray(starts, stops, wrap(x)), arrays + wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays]) + return lambda x: cls(starts, stops, wrap(x)), arrays def memo(function): memoname = "_memo_" + function.__name__ - def memofunction(self): - if not hasattr(self, memoname): - setattr(self, memoname, function(self)) - return getattr(self, memoname) + def memofunction(array): + wrap, (array,) = _unwrap_jagged(None, (array,)) + if not hasattr(array, memoname): + setattr(array, memoname, function(array)) + return wrap(getattr(array, memoname)) return memofunction class ROOTMethods(awkward.Methods): diff --git a/uproot_methods/classes/TLorentzVector.py b/uproot_methods/classes/TLorentzVector.py index faf57ed..ace6c37 100644 --- a/uproot_methods/classes/TLorentzVector.py +++ b/uproot_methods/classes/TLorentzVector.py @@ -75,11 +75,6 @@ def perp(self): def pt2(self): return self.p3.rho2 - @property - @uproot_methods.base.memo - def pt(self): - return self.p3.rho - @property def Et(self): return self.energy * self.pt / self.p @@ -92,19 +87,9 @@ def mag2(self): def mass2(self): return self.mag2 - @property - @uproot_methods.base.memo - def mass(self): - return self.mag - @property def mt2(self): return self.energy**2 - self.z**2 - - @property - @uproot_methods.base.memo - def phi(self): - return self.p3.phi @property def theta(self): @@ -193,6 +178,26 @@ def z(self): def t(self): return self["fE"] + @property + @uproot_methods.base.memo + def pt(self): + return awkward.util.numpy.sqrt(self.pt2) + + @property + @uproot_methods.base.memo + def eta(self): + return awkward.util.numpy.arcsinh(self.z / awkward.util.numpy.sqrt(self.x**2 + self.y**2)) + + @property + @uproot_methods.base.memo + def phi(self): + return self.p3.phi + + @property + @uproot_methods.base.memo + def mass(self): + return awkward.util.numpy.sqrt(self.mag2) + @property def mag(self): return awkward.util.numpy.sqrt(self.mag2) @@ -203,11 +208,6 @@ def mt(self): sign = awkward.util.numpy.sign(mt2) return awkward.util.numpy.sqrt(awkward.util.numpy.absolute(mt2)) * sign - @property - @uproot_methods.base.memo - def eta(self): - return awkward.util.numpy.arcsinh(self.z / awkward.util.numpy.sqrt(self.x**2 + self.y**2)) - @property def rapidity(self): return awkward.util.numpy.log((self.t + self.z) / (self.t - self.z)) / 2.0 @@ -370,6 +370,22 @@ def _vector(self, operator, vector, reverse=False): def _unary(self, operator): return TLorentzVector(operator(self.x), operator(self.y), operator(self.z), operator(self.t)) + @property + def pt(self): + return math.sqrt(self.pt2) + + @property + def eta(self): + return math.asinh(self.z / math.sqrt(self.x**2 + self.y**2)) + + @property + def phi(self): + return self.p3.phi + + @property + def mass(self): + return math.sqrt(self.mag2) + @property def mag(self): return math.sqrt(self.mag2) @@ -382,10 +398,6 @@ def mt(self): else: return -math.sqrt(out) - @property - def eta(self): - return math.asinh(self.z / math.sqrt(self.x**2 + self.y**2)) - @property def rapidity(self): return math.log((self.t + self.z) / (self.t - self.z)) / 2.0 @@ -559,54 +571,54 @@ def origin_like(cls, array): @classmethod def from_p3(cls, p3, t): - wrapmethods, wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((p3.x, p3.y, p3.z, t))) - return wrapmethods(cls(x, y, z, t)) + wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((p3.x, p3.y, p3.z, t))) + return wrap(cls(x, y, z, t)) @classmethod def from_cartesian(cls, x, y, z, t): - wrapmethods, wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, t))) - return wrapmethods(cls(x, y, z, t)) + wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, t))) + return wrap(cls(x, y, z, t)) @classmethod def from_spherical(cls, r, theta, phi, t): - wrapmethods, wrap, (r, theta, phi, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((r, theta, phi, t))) - return wrapmethods(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_spherical(r, theta, phi), t)) + wrap, (r, theta, phi, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((r, theta, phi, t))) + return wrap(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_spherical(r, theta, phi), t)) @classmethod def from_cylindrical(cls, rho, phi, z, t): - wrapmethods, wrap, (rho, phi, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z, t))) - return wrapmethods(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_cylindrical(rho, phi, z), t)) + wrap, (rho, phi, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z, t))) + return wrap(cls.from_p3(uproot_methods.classes.TVector3.TVector3Array.from_cylindrical(rho, phi, z), t)) @classmethod def from_xyzm(cls, x, y, z, m): - wrapmethods, wrap, (x, y, z, m) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, m))) - return wrapmethods(cls(x, y, z, awkward.util.numpy.sqrt(x*x + y*y + z*z + m*m*awkward.util.numpy.sign(m)))) + wrap, (x, y, z, m) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, m))) + return wrap(cls(x, y, z, awkward.util.numpy.sqrt(x*x + y*y + z*z + m*m*awkward.util.numpy.sign(m)))) @classmethod def from_ptetaphi(cls, pt, eta, phi, energy): - wrapmethods, wrap, (pt, eta, phi, energy) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, energy))) - out = wrapmethods(cls(pt * awkward.util.numpy.cos(phi), - pt * awkward.util.numpy.sin(phi), - pt * awkward.util.numpy.sinh(eta), - energy)) - out._memo_pt = wrap(pt) - out._memo_eta = wrap(eta) - out._memo_phi = wrap(phi) - return out + wrap, (pt, eta, phi, energy) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, energy))) + out = cls(pt * awkward.util.numpy.cos(phi), + pt * awkward.util.numpy.sin(phi), + pt * awkward.util.numpy.sinh(eta), + energy) + out._memo_pt = pt + out._memo_eta = eta + out._memo_phi = phi + return wrap(out) @classmethod def from_ptetaphim(cls, pt, eta, phi, mass): - wrapmethods, wrap, (pt, eta, phi, mass) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, mass))) + wrap, (pt, eta, phi, mass) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, mass))) x = pt * awkward.util.numpy.cos(phi) y = pt * awkward.util.numpy.sin(phi) z = pt * awkward.util.numpy.sinh(eta) p3 = uproot_methods.classes.TVector3.TVector3Array(x, y, z) - out = wrapmethods(cls.from_p3(p3, awkward.util.numpy.sqrt(x*x + y*y + z*z + mass*mass*awkward.util.numpy.sign(mass)))) - out._memo_pt = wrap(pt) - out._memo_eta = wrap(eta) - out._memo_phi = wrap(phi) - out._memo_mass = wrap(mass) - return out + out = cls.from_p3(p3, awkward.util.numpy.sqrt(x*x + y*y + z*z + mass*mass*awkward.util.numpy.sign(mass))) + out._memo_pt = pt + out._memo_eta = eta + out._memo_phi = phi + out._memo_mass = mass + return wrap(out) @property def x(self): diff --git a/uproot_methods/classes/TVector2.py b/uproot_methods/classes/TVector2.py index b8913f8..efa7e45 100644 --- a/uproot_methods/classes/TVector2.py +++ b/uproot_methods/classes/TVector2.py @@ -153,14 +153,14 @@ def origin_like(cls, array): @classmethod def from_cartesian(cls, x, y): - wrapmethods, wrap, (x, y) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y))) - return wrapmethods(cls(x, y)) + wrap, (x, y) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y))) + return wrap(cls(x, y)) @classmethod def from_polar(cls, rho, phi): - wrapmethods, wrap, (rho, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi))) - return wrapmethods(cls(rho * awkward.util.numpy.cos(phi), - rho * awkward.util.numpy.sin(phi))) + wrap, (rho, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi))) + return wrap(cls(rho * awkward.util.numpy.cos(phi), + rho * awkward.util.numpy.sin(phi))) @property def x(self): diff --git a/uproot_methods/classes/TVector3.py b/uproot_methods/classes/TVector3.py index b1f57eb..4406f5b 100644 --- a/uproot_methods/classes/TVector3.py +++ b/uproot_methods/classes/TVector3.py @@ -267,22 +267,22 @@ def origin_like(cls, array): @classmethod def from_cartesian(cls, x, y, z): - wrapmethods, wrap, (x, y, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z))) - return wrapmethods(cls(x, y, z)) + wrap, (x, y, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z))) + return wrap(cls(x, y, z)) @classmethod def from_spherical(cls, r, theta, phi): - wrapmethods, wrap, (r, theta, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((r, theta, phi))) - return wrapmethods(cls(r * awkward.util.numpy.sin(theta) * awkward.util.numpy.cos(phi), - r * awkward.util.numpy.sin(theta) * awkward.util.numpy.sin(phi), - r * awkward.util.numpy.cos(theta))) + wrap, (r, theta, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((r, theta, phi))) + return wrap(cls(r * awkward.util.numpy.sin(theta) * awkward.util.numpy.cos(phi), + r * awkward.util.numpy.sin(theta) * awkward.util.numpy.sin(phi), + r * awkward.util.numpy.cos(theta))) @classmethod def from_cylindrical(cls, rho, phi, z): - wrapmethods, wrap, (rho, phi, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z))) - return wrapmethods(cls(rho * awkward.util.numpy.cos(phi), - rho * awkward.util.numpy.sin(phi), - z)) + wrap, (rho, phi, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z))) + return wrap(cls(rho * awkward.util.numpy.cos(phi), + rho * awkward.util.numpy.sin(phi), + z)) @property def x(self): From 703866fa98de468551d7f5137257e56ce82b2564 Mon Sep 17 00:00:00 2001 From: Jim Pivarski Date: Wed, 12 Dec 2018 16:25:37 -0600 Subject: [PATCH 4/4] version 0.2.11 --- uproot_methods/version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/uproot_methods/version.py b/uproot_methods/version.py index 3ce216a..eccce97 100644 --- a/uproot_methods/version.py +++ b/uproot_methods/version.py @@ -30,7 +30,7 @@ import re -__version__ = "0.3.0" +__version__ = "0.2.11" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__))