diff --git a/uproot_methods/base.py b/uproot_methods/base.py index 9783b42..afed533 100644 --- a/uproot_methods/base.py +++ b/uproot_methods/base.py @@ -28,15 +28,33 @@ # 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 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: @@ -51,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): diff --git a/uproot_methods/version.py b/uproot_methods/version.py index 12b43c7..eccce97 100644 --- a/uproot_methods/version.py +++ b/uproot_methods/version.py @@ -30,7 +30,7 @@ import re -__version__ = "0.2.10" +__version__ = "0.2.11" version = __version__ version_info = tuple(re.split(r"[-\.]", __version__))