Merge pull request #92 from wiso/add_angle

add angle method to TVector3(Array)

tests/test_vector.py

@@ -83,6 +83,13 @@ def test_vector3(self):
         assert a == TVector3(104.4, 205.5, 0)
         a *= 10
         assert a == TVector3(1044, 2055, 0)
+        assert a.angle(a) == 0
+        assert a.angle(a * 3) == 0
+        self.assertAlmostEqual(a.angle(-2 * a), numpy.pi)
+        self.assertAlmostEqual(a.angle(a.rotatez(0.01)), 0.01)
+        self.assertAlmostEqual(a.angle(a.rotatez(-0.01)), 0.01)
+        assert a.angle(a * 0) == 0
+
 
     def test_vector3_array(self):
         a = TVector3Array(numpy.zeros(10), numpy.arange(10), numpy.zeros(10))
@@ -100,6 +107,11 @@ def test_vector3_array(self):
             self.assertAlmostEqual(aroti.y,-ai.y)
             self.assertAlmostEqual(aroti.z,ai.z)
 
+        numpy.testing.assert_almost_equal(a.angle(a), numpy.zeros_like(a))
+        numpy.testing.assert_almost_equal(a.angle(3 * a), numpy.zeros_like(a))
+        # first element is null vector, skip it, should return 0
+        numpy.testing.assert_almost_equal(a[1:].angle(-2 * a[1:]), numpy.ones_like(a[1:]) * numpy.pi)
+
     def test_vector3_jagged(self):
         TVector3Jagged = type("TVector3Jagged", (awkward.JaggedArray, uproot_methods.classes.TVector3.ArrayMethods), {})
         a = TVector3Jagged.fromoffsets([0, 3, 3, 5, 10], TVector3Array(numpy.zeros(10), numpy.arange(10), numpy.zeros(10)))

uproot_methods/classes/TVector3.py

@@ -222,6 +222,18 @@ def cross(self, other):
         x, y, z = self._cross(other)
         return TVector3(x, y, z)
 
+    def angle(self, other):
+        denominator = math.sqrt(self.dot(self) * other.dot(other))
+        if denominator == 0:
+            # one of the vector is null
+            return 0.
+        cos_angle = self.dot(other) / denominator
+        if cos_angle > 1:
+            cos_angle = 1
+        elif cos_angle < -1:
+            cos_angle = -1
+        return math.acos(cos_angle)
+
     @property
     def theta(self):
         return math.atan2(self.rho, self.z)