def eta(Vector3D):
"""Return the pseudorapidity. for some reason, the libary I used, only
allowed calculation of eta for lorentz vectors and not Vector3D as well, so I made my own"""
if abs(Vector3D.costheta()) < 1.:
return -0.5 * log( (1. - Vector3D.costheta())/(1. + Vector3D.costheta()) )
else:
return 10E10 if Vector3D.z > 0 else -10E10
def MM2_calculator(df, B_M_nominal=5279.5):
"""Main function to calculate COM values, input is df, and output is same
df but with added columns for the COM values"""
# set up flight related vectors
df['pv_vector'] = df.apply(lambda x: array([x.TwoBody_OWNPV_X, x.TwoBody_OWNPV_Y, x.TwoBody_OWNPV_Z]), axis=1)
df['sv_vector'] = df.apply(lambda x: array([x.TwoBody_ENDVERTEX_X, x.TwoBody_ENDVERTEX_Y, x.TwoBody_ENDVERTEX_Z]),
axis=1)
df['flight'] = df.apply(lambda x: x.sv_vector - x.pv_vector, axis=1)
df['tan_theta'] = df.apply(lambda x: mag(perp(x.flight) / x.flight[-1]), axis=1)
# this is setting the 4momentum of the B meson from kinematic info from Two Body vertex
df['p4B'] = df.apply(lambda x: LorentzVector(x.TwoBody_PX, x.TwoBody_PY, x.TwoBody_PZ, x.TwoBody_PE), axis=1)
# PT estimate based on reconstructed mass and flight vector
df['pt_est'] = df.apply(lambda x: (B_M_nominal / x.TwoBody_M) * x.tan_theta * x.TwoBody_PZ, axis=1)
# calculating the eta and phi of the flight vector
df['flight_eta'] = df.apply(lambda x: eta(Vector3D(x.flight[0], x.flight[1], x.flight[2]).unit()), axis=1)
df['flight_phi'] = df.apply(lambda x: Vector3D(x.flight[0], x.flight[1], x.flight[2]).unit().phi(), axis=1)
# estimated B candidate for this estimated momentum, measured flight direction and expected true B mass
df['p4B_est'] = df.apply(lambda x: my_SetPtEtaPhiM(x.pt_est, x.flight_eta, x.flight_phi, B_M_nominal), axis=1)
# estimating the boost needed to get to the B's rest frame
df['boost_est'] = df.apply(lambda x: boostvector(x.p4B_est), axis=1)
# calculating the missing mass^2 - this can go negative with resolution
df['mm2'] = df.apply(lambda x: (x.p4B_est - x.p4B).mass2, axis=1)
return df
def event_passes(self,momentum):
# obtain 3-momentum
p=Vector3D(momentum.px,momentum.py,momentum.pz)
# get position of
x=float(self.distance/p.z)*p
if type(x) is np.ndarray: x=Vector3D(x[0],x[1],x[2])
# check if it passes
if eval(self.selection): return True
else:return False
def twobody_decay(self, p0, m0, m1, m2, phi, costheta):
"""
function that decays p0 > p1 p2 and returns p1,p2
"""
#get axis of p0
zaxis=Vector3D(0,0,1)
rotaxis=zaxis.cross(p0.vector).unit()
rotangle=zaxis.angle(p0.vector)
#energy and momentum of p2 in the rest frame of p0
energy1 = (m0*m0+m1*m1-m2*m2)/(2.*m0)
energy2 = (m0*m0-m1*m1+m2*m2)/(2.*m0)
momentum1 = math.sqrt(energy1*energy1-m1*m1)
momentum2 = math.sqrt(energy2*energy2-m2*m2)
#4-momentum of p1 and p2 in the rest frame of p0
en1 = energy1
pz1 = momentum1 * costheta
py1 = momentum1 * math.sqrt(1.-costheta*costheta) * np.sin(phi)
px1 = momentum1 * math.sqrt(1.-costheta*costheta) * np.cos(phi)
p1=LorentzVector(-px1,-py1,-pz1,en1)
if rotangle!=0: p1=p1.rotate(rotangle,rotaxis)
en2 = energy2
pz2 = momentum2 * costheta
py2 = momentum2 * math.sqrt(1.-costheta*costheta) * np.sin(phi)
px2 = momentum2 * math.sqrt(1.-costheta*costheta) * np.cos(phi)
p2=LorentzVector(px2,py2,pz2,en2)
if rotangle!=0: p2=p2.rotate(rotangle,rotaxis)
#boost p2 in p0 restframe
p1_=p1.boost(-1.*p0.boostvector)
p2_=p2.boost(-1.*p0.boostvector)
return p1_,p2_
def __init__(self, point=Point3D(0, 0, 0), normal=Vector3D(0, 0, 1)):
if not isinstance(point, Point3D):
raise NotImplementedError("Plane3D: invalid ``point'' argument")
if not isinstance(normal, Vector3D):
raise NotImplementedError("Plane3D: invalid ``normal'' argument")
if 0 == normal.mag2:
raise ValueError("Plane3D: ``normal'' must not be zero!")
self.point = point.copy()
self.normal = normal.copy()
def __init__(self, point=Point3D(0, 0, 0), vector=Vector3D(0, 0, 1)):
if not isinstance(point, Point3D):
raise NotImplementedError("Line3D: invalid ``point'' argument")
if not isinstance(vector, Vector3D):
raise NotImplementedError("Line3D: invalid ``vector'' argument")
if 0 == vector.mag2:
raise ValueError("Line3D: ``vector'' must not be zero!")
self.point = point.copy()
self.direction = vector.copy()
def boostvector(self):
"""Return the spatial component divided by the time component."""
return Vector3D(self.x / self.t, self.y / self.t, self.z / self.t)
def test_invalid_Armenteros_Podolanski_variables():
d1 = Vector3D(1., 2., 3.)
d2 = Vector3D(-1., -2., -3.)
with pytest.raises(ValueError):
Armenteros_Podolanski_variables(d1, d2)
def test_valid_Armenteros_Podolanski_variables():
d1 = Vector3D(1., 2., 3.)
d2 = Vector3D(1., -2., 3.)
assert Armenteros_Podolanski_variables(d1, d2) == (2.0, 0.0)
def __init__(self, *args, **kwargs):
self._vct = Vector3D(*args, **kwargs)
>>> line2 = ...
>>> point1 = ...
>>> point2 = ...
>>> plane1 = ...
>>> plane2 = ...
>>> for o1 in ( line1, line2 , point1 , point2 , plane1 , plane2 ) :
... for o2 in ( line1, line2 , point1 , point2 , plane1 , plane2 ) :
... print distance ( o1 , o2 )
"""
return obj1.distance(obj2)
# ===========================================================================
if '__main__' == __name__:
v = Vector3D(1, 2, 3)
p = Point3D(1, 2, 3)
p += 0.5 * v
p -= 0.5 * v
p, p + v * 4, p - (p + v)
line = Line3D(p, v)
p in line
line = Line3D.from_points(p, Point3D(3, 2, 1))
p in line
plane = Plane3D(p, v)
line in plane
plane = Plane3D.from_points(p, Point3D(3, 2, 1), Point3D(0, 0, 3))