Skip to content
This repository has been archived by the owner on Jan 27, 2023. It is now read-only.

Commit

Permalink
Merge pull request #13 from nsmith-/fastlorentz
Browse files Browse the repository at this point in the history 
Faster LorentzVector stuff
  • Loading branch information
@jpivarski
jpivarski authored Dec 2, 2018
2 parents 158169a + 018b5f2 commit 2c689b8
Show file tree
Hide file tree
Showing 6 changed files with 60 additions and 39 deletions.
2 changes: 1 addition & 1 deletion tests/test_issues.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 11 additions & 3 deletions uproot_methods/base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
51 changes: 32 additions & 19 deletions uproot_methods/classes/TLorentzVector.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ def pt2(self):
return self.p3.rho2

@property
@uproot_methods.base.memo
def pt(self):
return self.p3.rho

Expand All @@ -92,6 +93,7 @@ def mass2(self):
return self.mag2

@property
@uproot_methods.base.memo
def mass(self):
return self.mag

Expand All @@ -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

Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down
10 changes: 5 additions & 5 deletions uproot_methods/classes/TVector2.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
20 changes: 10 additions & 10 deletions uproot_methods/classes/TVector3.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion uproot_methods/version.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@

import re

__version__ = "0.2.9"
__version__ = "0.2.10"
version = __version__
version_info = tuple(re.split(r"[-\.]", __version__))

Expand Down

0 comments on commit 2c689b8

Please sign in to comment.