Acessing additional attributes in numba

[jmduarte]

I’m trying to do a gen-matching between a H->bb decay and a large-radius jet. So I’d like to check 3 things:

1. dR(H, large-radius jet) < 0.8
2. dR(b_daughter_1, large-radius jet) < 0.8
3. dR(b_daughter_2, large-radius jet) < 0.8

for a positive match. I’m using numba + vector a la @jpivarski's excellent tutorial (https://github.com/jpivarski-talks/2021-07-06-pyhep-uproot-awkward-tutorial/blob/main/uproot-awkward-tutorial.ipynb):

@nb.njit
def match_fjet(higgses, bquarks, fjets, builder):
    for higgses_event, bquarks_event, fjets_event in zip(higgses, bquarks, fjets):
        builder.begin_list()
        for higgs in higgses_event:
            match_idx = -1
            bdaughters = []
            print(bquarks_event)
            for bquark in bquarks_event:
                print(bquark)
                # this line doesn't work
                # if bquark.m1 == higgs.idx:  bdaughters.append(bquark)
                bdaughters.append(bquark)
            for i, fjet in enumerate(fjets_event):
                dr_h = fjet.deltaR(higgs)
                dr_b0 = fjet.deltaR(bdaughters[0])
                dr_b1 = fjet.deltaR(bdaughters[1])
                if dr_h < FJET_DR and dr_b0 < FJET_DR and dr_b1 < FJET_DR:
                    match_idx = i
            if match_idx < 0:
                builder.append(None)
            else:
                builder.append(match_idx)
        builder.end_list()
    return builder

which works except for the fact that i can’t access the mother index m1 of the b’s or the index idx of the Higgs to check that it’s actually a daughter of the b. These additional attributes are there for the event-level array, but once I enter the for bquark in bquarks_event loop, bquark is now a pure vector.obj and the additional attributes are gone.

See the result of the print statements below:

[{pt: 191, eta: -0.146, phi: -0.31, mass: 4.7, pid: 5, ... pid: -5, m1: 4, idx: 10}]
vector.obj(pt=190.97007751464844, phi=-0.30980587005615234, eta=-0.14612193405628204, mass=4.699999809265137)

Is this expected? If so, how else can I access these additional attributes?

[jmduarte]

UPDATE. I realized I can do something like the following in order to access the additional attributes.

@nb.njit
def match_fjet(higgses, bquarks, fjets, builder):
    for higgses_event, bquarks_event, fjets_event in zip(higgses, bquarks, fjets):
        builder.begin_list()
        for higgs, higgs_idx in zip(higgses_event, higgses_event.idx):
            match_idx = -1
            bdaughters = []
            print(bquarks_event)
            for bquark, bquark_m1 in zip(bquarks_event, bquarks_event.m1):
                print(bquark)
                if bquark_m1 == higgs_idx:
                    bdaughters.append(bquark)
            for i, fjet in enumerate(fjets_event):
                dr_h = fjet.deltaR(higgs)
                dr_b0 = fjet.deltaR(bdaughters[0])
                dr_b1 = fjet.deltaR(bdaughters[1])
                if dr_h < FJET_DR and dr_b0 < FJET_DR and dr_b1 < FJET_DR:
                    match_idx = i
            if match_idx < 0:
                builder.append(None)
            else:
                builder.append(match_idx)
        builder.end_list()
    return builder