def from_square_dalitz_plot(self, mprimeac, thprimeac): """ sample: Given mprimeac and thprimeac, returns 2D tensor for (m2ab, m2bc). Make sure you don't pass in sqDP corner points as they lie outside phsp. """ m2AC = ( 0.25 * (self.maxac**0.5 * atfi.cos(math.pi * mprimeac) + self.maxac**0.5 - self.minac**0.5 * atfi.cos(math.pi * mprimeac) + self.minac**0.5) **2) m2AB = ( 0.5 * (-(m2AC**2) + m2AC * self.ma**2 + m2AC * self.mb**2 + m2AC * self.mc**2 + m2AC * self.md**2 - m2AC * atfi.sqrt( (m2AC * (m2AC - 2.0 * self.ma**2 - 2.0 * self.mc**2) + self.ma**4 - 2.0 * self.ma**2 * self.mc**2 + self.mc**4) / m2AC) * atfi.sqrt( (m2AC * (m2AC - 2.0 * self.mb**2 - 2.0 * self.md**2) + self.mb**4 - 2.0 * self.mb**2 * self.md**2 + self.md**4) / m2AC) * atfi.cos(math.pi * thprimeac) - self.ma**2 * self.mb**2 + self.ma**2 * self.md**2 + self.mb**2 * self.mc**2 - self.mc**2 * self.md**2) / m2AC) m2BC = self.msqsum - m2AC - m2AB return tf.stack([m2AB, m2BC], axis=1)
def model(x): # Get phase space variables cosThetaK = phsp.coordinate(x, 0) cosThetaL = phsp.coordinate(x, 1) phi = phsp.coordinate(x, 2) # Derived quantities sinThetaK = atfi.sqrt(1.0 - cosThetaK * cosThetaK) sinThetaL = atfi.sqrt(1.0 - cosThetaL * cosThetaL) sinTheta2K = (1.0 - cosThetaK * cosThetaK) sinTheta2L = (1.0 - cosThetaL * cosThetaL) sin2ThetaK = (2.0 * sinThetaK * cosThetaK) cos2ThetaL = (2.0 * cosThetaL * cosThetaL - 1.0) # Decay density pdf = (3.0 / 4.0) * (1.0 - FL) * sinTheta2K pdf += FL * cosThetaK * cosThetaK pdf += (1.0 / 4.0) * (1.0 - FL) * sin2ThetaK * cos2ThetaL pdf += (-1.0) * FL * cosThetaK * cosThetaK * cos2ThetaL pdf += (1.0 / 2.0) * ( 1.0 - FL) * AT2 * sinTheta2K * sinTheta2L * atfi.cos(2.0 * phi) pdf += S5 * sin2ThetaK * sinThetaL * atfi.cos(phi) return pdf
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 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 helicity_amplitude_3body(thetaR, phiR, thetaA, phiA, spinD, spinR, mu, lambdaR, lambdaA, lambdaB, lambdaC): """Calculate complex helicity amplitude for the 3-body decay D->ABC thetaR, phiR : polar and azimuthal angles of AB resonance in D rest frame thetaA, phiA : polar and azimuthal angles of A in AB rest frame spinD : D spin spinR : spin of the intermediate R resonance mu : D spin projection onto z axis lambdaR : R resonance helicity lambdaA : A helicity lambdaB : B helicity lambdaC : C helicity :param thetaR: :param phiR: :param thetaA: :param phiA: :param spinD: :param spinR: :param mu: :param lambdaR: :param lambdaA: :param lambdaB: :param lambdaC: """ lambda1 = lambdaR - lambdaC lambda2 = lambdaA - lambdaB ph = (mu - lambda1) / 2. * phiR + (lambdaR - lambda2) / 2. * phiA d_terms = wigner_small_d(thetaR, spinD, mu, lambda1) * \ wigner_small_d(thetaA, spinR, lambdaR, lambda2) h = atfi.complex(d_terms * atfi.cos(ph), d_terms * atfi.sin(ph)) return h
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 square_dalitz_plot_jacobian(self, sample): """ sample: [mAB^2, mBC^2] Return the jacobian determinant (|J|) of tranformation from dmAB^2*dmBC^2 -> |J|*dMpr*dThpr where Mpr, Thpr are defined in (AC) frame. """ mPrime = self.m_prime_ac(sample) thPrime = self.theta_prime_ac(sample) diff_AC = tf.cast( atfi.sqrt(self.maxac) - atfi.sqrt(self.minac), atfi.fptype()) mAC = atfi.const(0.5) * diff_AC * ( Const(1.0) + atfi.cos(atfi.pi() * mPrime)) + tf.cast( atfi.sqrt(self.minac), atfi.fptype()) mACSq = mAC * mAC eAcmsAC = (atfi.const(0.5) * (mACSq - tf.cast(self.mc2, atfi.fptype()) + tf.cast(self.ma2, atfi.fptype())) / mAC) eBcmsAC = (atfi.const(0.5) * (tf.cast(self.md, atfi.fptype())**2.0 - mACSq - tf.cast(self.mb2, atfi.fptype())) / mAC) pAcmsAC = atfi.sqrt(eAcmsAC**2.0 - tf.cast(self.ma2, atfi.fptype())) pBcmsAC = atfi.sqrt(eBcmsAC**2.0 - tf.cast(self.mb2, atfi.fptype())) deriv1 = Pi() * atfi.const(0.5) * diff_AC * atfi.sin( atfi.pi() * mPrime) deriv2 = Pi() * atfi.sin(atfi.pi() * thPrime) return atfi.const(4.0) * pAcmsAC * pBcmsAC * mAC * deriv1 * deriv2
def rotate_euler(v, phi, theta, psi): """Perform 3D rotation of the 3-vector v : vector to be rotated phi, theta, psi : Euler angles in Z-Y-Z convention :param v: :param phi: :param theta: :param psi: """ # rotate Z (phi) c1 = atfi.cos(phi) s1 = atfi.sin(phi) c2 = atfi.cos(theta) s2 = atfi.sin(theta) c3 = atfi.cos(psi) s3 = atfi.sin(psi) # rotate Y (theta) fzx2 = -s2 * c1 fzy2 = s2 * s1 fzz2 = c2 # rotate Z (psi) fxx3 = c3 * c2 * c1 - s3 * s1 fxy3 = -c3 * c2 * s1 - s3 * c1 fxz3 = c3 * s2 fyx3 = s3 * c2 * c1 + c3 * s1 fyy3 = -s3 * c2 * s1 + c3 * c1 fyz3 = s3 * s2 # Transform v vx = x_component(v) vy = y_component(v) vz = z_component(v) _vx = fxx3 * vx + fxy3 * vy + fxz3 * vz _vy = fyx3 * vx + fyy3 * vy + fyz3 * vz _vz = fzx2 * vx + fzy2 * vy + fzz2 * vz return vector(_vx, _vy, _vz)
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
def resonant_lass_lineshape(m2ab, m0, gamma0, a, r, ma, mb): """ LASS line shape, resonant part """ m = atfi.sqrt(m2ab) q0 = atfk.two_body_momentum(m0, ma, mb) q = atfk.two_body_momentum(m, ma, mb) cot_deltab = 1. / a / q + 1. / 2. * r * q phase = atfi.atan(1. / cot_deltab) width = gamma0 * q / m * m0 / q0 ampl = relativistic_breit_wigner(m2ab, m0, width) * atfi.complex( atfi.cos(phase), atfi.sin(phase)) * atfi.cast_complex( m2ab * gamma0 / q0) return ampl
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 normal_random(rnd1, rnd2): """ Normal distribution from two random numbers """ return atfi.sqrt(-2. * atfi.log(rnd1)) * atfi.cos(2. * math.pi * rnd2)