multisig

#!/usr/bin/env python

from multiprocessing import Process
from rootpy.io import root_open
from statstools.jobs import run_pool
from statstools import get_significance_workspace
import pickle
import os


class SigProcess(Process):
    def __init__(self, file, name, blind=True, profile='1', overwrite=False):
        super(SigProcess, self).__init__()
        self.file = file
        self.name = name
        self.blind = blind
        self.profile = profile
        self.overwrite = overwrite
    
    def run(self):
        if self.blind:
            pickle_name = os.path.splitext(self.file)[0] + '.pickle'
        else:
            pickle_name = os.path.splitext(self.file)[0] + '_profiled_mu{0}_unblind.pickle'.format(self.profile)
        if os.path.exists(pickle_name) and not self.overwrite:
            # abort
            return
        # get the significance of the workspace
        with root_open(self.file) as file:
            ws = file[self.name]
            if self.profile=='hat':
                ws.fit()
                poi = ws.obj('ModelConfig').GetParametersOfInterest().first()
                mu_profile_value = poi.getVal()
            else:
                mu_profile_value = float(self.profile)
            hist = get_significance_workspace(ws, 
                                              verbose=True,
                                              blind=self.blind,
                                              mu_profile_value=mu_profile_value)
            sig = hist[2].value
        # write the value into a pickle
        with open(pickle_name, 'w') as pickle_file:
            pickle.dump({self.name: sig}, pickle_file)


if __name__ == '__main__':
    from rootpy.extern.argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument('--jobs', type=int, default=-1)
    parser.add_argument('--name', default='combined')
    parser.add_argument('--overwrite', action='store_true', default=False)
    parser.add_argument('--unblind', action='store_false', dest='blind', default=True)
    parser.add_argument('--profile', choices=['0', '1', 'hat'], default='1')
    parser.add_argument('files', nargs='+')
    args = parser.parse_args()

    # find all root files
    root_files = []
    for file in args.files:
        if os.path.isdir(file):
            # walk below this path
            for dirpath, dirnames, filenames in os.walk(file):
                for filename in filenames:
                    if filename.endswith('.root'):
                        root_files.append(os.path.join(dirpath, filename))
        else:
            root_files.append(file)

    # define the workers
    workers = [SigProcess(file, args.name, blind=args.blind, profile=args.profile, overwrite=args.overwrite)
               for file in root_files]

    # run the pool
    run_pool(workers, n_jobs=args.jobs)