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)
def test_create_workspace(): ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES, EXPECTED_BKG_CAT) assert ws for cat in range(NCATEGORIES): for nproc in range(NPROCESS): np.testing.assert_allclose( ws.var("eff_cat%s_proc%s" % (format_index(cat), format_index(nproc))).getVal(), EFFICIENCIES[cat][nproc], ) all_nexp_cat = np.dot(EFFICIENCIES, NTRUE) + EXPECTED_BKG_CAT for cat, nexp_cat in zip(range(NCATEGORIES), all_nexp_cat): v = ws.obj('nexp_cat{cat}'.format(cat=format_index(cat))) assert (v) v1 = v.getVal() np.testing.assert_allclose(v1, nexp_cat) model_config = ws.obj('ModelConfig') obs = model_config.GetObservables() assert obs assert obs.getSize() == NCATEGORIES poi = model_config.GetParametersOfInterest() assert poi assert poi.getSize() == NPROCESS
def test_create_variable_matrix(): ws = ROOT.RooWorkspace() eff = np.arange(6).reshape(2, 3) r = countingworkspace.create_variables(ws, 'myeff_cat{cat}_proc{proc}', values=eff, index_names=('cat', 'proc')) assert r for r1, v1 in zip(r, eff): for r2, v2 in zip(r1, v1): assert r2.getVal() == v2 for y in range(2): for x in range(3): v = ws.var('myeff_cat{cat}_proc{proc}'.format( cat=format_index(y), proc=format_index(x))) assert v assert v.getVal() == eff[y][x] bins_proc = 'A', 'B', 'C' bins_cat = 'X', 'Y' countingworkspace.create_variables(ws, 'myeff2_cat{cat}_proc{proc}', values=eff, bins=[bins_cat, bins_proc], index_names=('cat', 'proc')) for icat, cat in enumerate(bins_cat): for iproc, proc in enumerate(bins_proc): v = ws.var('myeff2_cat{cat}_proc{proc}'.format(cat=cat, proc=proc)) assert v assert v.getVal() == eff[icat][iproc]
def test_create_workspace_systematics_nsignal_gen(): systematics_nsignal_gen = np.ones(NPROCESS) * 0.01 systematics_nsignal_gen[0] *= 2 ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES, EXPECTED_BKG_CAT, systematics_nsignal_gen=[{ 'name': 'lumi', 'values': systematics_nsignal_gen }]) assert ws for cat in range(NCATEGORIES): for nproc in range(NPROCESS): np.testing.assert_allclose( ws.var("eff_cat%s_proc%s" % (format_index(cat), format_index(nproc))).getVal(), EFFICIENCIES[cat][nproc], ) all_nexp_cat = np.dot(EFFICIENCIES, NTRUE) + EXPECTED_BKG_CAT for cat, nexp_cat in zip(range(NCATEGORIES), all_nexp_cat): v = ws.obj('nexp_cat{cat}'.format(cat=format_index(cat))) assert (v) v1 = v.getVal() np.testing.assert_allclose(v1, nexp_cat) ws.var('theta_lumi').setVal(1) all_nexp_cat = np.dot(EFFICIENCIES, NTRUE * (1. + systematics_nsignal_gen)) + EXPECTED_BKG_CAT for cat, nexp_cat in zip(range(NCATEGORIES), all_nexp_cat): v = ws.obj('nexp_cat{cat}'.format(cat=format_index(cat))) assert (v) v1 = v.getVal() np.testing.assert_allclose(v1, nexp_cat) ws.var('theta_lumi').setVal(2) all_nexp_cat = np.dot( EFFICIENCIES, NTRUE * (1. + 2 * systematics_nsignal_gen)) + EXPECTED_BKG_CAT for cat, nexp_cat in zip(range(NCATEGORIES), all_nexp_cat): v = ws.obj('nexp_cat{cat}'.format(cat=format_index(cat))) assert (v) v1 = v.getVal() np.testing.assert_allclose(v1, nexp_cat)
def test_add(): ws = ROOT.RooWorkspace() a = np.arange(10) b = np.arange(10) - 1.5 countingworkspace.create_variables(ws, 'a_{index0}', nbins=10, values=a) countingworkspace.create_variables(ws, 'b_{index0}', nbins=10, values=b) countingworkspace.add(ws, 'a_{index0}', 'b_{index0}', nvar=10) for i, c in enumerate(a + b): assert ws.obj('a_plus_b_%s' % format_index(i)).getVal() == c countingworkspace.add(ws, 'a_{index0}', 'b_{index0}', 'd_{index0}') for i, c in enumerate(a + b): assert ws.obj('d_%s' % format_index(i)).getVal() == c
def test_create_workspace_systematics_efficiencies(): systematics_efficiencies = np.ones_like(EFFICIENCIES) * 0.01 systematics_efficiencies[0] *= 2 ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES, EXPECTED_BKG_CAT, systematic_efficiencies=[{ 'name': 'lumi', 'values': systematics_efficiencies }]) assert ws for cat in range(NCATEGORIES): for nproc in range(NPROCESS): np.testing.assert_allclose( ws.var("eff_cat%s_proc%s" % (format_index(cat), format_index(nproc))).getVal(), EFFICIENCIES[cat][nproc], ) all_nexp_cat = np.dot(EFFICIENCIES, NTRUE) + EXPECTED_BKG_CAT for cat, nexp_cat in zip(range(NCATEGORIES), all_nexp_cat): v = ws.obj('nexp_cat{cat}'.format(cat=format_index(cat))) assert (v) v1 = v.getVal() np.testing.assert_allclose(v1, nexp_cat) ws.var('theta_lumi').setVal(2) all_nexp_cat = np.dot(EFFICIENCIES * (1 + 2 * systematics_efficiencies), NTRUE) + EXPECTED_BKG_CAT for cat, nexp_cat in zip(range(NCATEGORIES), all_nexp_cat): v = ws.obj('nexp_cat{cat}'.format(cat=format_index(cat))) assert (v) v1 = v.getVal() np.testing.assert_allclose(v1, nexp_cat) constrain = ws.pdf('constrain_syslumi') assert constrain
def test_toy_coverage(): ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES, EXPECTED_BKG_CAT) ntoys = 10 countingworkspace.utils.toy_study( ws, ntoys=ntoys, seed=42, plugins=[ countingworkspace.utils.ToyStudyError(save_asym=True), countingworkspace.utils.ToyStudyCoverage( ws.obj('ModelConfig').GetParametersOfInterest(), NTRUE, significance=1, output_var='isCoveredAll'), countingworkspace.utils.ToyStudyCoverage( ws.obj('ModelConfig').GetParametersOfInterest(), NTRUE, significance=2, output_var='isCoveredAll2sigma') ]) f = ROOT.TFile.Open('result_42.root') tree = f.Get("results") assert tree assert tree.GetEntries() == ntoys branches = [k.GetName() for k in tree.GetListOfBranches()] assert 'nll' in branches assert 'status' in branches assert 'isCoveredAll' in branches for nproc in range(NPROCESS): assert ("nsignal_gen_proc%s" % format_index(nproc)) in branches 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 h1sigma = ROOT.TH1F("h1sigma", "h1sigma", 2, 0, 2) h2sigma = ROOT.TH1F("h2sigma", "h2sigma", 2, 0, 2) tree.Draw("isCoveredAll>>h1sigma") tree.Draw("isCoveredAll2sigma>>h2sigma") assert (h1sigma.GetMean() <= h2sigma.GetMean())
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, plugins=[countingworkspace.utils.ToyStudyError(save_asym=True)]) f = ROOT.TFile.Open('result_42.root') tree = f.Get("results") assert tree assert tree.GetEntries() == ntoys branches = [k.GetName() for k in tree.GetListOfBranches()] assert 'nll' in branches assert 'status' in branches for nproc in range(NPROCESS): assert ("nsignal_gen_proc%s" % format_index(nproc)) in branches 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_create_efficiencies(): ws = ROOT.RooWorkspace() eff = np.arange(6).reshape(2, 3) countingworkspace.create_efficiencies(ws, eff, 'myeff_cat{cat}_proc{proc}') for y in range(2): for x in range(3): v = ws.var('myeff_cat{cat}_proc{proc}'.format( cat=format_index(y), proc=format_index(x))) assert v assert v.getVal() == eff[y][x] bins_proc = 'A', 'B', 'C' bins_cat = 'X', 'Y' countingworkspace.create_efficiencies(ws, eff, bins_proc=bins_proc, bins_cat=bins_cat) for icat, cat in enumerate(bins_cat): for iproc, proc in enumerate(bins_proc): v = ws.var('eff_cat{cat}_proc{proc}'.format(cat=cat, proc=proc)) assert v assert v.getVal() == eff[icat][iproc]
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))
def test_asimov_roostats(): ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES, EXPECTED_BKG_CAT) obs = ws.set('all_obs') pdf = ws.obj('model') assert (obs) assert (pdf) data_asimov = ROOT.RooStats.AsymptoticCalculator.GenerateAsimovData( pdf, obs) assert (data_asimov) assert (data_asimov.numEntries() == 1) for ivar, v in enumerate(iter_collection(data_asimov.get(0))): if not isinstance(v, ROOT.RooCategory): np.testing.assert_allclose( v.getVal(), ws.obj('nexp_cat%s' % format_index(ivar)).getVal())
def test_create_formula(): ws = ROOT.RooWorkspace() countingworkspace.create_variables(ws, 'a', values=10.) assert ws.var('a').getVal() == 10. countingworkspace.create_variables(ws, 'theta', values=0., ranges=(-5, 5)) assert ws.var('theta').getVal() == 0. assert ws.var('theta').getMin() == -5. assert ws.var('theta').getMax() == 5. countingworkspace.create_variables(ws, 'prod:X(a, b[20])') assert ws.var('b').getVal() == 20. assert ws.obj('X').getVal() == 10. * 20. NPROC = 4 xsvalues = np.arange(1, NPROC + 1) countingworkspace.create_variables(ws, 'xsec_{proc}', nbins=NPROC, values=xsvalues) countingworkspace.create_variables( ws, 'prod:ntrue_{proc}(lumi[100], xsec_{proc})', nbins=NPROC) for i, xs in enumerate(xsvalues): assert (ws.obj('ntrue_%s' % format_index(i)) == 100 * (i + 1))