diff --git a/src/coffea/nanoevents/methods/vector.py b/src/coffea/nanoevents/methods/vector.py index 99d09d84e..c71e1a054 100644 --- a/src/coffea/nanoevents/methods/vector.py +++ b/src/coffea/nanoevents/methods/vector.py @@ -65,10 +65,30 @@ def _mass2_kernel(t, x, y, z): numba.float64(numba.float64, numba.float64), ] ) -def _deltaphi_kernel(a, b): +def delta_phi(a, b): + """Compute difference in angle given two angles a and b + + Returns a value within [-pi, pi) + """ return (a - b + numpy.pi) % (2 * numpy.pi) - numpy.pi +@numba.vectorize( + [ + numba.float32(numba.float32, numba.float32, numba.float32, numba.float32), + numba.float64(numba.float64, numba.float64, numba.float64, numba.float64), + ] +) +def delta_r(eta1, phi1, eta2, phi2): + r"""Distance in (eta,phi) plane given two pairs of (eta,phi) + + :math:`\sqrt{\Delta\eta^2 + \Delta\phi^2}` + """ + deta = eta1 - eta2 + dphi = delta_phi(phi1, phi2) + return numpy.hypot(deta, dphi) + + behavior = {} @@ -201,7 +221,7 @@ def delta_phi(self, other): Returns a value within [-pi, pi) """ - return _deltaphi_kernel(self.phi, other.phi) + return delta_phi(self.phi, other.phi) def dot(self, other): """Compute the dot product of two vectors""" @@ -595,16 +615,14 @@ def multiply(self, other): def delta_r2(self, other): """Squared `delta_r`""" - deta = self.eta - other.eta - dphi = self.delta_phi(other) - return deta * deta + dphi * dphi + return delta_r(self.eta, self.phi, other.eta, other.phi) ** 2 def delta_r(self, other): r"""Distance between two Lorentz vectors in (eta,phi) plane :math:`\sqrt{\Delta\eta^2 + \Delta\phi^2}` """ - return numpy.hypot(self.eta - other.eta, self.delta_phi(other)) + return delta_r(self.eta, self.phi, other.eta, other.phi) @awkward.mixin_class_method(numpy.negative) def negative(self):