def test_create_workspace_raise():
with pytest.raises(ValueError):
create_workspace(NCATEGORIES - 1, NPROCESS, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
with pytest.raises(ValueError):
create_workspace(NCATEGORIES, NPROCESS + 1, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
def test_generate_and_fit_crossed():
ws_generate = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
ws_fit = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES / 2.,
EXPECTED_BKG_CAT)
assert (len(
list(
countingworkspace.utils.generate_and_fit(ws_generate,
ws_fit,
ntoys=10))) == 10)
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_asimov():
ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
data_asimov = countingworkspace.utils.generate_asimov(ws)
assert data_asimov
assert data_asimov.numEntries() == 1
d = data_asimov.get(0)
assert d.getSize() == NCATEGORIES
np.testing.assert_allclose([x.getVal() for x in iter_collection(d)],
np.dot(EFFICIENCIES, NTRUE) + EXPECTED_BKG_CAT)
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_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_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_fit_asimov():
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 = countingworkspace.utils.generate_asimov(ws)
pois = ws.obj('ModelConfig').GetParametersOfInterest()
assert pois
# not start the fit from the true values
for poi in iter_collection(pois):
poi.setVal(poi.getVal() * 1.1)
ws.Print()
fr = pdf.fitTo(data_asimov, ROOT.RooFit.Save())
assert (fr.status() == 0)
pois_fitted = fr.floatParsFinal()
for ntrue, poi_fitted in zip(NTRUE, iter_collection(pois_fitted)):
np.testing.assert_allclose(ntrue, poi_fitted.getVal(), rtol=0.002)
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_free_variables():
ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
free_variables = countingworkspace.utils.get_free_variables(ws)
assert free_variables.getSize() == NPROCESS
def test_create_workspace_luminosity():
# workspace where nsignal_gen[p] = xsec[p] * lumi
ws_with_lumi = ROOT.RooWorkspace()
ws_with_lumi.factory('lumi[%f]' % LUMI)
countingworkspace.create_variables(ws_with_lumi,
'xsec_{proc}',
nbins=NPROCESS,
values=XSECFID_X_BR_PRODUCTION_MODES,
ranges=[-1000, 10000])
create_workspace(
NCATEGORIES,
NPROCESS,
None,
EFFICIENCIES,
EXPECTED_BKG_CAT,
expression_nsignal_gen='prod:nsignal_gen_proc{proc}(lumi, xsec_{proc})',
ws=ws_with_lumi)
# workspace where nsignal_gen[p] = mu[p] * xsec[p] * lumi
ws_with_4mu = ROOT.RooWorkspace()
ws_with_4mu.factory('lumi[%f]' % LUMI)
countingworkspace.create_variables(ws_with_4mu,
'xsec_{proc}',
nbins=NPROCESS,
values=XSECFID_X_BR_PRODUCTION_MODES)
create_workspace(
NCATEGORIES,
NPROCESS,
None,
EFFICIENCIES,
EXPECTED_BKG_CAT,
expression_nsignal_gen=
'prod:nsignal_gen_proc{proc}(mu_{proc}[1, -4, 5], lumi, xsec_{proc})',
ws=ws_with_4mu)
# workspace where nsignal_gen[p] = mu * mu[p] * xsec[p] * lumi
# where true yield is created externally
ws_with_4mu_x_mu = ROOT.RooWorkspace()
ws_with_4mu_x_mu.factory('lumi[%f]' % LUMI)
countingworkspace.create_variables(ws_with_4mu_x_mu,
'xsec_{proc}',
nbins=NPROCESS,
values=XSECFID_X_BR_PRODUCTION_MODES)
countingworkspace.create_variables(
ws_with_4mu_x_mu,
'prod:nsignal_gen_proc{proc}(mu[1, -4, 5], mu_{proc}[1, -4, 5], lumi, xsec_{proc})',
nbins=NPROCESS)
create_workspace(NCATEGORIES,
NPROCESS,
None,
EFFICIENCIES,
EXPECTED_BKG_CAT,
expression_nsignal_gen='nsignal_gen_proc{proc}',
ws=ws_with_4mu_x_mu)
# same, but with names
ws_with_4mu_x_mu_names = ROOT.RooWorkspace()
ws_with_4mu_x_mu_names.factory('lumi[%f]' % LUMI)
countingworkspace.create_variables(ws_with_4mu_x_mu_names,
'xsec_{proc}',
bins=list(
map(format_index, range(NPROCESS))),
values=XSECFID_X_BR_PRODUCTION_MODES)
countingworkspace.create_variables(
ws_with_4mu_x_mu_names,
'prod:nsignal_gen_proc{proc}(mu[1, -4, 5], mu_{proc}[1, -4, 5], lumi, xsec_{proc})',
bins=list(map(format_index, range(NPROCESS))))
create_workspace(list(map(format_index, range(NCATEGORIES))),
list(map(format_index, range(NPROCESS))),
None,
EFFICIENCIES,
EXPECTED_BKG_CAT,
expression_nsignal_gen='nsignal_gen_proc{proc}',
ws=ws_with_4mu_x_mu_names)
# nominal workspace for reference
ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
all_vars = ws.allVars()
for v in iter_collection(all_vars):
v_lumi = ws_with_lumi.obj(v.GetName())
assert v_lumi
np.testing.assert_allclose(v.getVal(), v_lumi.getVal())
v_4mu = ws_with_4mu.obj(v.GetName())
assert v_4mu
np.testing.assert_allclose(v.getVal(), v_4mu.getVal())
v_4mu_x_mu = ws_with_4mu_x_mu.obj(v.GetName())
assert v_4mu_x_mu
np.testing.assert_allclose(v.getVal(), v_4mu_x_mu.getVal())
v_4mu_x_mu_names = ws_with_4mu_x_mu_names.obj(v.GetName())
assert v_4mu_x_mu_names
np.testing.assert_allclose(v.getVal(), v_4mu_x_mu_names.getVal())
all_f = ws.allFunctions()
for f in iter_collection(all_f):
f_lumi = ws_with_lumi.obj(f.GetName())
assert f_lumi
np.testing.assert_allclose(f.getVal(), f_lumi.getVal())
f_4mu = ws_with_4mu.obj(f.GetName())
assert f_4mu
np.testing.assert_allclose(f.getVal(), f_4mu.getVal())
f_4mu_x_mu = ws_with_4mu_x_mu.obj(f.GetName())
assert f_4mu_x_mu
np.testing.assert_allclose(f.getVal(), f_4mu_x_mu.getVal())
f_4mu_x_mu_names = ws_with_4mu_x_mu_names.obj(f.GetName())
assert f_4mu_x_mu_names
np.testing.assert_allclose(f.getVal(), f_4mu_x_mu_names.getVal())
all_pdf = ws.allPdfs()
for p in iter_collection(all_pdf):
p_lumi = ws_with_lumi.pdf(p.GetName())
assert p_lumi
np.testing.assert_allclose(p.getVal(), p_lumi.getVal())
p_4mu = ws_with_4mu.pdf(p.GetName())
assert p_4mu
np.testing.assert_allclose(p.getVal(), p_4mu.getVal())
p_4mu_x_mu = ws_with_4mu.pdf(p.GetName())
assert p_4mu_x_mu
np.testing.assert_allclose(p.getVal(), p_4mu_x_mu.getVal())
p_4mu_x_mu_names = ws_with_4mu_x_mu_names.pdf(p.GetName())
assert p_4mu_x_mu_names
np.testing.assert_allclose(p.getVal(), p_4mu_x_mu_names.getVal())
assert countingworkspace.utils.get_free_variables(ws).getSize() == \
countingworkspace.utils.get_free_variables(ws_with_lumi).getSize() == \
countingworkspace.utils.get_free_variables(ws_with_4mu).getSize() == \
countingworkspace.utils.get_free_variables(ws_with_4mu_x_mu).getSize() - 1 == \
countingworkspace.utils.get_free_variables(ws_with_4mu_x_mu_names).getSize() - 1 == \
NPROCESS
def test_fit_asimov_syst():
systematics_nsignal_gen = np.ones(NPROCESS) * 0.05
systematics_nsignal_gen[0] *= 2
systematics_nsignal_gen2 = np.ones(NPROCESS) * 0.06
systematics_nsignal_gen2[1] *= 2
ws = create_workspace(NCATEGORIES,
NPROCESS,
NTRUE,
EFFICIENCIES,
EXPECTED_BKG_CAT,
systematics_nsignal_gen=[{
'name':
'lumi',
'values':
systematics_nsignal_gen
}, {
'name':
'lumi2',
'values':
systematics_nsignal_gen2
}])
obs = ws.set('all_obs')
pdf = ws.obj('model')
assert obs
assert pdf
data_asimov = countingworkspace.utils.generate_asimov(ws)
pois = ws.obj('ModelConfig').GetParametersOfInterest()
assert pois
# not start the fit from the true values
for poi in iter_collection(pois):
poi.setVal(poi.getVal() * 1.1)
fr = pdf.fitTo(data_asimov, ROOT.RooFit.Save(), ROOT.RooFit.Hesse(True))
assert (fr.status() == 0)
pois_fitted = fr.floatParsFinal()
all_errors = []
for ntrue, poi_fitted in zip(NTRUE, iter_collection(pois_fitted)):
np.testing.assert_allclose(ntrue, poi_fitted.getVal(), rtol=0.002)
all_errors.append(poi_fitted.getError())
ws.loadSnapshot('initial')
for theta in iter_collection(ws.allVars().selectByName('theta*')):
theta.setVal(0)
theta.setConstant()
fr = pdf.fitTo(data_asimov, ROOT.RooFit.Save(), ROOT.RooFit.Hesse(True))
assert (fr.status() == 0)
pois_fitted = fr.floatParsFinal()
all_errors_stat = []
for ntrue, poi_fitted in zip(NTRUE, iter_collection(pois_fitted)):
np.testing.assert_allclose(ntrue, poi_fitted.getVal(), rtol=0.002)
all_errors_stat.append(poi_fitted.getError())
sys_only_errors = (
np.sqrt(np.array(all_errors)**2 - np.array(all_errors_stat)**2) /
NTRUE)
np.testing.assert_allclose(sys_only_errors,
np.sqrt(systematics_nsignal_gen**2 +
systematics_nsignal_gen2**2),
rtol=5E-2)
#!/bin/env python
from countingworkspace import create_workspace
import countingworkspace.utils
from countingworkspace.examples import NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES, EXPECTED_BKG_CAT
countingworkspace.utils.silence_roofit()
ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
countingworkspace.utils.toy_study(ws, 1000)
parser.add_argument('-s', '--selection', action='store', required=True, help='Selection to be used')
args = parser.parse_args()
########
# ATLAS Run 2 Luminosity
lumi = (36104.16 + 43593.8 + 58450.1) / 1000.0
combWS = ROOT.RooWorkspace("combWS","combWS")
combWS.factory('lumi[%f]' % lumi)
## add cross sections into workspace
ntrue = create_variables(combWS, 'xsec_{proc}',bins=SAMPLES,values=xsec_production_modes)
if args.selection == "nominal":
create_workspace(NCATEGORIES, SAMPLES, efficiencies=EFFICIENCIES, nexpected_bkg_cat = EXPECTED_BKG_CAT,expression_nsignal_gen='prod:nsignal_gen_proc{proc}(mu_{proc}[1, -4, 5], lumi, xsec_{proc})',ws=combWS)
if args.selection == "new":
create_workspace(newNCATEGORIES, SAMPLES, efficiencies=newEFFICIENCIES, nexpected_bkg_cat = newEXPECTED_BKG_CAT,expression_nsignal_gen='prod:nsignal_gen_proc{proc}(mu_{proc}[1, -4, 5], lumi, xsec_{proc})',ws=combWS)
combWS.set('all_exp').Print('V')
## set all the POIs to constant with the exception of ttH
for sample in SAMPLES:
if sample != "ttH":
combWS.obj("mu_"+sample).setConstant(True)
poi = ""
if args.fit == "mu":
def test_generate_toys():
ws = create_workspace(NCATEGORIES, NPROCESS, NTRUE, EFFICIENCIES,
EXPECTED_BKG_CAT)
toys = countingworkspace.utils.generate_toys(ws, 100)
assert toys.numEntries() == 100