Ejemplo n.º 1
0
def preprocessing(jet):
    jet = copy.deepcopy(jet)

    jet=jet.reshape(-1,4)
    n_consti=len(jet)

    # find the jet (eta, phi)
    center=jet.sum(axis=0)

    v_jet=LorentzVector(center[1], center[2], center[3], center[0])

    # centering parameters
    phi=v_jet.phi()
    bv = v_jet.boost_vector()
    bv.set_perp(0)    

    for i in range(n_consti):
        v = LorentzVector(jet[i,1], jet[i,2], jet[i,3], jet[i,0])
        v.rotate_z(-phi)
        v.boost(-bv)  
        jet[i, 0]=v[3] #e
        jet[i, 1]=v[0] #px
        jet[i, 2]=v[1] #py
        jet[i, 3]=v[2] #pz

    # rotating parameters
    weighted_phi=0
    weighted_eta=0
    for i in range(n_consti):
        if jet[i,0]<1e-10: # pass zero paddings
            continue
        v = LorentzVector(jet[i,1], jet[i,2], jet[i,3], jet[i,0])
        r=np.sqrt(v.phi()**2 + v.eta()**2)
        if r == 0: # in case there is only one component. In fact these data points should generally be invalid.
            continue
        weighted_phi += v.phi() * v.E()/r
        weighted_eta += v.eta() * v.E()/r
    #alpha = np.arctan2(weighted_phi, weighted_eta) # approximately align at eta
    alpha = np.arctan2(weighted_eta, weighted_phi) # approximately align at phi

    for i in range(n_consti):
        v = LorentzVector(jet[i,1], jet[i,2], jet[i,3], jet[i,0])
        #v.rotate_x(alpha) # approximately align at eta
        v.rotate_x(-alpha) # approximately align at phi

        jet[i, 0]=v[3]
        jet[i, 1]=v[0]
        jet[i, 2]=v[1]
        jet[i, 3]=v[2]

    #jet=jet.reshape(1,-1)
    jet=jet.ravel()
    return jet
Ejemplo n.º 2
0
def preprocessing(
    jet
):  # every entry would be a sequence of 4-vecs (E, px, py, pz) of jet constituents
    jet = copy.deepcopy(jet)
    jet = jet.reshape(-1, 4)
    n_consti = len(jet)

    # find the jet (eta, phi)
    center = jet.sum(axis=0)

    v_jet = LorentzVector(center[1], center[2], center[3], center[0])

    # centering
    phi = v_jet.phi()
    bv = v_jet.boost_vector()
    bv.set_perp(0)

    # rotating
    weighted_phi = 0
    weighted_eta = 0
    for i in range(n_consti):
        if jet[i, 0] < 1e-10:
            continue
        v = LorentzVector(jet[i, 1], jet[i, 2], jet[i, 3], jet[i, 0])
        r = np.sqrt(v.phi()**2 + v.eta()**2)
        weighted_phi += v.phi() * v.E() / r
        weighted_eta += v.eta() * v.E() / r
    alpha = -np.arctan2(weighted_phi, weighted_eta)

    for i in range(n_consti):
        v = LorentzVector(jet[i, 1], jet[i, 2], jet[i, 3], jet[i, 0])
        v.rotate_z(-phi)
        v.boost(-bv)
        v.rotate_x(alpha)

        jet[i, 0] = v[3]
        jet[i, 1] = v[0]
        jet[i, 2] = v[1]
        jet[i, 3] = v[2]

    jet = jet.reshape(1, -1)

    return jet
