Compare 2010 proton-proton collision CoM values with new method

boosting.py

# -*- coding: utf-8 -*-
from __future__ import print_function
from math import sqrt, sin
from array import array

import ROOT

# Proton mass in MeV
PROTON_MASS = 938.27

# 3.5 TeV is 3.5e6 MeV
ebeam = 3.5*10e6
beam1x = [
    207.7, 215.5, -322.3, 208.7, 209.1, -318.6, -323.3, -321.4, -320.2,
    -323.4, -334.9
]
beam2x = [
    -341.5, -341.0, 188.3, -332.0, -330.7, 187.4, 188.4, 187.9, 181.0,
    179.2, 204.7
]
# BUG: There's a typo in the old code: 17.4 is written is 17/4
beam1y = [
    23.1, 26.6, 15.3, 22.4, 25.1, 22.9, 24.0, 17.4, 17.2, 18.3,
    23.7
]
beam2y = [
    31.2, 33.5, 28.3, 31.2, 22.9, 27.1, 24.5, 21.7, 26.3, 22.3,
    25.6
]
fill_list = [
    1089, 1090, 1101, 1104, 1107, 1117, 1118, 1119, 1121,
    1122, 1134
]
# Map fills to beam crossing angles, converting angles to radians (from µrad)
beam1_fill_angles = {
    fill: (1e-6*beam1x[idx], 1e-6*beam1y[idx])
    for idx, fill in enumerate(fill_list)
}
beam2_fill_angles = {
    fill: (1e-6*beam2x[idx], 1e-6*beam2y[idx])
    for idx, fill in enumerate(fill_list)
}
run_to_fill = {}
for run in range(71474, 71495):
    run_to_fill[run] = 1089
for run in range(71525, 71531):
    run_to_fill[run] = 1090
for run in range(71803, 71817):
    run_to_fill[run] = 1101
for run in range(71857, 71888):
    run_to_fill[run] = 1104
for run in range(71919, 71932):
    run_to_fill[run] = 1107
for run in range(72216, 72223):
    run_to_fill[run] = 1117
for run in range(72251, 72255):
    run_to_fill[run] = 1118
for run in range(72266, 72270):
    run_to_fill[run] = 1119
for run in range(72309, 72311):
    run_to_fill[run] = 1121
for run in range(72325, 72332):
    run_to_fill[run] = 1122
for run in range(72908, 72910):
    run_to_fill[run] = 1134
# HACK: Some fills lack beam info. Use that of an adjacent fill (same
# polarity!)
# Fill 1109 (run 71958) looks like fill 1107:
for run in range(71958, 71959):
    run_to_fill[run] = 1107
# Fill 1112 (run 72020, 72021) looks like fill 1117:
for run in range(72020, 72022):
    run_to_fill[run] = 1117


def get_beams(ebeam, beam1_angles, beam2_angles):
    """Return the four-vectors of the two LHC proton beams.

    Keyword arguments:
    ebeam -- LHC beam energy in MeV
    beam1_angles -- 2-tuple of beam 1 crossing angles with the z-axis as
                    (horizontal angle, vertical angle) in radians
    beam2_angles -- 2-tuple of beam 2 crossing angles with the z-axis as
                    (horizontal angle, vertical angle) in radians
    """
    # p^2 = E^2 - m^2
    ptot = sqrt((ebeam*ebeam) - (PROTON_MASS*PROTON_MASS))
    # Compute the x, y, and z components of the beam unit vectors
    sx1 = sin(beam1_angles[0])
    sy1 = sin(beam1_angles[1])
    sx2 = sin(beam2_angles[0])
    sy2 = sin(beam2_angles[1])
    z1 = sqrt(1.0 - (sx1*sx1) - (sy1*sy1))
    z2 = sqrt(1.0 - (sx2*sx2) - (sy2*sy2))
    # TLorentzVector is initialised as (Px, Py, Pz, E)
    beam_1 = ROOT.TLorentzVector(ptot*sx1, ptot*sy1, ptot*z1, ebeam)
    beam_2 = ROOT.TLorentzVector(-ptot*sx2, -ptot*sy2, -ptot*z2, ebeam)
    return beam_1, beam_2


def boost_particle_to_com(particle, ebeam, beam1_angles, beam2_angles):
    """Boost particle to the proton-proton collision centre-of-mass frame.

    Returns a new TLorentzVector, representing the boosted particle.
    Keyword arguments:
    ebeam -- LHC beam energy in MeV
    beam1_angles -- 2-tuple of beam 1 crossing angles with the z-axis as
                    (horizontal angle, vertical angle) in radians
    beam2_angles -- 2-tuple of beam 2 crossing angles with the z-axis as
                    (horizontal angle, vertical angle) in radians
    """
    # Get the CoM frame
    beam1, beam2 = get_beams(ebeam, beam1_angles, beam2_angles)
    # SMASH!
    pp_frame = beam1 + beam2

    # If the crossing angles are not symmetric, the plane containing the beams
    # in the CoM frame is not coincident with the laboratory xz plane, so the
    # boosted particle must be rotated to preserve the definition of transverse
    # quantities like pT and rapidity
    zaxis = ROOT.TVector3(0, 0, 1)
    # Return the x, y and z components of the CoM boost
    boost = -pp_frame.BoostVector()
    beam1.Boost(boost)
    z_angle = -beam1.Angle(zaxis)
    pp_zaxis = zaxis.Cross(beam1.Vect())

    # Copy the input particle, boost it, and rotate it
    boosted = ROOT.TLorentzVector(particle)
    boosted.Boost(boost)
    boosted.Rotate(z_angle, pp_zaxis)
    return boosted


