features.py

def add_features(data):
    data['ht'] = data['jet1_pt'] + data['jet2_pt'] + data['jet3_pt'] + data['jet4_pt']
    data['meff'] = data['jet1_pt'] + data['jet2_pt'] + data['jet3_pt'] + data['jet4_pt'] + data['lepton_pt'] + data['mem_pt']
    data['rt'] = data['meff'] / data['m_wwbb']
    data['delta_m'] = np.abs(data['m_jjj'] - data['m_jlv'])
    
    variables = ['ht', 'meff', 'rt', 'delta_m']
    
    data['dr_j12'] = dr(data['jet1_eta'], data['jet2_eta'], data['jet1_phi'], data['jet2_phi'])
    data['dr_j34'] = dr(data['jet3_eta'], data['jet4_eta'], data['jet3_phi'], data['jet4_phi'])
    
    variables += ['dr_j12', 'dr_j34']
    
    for i in [1,2,3,4]:
        data['dr_lj%d' % i] = dr(data['jet%d_eta' % i], data['lepton_eta'], data['jet%d_phi' % i], data['lepton_phi'])
        data['dphi_met_j%d' % i] = np.abs(dphi(data['mem_phi'], data['jet%d_phi' % i]))

        variables += ['dr_lj%d' % i, 'dphi_met_j%d' % i]
        
    data['max_dr_lj'] = np.maximum(np.maximum(data['dr_lj1'], data['dr_lj2']),
                                   np.maximum(data['dr_lj3'], data['dr_lj4']))
    data['min_dr_lj'] = np.minimum(np.minimum(data['dr_lj1'], data['dr_lj2']),
                                   np.minimum(data['dr_lj3'], data['dr_lj4']))   

    data['max_dphi_met_j'] = np.maximum(np.maximum(data['dphi_met_j1'], data['dphi_met_j2']),
                                        np.maximum(data['dphi_met_j3'], data['dphi_met_j4']))
    data['min_dphi_met_j'] = np.minimum(np.minimum(data['dphi_met_j1'], data['dphi_met_j2']),
                                        np.minimum(data['dphi_met_j3'], data['dphi_met_j4']))

    variables += ['max_dr_lj', 'min_dr_lj', 'max_dphi_met_j', 'min_dphi_met_j']

    data['max_jet_pt'] = np.maximum(np.maximum(data['jet1_pt'], data['jet2_pt']),
                                    np.maximum(data['jet3_pt'], data['jet4_pt']))
    data['min_jet_pt'] = np.minimum(np.minimum(data['jet1_pt'], data['jet2_pt']),
                                    np.minimum(data['jet3_pt'], data['jet4_pt']))       
    
    variables += ['max_jet_pt', 'min_jet_pt']
    
    for obj in real_objects:
        data['%s_abs_eta' % obj] = np.abs(data['%s_eta' % obj])
        variables.append('%s_abs_eta' % obj)
        
        for obj2 in ['lepton']:
            if obj == obj2:
                continue
            
            dphi_name = "dphi_%s_%s" % (obj, obj2)
            deta_name = "deta_%s_%s" % (obj, obj2)
            
            data[dphi_name] = np.abs(dphi(data['%s_phi' % obj], data['%s_phi' % obj2]))
            data[deta_name] = np.abs(data['%s_eta' % obj] - data['%s_eta' % obj2])
            
            variables += [dphi_name, deta_name]
    
    return variables