Ejemplo n.º 3
0
def preprocess(jet, cluster, output="kt", regression=False, R_clustering=0.3):
    """
    preprocesses the data to make it usable by the recnn
    Preprocessing algorithm:
    1. j = the highest pt anti-kt jet (R=1)
    2. run kt (R=0.3) on the constituents c of j, resulting in subjets sj1, sj2, ..., sjN
    3. phi = sj1.phi(); for all c, do c.rotate_z(-phi)
    4. bv = sj1.boost_vector(); bv.set_perp(0); for all c, do c.boost(-bv)
    5. deltaz = sj1.pz - sj2.pz; deltay = sj1.py - sj2.py; alpha = -atan2(deltaz, deltay); for all c, do c.rotate_x(alpha)
    6. if sj3.pz < 0: for all c, do c.set_pz(-c.pz)
    7. finally recluster all transformed constituents c into a single jet
    """
    jet = copy.deepcopy(jet)
    constituents = jet["content"][jet["tree"][:, 0] == -1]
    if regression:
        genpt = jet["genpt"]

    ### run kt (R=0.3) on the constituents c of j, resulting in subjets sj1, sj2, ..., sjN ###
    subjets = cluster(constituents, R=R_clustering, jet_algorithm=0)
    oldeta = jet["eta"]
    oldpt = jet['pt']
    ### Rot phi ###
    # phi = sj1.phi()
    # for all c, do c.rotate_z(-phi)
    v = subjets[0][1][0]
    v = LorentzVector(v)

    phi = v.phi()

    for _, content, _, _ in subjets:
        for i in range(len(content)):
            v = LorentzVector(content[i][:4])
            v.rotate_z(-phi)
            content[i, 0] = v[0]
            content[i, 1] = v[1]
            content[i, 2] = v[2]
            content[i, 3] = v[3]

    ### boost ###
    # bv = sj1.boost_vector()
    # bv.set_perp(0)
    # for all c, do c.boost(-bv)
    v = subjets[0][1][0]
    v = LorentzVector(v)
    bv = v.boost_vector()
    bv.set_perp(0)
    for _, content, _, _ in subjets:
        for i in range(len(content)):
            v = LorentzVector(content[i][:4])
            v.boost(-bv)
            content[i, 0] = v[0]
            content[i, 1] = v[1]
            content[i, 2] = v[2]
            content[i, 3] = v[3]

    ### Rot alpha ###
    # deltaz = sj1.pz - sj2.pz
    # deltay = sj1.py - sj2.py
    # alpha = -atan2(deltaz, deltay)
    # for all c, do c.rotate_x(alpha)
    if len(subjets) >= 2:
        deltaz = subjets[0][1][0, 2] - subjets[1][1][0, 2]
        deltay = subjets[0][1][0, 1] - subjets[1][1][0, 1]
        alpha = -np.arctan2(deltaz, deltay)
        for _, content, _, _ in subjets:
            for i in range(len(content)):
                v = LorentzVector(content[i][:4])
                v.rotate_x(alpha)
                content[i, 0] = v[0]
                content[i, 1] = v[1]
                content[i, 2] = v[2]
                content[i, 3] = v[3]

    ### flip if necessary ###
    # if sj3.pz < 0: for all c, do c.set_pz(-c.pz)
    if len(subjets) >= 3 and subjets[2][1][0, 2] < 0:
        for _, content, _, _ in subjets:
            for i in range(len(content)):
                content[i, 2] *= -1.0

    ### finally recluster all transformed constituents c into a single jet ###
    constituents = []

    for tree, content, _, _ in subjets:
        constituents.append(content[tree[:, 0] == -1])

    constituents = np.vstack(constituents)

    if output == "anti-kt":
        subjets = cluster(constituents, R=100., jet_algorithm=1)
    elif output == "kt":
        subjets = cluster(constituents, R=100., jet_algorithm=0)
    elif output == "cambridge":
        subjets = cluster(constituents, R=100., jet_algorithm=2)
    else:
        raise

    jet["tree"] = subjets[0][0]
    jet["content"] = subjets[0][1]
    v = LorentzVector(jet["content"][0])
    jet["phi"] = v.phi()
    jet["eta"] = v.eta()
    jet["energy"] = v.E()
    jet["mass"] = v.m()
    jet["pt"] = v.pt()
    jet["root_id"] = 0
    jet['oldeta'] = oldeta
    jet['oldpt'] = oldpt
    if regression:
        jet["genpt"] = genpt
    return (jet)