Merge pull request #12 from scikit-hep/issue-override-jaggedarray-sum

Issue override jaggedarray sum

uproot_methods/classes/TLorentzVector.py

@@ -276,9 +276,15 @@ def rotate_euler(self, phi=0, theta=0, psi=0):
     def islightlike(self, tolerance=1e-10):
         return awkward.util.numpy.absolute(self.mag2) < tolerance
 
+    def sum(self):
+        if isinstance(self, awkward.JaggedArray):
+            return TLorentzVectorArray.from_cartesian(self.x.sum(), self.y.sum(), self.z.sum(), self.t.sum())
+        else:
+            return TLorentzVector(self.x.sum(), self.y.sum(), self.z.sum(), self.t.sum())
+
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         if method != "__call__":
-            raise NotImplemented
+            return NotImplemented
 
         inputs = list(inputs)
         for i in range(len(inputs)):

uproot_methods/classes/TVector2.py

@@ -59,6 +59,12 @@ def x(self):
     def y(self):
         return self["fY"]
 
+    def sum(self):
+        if isinstance(self, awkward.JaggedArray):
+            return TVector2Array.from_cartesian(self.x.sum(), self.y.sum())
+        else:
+            return TVector2(self.x.sum(), self.y.sum())
+
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         if method != "__call__":
             return NotImplemented

uproot_methods/classes/TVector3.py

@@ -150,9 +150,15 @@ def rotate_euler(self, phi=0, theta=0, psi=0):
         out["fZ"] = z
         return out
 
+    def sum(self):
+        if isinstance(self, awkward.JaggedArray):
+            return TVector3Array.from_cartesian(self.x.sum(), self.y.sum(), self.z.sum())
+        else:
+            return TVector3(self.x.sum(), self.y.sum(), self.z.sum())
+
     def __array_ufunc__(self, ufunc, method, *inputs, **kwargs):
         if method != "__call__":
-            raise NotImplemented
+            return NotImplemented
 
         inputs = list(inputs)
         for i in range(len(inputs)):

uproot_methods/version.py

@@ -30,7 +30,7 @@
 
 import re
 
-__version__ = "0.2.8"
+__version__ = "0.2.9"
 version = __version__
 version_info = tuple(re.split(r"[-\.]", __version__))