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))