diff --git a/tests/test_vector.py b/tests/test_vector.py index 5f2bdd0..94d9436 100644 --- a/tests/test_vector.py +++ b/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))) diff --git a/uproot_methods/classes/TVector3.py b/uproot_methods/classes/TVector3.py index 89c1320..15a6c1b 100644 --- a/uproot_methods/classes/TVector3.py +++ b/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)