def boosting():
    """Tests get_boost_vector by comparing results with 2010."""
    f = ROOT.TFile((
        '/afs/cern.ch/work/a/apearce/CharmProduction2010'
        '/DpToKKpi/ocptftuple_data_dsdplus.root'
    ))
    t = f.Get('Ds2piPhiTuple/DecayTree')
    entries = t.GetEntries()
    print('Got ntuple with {0} entries'.format(entries))

    # Set up branches
    ds_px = array('f', [0.0])
    ds_py = array('f', [0.0])
    ds_pz = array('f', [0.0])
    ds_pe = array('f', [0.0])
    ds_m = array('f', [0.0])
    ds_y_lab = array('f', [0.0])
    ds_pt_lab = array('f', [0.0])
    ds_y_cm = array('f', [0.0])
    ds_pt_cm = array('f', [0.0])
    run_number = array('i', [0])
    t.SetBranchAddress('D_s_PX', ds_px)
    t.SetBranchAddress('D_s_PY', ds_py)
    t.SetBranchAddress('D_s_PZ', ds_pz)
    t.SetBranchAddress('D_s_PE', ds_pe)
    t.SetBranchAddress('D_s_M', ds_m)
    t.SetBranchAddress('ptlab', ds_pt_lab)
    t.SetBranchAddress('ylab', ds_y_lab)
    t.SetBranchAddress('ptcm', ds_pt_cm)
    t.SetBranchAddress('ycm', ds_y_cm)
    t.SetBranchAddress('runNumber', run_number)

    # Set up output file and tree
    fout = ROOT.TFile('boosting.root', 'recreate')
    tout = ROOT.TTree('DecayTree', 'DecayTree')
    ds_pt_cm_new = array('f', [0.0])
    ds_y_cm_new = array('f', [0.0])
    ds_pt_cm_diff = array('f', [0.0])
    ds_y_cm_diff = array('f', [0.0])
    tout.Branch('D_s_PX', ds_px, 'D_s_PX/F')
    tout.Branch('D_s_PY', ds_py, 'D_s_PY/F')
    tout.Branch('D_s_PZ', ds_pz, 'D_s_PZ/F')
    tout.Branch('D_s_PE', ds_pe, 'D_s_PE/F')
    tout.Branch('D_s_M', ds_m, 'D_s_M/F')
    tout.Branch('D_s_PT_LAB', ds_pt_lab, 'D_s_PT_LAB/F')
    tout.Branch('D_s_Y_LAB', ds_y_lab, 'D_s_Y_LAB/F')
    tout.Branch('D_s_PT_CM', ds_pt_cm, 'D_s_PT_CM/F')
    tout.Branch('D_s_Y_CM', ds_y_cm, 'D_s_Y_CM/F')
    tout.Branch('D_s_PT_CM_NEW', ds_pt_cm_new, 'D_s_PT_CM_NEW/F')
    tout.Branch('D_s_Y_CM_NEW', ds_y_cm_new, 'D_s_Y_CM_NEW/F')
    tout.Branch('D_s_PT_CM_DIFF', ds_pt_cm_diff, 'D_s_PT_CM_DIFF/F')
    tout.Branch('D_s_Y_CM_DIFF', ds_y_cm_diff, 'D_s_Y_CM_DIFF/F')
    tout.Branch('run_number', run_number, 'run_number/I')

    ds = ROOT.TLorentzVector()
    for entry in range(1, entries + 1):
        t.GetEntry(entry)
        ds.SetPxPyPzE(ds_px[0], ds_py[0], ds_pz[0], ds_pe[0])
        fill = run_to_fill[run_number[0]]
        b1 = beam1_fill_angles[fill]
        b2 = beam2_fill_angles[fill]
        boosted = boost_particle_to_com(ds, ebeam, b1, b2)
        ds_pt_cm_new[0] = boosted.Pt()
        ds_y_cm_new[0] = boosted.Rapidity()
        ds_pt_cm_diff[0] = ds_pt_cm_new[0] - ds_pt_cm[0]
        ds_y_cm_diff[0] = ds_y_cm_new[0] - ds_y_cm[0]

        tout.Fill()

    fout.Write()
    fout.Close()

    f.Close()


if __name__ == '__main__':
    # Shut up, ROOT
    ROOT.gROOT.SetBatch(True)
    boosting()