def test_sweights_constructor():
mass, p, loss, Nsig, Nbkg, sig_p, bkg_p = get_data_and_loss()
with pytest.raises(ValueError):
compute_sweights("model", mass)
with pytest.raises(ValueError):
compute_sweights(loss.model[0].get_models()[0], mass)
def test_sweights():
minimizer = Minuit()
mass, p, loss, Nsig, Nbkg, sig_p, bkg_p = get_data_and_loss()
with pytest.raises(ModelNotFittedToData):
compute_sweights(loss.model[0], mass)
minimizer.minimize(loss)
model = loss.model[0]
assert is_sum_of_extended_pdfs(model)
yields = [Nsig, Nbkg]
sweights = compute_sweights(loss.model[0], mass)
assert np.allclose(
[np.sum(sweights[y]) / get_value(y.value()) for y in yields], 1.0)
nbins = 30
hist_conf = dict(bins=nbins, range=[0, 10])
hist_sig_true_p, _ = np.histogram(sig_p, **hist_conf)
sel = hist_sig_true_p != 0
hist_sig_true_p = hist_sig_true_p[sel]
hist_sig_sweights_p = np.histogram(p, weights=sweights[Nsig],
**hist_conf)[0][sel]
assert chisquare(hist_sig_sweights_p, hist_sig_true_p)[-1] < 0.01
hist_bkg_true_p, _ = np.histogram(bkg_p, **hist_conf)
sel = hist_bkg_true_p != 0
hist_bkg_true_p = hist_bkg_true_p[sel]
hist_bkg_sweights_p = np.histogram(p, weights=sweights[Nbkg],
**hist_conf)[0][sel]
assert chisquare(hist_bkg_sweights_p, hist_bkg_true_p)[-1] < 0.01
with pytest.warns(AboveToleranceWarning):
compute_sweights(
loss.model[0],
np.concatenate([mass, np.random.normal(0.8, 0.1, 100)]))
with pytest.raises(ModelNotFittedToData):
compute_sweights(
loss.model[0],
np.concatenate([mass, np.random.normal(0.8, 0.1, 1000)]))
xlabel=args.xLabel,
ax1_ylabel=args.yLabel,
data_lbl=args.dataLabel,
title='Before fit',
bins=args.bins)
if args.noFit:
sys.exit(0)
# Now do the fit
fit_result, fit_nll = fit(obs, fit_var, fit_model)
print('Fit result:\n', fit_result, sep='')
# Compute sWeights
print('Compute sWeights...')
sweights = compute_sweights(fit_model, fit_var)
# Dump result
print('Dump result...')
ensure_dir(args.output, is_file=False)
ntp_path = f'{args.output}/fit.root' if not args.outputNtp \
else args.outputNtp
ntp = recreate(ntp_path)
ntp_brs['sw_sig'] = sweights[fit_params['yld_sig']]
ntp_brs['sw_bkg'] = sweights[fit_params['yld_bkg']]
ntp['tree'] = ntp_brs
with open(f'{args.output}/params.yml', 'w') as f:
params_formatted = {
k.name: filter_dict(v)