harvest.py

#!/usr/bin/env python

import json
import glob
import os
import re
import click

# This is only needed for Python 2/3 compatibility
def ensure_dirs(path):
    try:
        os.makedirs(path, exist_ok=True)
    except TypeError:
        if not os.path.exists(path):
            os.makedirs(path)


def make_harvest_from_result(result, masses):
    return {
        "CLs": result["CLs_obs"],
        "CLsexp": result["CLs_exp"][2],
        "clsd1s": result["CLs_exp"][1],
        "clsd2s": result["CLs_exp"][0],
        "clsu1s": result["CLs_exp"][3],
        "clsu2s": result["CLs_exp"][4],
        "covqual": 3,
        "dodgycov": 0,
        "excludedXsec": -999007,
        "expectedUpperLimit": -1,
        "expectedUpperLimitMinus1Sig": -1,
        "expectedUpperLimitMinus2Sig": -1,
        "expectedUpperLimitPlus1Sig": -1,
        "expectedUpperLimitPlus2Sig": -1,
        "fID": -1,
        "failedcov": 0,
        "failedfit": 0,
        "failedp0": 0,
        "failedstatus": 0,
        "fitstatus": 0,
        "mn1": masses[2],
        "mn2": masses[1],
        "mode": -1,
        "msb": masses[0],
        "nexp": -1,
        "nofit": 0,
        "p0": 0,
        "p0d1s": -1,
        "p0d2s": -1,
        "p0exp": -1,
        "p0u1s": -1,
        "p0u2s": -1,
        "p1": 0,
        "seed": 0,
        "sigma0": -1,
        "sigma1": -1,
        "upperLimit": -1,
        "upperLimitEstimatedError": -1,
        "xsec": -999007,
    }


@click.command()
@click.argument("master_input_dir", default="data/inputs")
@click.argument("master_output_dir", default="data/outputs")
def main(master_input_dir, master_output_dir):
    pattern = re.compile("sbottom_(\d+)_(\d+)_(\d+)")
    for kind in ["root", "pyhf"]:
        ensure_dirs(os.path.join(master_output_dir, "harvest", kind))
        for region in ["A", "B", "C"]:
            harvest = []

            thedir = os.path.join(
                master_output_dir,
                "results/{kind}/Region{region}".format(kind=kind, region=region),
            )

            files = "{master_output_dir}/{kind}_result_sbottom_*_*_*.json".format(
                master_output_dir=thedir, kind=kind
            )
            print(files)
            for fname in glob.glob(files):
                # print(fname)
                result = json.load(open(fname))
                m = pattern.search(fname)
                masses = list(map(int, m.groups()))
                # only use 60 GeV
                if masses[2] != 60:
                    continue
                harvest.append(make_harvest_from_result(result, masses))
            json.dump(
                harvest,
                open(
                    os.path.join(
                        master_output_dir,
                        "harvest",
                        kind,
                        "Region{region}.json".format(region=region),
                    ),
                    "w+",
                ),
                sort_keys=True,
                indent=2,
            )


if __name__ == "__main__":
    main()