Shortcut multiply and divide for PtEtaPhiMassArrayMethods

tests/test_vector.py

@@ -134,6 +134,23 @@ def test_lorentzvector_array(self):
         assert (a + TLorentzVector(1000, 2000, 0, 0) == TLorentzVectorArray(numpy.full(10, 1000), numpy.arange(2000, 2010), numpy.zeros(10), numpy.zeros(10))).tolist() == [True, True, True, True, True, True, True, True, True, True]
         assert (a**2).tolist() == [0.0, -1.0, -4.0, -9.0, -16.0, -25.0, -36.0, -49.0, -64.0, -81.0]
 
+    def test_ptetaphim_array(self):
+        a = TLorentzVectorArray.from_ptetaphim(
+            numpy.full(5, 20.),
+            numpy.linspace(-5, 5, 5),
+            numpy.linspace(-numpy.pi, numpy.pi, 6)[:-1],
+            numpy.linspace(0, 20., 5),
+        )
+        assert (a * 5).tolist() == [
+            PtEtaPhiMassLorentzVector(pt=100, eta=-5,   phi=-numpy.pi + 0*numpy.pi/5, mass=0),
+            PtEtaPhiMassLorentzVector(pt=100, eta=-2.5, phi=-numpy.pi + 2*numpy.pi/5, mass=25),
+            PtEtaPhiMassLorentzVector(pt=100, eta=0,    phi=-numpy.pi + 4*numpy.pi/5, mass=50),
+            PtEtaPhiMassLorentzVector(pt=100, eta=2.5,  phi=-numpy.pi + 6*numpy.pi/5, mass=75),
+            PtEtaPhiMassLorentzVector(pt=100, eta=5,    phi=-numpy.pi + 8*numpy.pi/5, mass=100)
+        ]
+        assert a.sum().p < 1e-10
+        assert a.sum().mass == numpy.hypot(20 * numpy.cosh(a.eta), a.mass).sum()
+
     def test_lorentzvector_jagged(self):
         TLorentzVectorJagged = type("TLorentzVectorJagged", (awkward.JaggedArray, uproot_methods.classes.TLorentzVector.ArrayMethods), {})
         a = TLorentzVectorJagged.fromoffsets([0, 3, 3, 5, 10], TLorentzVectorArray(numpy.zeros(10), numpy.arange(10), numpy.zeros(10), numpy.zeros(10)))

uproot_methods/__init__.py

@@ -4,7 +4,7 @@
 
 from uproot_methods.classes.TVector2 import TVector2, TVector2Array
 from uproot_methods.classes.TVector3 import TVector3, TVector3Array
-from uproot_methods.classes.TLorentzVector import TLorentzVector, TLorentzVectorArray
+from uproot_methods.classes.TLorentzVector import TLorentzVector, TLorentzVectorArray, PtEtaPhiMassLorentzVector
 
 # convenient access to the version number
 from uproot_methods.version import __version__

uproot_methods/classes/TLorentzVector.py

@@ -405,6 +405,38 @@ def p(self):
     def p2(self):
         return self.p**2
 
+    def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
+        if "out" in kwargs:
+            raise NotImplementedError("in-place operations not supported")
+
+        if method != "__call__":
+            return NotImplemented
+
+        inputs = list(inputs)
+        for i in range(len(inputs)):
+            if isinstance(inputs[i], self.awkward.numpy.ndarray) and inputs[i].dtype == self.awkward.numpy.dtype(object) and len(inputs[i]) > 0:
+                idarray = self.awkward.numpy.frombuffer(inputs[i], dtype=self.awkward.numpy.uintp)
+                if (idarray == idarray[0]).all():
+                    inputs[i] = inputs[i][0]
+
+        if ufunc is self.awkward.numpy.multiply or ufunc is self.awkward.numpy.divide:
+            if sum(isinstance(x, PtEtaPhiMassArrayMethods) for x in inputs) > 1:
+                raise ValueError("cannot multiply or divide two PtEtaPhiMassArrayMethods")
+            this_input = None
+            for i in range(len(inputs)):
+                if isinstance(inputs[i], PtEtaPhiMassArrayMethods) and not isinstance(inputs[i], self.awkward.JaggedArray) and this_input is None:
+                    this_input = inputs[i]
+                    inputs[i] = self.awkward.Table(fPt=inputs[i]['fPt'], fMass=inputs[i]['fMass'])
+
+            out = super(PtEtaPhiMassArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs)
+            if this_input is not None:
+                out['fEta'] = this_input['fEta']
+                out['fPhi'] = this_input['fPhi']
+            return out
+
+        else:
+            return super(PtEtaPhiMassArrayMethods, self).__array_ufunc__(ufunc, method, *inputs, **kwargs)
+
 PtEtaPhiMassJaggedArrayMethods = PtEtaPhiMassArrayMethods.mixin(PtEtaPhiMassArrayMethods, awkward.JaggedArray)
 
 class Methods(Common, uproot_methods.base.ROOTMethods):
@@ -434,7 +466,7 @@ def _to_cartesian(self):
         return TLorentzVector(self.x,self.y,self.z,self.t)
 
     def __repr__(self):
-        return "TLorentzVector({0:.5g}, {1:.5g}, {2:.5g}, {3:.5g})".format(self._fP._fX, self._fP._fY, self._fP._fZ, self._fE)
+        return "TLorentzVector(x={0:.5g}, y={1:.5g}, z={2:.5g}, t={3:.5g})".format(self._fP._fX, self._fP._fY, self._fP._fZ, self._fE)
 
     def __str__(self):
         return repr(self)
@@ -706,7 +738,7 @@ def t(self):
         return self.awkward.numpy.sqrt(x*x + y*y + z*z + mass*mass*self.awkward.numpy.sign(mass))
 
     def __repr__(self):
-        return "TLorentzVector({0:.5g}, {1:.5g}, {2:.5g}, {3:.5g})".format(self._fPt, self._fEta, self._fPhi, self._fMass)
+        return "PtEtaPhiMassLorentzVector(pt={0:.5g}, eta={1:.5g}, phi={2:.5g}, mass={3:.5g})".format(self._fPt, self._fEta, self._fPhi, self._fMass)
 
 class PtEtaPhiMassLorentzVectorArray(PtEtaPhiMassArrayMethods, uproot_methods.base.ROOTMethods.awkward.ObjectArray):
     def __init__(self, pt, eta, phi, mass):