diff --git a/countingworkspace/countingworkspace.py b/countingworkspace/countingworkspace.py index c7c6774..29cbb57 100644 --- a/countingworkspace/countingworkspace.py +++ b/countingworkspace/countingworkspace.py @@ -185,13 +185,14 @@ def create_model(ws, categories, processes, expression_nexpected_cat='nexp_cat{cat}', expression_nobs_cat='nobs_cat{cat}', expression_model_cat='model_cat{cat}', + factory_model_cat='poisson', + expression_nobs_err_cat=None, expression_model='model'): if type(categories) is int: categories = string_range(categories) if type(processes) is int: processes = string_range(processes) - all_poissons = [] all_exp = ROOT.RooArgSet() for cat in categories: s = ','.join([expression_nexpected_signal_cat_proc.format(cat=cat, proc=proc) for proc in processes]) @@ -199,13 +200,32 @@ def create_model(ws, categories, processes, s += ',' + expression_nexpected_bkg_cat.format(cat=cat) var_expected = ws.factory('sum:{expression_nexpected_cat}({s})'.format(expression_nexpected_cat=expression_nexpected_cat.format(cat=cat), s=s)) all_exp.add(var_expected) - model = 'Poisson:{model_cat}({nobs_cat}, {nexp_cat})'.format(model_cat=expression_model_cat, - nobs_cat=expression_nobs_cat, - nexp_cat=expression_nexpected_cat) - all_poissons.append(str(ws.factory(model.format(cat=cat)).getTitle())) + if factory_model_cat == 'gaussian' and expression_nobs_err_cat is None: + raise NotImplementedError('not implemented due to a ROOT bug https://sft.its.cern.ch/jira/browse/ROOT-10069') + expression_nobs_err_cat = 'nobs_err_cat{cat}' + for cat in categories: + ws.factory('expr:{nobs_err_cat}("sqrt(@0)", {nexp_cat})'.format(nobs_err_cat=expression_nobs_err_cat.format(cat=cat), + nexp_cat=expression_nexpected_cat.format(cat=cat))) + + all_pdfs = [] + for cat in categories: + + if factory_model_cat == 'poisson': + model = 'Poisson:{model_cat}({nobs_cat}, {nexp_cat})'.format(model_cat=expression_model_cat, + nobs_cat=expression_nobs_cat, + nexp_cat=expression_nexpected_cat) + elif factory_model_cat == 'gaussian': + model = 'RooGaussian:{model_cat}({nobs_cat}, {nexp_cat}, {nobs_err_cat})'.format(model_cat=expression_model_cat, + nobs_cat=expression_nobs_cat, + nexp_cat=expression_nexpected_cat, + nobs_err_cat=expression_nobs_err_cat) + + else: + model = factory_model_cat + all_pdfs.append(str(ws.factory(model.format(cat=cat)).getTitle())) ws.defineSet('all_exp', all_exp) - ws.factory('PROD:%s(%s)' % (expression_model, ','.join(all_poissons))) + ws.factory('PROD:%s(%s)' % (expression_model, ','.join(all_pdfs))) def dot(ws, var1, var2, name=None, nvar=None, operation='prod'): @@ -249,6 +269,8 @@ def create_workspace(categories, processes, expression_efficiency='eff_cat{cat}_proc{proc}', expression_efficiency_with_sys='eff_cat{cat}_proc{proc}_with_sys', expression_nexpected_bkg_cat='nexp_bkg_cat{cat}', + factory_model_cat='poisson', + expression_nobs_err_cat=None, name_constrain='constrain_sys{sysname}', ws=None): if type(categories) is int: @@ -326,10 +348,14 @@ def create_workspace(categories, processes, ws.factory('PROD:prod_constrains(%s)' % ','.join([v.GetName() for v in utils.iter_collection(all_constrains)])) if sysnames: - create_model(ws, categories, processes, expression_model='model_nosys') + create_model(ws, categories, processes, + factory_model_cat=factory_model_cat, expression_nobs_err_cat=expression_nobs_err_cat, + expression_model='model_nosys') ws.factory('PROD:model(model_nosys, prod_constrains)') else: - create_model(ws, categories, processes) + create_model(ws, categories, processes, + factory_model_cat=factory_model_cat, + expression_nobs_err_cat=expression_nobs_err_cat) ws.saveSnapshot('initial', ws.allVars()) model_config = ROOT.RooStats.ModelConfig('ModelConfig', ws) diff --git a/tests/test_all.py b/tests/test_all.py index aa9ea6d..ace61b1 100644 --- a/tests/test_all.py +++ b/tests/test_all.py @@ -4,7 +4,7 @@ import numpy as np import pytest import countingworkspace.utils -from countingworkspace.utils import iter_collection +from countingworkspace.utils import iter_collection, get_free_variables import ROOT @@ -565,12 +565,12 @@ def test_generate_and_fit_crossed(): def test_toy_study(): ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES, EXPECTED_BKG_CAT) - NTOYS = 10 - countingworkspace.utils.toy_study(ws, ntoys=NTOYS, seed=42) + ntoys = 10 + countingworkspace.utils.toy_study(ws, ntoys=ntoys, seed=42) f = ROOT.TFile.Open('result_42.root') tree = f.Get("results") assert tree - assert tree.GetEntries() == NTOYS + assert tree.GetEntries() == ntoys branches = [k.GetName() for k in tree.GetListOfBranches()] assert 'nll' in branches assert 'status' in branches @@ -579,3 +579,70 @@ def test_toy_study(): assert ("nsignal_gen_proc%s_error" % format_index(nproc)) in branches assert ("nsignal_gen_proc%s_error_up" % format_index(nproc)) in branches assert ("nsignal_gen_proc%s_error_down" % format_index(nproc)) in branches + + +def test_trivial_unfolding(): + ncat = 4 + nprocess = 4 + efficiencies = np.eye(4) + ntrue = np.array([10, 20, 30, 40]) + bkg = np.zeros(4) + ws = create_workspace(ncat, nprocess, ntrue, efficiencies, bkg) + data_asimov = countingworkspace.utils.generate_asimov(ws) + pdf = ws.obj('model') + fr = pdf.fitTo(data_asimov, ROOT.RooFit.Save()) + assert(fr.status() == 0) + + for nproc, t in enumerate(ntrue): + var_name = "nsignal_gen_proc%s" % format_index(nproc) + assert ws.obj(var_name).getVal() == pytest.approx(t) + assert ws.obj(var_name).getError() == pytest.approx(np.sqrt(t)) + + +@pytest.mark.skip +def test_trivial_unfolding_gaus(): + ncat = 4 + nprocess = 4 + efficiencies = np.eye(4) + ntrue = np.array([10, 20, 30, 40]) + bkg = np.zeros(4) + ws = create_workspace(ncat, nprocess, ntrue, efficiencies, bkg, factory_model_cat='gaussian') + + data_asimov = countingworkspace.utils.generate_asimov(ws) + pdf = ws.obj('model') + fr = pdf.fitTo(data_asimov, ROOT.RooFit.Save()) + assert(fr.status() == 0) + + for nproc, t in enumerate(ntrue): + var_name = "nsignal_gen_proc%s" % format_index(nproc) + assert ws.obj(var_name).getVal() == pytest.approx(t) + assert ws.obj(var_name).getError() == pytest.approx(np.sqrt(t)) + + +def test_trivial_unfolding_gaus_fixed_error(): + ncat = 4 + nprocess = 4 + efficiencies = np.eye(4) + ntrue = np.array([10, 20, 30, 40]) + bkg = np.zeros(4) + ws = ROOT.RooWorkspace() + expected_errors = [5., 10., 1., 0.001] + for icat, error in enumerate(expected_errors): + ws.factory('error_cat%s[%f]' % (format_index(icat), error)) + ws = create_workspace(ncat, nprocess, ntrue, efficiencies, bkg, + factory_model_cat='gaussian', + expression_nobs_err_cat='error_cat{cat}', + ws=ws + ) + + data_asimov = countingworkspace.utils.generate_asimov(ws) + pdf = ws.obj('model') + fr = pdf.fitTo(data_asimov, ROOT.RooFit.Save()) + assert(fr.status() == 0) + + + for nproc, t in enumerate(ntrue): + var_name = "nsignal_gen_proc%s" % format_index(nproc) + expected_error = ws.obj('error_cat%s' % format_index(nproc)).getVal() + assert ws.obj(var_name).getVal() == pytest.approx(t, rel=0.1) + assert ws.obj(var_name).getError() == pytest.approx(expected_error, rel=0.1)