diff --git a/tests/test_issues.py b/tests/test_issues.py index 58b460e..9cef285 100644 --- a/tests/test_issues.py +++ b/tests/test_issues.py @@ -42,7 +42,7 @@ def runTest(self): def test_issue10(self): p4 = TLorentzVectorArray.from_ptetaphim(awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]]), awkward.JaggedArray.fromiter([[1.0]])) - assert p4.mass.tolist() == [[0.9999999999999999]] + assert p4.mass.tolist() == [[1.0]] assert p4[0].mass.tolist() == [0.9999999999999999] assert p4[0][0].mass == 0.9999999999999999 assert type(p4.mass) is awkward.JaggedArray diff --git a/uproot_methods/base.py b/uproot_methods/base.py index d81bb3e..9783b42 100644 --- a/uproot_methods/base.py +++ b/uproot_methods/base.py @@ -51,12 +51,20 @@ def _normalize_arrays(arrays): def _unwrap_jagged(ArrayMethods, arrays): if not isinstance(arrays[0], awkward.JaggedArray): - return lambda x: x, arrays + return lambda x: x, lambda x: x, arrays else: JaggedArrayMethods = ArrayMethods.mixin(ArrayMethods, awkward.JaggedArray) starts, stops = arrays[0].starts, arrays[0].stops - wrap, arrays = _unwrap_jagged(ArrayMethods, [x.content for x in arrays]) - return lambda x: JaggedArrayMethods(starts, stops, wrap(x)), arrays + 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 + +def memo(function): + memoname = "_memo_" + function.__name__ + def memofunction(self): + if not hasattr(self, memoname): + setattr(self, memoname, function(self)) + return getattr(self, memoname) + return memofunction class ROOTMethods(awkward.Methods): _arraymethods = None diff --git a/uproot_methods/classes/TLorentzVector.py b/uproot_methods/classes/TLorentzVector.py index 60bb610..faf57ed 100644 --- a/uproot_methods/classes/TLorentzVector.py +++ b/uproot_methods/classes/TLorentzVector.py @@ -76,6 +76,7 @@ def pt2(self): return self.p3.rho2 @property + @uproot_methods.base.memo def pt(self): return self.p3.rho @@ -92,6 +93,7 @@ def mass2(self): return self.mag2 @property + @uproot_methods.base.memo def mass(self): return self.mag @@ -100,6 +102,7 @@ def mt2(self): return self.energy**2 - self.z**2 @property + @uproot_methods.base.memo def phi(self): return self.p3.phi @@ -201,8 +204,9 @@ def mt(self): return awkward.util.numpy.sqrt(awkward.util.numpy.absolute(mt2)) * sign @property + @uproot_methods.base.memo def eta(self): - return -awkward.util.numpy.log((1.0 - awkward.util.numpy.cos(self.theta)) / (1.0 + awkward.util.numpy.cos(self.theta))) / 2.0 + return awkward.util.numpy.arcsinh(self.z / awkward.util.numpy.sqrt(self.x**2 + self.y**2)) @property def rapidity(self): @@ -380,7 +384,7 @@ def mt(self): @property def eta(self): - return -math.log((1.0 - math.cos(self.theta)) / (1.0 + math.cos(self.theta)))/2.0 + return math.asinh(self.z / math.sqrt(self.x**2 + self.y**2)) @property def rapidity(self): @@ -555,45 +559,54 @@ def origin_like(cls, array): @classmethod def from_p3(cls, p3, 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)) + 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)) @classmethod def from_cartesian(cls, x, y, z, t): - wrap, (x, y, z, t) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z, t))) - return wrap(cls(x, y, z, t)) + 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)) @classmethod def from_spherical(cls, 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)) + 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)) @classmethod def from_cylindrical(cls, 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)) + 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)) @classmethod def from_xyzm(cls, x, y, z, 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)))) + 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)))) @classmethod def from_ptetaphi(cls, pt, eta, phi, energy): - wrap, (pt, eta, phi, energy) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, energy))) - return wrap(cls(pt * awkward.util.numpy.cos(phi), - pt * awkward.util.numpy.sin(phi), - pt * awkward.util.numpy.sinh(eta), - energy)) + 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 @classmethod def from_ptetaphim(cls, pt, eta, phi, mass): - wrap, (pt, eta, phi, mass) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, mass))) + wrapmethods, wrap, (pt, eta, phi, mass) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((pt, eta, phi, mass))) x = pt * awkward.util.numpy.cos(phi) y = pt * awkward.util.numpy.sin(phi) z = pt * awkward.util.numpy.sinh(eta) p3 = uproot_methods.classes.TVector3.TVector3Array(x, y, z) - return wrap(cls.from_p3(p3, awkward.util.numpy.sqrt(x*x + y*y + z*z + mass*mass*awkward.util.numpy.sign(mass)))) + 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 @property def x(self): diff --git a/uproot_methods/classes/TVector2.py b/uproot_methods/classes/TVector2.py index efa7e45..b8913f8 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): - wrap, (x, y) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y))) - return wrap(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)) @classmethod def from_polar(cls, rho, phi): - wrap, (rho, phi) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi))) - return wrap(cls(rho * awkward.util.numpy.cos(phi), - rho * awkward.util.numpy.sin(phi))) + 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))) @property def x(self): diff --git a/uproot_methods/classes/TVector3.py b/uproot_methods/classes/TVector3.py index 4406f5b..b1f57eb 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): - wrap, (x, y, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((x, y, z))) - return wrap(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)) @classmethod def from_spherical(cls, r, theta, phi): - 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))) + 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))) @classmethod def from_cylindrical(cls, rho, phi, z): - wrap, (rho, phi, z) = uproot_methods.base._unwrap_jagged(ArrayMethods, uproot_methods.base._normalize_arrays((rho, phi, z))) - return wrap(cls(rho * awkward.util.numpy.cos(phi), - rho * awkward.util.numpy.sin(phi), - z)) + 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)) @property def x(self): diff --git a/uproot_methods/version.py b/uproot_methods/version.py index 8e1f5ed..12b43c7 100644 --- a/uproot_methods/version.py +++ b/uproot_methods/version.py @@ -30,7 +30,7 @@ import re -__version__ = "0.2.9" +__version__ = "0.2.10" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__))