funcx_toys.py

from time import sleep
import pyhf

from funcx.sdk.client import FuncXClient

pyhf_endpoint = 'a727e996-7836-4bec-9fa2-44ebf7ca5302'

fxc = FuncXClient()
fxc.max_requests = 200


def prepare_workspace(nSig, errSig, nBkg, errBkgUp, errBkgDown):
    import pyhf

    spec = {
        "channels": [
            {
                "name": "SR_combined",
                "samples": [
                    {
                        "name": "signal",
                        "data": [nSig],
                        "modifiers": [
                            {"name": "mu", "type": "normfactor", "data": None},
                            #                              {"name": "uncorr_siguncrt", "type": "shapesys", "data": [0.1]}
                        ],
                    },
                    {
                        "name": "background",
                        "data": [nBkg],
                        "modifiers": [
                            {
                                "name": "uncorr_bkguncrt",
                                "type": "normsys",
                                "data": {"hi": errBkgUp, "lo": errBkgDown},
                            }
                        ],
                    },
                ],
            }
        ],
        "observations": [{"name": "SR_combined", "data": [0.0]}],
        "measurements": [
            {"name": "Measurement", "config": {"poi": "mu", "parameters": []}}
        ],
        "version": "1.0.0",
    }
    ws = pyhf.Workspace(spec)
    return ws


prepare_func = fxc.register_function(prepare_workspace)


def qmu_tilde(kind, poi_test, sample, model, signal_pars, par_bounds, fixed_params):
    import pyhf
    import time

    tick = time.time()
    return {
        'kind': kind,
        'teststat': float(
            pyhf.infer.test_statistics.qmu_tilde(
                poi_test, sample, model, signal_pars, par_bounds, fixed_params
            )
        ),
        'Fit-Time': time.time() - tick,
    }


teststat_func = fxc.register_function(qmu_tilde)


# configs for workspace spec
nSig = 4.166929245
errSig = 4.166929245
nBkg = 0.11
errBkgUp = 0.20
errBkgDown = 0.11

prepare_task = fxc.run(
    nSig,
    errSig,
    nBkg,
    errBkgUp,
    errBkgDown,
    endpoint_id=pyhf_endpoint,
    function_id=prepare_func,
)

workspace = None
while not workspace:
    try:
        workspace = fxc.get_result(prepare_task)
    except Exception as e:
        print("prepare ", e)
        sleep(15)

model = workspace.model()
data = workspace.data(model)

print("--------------------")
print(workspace)
print(model)
print(data)

NUM_TOYS = 50
POI_TEST = 1.0

# Copied from toy calculator
sample_shape = (NUM_TOYS,)

signal_pars = model.config.suggested_init()
signal_pars[model.config.poi_index] = POI_TEST
signal_pdf = model.make_pdf(pyhf.tensorlib.astensor(signal_pars))
signal_sample = signal_pdf.sample(sample_shape)

bkg_pars = model.config.suggested_init()
bkg_pars[model.config.poi_index] = 0.0
bkg_pdf = model.make_pdf(pyhf.tensorlib.astensor(bkg_pars))
bkg_sample = bkg_pdf.sample(sample_shape)

tasks = {}
for idx, sample in enumerate(signal_sample):
    task_id = fxc.run(
        'signal',
        POI_TEST,
        sample,
        model,
        signal_pars,
        model.config.suggested_bounds(),
        model.config.suggested_fixed(),
        endpoint_id=pyhf_endpoint,
        function_id=teststat_func,
    )
    tasks[f'signal_{idx}'] = {"id": task_id, "result": None}

for idx, sample in enumerate(bkg_sample):
    task_id = fxc.run(
        'background',
        POI_TEST,
        sample,
        model,
        signal_pars,
        model.config.suggested_bounds(),
        model.config.suggested_fixed(),
        endpoint_id=pyhf_endpoint,
        function_id=teststat_func,
    )
    tasks[f'background_{idx}'] = {"id": task_id, "result": None}


def count_complete(l):
    return len(list(filter(lambda e: e['result'], l)))


while count_complete(tasks.values()) < NUM_TOYS:
    for task in tasks.keys():
        if not tasks[task]['result']:
            try:
                result = fxc.get_result(tasks[task]['id'])
                print(
                    f"Task {task} complete, there are {count_complete(tasks.values())} results now"
                )
                tasks[task]['result'] = result
            except Exception as e:
                print(e)
                sleep(15)

print("--------------------")
print(tasks.values())