top10

#!/usr/bin/env python

# create an ntuple containing trees for data with classifier scores
# and event weights

from mva.cmd import get_parser

args = get_parser(actions=False).parse_args()

from mva.analysis import get_analysis

analysis = get_analysis(args)

from rootpy.io import root_open
import numpy as np
from numpy.lib import recfunctions
from root_numpy import array2tree


out = root_open('hhntup_top10.root', 'recreate')

fields = [
    'RunNumber', 'EventNumber', 'actualIntPerXing', 'averageIntPerXing', 'number_of_good_vertices',
    'mmc1_mass', 'mass_jet1_jet2', 'dEta_jets', 'numJets',
    'MET_et', 'MET_phi', 'MET_ety', 'MET_etx', 'MET_bisecting',
    'jet1_pt', 'jet1_eta', 'jet1_phi', 'jet1_m', 'jet1_et', 'jet1_e', 'jet1_p',
    'jet2_pt', 'jet2_eta', 'jet2_phi', 'jet2_m', 'jet2_et', 'jet2_e', 'jet2_p',
    'tau1_pt', 'tau1_eta', 'tau1_phi', 'tau1_m', 'tau1_et', 'tau1_e', 'tau1_p', 'tau1_charge', 'tau1_numTrack',
    'tau2_pt', 'tau2_eta', 'tau2_phi', 'tau2_m', 'tau2_et', 'tau2_e', 'tau2_p', 'tau2_charge', 'tau2_numTrack',
]

target_region = args.target_region
data = analysis.data

for category in analysis.iter_categories(
        args.categories, args.controls, names=args.category_names):

    if category.analysis_control:
        continue
    
    clf = analysis.get_clf(category, load=True)
        
    # get the record array
    rec = data.merged_records(
        category, target_region, fields=fields, include_weight=False)
    # get the scores
    scores, _ = data.scores(
        clf, category, target_region,
        systematics=False)

    # append the scores field
    rec = recfunctions.rec_append_fields(rec,
        data=scores,
        names='score',
        dtypes='f4')

    # sort by the score and pick up the top 10
    rec.sort(order=['score'])
    rec = rec[-10:]
    # reverse so highest score is first
    rec = rec[::-1]

    out.cd()
    output_name = 'category_{0}'.format(
        category.name).replace('.', '_')
    outtree = array2tree(rec) 
    outtree.Write(output_name)

out.Close()