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 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 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, 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 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 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 get_gen_sample(sample='mu'): den_columns = [] den_columns += ['Lambda_b0_TRUEP_E'] den_columns += ['Lambda_b0_TRUEP_X'] den_columns += ['Lambda_b0_TRUEP_Y'] den_columns += ['Lambda_b0_TRUEP_Z'] den_columns += ['Lambda_cplus_TRUEP_E'] den_columns += ['Lambda_cplus_TRUEP_X'] den_columns += ['Lambda_cplus_TRUEP_Y'] den_columns += ['Lambda_cplus_TRUEP_Z'] den_columns += [sample + 'minus_TRUEP_E'] den_columns += [sample + 'minus_TRUEP_X'] den_columns += [sample + 'minus_TRUEP_Y'] den_columns += [sample + 'minus_TRUEP_Z'] den_columns += ['nu_' + sample + '~_TRUEP_E'] den_columns += ['nu_' + sample + '~_TRUEP_X'] den_columns += ['nu_' + sample + '~_TRUEP_Y'] den_columns += ['nu_' + sample + '~_TRUEP_Z'] den_fname = '~/LbToLclnu_RunTwo/Selection/PID/FFs/GenMC/Lc' + sample.capitalize( ) + 'Nu_gen.root' df_den = rpd.read_root(den_fname, columns=den_columns, key='MCDecayTreeTuple/MCDecayTree') PLc_lab = atfk.lorentz_vector( atfk.vector(df_den['Lambda_cplus_TRUEP_X'], df_den['Lambda_cplus_TRUEP_Y'], df_den['Lambda_cplus_TRUEP_Z']), df_den['Lambda_cplus_TRUEP_E']) Pl_lab = atfk.lorentz_vector( atfk.vector(df_den[sample + 'minus_TRUEP_X'], df_den[sample + 'minus_TRUEP_Y'], df_den[sample + 'minus_TRUEP_Z']), df_den[sample + 'minus_TRUEP_E']) PNu_lab = atfk.lorentz_vector( atfk.vector(df_den["nu_" + sample + "~_TRUEP_X"], df_den["nu_" + sample + "~_TRUEP_Y"], df_den["nu_" + sample + "~_TRUEP_Z"]), df_den["nu_" + sample + "~_TRUEP_E"]) PLb_lab = PLc_lab + Pl_lab + PNu_lab df_den['Lb_True_Q2'], df_den['Lb_True_Costhetal'] = get_phasespace_vars( PLb_lab, PLc_lab, Pl_lab) return df_den[['Lb_True_Q2', 'Lb_True_Costhetal']].to_numpy(),
def generate_4momenta(rnd, meanpt, ptcut, m): """ Generate random 4-momenta according to specified mean Pt, minimum Pt, and mass of the particle, flat in eta and phi """ pt = generate_pt(rnd[:, 0], meanpt, ptcut, atfi.const(200.)) # Pt in GeV eta = generate_eta(rnd[:, 1]) # Eta phi = generate_phi(rnd[:, 2]) # Phi theta = 2. * atfi.atan(atfi.exp(-eta)) p = pt / atfi.sin(theta) # Full momentum e = atfi.sqrt(p**2 + m**2) # Energy px = p * atfi.sin(theta) * atfi.sin(phi) py = p * atfi.sin(theta) * atfi.cos(phi) pz = p * atfi.cos(theta) return atfk.lorentz_vector(atfk.vector(px, py, pz), e)
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))