class MLFit :
def __init__(self, plot_dire, condition):
self.plot_dire = plot_dire
self.condition = condition
self.mjj = RooRealVar("ZJJMass", "ZJJMass", 50., 7000.)
#self.costhetastar = RooRealVar("costhetastar", "costhetastar", -1., 1.)
self.weight = RooRealVar("weight", "weight", -100., 100.)
self.isSignal = RooCategory("isSignal", "isSignal")
self.isSignal.defineType("signal", 1);
self.isSignal.defineType("background", 0);
self.ras = RooArgSet(self.mjj, self.weight)
self.ds = RooDataSet("ds"+condition, "ds"+condition, self.ras, "weight")
#self.mu = RooRealVar("mu", "mu", 90., 80., 100.)
#self.widthL = RooRealVar("widthL", "widthL", 15., 2., 30.)
#self.widthR = RooRealVar("widthR", "widthR", 4., 2., 15.)
#self.sigmass = RooBifurGauss("bifurgauss", "bifurgauss", self.mjj, self.mu, self.widthL, self.widthR)
self.c0 = RooRealVar("c0", "c0", -100., 100.)
self.bkgmass = RooExponential("expo", "expo", self.mjj, self.c0)
#self.nsig = RooRealVar("nsig", "nsig", 100, 0, 200)
#self.nbkg = RooRealVar("nbkg", "nbkg", 100, 0, 200)
#self.components = RooArgList(self.sigmass, self.bkgmass)
#self.coefficients = RooArgList(self.nsig, self.nbkg)
self.modelmass = self.bkgmass #RooAddPdf("massmodel", "massmodel", self.components, self.coefficients)
def addToDataset(self, event, isSignal):
if not eval(self.condition):
return
self.mjj = event.ZJJMass
print self.mjj
#self.costhetastar = event.costhetastar
self.weight = event.weight
#self.isSignal = isSignal
self.ds.fill()
def fit(self):
print "nentries", self.ds.numEntries()
#for i in range(self.ds.numEntries()):
# argset = self.ds.get(i)
# argset.Dump()
fitresult = self.modelmass.fitTo(self.ds, RooFit.Save(True), RooFit.Extended(), RooFit.PrintLevel(3), RooFit.Strategy(2)) #, RooFit.SumW2Error(True))
def reweigh(target, target_cat, source, source_cat, binning = None):
cats = {}
for cat, data, name in [(target_cat, target, 'target_cat'), (source_cat, source, 'source_cat')]:
if hasattr(cat, '_target_'):
cat = cat._target_()
elif type(cat) == str:
cat_name = cat
cat = data.get().find(cat_name)
if not cat:
data.Print()
raise RuntimeError('observable or category %s is not present in data' % cat_name)
if not any([isinstance(cat, t) for t in [RooCategory, RooRealVar]]):
raise RuntimeError('category must be either a RooRealVar or a RooCategory')
if isinstance(cat, RooRealVar):
assert(binning)
if type(binning) == array:
binning = RooBinning(len(binning) - 1, binning)
binning.SetName('reweigh')
cat.setBinning(binning, 'reweigh')
cat_name = cat.GetName() + '_cat_' + data.GetName()
test_cat = data.get().find(cat_name)
if not test_cat:
cat = BinningCategory(Name = cat_name, Observable = cat,
Binning = binning, Data = data, Fundamental = True)
else:
cat = test_cat
cats[name] = cat
target_cat = cats['target_cat']
source_cat = cats['source_cat']
print target_cat
print source_cat
target_bins = dict([(ct.getVal(), ct.GetName()) for ct in target_cat])
source_bins = dict([(ct.getVal(), ct.GetName()) for ct in source_cat])
print target_bins
print source_bins
target_table = target.table(target_cat)
source_table = source.table(source_cat)
target_table.Print('v')
source_table.Print('v')
from collections import defaultdict
reweigh_weights = {}
for i, l in sorted(target_bins.iteritems()):
sf = source_table.getFrac(source_bins[i])
if sf == 0:
w = 0
else:
try:
w = sf / target_table.getFrac(l)
except ZeroDivisionError:
print 'Warning bin %s in wpv_data is 0, setting weight to 0' % l
w = 0.
reweigh_weights[i] = w
# RooFit infinity
from ROOT import RooNumber
RooInf = RooNumber.infinity()
from RooFitWrappers import RealVar
weight_var = RealVar('reweigh_var', MinMax = (RooInf, RooInf))
from ROOT import RooDataSet
data_name = target.GetName() + 'weight_data'
from ROOT import RooArgSet
weight_data = RooDataSet(data_name, data_name, RooArgSet(weight_var))
weight_var = weight_data.get().find(weight_var.GetName())
for i in range(target.numEntries()):
r = target.get(i)
n = target_cat.getIndex()
w = reweigh_weights[n]
weight_var.setVal(target.weight() * w)
weight_data.fill()
target.merge(weight_data)
target = RooDataSet(target.GetName(), target.GetTitle(), target,
target.get(), '', weight_var.GetName())
return target, reweigh_weights
