def test_vectors_Lorentz_constructors():
with pytest.raises(ValueError):
LorentzVector(*['str1', 'str2', 'str3', 'str4'])
#
lv1 = LorentzVector()
assert str(lv1) == str(np.array([0., 0., 0., 0.]))
lv2 = LorentzVector(1., 1., 1., 1.)
assert str(lv2) == str(np.array([1., 1., 1., 1.]))
def test_lorentz_vectors_properties_again():
p1 = LorentzVector(5., 5., 10., 5)
assert p1.x == 5.
assert p1.y == 5.
assert p1.z == 10.
assert p1.t == 5.
assert p1.mag == approx(-11.180340)
assert p1.mag2 == approx(-125)
assert p1.pt == sqrt(p1.x**2 + p1.y**2)
assert p1.p == sqrt(p1.x**2 + p1.y**2 + p1.z**2)
assert p1.p == sqrt(p1.pt**2 + p1.z**2)
assert p1.e == sqrt(p1.mag2 + p1.p**2)
assert p1.beta == p1.p / p1.e
def test_lorentz_vectors_boosting():
p3 = LorentzVector(5., 5., 10., 20)
assert p3.x == 5.
assert p3.y == 5.
assert p3.z == 10.
assert p3.e == 20.
assert p3.mag == sqrt(p3.e**2 - p3.p**2)
assert p3.beta == p3.p / p3.e
p4 = LorentzVector.from_pt_eta_phi_m(10., 3.5, pi / 3, 5.)
assert p4.pt == approx(10.)
assert p4.eta == approx(3.5)
assert p4.phi == approx(pi / 3)
assert p4.mag == approx(5.)
p5 = LorentzVector.from_pt_eta_phi(10., 3.9, -2 * (pi / 3), 20.)
assert p5.pt == approx(10.)
assert p5.eta == approx(3.9)
assert p5.phi == approx(-2 * (pi / 3))
assert p5.e == approx(20.)
assert p5.eta - p4.eta == approx(0.4)
assert p5.phi - p4.phi == approx(-pi)
assert p5.delta_r(p4) == approx(sqrt(0.4**2 + pi**2))
p6 = LorentzVector.from_pt_eta_phi(10., 3.9, -pi, 20.)
assert np.mod(p6.phi - p4.phi + np.pi, np.pi * 2) - np.pi == approx(
2 * (pi / 3))
assert np.mod(p4.phi - p6.phi + np.pi, np.pi * 2) - np.pi == approx(
-2 * (pi / 3))
p7 = LorentzVector(5., 5., 5., 0.)
assert p7.beta == np.inf
assert np.isnan(p7.gamma), "Gamma of the photons is +inf"
# assert p7.p == p7.e, "Momentum = Energy for photons"
assert np.isnan(p7.rapidity)
p8 = LorentzVector.from_pt_eta_phi(10., 2E2, -2 * (pi / 3),
20.) #sinh(eta) diverge quickly
assert p8.theta == approx(0.0)
assert p8.eta > 1E10
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_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_lorentz_vectors_errors():
with pytest.raises(AttributeError):
LorentzVector.boost(LorentzVector(), 1)
with pytest.raises(AttributeError):
LorentzVector.boost(LorentzVector(), [1, 2])
with pytest.raises(TypeError):
LorentzVector.boost(LorentzVector(), 1, 2, 3, 4)
with pytest.raises(TypeError):
LorentzVector.boost(LorentzVector(), 0, 1, 'a')
with pytest.raises(AttributeError):
LorentzVector.boost(LorentzVector(), ['a', 'b', 3])
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_vectors_Lorentz_operators():
with pytest.raises(TypeError):
LorentzVector.__iadd__(LorentzVector(), "a")
with pytest.raises(TypeError):
LorentzVector.__isub__(LorentzVector(), "a")
#
lv1 = LorentzVector(0., 0., 0., 0.)
lv2 = LorentzVector(1., 1., 1., 0.)
lv3 = LorentzVector(2., 2., 2., 1.)
lv4 = LorentzVector(3., 3., 3., 1.)
lv5 = LorentzVector(1., 1., 1., 6.)
assert np.all(lv1 == 0.)
assert_allclose(lv1 + lv2, LorentzVector(1., 1., 1., 0.))
assert_allclose(lv1 - lv2, LorentzVector(-1., -1., -1., 0.))
assert_allclose(lv1 * 2., LorentzVector(0., 0., 0., 0.))
assert_allclose(lv2 * 2., LorentzVector(2., 2., 2., 0.))
assert_allclose(2. * lv1, LorentzVector(0., 0., 0., 0.))
assert_allclose(2. * lv2, LorentzVector(2., 2., 2., 0.))
assert lv2.dot(lv1) == 0.
assert lv1.dot(lv2) == 0.
assert lv3.dot(lv4) == lv4.dot(lv3)
assert lv3.dot(lv4) == -17.
assert_allclose(lv3 / 2., LorentzVector(1., 1., 1., 0.5))
assert_allclose(lv4 / 3., LorentzVector(1., 1., 1., 1. / 3))
assert lv3.dot(lv5) == 0.0
lv1 *= 2.
lv2 *= 2.
assert_allclose(lv1, LorentzVector(0., 0., 0., 0.))
assert_allclose(lv2, LorentzVector(2., 2., 2., 0.))
lv1 /= 2.
lv2 /= 2.
assert_allclose(lv1, LorentzVector(0., 0., 0., 0.))
assert_allclose(lv2, LorentzVector(1., 1., 1., 0.))
assert_allclose(lv2, lv2)
def test_values(xyz, t):
np.seterr(divide='ignore', invalid='ignore')
x, y, z = xyz
rvec = ROOT.TLorentzVector(x, y, z, t)
svec = LorentzVector(x, y, z, t)
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 r(self):
# def in_basis(self, xhat, yhat, zhat):
# def rotate_axis(self, axis, angle):
# def rotate_euler(self, phi=0, theta=0, psi=0):
# def from_spherical_coords(cls, r, theta, phi):
# def from_cylindrical_coords(cls, rho, phi, z):
# def __new__(cls, x=0, y=0, z=0, t=0, dtype=np.double):
# def from_pt_eta_phi(cls, pt, eta, phi, t):
# def from_pt_eta_phi_m(cls, pt, eta, phi, m):
# def vect(self):
# def vect(self, obj):
# def p(self):
assert svec.p() == approx(rvec.P())
# def e(self):
# def eta(self):
# def gamma(self):
# def beta(self):
# def boost_vector(self):
assert_same(svec.boost_vector(), rvec.BoostVector())
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())