def spin_rotation_angle(pa, pb, pc, bachelor=2): """Calculate the angle between two spin-quantisation axes for the 3-body D->ABC decay aligned along the particle B and particle A. pa, pb, pc : 4-momenta of the final-state particles bachelor : index of the "bachelor" particle (0=A, 1=B, or 2=C) :param pa: :param pb: :param pc: :param bachelor: (Default value = 2) """ if bachelor == 2: return atfi.const(0.) pboost = lorentz_vector( -spatial_components(pb) / scalar(time_component(pb)), time_component(pb)) if bachelor == 0: pa1 = spatial_components(lorentz_boost(pa, pboost)) pc1 = spatial_components(lorentz_boost(pc, pboost)) return atfi.acos(scalar_product(pa1, pc1) / norm(pa1) / norm(pc1)) if bachelor == 1: pac = pa + pc pac1 = spatial_components(lorentz_boost(pac, pboost)) pa1 = spatial_components(lorentz_boost(pa, pboost)) return atfi.acos(scalar_product(pac1, pa1) / norm(pac1) / norm(pa1)) return None
def final_state_momenta(self, x): """ Return final state momenta p(A1), p(A2), p(B1), p(B2) for the decay defined by the phase space vector x. The momenta are calculated in the D rest frame. """ ma1a2 = self.m_a1a2(x) mb1b2 = self.m_b1b2(x) ctha = self.cos_helicity_a(x) cthb = self.cos_helicity_b(x) phi = self.phi(x) p0 = atfk.two_body_momentum(self.md, ma1a2, mb1b2) pA = atfk.two_body_momentum(ma1a2, self.ma1, self.ma2) pB = atfk.two_body_momentum(mb1b2, self.mb1, self.mb2) zeros = atfi.zeros(pA) p3A = atfk.rotate_euler(atfk.vector(zeros, zeros, pA), zeros, atfi.acos(ctha), zeros) p3B = atfk.rotate_euler(atfk.vector(zeros, zeros, pB), zeros, atfi.acos(cthb), phi) ea = atfi.sqrt(p0 ** 2 + ma1a2 ** 2) eb = atfi.sqrt(p0 ** 2 + mb1b2 ** 2) v0a = atfk.vector(zeros, zeros, p0 / ea) v0b = atfk.vector(zeros, zeros, -p0 / eb) p4A1 = atfk.lorentz_boost(atfk.lorentz_vector(p3A, atfi.sqrt(self.ma1 ** 2 + pA ** 2)), v0a) p4A2 = atfk.lorentz_boost(atfk.lorentz_vector(-p3A, atfi.sqrt(self.ma2 ** 2 + pA ** 2)), v0a) p4B1 = atfk.lorentz_boost(atfk.lorentz_vector(p3B, atfi.sqrt(self.mb1 ** 2 + pB ** 2)), v0b) p4B2 = atfk.lorentz_boost(atfk.lorentz_vector(-p3B, atfi.sqrt(self.mb2 ** 2 + pB ** 2)), v0b) return (p4A1, p4A2, p4B1, p4B2)
def helicity_angles_3body(pa, pb, pc): """Calculate 4 helicity angles for the 3-body D->ABC decay defined as: theta_r, phi_r : polar and azimuthal angles of the AB resonance in the D rest frame theta_a, phi_a : polar and azimuthal angles of the A in AB rest frame :param pa: :param pb: :param pc: """ theta_r = atfi.acos(-z_component(pc) / norm(spatial_components(pc))) phi_r = atfi.atan2(-y_component(pc), -x_component(pc)) pa_prime = lorentz_vector( rotate_euler(spatial_components(pa), -phi_r, atfi.pi() - theta_r, phi_r), time_component(pa)) pb_prime = lorentz_vector( rotate_euler(spatial_components(pb), -phi_r, atfi.pi() - theta_r, phi_r), time_component(pb)) w = time_component(pa) + time_component(pb) pab = lorentz_vector(-(pa_prime + pb_prime) / scalar(w), w) pa_prime2 = lorentz_boost(pa_prime, pab) theta_a = atfi.acos( z_component(pa_prime2) / norm(spatial_components(pa_prime2))) phi_a = atfi.atan2(y_component(pa_prime2), x_component(pa_prime2)) return (theta_r, phi_r, theta_a, phi_a)
def final_state_momenta(data): # Obtain the vectors of angles from the input tensor using the functions # provided by phasespace object cos_theta_jpsi = phsp.cos_theta1(data) cos_theta_phi = phsp.cos_theta2(data) phi = phsp.phi(data) # Rest-frame momentum of two-body Bs->Jpsi phi decay p0 = atfk.two_body_momentum(mb, mjpsi, mphi) # Rest-frame momentum of two-body Jpsi->mu mu decay pjpsi = atfk.two_body_momentum(mjpsi, mmu, mmu) # Rest-frame momentum of two-body phi->K K decay pphi = atfk.two_body_momentum(mphi, mk, mk) # Vectors of zeros and ones of the same size as the data sample # (needed to use constant values that do not depend on the event) zeros = atfi.zeros(phi) ones = atfi.ones(phi) # 3-vectors of Jpsi->mumu and phi->KK decays (in the corresponding rest frames), # rotated by the helicity angles p3jpsi = atfk.rotate_euler( atfk.vector(zeros, zeros, pjpsi * ones), zeros, atfi.acos(cos_theta_jpsi), zeros ) p3phi = atfk.rotate_euler( atfk.vector(zeros, zeros, pphi * ones), zeros, atfi.acos(cos_theta_phi), phi ) ejpsi = atfi.sqrt(p0 ** 2 + mjpsi ** 2) # Energy of Jpsi in Bs rest frame ephi = atfi.sqrt(p0 ** 2 + mphi ** 2) # Energy of phi in Bs rest frame v0jpsi = atfk.vector( zeros, zeros, p0 / ejpsi * ones ) # 3-vector of Jpsi in Bs rest frame v0phi = atfk.vector( zeros, zeros, -p0 / ephi * ones ) # 3-vector of phi in Bs rest frame # Boost momenta of final-state particles into Bs rest frame p4mu1 = atfk.lorentz_boost( atfk.lorentz_vector(p3jpsi, atfi.sqrt(mmu ** 2 + pjpsi ** 2) * ones), v0jpsi ) p4mu2 = atfk.lorentz_boost( atfk.lorentz_vector(-p3jpsi, atfi.sqrt(mmu ** 2 + pjpsi ** 2) * ones), v0jpsi ) p4k1 = atfk.lorentz_boost( atfk.lorentz_vector(p3phi, atfi.sqrt(mk ** 2 + pphi ** 2) * ones), v0phi ) p4k2 = atfk.lorentz_boost( atfk.lorentz_vector(-p3phi, atfi.sqrt(mk ** 2 + pphi ** 2) * ones), v0phi ) return (p4mu1, p4mu2, p4k1, p4k2)
def generate_rotation_and_boost(moms, minit, meanpt, ptcut, rnd): """ Generate 4-momenta of final state products according to 3-body phase space distribution moms - initial particle momenta (in the rest frame) meanpt - mean Pt of the initial particle ptcut - miminum Pt of the initial particle rnd - Auxiliary random tensor """ pt = generate_pt(rnd[:, 0], meanpt, ptcut, 200.) # Pt in GeV eta = generate_eta(rnd[:, 1]) # Eta phi = generate_phi(rnd[:, 2]) # Phi theta = 2. * atfi.atan(atfi.exp(-eta)) # Theta angle p = pt / atfi.sin(theta) # Full momentum e = atfi.sqrt(p**2 + minit**2) # Energy px = p * atfi.sin(theta) * atfi.sin(phi) # 3-momentum of initial particle py = p * atfi.sin(theta) * atfi.cos(phi) pz = p * atfi.cos(theta) p4 = atfk.lorentz_vector(atfk.vector(px, py, pz), e) # 4-momentum of initial particle rotphi = uniform_random(rnd[:, 3], 0., 2 * atfi.pi()) rotpsi = uniform_random(rnd[:, 4], 0., 2 * atfi.pi()) rottheta = atfi.acos(uniform_random(rnd[:, 5], -1, 1.)) moms2 = [] for m in moms: m1 = atfk.rotate_lorentz_vector(m, rotphi, rottheta, rotpsi) moms2 += [atfk.boost_from_rest(m1, p4)] return moms2
def m_prime_bc(self, sample): """ Square Dalitz plot variable m' """ mbc = atfi.sqrt(self.m2bc(sample)) return atfi.acos(2 * (mbc - math.sqrt(self.minbc)) / (math.sqrt(self.maxbc) - math.sqrt(self.minbc)) - 1.) / math.pi
def m_prime_ab(self, sample): """ Square Dalitz plot variable m' """ mab = atfi.sqrt(self.m2ab(sample)) return atfi.acos(2 * (mab - math.sqrt(self.minab)) / (math.sqrt(self.maxab) - math.sqrt(self.minab)) - 1.) / math.pi
def m_prime_ac(self, sample): """ Square Dalitz plot variable m' """ mac = atfi.sqrt(self.m2ac(sample)) return ( atfi.acos(2 * (mac - math.sqrt(self.minac)) / (math.sqrt(self.maxac) - math.sqrt(self.minac)) - 1.0) / math.pi)
def helicity_angles_4body(pa, pb, pc, pd): """Calculate 4 helicity angles for the 4-body E->ABCD decay defined as: theta_ab, phi_ab : polar and azimuthal angles of the AB resonance in the E rest frame theta_cd, phi_cd : polar and azimuthal angles of the CD resonance in the E rest frame theta_ac, phi_ac : polar and azimuthal angles of the AC resonance in the E rest frame theta_bd, phi_bd : polar and azimuthal angles of the BD resonance in the E rest frame theta_ad, phi_ad : polar and azimuthal angles of the AD resonance in the E rest frame theta_bc, phi_bc : polar and azimuthal angles of the BC resonance in the E rest frame phi_ab_cd : azimuthal angle between AB and CD phi_ac_bd : azimuthal angle between AC and BD phi_ad_bc : azimuthal angle between AD and BC :param pa: :param pb: :param pc: :param pd: """ theta_r = atfi.acos(-z_component(pc) / norm(spatial_components(pc))) phi_r = atfi.atan2(-y_component(pc), -x_component(pc)) pa_prime = lorentz_vector( rotate_euler(spatial_components(pa), -phi_r, atfi.pi() - theta_r, phi_r), time_component(pa), ) pb_prime = lorentz_vector( rotate_euler(spatial_components(pb), -phi_r, atfi.pi() - theta_r, phi_r), time_component(pb), ) w = time_component(pa) + time_component(pb) pab = lorentz_vector(-(pa_prime + pb_prime) / scalar(w), w) pa_prime2 = lorentz_boost(pa_prime, pab) theta_a = atfi.acos( z_component(pa_prime2) / norm(spatial_components(pa_prime2))) phi_a = atfi.atan2(y_component(pa_prime2), x_component(pa_prime2)) return (theta_r, phi_r, theta_a, phi_a)
def spherical_angles(pb): """theta, phi : polar and azimuthal angles of the vector pb :param pb: """ z1 = unit_vector( spatial_components(pb)) # New z-axis is in the direction of pb theta = atfi.acos(z_component(z1)) # Helicity angle phi = atfi.atan2(y_component(pb), x_component(pb)) # phi angle return (theta, phi)
def spherical_angles(pb): """ Return polar and azimuthal angles of the 3-vector or Lorentz vector :param pb: Input vector :returns: tuple (theta, phi), where theta is the polar and phi the azimuthal angles """ z1 = unit_vector( spatial_components(pb)) # Unit vector in the direction of pb theta = atfi.acos(z_component(z1)) # Helicity angle phi = atfi.atan2(y_component(pb), x_component(pb)) # phi angle return (theta, phi)
def euler_angles(x1, y1, z1, x2, y2, z2): """Calculate Euler angles (phi, theta, psi in the ZYZ convention) which transform the coordinate basis (x1, y1, z1) to the basis (x2, y2, z2). Both x1,y1,z1 and x2,y2,z2 are assumed to be orthonormal and right-handed. :param x1: :param y1: :param z1: :param x2: :param y2: :param z2: """ theta = atfi.acos(scalar_product(z1, z2)) phi = atfi.atan2(scalar_product(z1, y2), scalar_product(z1, x2)) psi = atfi.atan2(scalar_product(y1, z2), scalar_product(x1, z2)) return (phi, theta, psi)
def final_state_momenta(self, m2ab, m2bc, costhetaa, phia, phibc): """ Calculate 4-momenta of final state tracks in the 5D phase space m2ab, m2bc : invariant masses of AB and BC combinations (cos)thetaa, phia : direction angles of the particle A in the D reference frame phibc : angle of BC plane wrt. polarisation plane z x p_a """ thetaa = atfi.acos(costhetaa) m2ac = self.msqsum - m2ab - m2bc # Magnitude of the momenta p_a = atfk.two_body_momentum(self.md, self.ma, atfi.sqrt(m2bc)) p_b = atfk.two_body_momentum(self.md, self.mb, atfi.sqrt(m2ac)) p_c = atfk.two_body_momentum(self.md, self.mc, atfi.sqrt(m2ab)) cos_theta_b = (p_a * p_a + p_b * p_b - p_c * p_c) / (2. * p_a * p_b) cos_theta_c = (p_a * p_a + p_c * p_c - p_b * p_b) / (2. * p_a * p_c) # Fix momenta with p3a oriented in z (quantisation axis) direction p3a = atfk.vector(atfi.zeros(p_a), atfi.zeros(p_a), p_a) p3b = atfk.vector(p_b * Sqrt(1. - cos_theta_b**2), atfi.zeros(p_b), -p_b * cos_theta_b) p3c = atfk.vector(-p_c * Sqrt(1. - cos_theta_c**2), atfi.zeros(p_c), -p_c * cos_theta_c) # rotate vectors to have p3a with thetaa as polar helicity angle p3a = atfk.rotate_euler(p3a, atfi.const(0.), thetaa, atfi.const(0.)) p3b = atfk.rotate_euler(p3b, atfi.const(0.), thetaa, atfi.const(0.)) p3c = atfk.rotate_euler(p3c, atfi.const(0.), thetaa, atfi.const(0.)) # rotate vectors to have p3a with phia as azimuthal helicity angle p3a = atfk.rotate_euler(p3a, phia, atfi.const(0.), atfi.const(0.)) p3b = atfk.rotate_euler(p3b, phia, atfi.const(0.), atfi.const(0.)) p3c = atfk.rotate_euler(p3c, phia, atfi.const(0.), atfi.const(0.)) # rotate BC plane to have phibc as angle with the polarization plane p3b = atfk.rotate(p3b, phibc, p3a) p3c = atfk.rotate(p3c, phibc, p3a) # Define 4-vectors p4a = atfk.lorentz_vector(p3a, atfi.sqrt(p_a**2 + self.ma2)) p4b = atfk.lorentz_vector(p3b, atfi.sqrt(p_b**2 + self.mb2)) p4c = atfk.lorentz_vector(p3c, atfi.sqrt(p_c**2 + self.mc2)) return (p4a, p4b, p4c)
def random_rotation_and_boost(moms, rnd): """ Apply random boost and rotation to the list of 4-vectors moms : list of 4-vectors rnd : random array of shape (N, 6), where N is the length of 4-vector array """ pt = -5.0 * atfi.log( rnd[:, 0] ) # Random pT, exponential distribution with mean 5 GeV eta = rnd[:, 1] * 3.0 + 2.0 # Uniform distribution in pseudorapidity eta phi = rnd[:, 2] * 2.0 * atfi.pi() # Uniform distribution in phi theta = 2.0 * atfi.atan(atfi.exp(-eta)) # Theta angle is a function of eta p = pt / atfi.sin(theta) # Full momentum e = atfi.sqrt(p ** 2 + mb ** 2) # Energy of the Bs px = ( p * atfi.sin(theta) * atfi.sin(phi) ) # 3-momentum of initial particle (Bs meson) py = p * atfi.sin(theta) * atfi.cos(phi) pz = p * atfi.cos(theta) boost = atfk.lorentz_vector( atfk.vector(px, py, pz), e ) # Boost vector of the Bs meson rot_theta = atfi.acos( rnd[:, 3] * 2.0 - 1.0 ) # Random Euler rotation angles for Bs decay rot_phi = rnd[:, 4] * 2.0 * atfi.pi() rot_psi = rnd[:, 5] * 2.0 * atfi.pi() # Apply rotation and boost to the momenta in input list moms1 = [] for m in moms: m1 = atfk.rotate_lorentz_vector(m, rot_phi, rot_theta, rot_psi) moms1 += [atfk.boost_from_rest(m1, boost)] return moms1
import sys import tensorflow as tf sys.path.append("../") import amplitf.interface as atfi import amplitf.kinematics as atfk atfi.set_seed(2) rndvec = tf.random.uniform([32, 3], dtype=atfi.fptype()) v = rndvec[:, 0] th = atfi.acos(rndvec[:, 1]) phi = (rndvec[:, 2] * 2 - 1) * atfi.pi() p = atfk.lorentz_vector( atfk.vector(atfi.zeros(v), atfi.zeros(v), atfi.zeros(v)), atfi.ones(v)) bp = atfk.lorentz_boost( p, atfk.rotate_euler(atfk.vector(v, atfi.zeros(v), atfi.zeros(v)), th, phi, atfi.zeros(v))) print(bp) print(atfk.mass(bp))
def theta_prime_bc(self, sample): """ Square Dalitz plot variable theta' """ return atfi.acos(-self.cos_helicity_bc(sample)) / math.pi