def test_lorentz_vectors_properties():
lv0 = LorentzVector()
lv1 = LorentzVector(1., 1., 1., 1.)
lv2 = LorentzVector(1., 1., 1., 2.)
assert np.all(lv1.boostp3 == Vector3D(1., 1., 1.))
assert lv1.p == np.sqrt(3.)
assert lv1.e == 1.
assert lv1.beta == sqrt(3.)
assert np.all(lv2.boostp3 == Vector3D(0.5, 0.5, 0.5))
lv3 = LorentzVector(0., 0., 1., 0.)
beta = 0.05
gamma = 1 / sqrt(1 - beta**2)
lv4 = lv3.boost(Vector3D(0, 0, beta))
assert lv4.z == lv3.z * gamma, "Check length contraction"
assert lv4.x == lv3.x
assert lv4.y == lv3.x
assert lv4.t == -gamma * (lv3.t - beta * lv3.z) # Added -
assert_allclose(lv4,
LorentzVector(lv3.x, lv3.y, lv3.z * gamma,
-gamma * (lv3.t - beta * lv3.z))) # Added -
lv5 = LorentzVector(0., 0., 0., 1.)
assert lv5.beta == 0.
assert lv5.gamma == 1.
assert_allclose(lv5.boost(Vector3D(0, 0, 0)), lv5)
lv6 = lv5.boost(Vector3D(0, beta, 0))
assert lv6.x == lv5.x
assert lv6.z == lv5.z
assert lv6.t == lv5.t * gamma, "Check time dilation"
assert lv6.y == -gamma * (lv5.y - beta * lv5.t)
assert_allclose(
lv6,
LorentzVector(lv5.x, -gamma * (lv5.y - beta * lv5.t), lv5.z,
lv5.t * gamma))
def test_from_iterable():
iterable = [1, 2, 3]
vec = Vector3D(*iterable)
assert np.all(vec == Vector3D(1, 2, 3))
with pytest.raises(TypeError):
iterable = [1, 2, 3, 4]
Vector3D(*iterable)
def test_default_const_iterable():
several_vec = Vector3D(1, 2, [3, 4, 5])
assert several_vec.shape == (3, 3)
several_vec = Vector3D([1, 2, 3], [2, 3, 4], [3, 4, 5])
assert several_vec.shape == (3, 3)
with pytest.raises(ValueError):
Vector3D([1, 2, 3], [2, 3], [3, 4, 5])
def test_vectors_3D_rotations():
v1 = Vector3D.from_cylindrical_coords(1., 0., 0.)
assert v1.phi == 0.0
assert v1.rotate_axis(Vector3D.Z, pi / 2).phi == pi / 2
assert_allclose(v1.rotate_axis(Vector3D.Z, pi / 2),
Vector3D(0., 1., 0.),
atol=.0000001)
assert v1.rotate_axis(Vector3D.Y, pi).phi == pi
assert v1.rotate_axis(Vector3D.Y, -pi).phi == pi
assert_allclose(v1.rotate_axis(Vector3D.Y, pi),
Vector3D(-1., 0., 0.),
atol=.0000001)
assert_allclose(v1.rotate_axis(Vector3D.Y, -pi),
Vector3D(-1., 0., 0.),
atol=.0000001)
assert v1.rotate_axis(Vector3D.X, pi).phi == 0.
assert_allclose(v1.rotate_axis(Vector3D.X, pi),
Vector3D(1., 0., 0.),
atol=.0000001)
v2 = Vector3D.from_spherical_coords(1.0, pi / 2, pi / 2)
assert v2.phi == pi / 2
assert v2.theta == pi / 2
assert v2.rotate_axis(Vector3D.X, pi).phi == -pi / 2
assert v2.rotate_axis(Vector3D.X, pi).theta == pi / 2
assert_allclose(v2.rotate_axis(Vector3D.X, pi),
Vector3D(0., -1., 0.),
atol=.0000001)
v3 = Vector3D.from_spherical_coords(1.0, pi / 4, pi / 4)
angle = v2.angle(v3)
axis = v2.cross(v3)
assert_allclose(v2.rotate_axis(axis, angle), v3, atol=.0000001)
assert_allclose(v2.rotate_axis(-1. * axis, -angle), v3, atol=.0000001)
def test_vectors_Lorentz_rotations():
v1 = Vector3D.from_cylindrical_coords(1., 0., 0.)
v2 = Vector3D.from_spherical_coords(1.0, pi / 2, pi / 2)
v3 = Vector3D.from_spherical_coords(1.0, pi / 4, pi / 4)
angle = v2.angle(v3)
axis = v2.cross(v3)
with pytest.raises(AttributeError):
LorentzVector.rotate_axis(LorentzVector(), pi, 1)
with pytest.raises(AttributeError):
LorentzVector.rotate_axis(LorentzVector(), pi, [1, 2])
with pytest.raises(TypeError):
LorentzVector.rotate_axis(LorentzVector(), pi, 1, 2, 3, 4)
with pytest.raises(TypeError):
LorentzVector.rotate_axis(LorentzVector(), pi, 0, 1, 'a')
with pytest.raises(AttributeError):
LorentzVector.rotate_axis(LorentzVector(), pi, ['a', 'b', 3])
#
lv1 = LorentzVector(v1[0], v1[1], v1[2], 1.)
assert lv1.phi == 0.
assert lv1.rotate_axis(Vector3D.Z, pi / 2).phi == pi / 2
assert_allclose(lv1.rotate_axis(Vector3D.Z, pi / 2),
LorentzVector(0., 1., 0., 1.))
assert lv1.rotate_axis(Vector3D.Y, pi).phi == pi
assert lv1.rotate_axis(Vector3D.Y, -pi).phi == pi
assert_allclose(lv1.rotate_axis(Vector3D.Y, pi),
LorentzVector(-1., 0., 0., 1.))
assert_allclose(lv1.rotate_axis(Vector3D.Y, -pi),
LorentzVector(-1., 0., 0., 1.))
assert lv1.rotate_axis(Vector3D.X, pi).phi == 0.
assert_allclose(lv1.rotate_axis(Vector3D.X, pi),
LorentzVector(1., 0., 0., 1.))
lv2 = LorentzVector(v2[0], v2[1], v2[2], 2.0)
assert lv2.phi == pi / 2
assert lv2.theta == pi / 2
assert lv2.rotate_axis(Vector3D.X, pi).phi == -pi / 2
assert lv2.rotate_axis(Vector3D.X, pi).theta == pi / 2
assert_allclose(lv2.rotate_axis(Vector3D.X, pi),
LorentzVector(0., -1., 0., 2.0))
lv3 = LorentzVector(v3[0], v3[1], v3[2], 2.0)
assert_allclose(lv2.rotate_axis(axis, angle), lv3, atol=.0000001)
assert_allclose(lv2.rotate_axis(-1. * axis, -angle), lv3, atol=.0000001)
def test_values(xyz):
np.seterr(divide='ignore', invalid='ignore')
x, y, z = xyz
rvec = ROOT.TVector3(x, y, z)
svec = Vector3D(x, y, z)
assert_same(svec, rvec)
# assert svec.cos_theta() == approx(rvec.CosTheta()) # Not implemented
# Adding all the def's in LorentzVectors here, along with tests if applicable
# def origin(cls):
# def from_pandas(cls, pd_dataframe):
# def from_vector(cls, other):
# def dot(self, other):
# def mag(self):
assert svec.mag() == approx(rvec.Mag())
# def mag2(self):
assert svec.mag2() == approx(rvec.Mag2())
# def unit(self, inplace=False):
# def T(self):
# def T(self, val):
# def to_pd(self):
# def angle(self, other, normal=None):
# def __array_finalize__(self, obj):
# def __array_wrap__(self, out_arr, context=None):
# def __getitem__(self, item):
# def __setitem__(self, item, value):
# def dims(self):
# def _repr_html_(self):
# def __new__(cls, x=0, y=0, dtype=np.double):
# def phi(self):
assert svec.phi() == approx(rvec.Phi())
# def rho(self):
# assert svec.rho() == approx(rvec.Rho()) # Fails
# def angle(self, other):
# def pt2(self):
# def pt(self):
# def __new__(cls, x=0, y=0, z=0, dtype=np.double):
# def cross(self, other):
# def theta(self):
assert svec.theta() == approx(rvec.Theta())
def test_3Dvectors_properties():
v0 = Vector3D()
v1, v2 = Vector3D(1., 1., 1.), Vector3D(2., 2., 2.)
v3, v4 = Vector3D(-1., -1., -1.), Vector3D(-2., -2., -2.)
v5, v6 = Vector3D(1., 1., 0.), Vector3D(0., 0., 2.)
assert np.all(v0 == [0, 0, 0])
assert np.all(v1 == [1, 1, 1])
assert v1.mag2 == 3.
assert v1.mag == sqrt(3.)
assert v2.mag2 == 12.
assert v2.mag == sqrt(12.)
assert v1.unit.mag == 1.
v7 = v1.unit
assert np.all(v1.unit == v7.unit)
def to_lv(rvec):
return Vector3D(rvec.X(), rvec.Y(), rvec.Z())
def test_containers_properties():
with pytest.raises(IndexError):
Vector3D.__setitem__(Vector3D(), 3, 1.)
with pytest.raises(IndexError):
Vector3D.__getitem__(Vector3D(), 3)
#
v1 = Vector3D()
v1.x = 5.
v1.y = -5.
v1.z = 10
assert np.all(v1 == Vector3D(5., -5., 10.))
v1[:] = (-5., 5., 10.)
assert np.all(v1 == Vector3D(-5., 5., 10.))
v1[0] = 1.
assert np.all(v1 == Vector3D(1., 5., 10.))
v1[1] = 1.
assert np.all(v1 == Vector3D(1., 1., 10.))
v1[2] = 1.
assert np.all(v1 == Vector3D(1., 1., 1.))
assert v1[0] == 1.
assert v1[1] == 1.
assert v1[2] == 1.
assert len(v1) == 3.
assert v1.tolist() == [1., 1., 1.]
assert list(v1) == [1., 1., 1.]
assert [v for v in v1] == [1., 1., 1.]
#
with pytest.raises(IndexError):
LorentzVector.__setitem__(LorentzVector(), 4, 1.)
with pytest.raises(IndexError):
LorentzVector.__getitem__(LorentzVector(), 4)
#
lv1 = LorentzVector()
lv1.x = 5.
lv1.y = -5.
lv1.z = 10
lv1.t = 2.
assert np.all(lv1 == LorentzVector(5., -5., 10., 2.))
lv1.x = 5.
lv1.y = 5.
lv1.z = -10
lv1.t = 2.
assert np.all(lv1 == LorentzVector(5., 5., -10., 2.))
lv1[:] = (-5., 5., 10., -2.)
assert np.all(lv1 == LorentzVector(-5., 5., 10., -2.))
lv1[0] = 1.
assert np.all(lv1 == LorentzVector(1., 5., 10., -2.))
lv1[1] = 1.
assert np.all(lv1 == LorentzVector(1., 1., 10., -2.))
lv1[2] = 1.
assert np.all(lv1 == LorentzVector(1., 1., 1., -2.))
lv1[3] = 1.
assert np.all(lv1 == LorentzVector(1., 1., 1., 1.))
assert lv1[0] == 1.
assert lv1[1] == 1.
assert lv1[2] == 1.
assert lv1[3] == 1.
assert len(lv1) == 4.
assert list(lv1) == [1., 1., 1., 1.]
def test_vectors_3D_constructors():
v1 = Vector3D()
assert str(v1) == str(np.array([0., 0., 0.]))
assert str(v1) == str(Vector3D.origin())
v2 = Vector3D(1., 1., 1.)
assert str(v2) == str(np.array([1., 1., 1.]))
v3 = Vector3D.from_vector(v1)
assert str(v3) == str(np.array([0., 0., 0.]))
v4 = Vector3D(*[1.0, 1.0, 1.0])
assert str(v4) == str(np.array([1., 1., 1.]))
v5 = Vector3D.from_cylindrical_coords(1., 0., 1.)
assert np.all(v5 == Vector3D(1., 0., 1.))
assert v5.rho == 1.
v6 = Vector3D.from_cylindrical_coords(0.5, pi / 2, 0.)
assert_allclose(v6, Vector3D(0., 0.5, 0.))
v7 = Vector3D.from_spherical_coords(1.0, 0., 0.)
assert np.all(v7 == Vector3D(0., 0., 1.))
assert v7.r == 1.
assert v7.theta == 0.
v8 = Vector3D.from_spherical_coords(1.0, 0., pi / 2)
assert_allclose(v8, Vector3D(0., 0., 1.))
v9 = Vector3D.from_spherical_coords(2.0, pi / 2, pi / 4)
assert_allclose(v9, Vector3D(sqrt(2), sqrt(2), 0.))
def test_vectors_3D_operators():
v1 = Vector3D(0., 0., 0.)
v2 = Vector3D(1., 1., 1.)
v3 = Vector3D(2., 2., 2.)
v4 = Vector3D(3., 3., 3.)
v5 = Vector3D(1., 2., 3.)
assert np.all(v1 == 0.)
assert_allclose(v1 + v2, Vector3D(1., 1., 1.))
assert_allclose(v1 - v2, Vector3D(-1., -1., -1.))
assert_allclose(v1 - v2, -1. * Vector3D(1., 1., 1.))
assert_allclose(v1 * 2., Vector3D(0., 0., 0.))
assert_allclose(v2 * 2., Vector3D(2., 2., 2.))
assert_allclose(2. * v1, Vector3D(0., 0., 0.))
assert_allclose(2. * v2, Vector3D(2., 2., 2.))
assert v1.dot(v2) == 0.
assert v2.dot(v1) == 0.
assert v4.dot(v3) == v3.dot(v4)
assert v3.dot(v4) == 18.
assert_allclose(v3 / 2., v2)
assert_allclose(v4 / 3., v2)
v1 *= 2.
v2 *= 2.
assert np.all(v1 == Vector3D(0., 0., 0.))
assert np.all(v2 == Vector3D(2., 2., 2.))
v1 /= 2.
v2 /= 2.
assert np.all(v1 == Vector3D(0., 0., 0.))
assert np.all(v2 == Vector3D(1., 1., 1.))
assert (v1 + v2).dot(v3) == 6., "Check operations combination"
assert (v2 - v1).dot(v3) == 6., "Check operations combination"
assert v3.dot((v1 - v2)) == -6., "Check operations combination"
assert 18. / (v3.dot(v4)) == 1., "Check operations combination"
assert_allclose(v4 / (v3.dot(v2)), Vector3D(0.5, 0.5, 0.5))
assert_allclose(v2.cross(v2), Vector3D(0., 0., 0.))
assert_allclose(v2.cross(v5), Vector3D(1., -2., 1.))
assert_allclose(v5.cross(v2), -1 * v2.cross(v5))
assert_allclose(v2, v2)
def to_lv(rvec):
if isinstance(rvec, ROOT.TLorentzVector):
return LorentzVector(rvec.X(), rvec.Y(), rvec.Z(), rvec.T())
else:
return Vector3D(rvec.X(), rvec.Y(), rvec.Z())