def final_state_momenta(self, m2ab, m2bc): """ Calculate 4-momenta of final state tracks in a certain reference frame (decay is in x-z plane, particle A moves along z axis) m2ab, m2bc : invariant masses of AB and BC combinations """ m2ac = self.msqsum - m2ab - m2bc 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) p4a = atfk.lorentz_vector( atfk.vector(atfi.zeros(p_a), atfi.zeros(p_a), p_a), atfi.sqrt(p_a**2 + self.ma2)) p4b = atfk.lorentz_vector( atfk.vector(p_b * atfi.sqrt(1. - cos_theta_b**2), atfi.zeros(p_b), -p_b * cos_theta_b), atfi.sqrt(p_b**2 + self.mb2)) p4c = atfk.lorentz_vector( atfk.vector(-p_c * atfi.sqrt(1. - cos_theta_c**2), atfi.zeros(p_c), -p_c * cos_theta_c), atfi.sqrt(p_c**2 + self.mc2)) return (p4a, p4b, p4c)
def rotate(v, angle, axis): """rotate vector around an arbitrary axis, from ROOT implementation :param v: :param angle: :param axis: """ if (angle != atfi.zeros(angle)): ll = norm(axis) if (ll == atfi.zeros(ll)): sys.exit('ERROR in rotate: rotation axis is zero') else: sa = atfi.sin(angle) ca = atfi.cos(angle) dx = x_component(axis) / ll dy = y_component(axis) / ll dz = z_component(axis) / ll vx = x_component(v) vy = y_component(v) vz = z_component(v) _vx = (ca+(1-ca)*dx*dx)*vx + ((1-ca)*dx*dy-sa*dz) * \ vy + ((1-ca)*dx*dz+sa*dy)*vz _vy = ((1-ca)*dy*dx+sa*dz)*vx + (ca+(1-ca)*dy*dy) * \ vy + ((1-ca)*dy*dz-sa*dx)*vz _vz = ((1-ca)*dz*dx-sa*dy)*vx + \ ((1-ca)*dz*dy+sa*dx)*vy + (ca+(1-ca)*dz*dz)*vz return vector(_vx, _vy, _vz) else: return v
def generate_rho(cuts, rnd): """ Generate random combinations of rho -> pi pi and a kaon track """ meanrhopt = cuts[5] meankpt = cuts[0] ptcut = cuts[2] mrhogen = breit_wigner_random(rnd[:, 0], mrho, wrho) ones = atfi.ones(rnd[:, 0]) p = atfk.two_body_momentum(mrhogen, mpi * ones, mpi * ones) zeros = atfi.zeros(p) mom = [ atfk.lorentz_vector(atfk.vector(p, zeros, zeros), atfi.sqrt(p**2 + mpi**2)), atfk.lorentz_vector(-atfk.vector(p, zeros, zeros), atfi.sqrt(p**2 + mpi**2)) ] mom = generate_rotation_and_boost(mom, mrhogen, meanrhopt, ptcut, rnd[:, 2:8]) p4k = generate_4momenta(rnd[:, 8:11], meankpt, ptcut, atfi.const(mk)) p4pi1 = mom[0] p4pi2 = mom[1] return p4k, p4pi1, p4pi2
def generate_kstar(cuts, rnd): """ Generate random combinations of Kstar and a pion track """ meankstarpt = cuts[4] meanpipt = cuts[1] ptcut = cuts[2] mkstargen = breit_wigner_random(rnd[:, 0], mkstar, wkstar) ones = atfi.ones(rnd[:, 0]) p = atfk.two_body_momentum(mkstargen, mk * ones, mpi * ones) zeros = atfi.zeros(p) mom = [ atfk.lorentz_vector(atfk.vector(p, zeros, zeros), atfi.sqrt(p**2 + mk**2)), atfk.lorentz_vector(-atfk.vector(p, zeros, zeros), atfi.sqrt(p**2 + mpi**2)) ] mom = generate_rotation_and_boost(mom, mkstargen, meankstarpt, ptcut, rnd[:, 2:8]) p4k = mom[0] p4pi1 = mom[1] p4pi2 = generate_4momenta(rnd[:, 8:11], meanpipt, ptcut, atfi.const(mpi)) return p4k, p4pi1, p4pi2
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(Vector(zeros, zeros, pA), zeros, Acos(ctha), zeros) p3B = atfk.rotate_euler(Vector(zeros, zeros, pB), zeros, 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 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 axes_before_rotation(pb): """Calculate old (before rotation) axes in the frame aligned with the momentum vector pb :param pb: """ z1 = unit_vector( spatial_components(pb)) # New z-axis is in the direction of pb eb = time_component(pb) z0 = vector(atfi.zeros(eb), atfi.zeros(eb), atfi.ones(eb)) # Old z-axis vector x0 = vector(atfi.ones(eb), atfi.zeros(eb), atfi.zeros(eb)) # Old x-axis vector sp = scalar_product(z1, z0) a0 = z0 - z1 * scalar(sp) # vector in z-pb plane perpendicular to z0 x1 = tf.where(tf.equal(sp, 1.0), x0, -unit_vector(a0)) y1 = vector_product(z1, x1) # New y-axis x = vector(x_component(x1), x_component(y1), x_component(z1)) y = vector(y_component(x1), y_component(y1), y_component(z1)) z = vector(z_component(x1), z_component(y1), z_component(z1)) return (x, y, z)
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 four_momenta_from_helicity_angles(md, ma, mb, theta, phi): """Calculate the four-momenta of the decay products in D->AB in the rest frame of D md: mass of D ma: mass of A mb: mass of B theta: angle between A momentum in D rest frame and D momentum in its helicity frame phi: angle of plane formed by A & B in D helicity frame :param md: :param ma: :param mb: :param theta: :param phi: """ # Calculate magnitude of momentum in D rest frame p = two_body_momentum(md, ma, mb) # Calculate energy in D rest frame Ea = atfi.sqrt(p**2 + ma**2) Eb = atfi.sqrt(p**2 + mb**2) # Construct four-momenta with A aligned with D in D helicity frame Pa = lorentz_vector(vector(atfi.zeros(p), atfi.zeros(p), p), Ea) Pb = lorentz_vector(vector(atfi.zeros(p), atfi.zeros(p), -p), Eb) # rotate four-momenta Pa = rotate_lorentz_vector(Pa, atfi.zeros(phi), -theta, -phi) Pb = rotate_lorentz_vector(Pb, atfi.zeros(phi), -theta, -phi) return Pa, Pb
def rotate(v, angle, axis): """ Rotate vector around an arbitrary axis. Uses ROOT implementation. :param v: Input 3-vector :param angle: Rotation angle :param axis: 3-vector defining rotation axis :returns: Rotated vector """ if angle != atfi.zeros(angle): ll = norm(axis) if ll == atfi.zeros(ll): sys.exit("ERROR in rotate: rotation axis is zero") else: sa = atfi.sin(angle) ca = atfi.cos(angle) dx = x_component(axis) / ll dy = y_component(axis) / ll dz = z_component(axis) / ll vx = x_component(v) vy = y_component(v) vz = z_component(v) _vx = ((ca + (1 - ca) * dx * dx) * vx + ((1 - ca) * dx * dy - sa * dz) * vy + ((1 - ca) * dx * dz + sa * dy) * vz) _vy = (((1 - ca) * dy * dx + sa * dz) * vx + (ca + (1 - ca) * dy * dy) * vy + ((1 - ca) * dy * dz - sa * dx) * vz) _vz = (((1 - ca) * dz * dx - sa * dy) * vx + ((1 - ca) * dz * dy + sa * dx) * vy + (ca + (1 - ca) * dz * dz) * vz) return vector(_vx, _vy, _vz) else: return v
def axes_after_rotation(pb, oldaxes=None): """Calculate new (rotated) axes aligned with the momentum vector pb :param pb: :param oldaxes: (Default value = None) """ z1 = unit_vector( spatial_components(pb)) # New z-axis is in the direction of pb eb = time_component(pb) zeros = atfi.zeros(eb) ones = atfi.ones(eb) # Old z-axis vector z0 = vector(zeros, zeros, ones) if oldaxes == None else oldaxes[2] # Old x-axis vector x0 = vector(ones, zeros, zeros) if oldaxes == None else oldaxes[0] sp = scalar_product(z1, z0) a0 = z0 - z1 * scalar(sp) # vector in z-pb plane perpendicular to z0 x1 = tf.where(scalar(tf.equal(sp, 1.)), x0, -unit_vector(a0)) y1 = vector_product(z1, x1) # New y-axis return (x1, y1, z1)
def rotation_and_boost(ps, pb): """rotate and boost all momenta from the list ps to the rest frame of pb After the rotation, the coordinate system is defined as: z axis: direction of pb y axis: perpendicular to the plane formed by the old z and pb x axis: [y,z] ps : list of Lorentz vectors to rotate and boost pb : Lorentz vector defining the new frame :param ps: :param pb: :returns: list of transformed Lorentz vectors :rtype: ps1 """ newaxes = axes_after_rotation(pb) eb = time_component(pb) zeros = atfi.zeros(eb) # Boost vector in the rotated coordinates along z axis boost = vector(zeros, zeros, -norm(spatial_components(pb)) / eb) return nested_rotation_and_boost(ps, newaxes, boost)
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))