def gendata(filename, q2Bin, nSample):
time_start = time.time()
obs = getobs()
file = TFile.Open(filename)
if not file:
print('File not found: ', filename)
return
wspname = 'ws_b{}p0'.format(q2Bin)
wsp = file.Get(wspname)
if not wsp:
print('Workspace not found in file: ', filename)
return
ctK = wsp.var('ctK')
ctL = wsp.var('ctL')
phi = wsp.var('phi')
rand = wsp.var("rand")
if (not ctK) or (not ctL) or (not phi):
print('Variables not found in file: ', filename)
return
vars = RooArgList(ctK, ctL, phi)
datasetname = 'data_genDen_ev_b{}'.format(q2Bin)
fullData = wsp.data(datasetname)
if not fullData:
print('DataSet ', datasetname, ' not found in file: ', filename)
return
data = RooDataSet()
if nSample == 0:
data = fullData
elif nSample <= 10000:
data = fullData.reduce(
RooArgSet(vars), 'rand > {:1.6f} && rand < {:1.6f}'.format(
(nSample - 1) * dataStat[q2Bin] / genStat[q2Bin],
nSample * dataStat[q2Bin] / genStat[q2Bin]))
else:
data = fullData.reduce(
RooArgSet(vars), 'rand > {:.6f} && rand < {:.6f}'.format(
(nSample % 10000 - 1) * d16Stat[q2Bin] / genStat[q2Bin],
(nSample % 10000) * d16Stat[q2Bin] / genStat[q2Bin]))
print('Fit ', data.numEntries(), ' events')
###convert the final dataset to pd.dataframe
datalist = []
for i in range(0, data.numEntries()):
ctKvalue = data.get(i).getRealValue('ctK')
ctLvalue = data.get(i).getRealValue('ctL')
phivalue = data.get(i).getRealValue('phi')
datalist.append([ctKvalue, ctLvalue, phivalue])
#print ('lenth of datalist ', len(datalist))
finaldata_df = pd.DataFrame(datalist, columns=['costk', 'costl', 'phi'])
finaldata = zfit.Data.from_pandas(finaldata_df, obs=obs)
print('generate zfit data from RooDataSet successfully')
time_end = time.time()
time_c = time_end - time_start
print('time used in fit data import: ', time_c, 's')
return finaldata
def main1():
m = RooRealVar('evt.m', 'evt.m', 1.92, 2.02)
mode = RooCategory('evt.mode', 'evt.mode')
mode.defineType('phipi', 0)
aset = RooArgSet('aset')
aset.add(m)
aset.add(mode)
tuplist = tupleList('genmc', 7)
dst, _, f = getTree('/'.join([tuplePath(), tuplist[-1]]))
# dst.Print()
# dst.Show(0)
# for evt in dst:
# print('Hello!')
# print(evt.evt.m)
# break
ds = RooDataSet('ds', 'ds', dst, aset)
print(ds.numEntries())
mean = RooRealVar('mean', 'mean', 1.96, 1.92, 2.0)
width = RooRealVar('width', 'width', 0.006, 0.001, 0.010)
pdf = RooGaussian('pdf', 'pdf', m, mean, width)
pdf.fitTo(ds, Verbose(), Timer(True))
makePlot(m, ds, pdf)
raw_input("Press Enter to continue...")
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))
massmumu = RooRealVar("mumuM", "mumuM", 2.5, 3.5)
cutFormula = RooFormulaVar("cutFormula", "cutFormula", "xHlt!=8.0",
RooArgList(hlt))
# In[9]:
alldata = RooDataSet("alldata", "alldata", xTuple,
RooArgSet(masskk, mass, lxy, hlt,
massmumu)) #,cutFormula)
datasetfile = TFile("xMassDataset.root", "RECREATE")
datasetfile.cd()
alldata.Write()
# In[10]:
alldata.numEntries()
# In[ ]:
#xb->setRange("alt","x_coarse_bin1,x_coarse_bin3,x_coarse_bin5,x_coarse_bin7,x_coarse_bin9") ;
b0dataNonPrompt = ((
alldata.reduce('xHlt!=8')).reduce('xM>5.2')).reduce("xL>3.0")
# In[ ]:
b0dataNonPromptMass = b0dataNonPrompt.reduce(SelectVars(RooArgSet(mass)))
b0dataNonPrompt.numEntries()
# In[ ]:
c = TCanvas("canvas", "canvas", 1200, 800)
def predict_and_plot(self, x):
"""
Do the same as predict(), but add plots
Return -logL ratio, for external plotting
"""
rcParams['xtick.major.pad'] = 12
rcParams['ytick.major.pad'] = 12
roodata = RooDataSet('data', 'data', RooArgSet(self.phistar))
for xval in x:
self.phistar.setVal(xval)
roodata.add(RooArgSet(self.phistar))
theta = RooRealVar('theta', 'theta', self.theta_min, self.theta_max)
model = RooAddPdf('model', 'model',
RooArgList(self.pdfs['A'], self.pdfs['H']),
RooArgList(theta))
#with stdout_redirected_to():
print '\n\nPHISTAR FIT'
model.fitTo(roodata)
fitted_theta = theta.getValV()
# Histogram binning for data points
nbins = 10
xvals = np.linspace(0, 2*pi, 200)
yvals_H = []
yvals_A = []
# Get points for pdf curves
for xval in xvals:
self.phistar.setVal(xval)
yvals_H.append(self.pdfs['H'].getValV(RooArgSet(self.phistar)))
yvals_A.append(self.pdfs['A'].getValV(RooArgSet(self.phistar)))
yvals_H = np.array(yvals_H)
yvals_A = np.array(yvals_A)
# Plot pdfs by themselves
#print 'integral H =', np.trapz(yvals_H, dx=1.0/200.0)
fig = plt.figure()
plt.plot(xvals, yvals_H, color=self.colors['H'], label=r'$p_{H}(\varphi^{*})$')
plt.plot(xvals, yvals_A, color=self.colors['A'], label=r'$p_{A}(\varphi^{*})$')
plt.fill_between(xvals, 0, yvals_H, color=self.colors['H'], alpha=0.2)
plt.fill_between(xvals, 0, yvals_A, color=self.colors['A'], alpha=0.2)
plt.xlim([0, 2*pi])
plt.ylim([0, 0.295])
plt.xlabel(r'$\varphi^*$')
plt.ylabel('Probability density')
plt.legend(loc='upper right')
plt.tight_layout()
fig.show()
fig.savefig('phistar_pdfs.pdf')
# Scale to event yield
yvals_H *= roodata.numEntries()*(1-fitted_theta)/float(nbins)*2*pi
yvals_A *= roodata.numEntries()*(fitted_theta)/float(nbins)*2*pi
yvals_sum = yvals_H + yvals_A
# Make plot
fig, ax = plt.subplots(1)
histentries, binedges = np.histogram(x, bins=nbins, range=(0, 2*pi))
bincenters = (binedges[:-1] + binedges[1:])/2.0
yerr = np.sqrt(histentries)
plt.errorbar(bincenters, histentries, xerr=np.diff(binedges)*0.5, yerr=yerr, linestyle='None', ecolor='black', label='Data')
plt.plot(xvals, yvals_H, color=self.colors['H'], label=r'$p_{H}(\varphi^{*})$')
plt.plot(xvals, yvals_A, color=self.colors['A'], label=r'$p_{A}(\varphi^{*})$')
plt.plot(xvals, yvals_sum, color=self.colors['model'], label=r'$p(\varphi^{*} \,|\, \alpha = %.2f)$' % fitted_theta)
plt.fill_between(xvals, 0, yvals_H, color=self.colors['H'], alpha=0.2)
plt.fill_between(xvals, 0, yvals_A, color=self.colors['A'], alpha=0.2)
# Set correct legend order
handles, labels = ax.get_legend_handles_labels()
handles = [handles[3], handles[2], handles[0], handles[1]]
labels = [labels[3], labels[2], labels[0], labels[1]]
ax.legend(handles, labels, loc='upper right')
plt.xlabel(r'$\varphi^{*}$')
plt.ylabel('Events / %.2f' % ((2*pi)/nbins))
axes = plt.gca()
axes.set_xlim([0, 2*pi])
axes.set_ylim([0, max(histentries)*1.8])
plt.tight_layout()
fig.show()
fig.savefig('phistar_fit.pdf')
# Create likelihood curve
logl = model.createNLL(roodata)
# Extract curve
xvals = np.linspace(0, 1, 200)
yvals = []
ymin = 999.
for xval in xvals:
theta.setVal(xval)
yvals.append(logl.getValV())
if yvals[-1] < ymin:
ymin = yvals[-1]
# Shift minimum to zero
yvals = np.array(yvals)
yvals -= ymin
# Return points for the NLL curve
return xvals, yvals
class DilutionToy(Toy):
def __init__(self):
Toy.__init__(self)
self._gen_params = []
def run(self, **kwargs):
from ROOT import RooArgSet
__check_req_kw__("Observables", kwargs)
__check_req_kw__("Pdf", kwargs)
__check_req_kw__("Sigmat", kwargs)
__check_req_kw__("Time", kwargs)
__check_req_kw__("SigmaGen", kwargs)
sigma_gen = kwargs.pop("SigmaGen")
observables = kwargs.pop("Observables")
obs_set = RooArgSet(*observables)
pdf = kwargs.pop("Pdf")
sigmat = kwargs.pop("Sigmat")
time = kwargs.pop("Time")
gen_obs_set = RooArgSet(*observables)
# Make another ArgSet to put the fit results in
result_params = RooArgSet("result_params")
from P2VV.RooFitWrappers import RealVar
da = RealVar("da", Observable=True, MinMax=(0.01, 1.1))
dft = RealVar("dft", Observable=True, MinMax=(0.01, 1.1))
result_params.add(da._target_())
result_params.add(dft._target_())
transform = self.transform()
if transform:
trans_params = transform.gen_params(gen_obs_set)
self._gen_params.extend(trans_params)
for p in trans_params:
result_params.add(p)
# Some extra numbers of interest
from ROOT import RooRealVar
seed = RooRealVar("seed", "random seed", 0.0)
result_params.add(seed)
# The dataset to store the results
from ROOT import RooDataSet
self._data = RooDataSet("result_data", "result_data", result_params)
data_params = self._data.get()
from ROOT import RooRandom
import struct, os
# Reset pdf parameters to initial values. Note: this does not reset the estimated errors...
args = dict(NumEvents=self.options().nevents)
if "ProtoData" in kwargs:
args["ProtoData"] = kwargs.pop("ProtoData")
spec = pdf.prepareMultiGen(obs_set, **args)
while self._data.numEntries() < self.options().ntoys:
# Get a good random seed, set it and store it
s = struct.unpack("I", os.urandom(4))[0]
RooRandom.randomGenerator().SetSeed(s)
seed.setVal(s)
data = pdf.generate(spec)
if self.transform():
old_data = data
data = self.transform()(old_data)
if not data:
transform.set_params(data_params)
self._data.add(data_params)
continue
from P2VV import Dilution
d_ft = Dilution.dilution_ft(data, time, t_range=2, quiet=True)
d_a = Dilution.signal_dilution_dg(data, sigmat, *sigma_gen)
da.setVal(d_a[0])
da.setError(d_a[1] if d_a[1] != None else 0.0)
dft.setVal(d_ft[0])
dft.setError(d_ft[1] if d_ft[1] != None else 0.0)
if transform:
transform.set_params(data_params)
self._data.add(result_params)
return self.data()
def gen_params(self):
return self._gen_params
class FitToy(Toy):
def __init__(self):
Toy.__init__(self)
def run(self, **kwargs):
from ROOT import RooArgSet
__check_req_kw__("Observables", kwargs)
__check_req_kw__("Pdf", kwargs)
observables = kwargs.pop("Observables")
obs_set = RooArgSet(*observables)
pdf = kwargs.pop("Pdf")
genPdf = kwargs.pop("GenPdf", pdf)
gen_obs_set = RooArgSet()
for o in list(observables) + list(genPdf.ConditionalObservables()):
gen_obs_set.add(o._target_())
gen_pdf_params = genPdf.getParameters(gen_obs_set).snapshot(True)
genPdf = genPdf.clone(genPdf.GetName() + "_toy_clone")
genPdf.recursiveRedirectServers(gen_pdf_params)
fit_obs_set = RooArgSet()
for o in list(observables) + list(pdf.ConditionalObservables()):
fit_obs_set.add(o._target_())
params = pdf.getParameters(fit_obs_set)
pdf_params = RooArgSet()
for p in params:
if p.isConstant():
continue
pdf_params.add(p)
## for param in pdf_params:
## if param.GetName() not in ['Gamma', 'dGamma']:
## param.setConstant()
self._gen_params = pdf_params.snapshot(True)
# Make another ArgSet to put the fit results in
result_params = RooArgSet(pdf_params, "result_params")
transform = self.transform()
if transform:
trans_params = transform.gen_params(gen_obs_set)
for p in trans_params:
result_params.add(p)
# Some extra numbers of interest
from ROOT import RooRealVar
NLL = RooRealVar("NLL", "-log(Likelihood)", 1.0)
ngen = RooRealVar("ngen", "number of generated events", self.options().nevents)
seed = RooRealVar("seed", "random seed", 0.0)
from ROOT import RooCategory
status = RooCategory("status", "fit status")
status.defineType("success", 0)
status.defineType("one", 1)
status.defineType("two", 2)
status.defineType("three", 3)
status.defineType("other", 4)
result_params.add(status)
result_params.add(NLL)
result_params.add(ngen)
result_params.add(seed)
# The dataset to store the results
from ROOT import RooDataSet
self._data = RooDataSet("result_data", "result_data", result_params)
data_params = self._data.get()
from ROOT import RooRandom
import struct, os
while self._data.numEntries() < self.options().ntoys:
# Get a good random seed, set it and store it
s = struct.unpack("I", os.urandom(4))[0]
RooRandom.randomGenerator().SetSeed(s)
seed.setVal(s)
# Reset pdf parameters to initial values. Note: this does not reset the estimated errors...
pdf_params.assignValueOnly(self.gen_params())
args = dict(NumEvents=self.options().nevents)
if "ProtoData" in kwargs:
args["ProtoData"] = kwargs.pop("ProtoData")
genPdf.getParameters(obs_set).assignValueOnly(gen_pdf_params)
data = genPdf.generate(obs_set, **args)
if transform:
data = transform(data)
if not data:
# Transform has failed
transform.set_params(data_params)
self._data.add(data_params)
continue
if data.isWeighted() and "SumW2Error" not in self.fit_opts():
self.fit_opts()["SumW2Error"] = False
j = 0
while j < 4:
fit_result = pdf.fitTo(data, NumCPU=self.options().ncpu, **(self.fit_opts()))
if fit_result.status() == 0:
fit_result.Print()
break
j += 1
if fit_result.status() != 0:
print "Fit result status = %s" % fit_result.status()
NLL.setVal(fit_result.minNll())
if fit_result.status() < 4:
status.setIndex(fit_result.status())
else:
status.setIndex(4)
for result_param in result_params:
data_param = data_params.find(result_param.GetName())
if isinstance(result_param, RooCategory):
data_param.setIndex(result_param.getIndex())
else:
data_param.setVal(result_param.getVal())
# This sets a symmetric error, but since we don't run Minos, that's ok
data_param.setError(result_param.getError())
if transform:
transform.set_params(data_params)
self._data.add(data_params)
return self.data()
def gen_params(self):
return self._gen_params
binning = 400
massbins = (phimax - phimean)/0.005
dimuonditrk_m_rf_c = RooRealVar("dimuonditrk_m_rf_c","M(#mu#muKK)[GeV]",4.0,6.0)
dimuonditrk_m_rf_c.setBins(binning)
masskk = RooRealVar("ditrak_m","M(KK) [GeV]",phimin,phimax);
masskk.setBins(int(200))
dimuonditrk_ctauPV = RooRealVar("dimuonditrk_ctauPV","dimuonditrk_ctauPV",-1000.0,1000.0)
dimuon_pt = RooRealVar("dimuon_pt","dimuon_pt",0.0,1000.0)
dimuonditrk_ctauErrPV = RooRealVar("dimuonditrk_ctauErrPV","dimuonditrk_ctauErrPV",-1000.0,1000.0)
ditrak_pt = RooRealVar("ditrak_pt","ditrak_pt",0.0,1000.0)
theSet = RooArgSet(masskk,dimuonditrk_m_rf_c,dimuonditrk_ctauPV,dimuonditrk_ctauErrPV,dimuon_pt,ditrak_pt)
splotData = RooDataSet("alldata","alldata",xTuple,theSet)
#
print "Tree entries %d"%(splotData.numEntries())
a0 = RooRealVar("a_{0}","a0",3.92267e-01,-10.,10.)
a1 = RooRealVar("a_{1}","a1",-1.14362e-01,-5.0,5.0)
a2 = RooRealVar("a_{2}","a2",2.43201e-02,-2.,2.)
a3 = RooRealVar("a_{3}","a3",-2.42999e-02,-0.5,0.5)
a4 = RooRealVar("a_{4}","a4",0.01,-0.2,0.2)
a5 = RooRealVar("a_{5}","a5",0.0,-0.025,0.05)
a6 = RooRealVar("a6","a6",0.0,-0.001,0.001)
aset = RooArgList(a0,a1,a2,a3,a4)#,a5)
''' 2018 low cuts fit
1 #Gamma 6.41369e-03 4.22079e-05 2.09351e-03 6.48765e+00
#Iteartion over the workspce files
for i, f_wsp in enumerate(wlist):
print "Processing workspace "+str(i)+"/"+str(len(wlist))
ws_file = TFile(f_wsp, "read")
wsp = ws_file.Get("wspace")
ws_file.Close()
dataset_RS = wsp.data("dataset_RS")
dataset_COMB_OS = wsp.data("dataset_COMB_OS")
dataset_COMB_OS_dtb = RooDataSet("dataset_COMB_OS_dtb", "Combinatorial shape",dataset_COMB_OS, varset_comb,combination_cut)
Runs = {}
#Getting list of runs in the workspace
start = datetime.now()
for i in range(dataset_COMB_OS_dtb.numEntries()):
i_runNumber = dataset_COMB_OS_dtb.get(i).getRealValue("runNumber")
if i_runNumber not in Runs:
Runs[i_runNumber]=[]
run_indexing_time = datetime.now()-start
start = datetime.now()
#Iterating over runs
for i, i_run in enumerate(Runs):
#For each run we create a lists of B and D* candidates and match them later using the Fill_shape() function
signals = {}
combs = []
i_dataset_COMB_OS_dtb = RooDataSet("i_dataset_COMB_OS_dtb", "Combinatorial shape",dataset_COMB_OS_dtb, varset_comb,"runNumber=="+str(i_run))
i_dataset_RS_dtb = RooDataSet("i_dataset_RS_dtb", "Signal shape",dataset_RS, varset, preselection_cut+" && ( runNumber=="+str(i_run)+" )")
for s_i in range(i_dataset_RS_dtb.numEntries()):
signals[str(i_dataset_RS_dtb.get(s_i).getRealValue("LOG_D0_IPCHI2_OWNPV"))+"__"+str(i_dataset_RS_dtb.get(s_i).getRealValue("D0_M"))] = [i_dataset_RS_dtb.get(s_i).getRealValue("LOG_D0_IPCHI2_OWNPV"),i_dataset_RS_dtb.get(s_i).getRealValue("D0_M"), i_dataset_RS_dtb.get(s_i).getRealValue("Dst_DTF_D0_CTAU"), i_dataset_RS_dtb.get(s_i).getRealValue("DTF_D0sPi_M")]
if args.prompt:
mmttData = mmttData.reduce("lxysig < 2.0")
region = "_promt_"
if args.ptcuts is not None:
trakData = trakData.reduce("trigp_pT > " + str(args.ptcuts))
trakData = trakData.reduce("trign_pT > " + str(args.ptcuts))
cuts += "pt_" + str(args.ptcuts) + "_"
#### #### Plotting variables
#### TrakTrak Data
if args.noplot:
print("TrakTrak data plotting . . .")
print("All : " + str(trakData.numEntries()))
ttFrame = tt_mass.frame()
trakData.plotOn(ttFrame)
ttFrame.Draw()
c.SaveAs("tt_mass" + cuts + ".png")
tt_trig_Frame = tt_mass_trig.frame()
trakData.plotOn(tt_trig_Frame)
tt_trig_Frame.Draw()
c.SaveAs("tt_mass" + cuts + "trigger.png")
#### MuMu Data
print("MuMu data plotting . . .")
print("All : " + str(mumuData.numEntries()))
mumuFrame = mm_mass.frame()
obsSet.add(wVar)
import os, subprocess, tempfile
tmp = tempfile.NamedTemporaryFile()
tmpFileName = os.path.realpath(tmp.name)
tmp.close()
from ROOT import RooDataSet, TObject, gROOT
nEvTot = 0
nEvDS = nTupleIn.GetEntries() / numDataSets
for it in range(numDataSets) :
startEv = it * nEvDS
nEv = nEvDS if it < numDataSets - 1 else nTupleIn.GetEntries() - ( numDataSets - 1 ) * nEvDS
tmpFile = TFile.Open( tmpFileName, 'RECREATE' )
nTupleSplit = nTupleIn.CopyTree( selection, '', nEv, startEv )
if weightName :
dataOut = RooDataSet( protoData.GetName(), protoData.GetTitle(), obsSet, Import = nTupleSplit, WeightVar = ( weightName, True ) )
else :
dataOut = RooDataSet( protoData.GetName(), protoData.GetTitle(), obsSet, Import = nTupleSplit )
tmpFile.Close()
print 'produced dataset:'
dataOut.Print()
nEvTot += dataOut.numEntries()
dataFileOut = TFile.Open( dataFilePathOut % it, 'RECREATE' )
dataFileOut.Append(dataOut)
dataFileOut.Write( dataFilePathOut % it, TObject.kOverwrite )
dataFileOut.Close()
os.remove(tmpFileName)
print 'total number of events: %d' % nEvTot
if args.prompt:
theData = theData.reduce("xL < 1.5")
region = "_promt_"
if args.ptcuts is not None:
theData = theData.reduce("trigp_pT > " + str(args.ptcuts))
theData = theData.reduce("trign_pT > " + str(args.ptcuts))
cuts += "P_t_" + str(args.ptcuts) + "_"
#### #### Plotting variables
#### TrakTrak Data
if args.noplot:
print("TrakTrak data plotting . . .")
print("All : " + str(theData.numEntries()))
ttFrame = tt_mass.frame()
theData.plotOn(ttFrame)
ttFrame.Draw()
c.SaveAs("tt_mass" + cuts + ".png")
#### MuMu Data
print("MuMu data plotting . . .")
print("All : " + str(theData.numEntries()))
mumuFrame = mm_mass.frame()
theData.plotOn(mumuFrame)
mumuFrame.Draw()
c.SaveAs("mm_mass.png")
#### X Data
elif config['norm'] is 'kpi':
pipiFile = TFile(config['kpiFile'],"OPEN")
normTree = pipiFile.Get("subTree")
pipiUnbDataSet = RooDataSet("pipiUnbDataSet", "pipiUnbDataSet", normTree, w.set("argsPreCutKPi"), "Del_Mass>139.4 && RAND < %f" % (config["normScale"]))
pipiUnbDataSet.Print()
pipiUnbPruneDataSet = RooDataSet("pipiUnbPruneDataSet", "pipiUnbPruneDataSet", pipiUnbDataSet, w.set("argsBasic"), "D0_Mass>1800 && D0_Mass<1920")
pipiUnbPruneDataSet.Print()
if config['mode'] is "toy":
print "Setting random seed to 1"
RooRandom.randomGenerator().SetSeed(1)
print "Generating dataset 1 ..."
bdt1UnbToyDataSet = w.obj("BDT1_Comb").generate(w.set("argsBasic"),int(round(toyRatio*bdt1UnbPruneDataSet.numEntries())))
print "Generating dataset 2 ..."
bdt2UnbToyDataSet = w.obj("BDT2_Comb").generate(w.set("argsBasic"),int(round(toyRatio*bdt2UnbPruneDataSet.numEntries())))
print "Generating dataset 3 ..."
bdt3UnbToyDataSet = w.obj("BDT3_Comb").generate(w.set("argsBasic"),int(round(toyRatio*bdt3UnbPruneDataSet.numEntries())))
fullUnbDataSet = RooDataSet("fullUnbDataSet",
"fullUnbDataSet",
w.set("argsBasic"),
RooFit.Index(w.cat("DataSet")),
RooFit.Import("BDT1",bdt1UnbToyDataSet),
RooFit.Import("BDT2",bdt2UnbToyDataSet),
RooFit.Import("BDT3",bdt3UnbToyDataSet),
RooFit.Import("Norm",pipiUnbPruneDataSet))
#lab2_logIP = dataset.addColumn(RooFormulaVar("log(lab2_IP_OWNPV)", "log(lab2_IP_OWNPV)", RooArgList(lab2_IP_OWNPV)))
#lab2_logIP.setMin(-10.)
#lab2_logIP.setMax( 5.)
if is_MC and not desc == desc_Prompt_MC:
pull_formula = "(lab0_LifetimeFit_ctau0/{}-lab0_TRUETAU/{})/(lab0_LifetimeFit_ctauErr0/{})".format(speedOfLight_mm_fs, ns_to_fs, speedOfLight_mm_fs)
args = [lab0_LifetimeFit_ctau0, lab0_TRUETAU, lab0_LifetimeFit_ctauErr0]
else:
pull_formula = "lab0_LifetimeFit_ctau0/lab0_LifetimeFit_ctauErr0"
args = [lab0_LifetimeFit_ctau0, lab0_LifetimeFit_ctauErr0]
tauOverTauErr = dataset.addColumn(RooFormulaVar("lab0_LifetimeFit_TAU_PULL", "lab0_LifetimeFit_TAU_PULL", pull_formula, RooArgList(*args)))
tauOverTauErr.setMin(-8.)
tauOverTauErr.setMax( 8.)
nentries = dataset.numEntries()
print('Number of entries: {}'.format(nentries))
# Get the MC from EOS
"""
f_MC = ROOT.TFile.Open("root://eoslhcb.cern.ch//eos/lhcb/user/l/lbel/TD_DsK3fb_MC/B2DX_MC_Bs_Dspi_KKpi_PTR.root", "READ")
t_MC = f_MC.Get("DecayTree")
mc_data = RooDataSet("mc", "mc", t_MC, varSetMC)
# Add difference between measured and true lifetime as a variable
taudiff = mc_data.addColumn(RooFormulaVar("lab0_LifetimeFit_TAU_DIFF", "lab0_LifetimeFit_TAU_DIFF", "lab0_LifetimeFit_TAU-(lab0_TRUETAU/{})".format(ns_to_fs), RooArgList(lab0_LifetimeFit_TAU, lab0_TRUETAU)))
taudiff.setUnit("fs")
taudiff.setMin(-1.)
taudiff.setMax( 1.)
"""
# Create the RooWorkspace
def predict_and_plot(self, x):
"""
Do the same as predict(), but add plots
Return -logL ratio, for external plotting
"""
rcParams['xtick.major.pad'] = 12
rcParams['ytick.major.pad'] = 12
xs = self.scaler.transform(x)
preds = self.model.predict(xs)[:, 1]
roodata = RooDataSet('data', 'data', RooArgSet(self.roopred))
for pred in preds:
self.roopred.setVal(pred)
roodata.add(RooArgSet(self.roopred))
theta = RooRealVar('theta', 'theta', 0.5, self.theta_min, self.theta_max)
model = RooAddPdf('model', 'model',
RooArgList(self.pdfs['A'], self.pdfs['H']),
RooArgList(theta))
#with stdout_redirected_to():
print '\n\nNEURAL NETWORK FIT'
res = model.fitTo(roodata, PrintLevel(10))
fitted_theta = theta.getValV()
# Histogram binning for data points
nbins = 14
xvals = np.linspace(0, 1, 300)
yvals_H = []
yvals_A = []
# Get points for pdf curves
for xval in xvals:
self.roopred.setVal(xval)
yvals_H.append(self.pdfs['H'].getValV(RooArgSet(self.roopred)))
yvals_A.append(self.pdfs['A'].getValV(RooArgSet(self.roopred)))
yvals_H = np.array(yvals_H)
yvals_A = np.array(yvals_A)
# Plot pdfs by themselves
fig = plt.figure()
plt.plot(xvals, yvals_H, color=self.colors["H"], label=r'$p_{H}(y)$')
plt.plot(xvals, yvals_A, color=self.colors["A"], label=r'$p_{A}(y)$')
plt.fill_between(xvals, 0, yvals_H, color=self.colors["H"], alpha=0.2)
plt.fill_between(xvals, 0, yvals_A, color=self.colors["A"], alpha=0.2)
plt.xlim([0, 1])
plt.ylim([0, 11])
plt.xlabel(r'Network output ($y$)')
plt.ylabel('Probability density')
plt.legend(loc='upper right')
plt.tight_layout()
fig.show()
fig.savefig('mle_nn_pdfs.pdf')
# Scale to event yield
yvals_H *= roodata.numEntries()*(1-fitted_theta)/float(nbins)
yvals_A *= roodata.numEntries()*(fitted_theta)/float(nbins)
yvals_sum = yvals_H + yvals_A
# Make plot of fit to data
fig, ax = plt.subplots(1)
histentries, binedges = np.histogram(preds, bins=nbins, range=(0, 1))
bincenters = (binedges[:-1] + binedges[1:])/2.0
yerr = np.sqrt(histentries)
plt.errorbar(bincenters, histentries, xerr=np.diff(binedges)*0.5, yerr=yerr, linestyle='None', ecolor='black', label='Data')
plt.plot(xvals, yvals_H, color=self.colors["H"], label=r'$p_{H}(y)$')
plt.plot(xvals, yvals_A, color=self.colors["A"], label=r'$p_{A}(y)$')
plt.plot(xvals, yvals_sum, color=self.colors["model"], label=r'$p(y \,|\, \alpha = %.2f)$' % fitted_theta)
plt.fill_between(xvals, 0, yvals_H, color=self.colors["H"], alpha=0.2)
plt.fill_between(xvals, 0, yvals_A, color=self.colors["A"], alpha=0.2)
# Set correct legend order
handles, labels = ax.get_legend_handles_labels()
handles = [handles[3], handles[2], handles[0], handles[1]]
labels = [labels[3], labels[2], labels[0], labels[1]]
ax.legend(handles, labels, loc='upper right')
plt.xlabel(r'Network output ($y$)')
plt.ylabel('Events / %.2f' % (1.0/nbins))
axes = plt.gca()
axes.set_xlim([0, 1])
axes.set_ylim([0, max(histentries)*2.3])
plt.tight_layout()
fig.show()
fig.savefig('mle_nn_fit.pdf')
# Create likelihood curve
logl = model.createNLL(roodata)
xvals = np.linspace(0, 1, 200)
yvals = []
ymin = 999.
for xval in xvals:
theta.setVal(xval)
yvals.append(logl.getValV())
if yvals[-1] < ymin:
ymin = yvals[-1]
yvals = np.array(yvals)
yvals -= ymin
# Return points for the NLL curve
return xvals, yvals
masskk = RooRealVar("phiM","phiM",phimin,phimax)
massmumu = RooRealVar("jPsiM","jPsiM",2.5,3.5)
lxy = RooRealVar("xL","l(xy)",0.0,10000.)
hlt = RooRealVar("xHlt","xHlt",0.0,20.0)
# In[8]:
alldata = RooDataSet("alldata","alldata",xTree,RooArgSet(masskk,mass,massmumu))
# In[9]:
alldata.numEntries()
# In[10]:
alldata.numEntries()
c = TCanvas("canvas","canvas",1200,800)
massFrame = mass.frame()
alldata.plotOn(massFrame)
massFrame.Draw()
c.SaveAs("plots/testmass.png")
massKKFrame = masskk.frame(Range(phimin,phimax))
def fit_mass(data,
column,
x,
sig_pdf=None,
bkg_pdf=None,
n_sig=None,
n_bkg=None,
blind=False,
nll_profile=False,
second_storage=None,
log_plot=False,
pulls=True,
sPlot=False,
bkg_in_region=False,
importance=3,
plot_importance=3):
"""Fit a given pdf to a variable distribution
Parameter
---------
data : |hepds_type|
The data containing the variable to fit to
column : str
The name of the column to fit the pdf to
sig_pdf : RooFit pdf
The signal Probability Density Function. The variable to fit to has
to be named 'x'.
bkg_pdf : RooFit pdf
The background Probability Density Function. The variable to fit to has
to be named 'x'.
n_sig : None or numeric
The number of signals in the data. If it should be fitted, use None.
n_bkg : None or numeric
The number of background events in the data.
If it should be fitted, use None.
blind : boolean or tuple(numberic, numberic)
If False, the data is fitted. If a tuple is provided, the values are
used as the lower (the first value) and the upper (the second value)
limit of a blinding region, which will be omitted in plots.
Additionally, no true number of signal will be returned but only fake.
nll_profile : boolean
If True, a Negative Log-Likelihood Profile will be generated. Does not
work with blind fits.
second_storage : |hepds_type|
A second data-storage that will be concatenated with the first one.
importance : |importance_type|
|importance_docstring|
plot_importance : |plot_importance_type|
|plot_importance_docstring|
Return
------
tuple(numerical, numerical)
Return the number of signals and the number of backgrounds in the
signal-region. If a blind fit is performed, the signal will be a fake
number. If no number of background events is required, -999 will be
returned.
"""
if not (isinstance(column, str) or len(column) == 1):
raise ValueError("Fitting to several columns " + str(column) +
" not supported.")
if type(sig_pdf) == type(bkg_pdf) == None:
raise ValueError("sig_pdf and bkg_pdf are both None-> no fit possible")
if blind is not False:
lower_blind, upper_blind = blind
blind = True
n_bkg_below_sig = -999
# create data
data_name = data.name
data_array, _t1, _t2 = data.make_dataset(second_storage, columns=column)
del _t1, _t2
# double crystalball variables
min_x, max_x = min(data_array[column]), max(data_array[column])
# x = RooRealVar("x", "x variable", min_x, max_x)
# create data
data_array = np.array([i[0] for i in data_array.as_matrix()])
data_array.dtype = [('x', np.float64)]
tree1 = array2tree(data_array, "x")
data = RooDataSet("data", "Data", RooArgSet(x), RooFit.Import(tree1))
# # TODO: export somewhere? does not need to be defined inside...
# mean = RooRealVar("mean", "Mean of Double CB PDF", 5280, 5100, 5600)#, 5300, 5500)
# sigma = RooRealVar("sigma", "Sigma of Double CB PDF", 40, 0.001, 200)
# alpha_0 = RooRealVar("alpha_0", "alpha_0 of one side", 5.715)#, 0, 150)
# alpha_1 = RooRealVar("alpha_1", "alpha_1 of other side", -4.019)#, -200, 0.)
# lambda_0 = RooRealVar("lambda_0", "Exponent of one side", 3.42)#, 0, 150)
# lambda_1 = RooRealVar("lambda_1", "Exponent of other side", 3.7914)#, 0, 500)
#
# # TODO: export somewhere? pdf construction
# frac = RooRealVar("frac", "Fraction of crystal ball pdfs", 0.479, 0.01, 0.99)
#
# crystalball1 = RooCBShape("crystallball1", "First CrystalBall PDF", x,
# mean, sigma, alpha_0, lambda_0)
# crystalball2 = RooCBShape("crystallball2", "Second CrystalBall PDF", x,
# mean, sigma, alpha_1, lambda_1)
# doubleCB = RooAddPdf("doubleCB", "Double CrystalBall PDF",
# crystalball1, crystalball2, frac)
# n_sig = RooRealVar("n_sig", "Number of signals events", 10000, 0, 1000000)
# test input
if n_sig == n_bkg == 0:
raise ValueError("n_sig as well as n_bkg is 0...")
if n_bkg is None:
n_bkg = RooRealVar("n_bkg", "Number of background events", 10000, 0,
500000)
elif n_bkg >= 0:
n_bkg = RooRealVar("n_bkg", "Number of background events", int(n_bkg))
else:
raise ValueError("n_bkg is not >= 0 or None")
if n_sig is None:
n_sig = RooRealVar("n_sig", "Number of signal events", 1050, 0, 200000)
# START BLINDING
blind_cat = RooCategory("blind_cat", "blind state category")
blind_cat.defineType("unblind", 0)
blind_cat.defineType("blind", 1)
if blind:
blind_cat.setLabel("blind")
blind_n_sig = RooUnblindPrecision("blind_n_sig",
"blind number of signals",
"wasistdas", n_sig.getVal(),
10000, n_sig, blind_cat)
else:
# blind_cat.setLabel("unblind")
blind_n_sig = n_sig
print "n_sig value " + str(n_sig.getVal())
# raw_input("blind value " + str(blind_n_sig.getVal()))
# n_sig = blind_n_sig
# END BLINDING
elif n_sig >= 0:
n_sig = RooRealVar("n_sig", "Number of signal events", int(n_sig))
else:
raise ValueError("n_sig is not >= 0")
# if not blind:
# blind_n_sig = n_sig
# # create bkg-pdf
# lambda_exp = RooRealVar("lambda_exp", "lambda exp pdf bkg", -0.00025, -1., 1.)
# bkg_pdf = RooExponential("bkg_pdf", "Background PDF exp", x, lambda_exp)
if blind:
comb_pdf = RooAddPdf("comb_pdf", "Combined DoubleCB and bkg PDF",
RooArgList(sig_pdf, bkg_pdf),
RooArgList(blind_n_sig, n_bkg))
else:
comb_pdf = RooAddPdf("comb_pdf", "Combined DoubleCB and bkg PDF",
RooArgList(sig_pdf, bkg_pdf),
RooArgList(n_sig, n_bkg))
# create test dataset
# mean_gauss = RooRealVar("mean_gauss", "Mean of Gaussian", 5553, -10000, 10000)
# sigma_gauss = RooRealVar("sigma_gauss", "Width of Gaussian", 20, 0.0001, 300)
# gauss1 = RooGaussian("gauss1", "Gaussian test dist", x, mean_gauss, sigma_gauss)
# lambda_data = RooRealVar("lambda_data", "lambda exp data", -.002)
# exp_data = RooExponential("exp_data", "data example exp", x, lambda_data)
# frac_data = RooRealVar("frac_data", "Fraction PDF of data", 0.15)
#
# data_pdf = RooAddPdf("data_pdf", "Data PDF", gauss1, exp_data, frac_data)
# data = data_pdf.generate(RooArgSet(x), 30000)
# data.printValue()
# xframe = x.frame()
# data_pdf.plotOn(xframe)
# print "n_cpu:", meta_config.get_n_cpu()
# input("test")
# comb_pdf.fitTo(data, RooFit.Extended(ROOT.kTRUE), RooFit.NumCPU(meta_config.get_n_cpu()))
# HACK to get 8 cores in testing
c5 = TCanvas("c5", "RooFit pdf not fit vs " + data_name)
c5.cd()
x_frame1 = x.frame()
# data.plotOn(x_frame1)
# comb_pdf.pdfList()[1].plotOn(x_frame1)
if __name__ == "__main__":
n_cpu = 8
else:
n_cpu = meta_config.get_n_cpu()
print "n_cpu = ", n_cpu
# HACK
# n_cpu = 8
result_fit = comb_pdf.fitTo(data, RooFit.Minos(ROOT.kTRUE),
RooFit.Extended(ROOT.kTRUE),
RooFit.NumCPU(n_cpu))
# HACK end
if bkg_in_region:
x.setRange("signal", bkg_in_region[0], bkg_in_region[1])
bkg_pdf_fitted = comb_pdf.pdfList()[1]
int_argset = RooArgSet(x)
# int_argset = x
# int_argset.setRange("signal", bkg_in_region[0], bkg_in_region[1])
integral = bkg_pdf_fitted.createIntegral(int_argset,
RooFit.NormSet(int_argset),
RooFit.Range("signal"))
bkg_cdf = bkg_pdf_fitted.createCdf(int_argset, RooFit.Range("signal"))
bkg_cdf.plotOn(x_frame1)
# integral.plotOn(x_frame1)
n_bkg_below_sig = integral.getVal(int_argset) * n_bkg.getVal()
x_frame1.Draw()
if plot_importance >= 3:
c2 = TCanvas("c2", "RooFit pdf fit vs " + data_name)
c2.cd()
x_frame = x.frame()
# if log_plot:
# c2.SetLogy()
# x_frame.SetTitle("RooFit pdf vs " + data_name)
x_frame.SetTitle(data_name)
if pulls:
pad_data = ROOT.TPad("pad_data", "Pad with data and fit", 0, 0.33,
1, 1)
pad_pulls = ROOT.TPad("pad_pulls", "Pad with data and fit", 0, 0,
1, 0.33)
pad_data.SetBottomMargin(0.00001)
pad_data.SetBorderMode(0)
if log_plot:
pad_data.SetLogy()
pad_pulls.SetTopMargin(0.00001)
pad_pulls.SetBottomMargin(0.2)
pad_pulls.SetBorderMode(0)
pad_data.Draw()
pad_pulls.Draw()
pad_data.cd()
else:
if log_plot:
c2.SetLogy()
if blind:
# HACK
column = 'x'
# END HACK
x.setRange("lower", min_x, lower_blind)
x.setRange("upper", upper_blind, max_x)
range_str = "lower,upper"
lower_cut_str = str(
min_x) + "<=" + column + "&&" + column + "<=" + str(lower_blind)
upper_cut_str = str(
upper_blind) + "<=" + column + "&&" + column + "<=" + str(max_x)
sideband_cut_str = "(" + lower_cut_str + ")" + "||" + "(" + upper_cut_str + ")"
n_entries = data.reduce(
sideband_cut_str).numEntries() / data.numEntries()
# raw_input("n_entries: " + str(n_entries))
if plot_importance >= 3:
data.plotOn(x_frame, RooFit.CutRange(range_str),
RooFit.NormRange(range_str))
comb_pdf.plotOn(
x_frame, RooFit.Range(range_str),
RooFit.Normalization(n_entries, RooAbsReal.Relative),
RooFit.NormRange(range_str))
if pulls:
# pull_hist(pull_frame=x_frame, pad_data=pad_data, pad_pulls=pad_pulls)
x_frame_pullhist = x_frame.pullHist()
else:
if plot_importance >= 3:
data.plotOn(x_frame)
comb_pdf.plotOn(x_frame)
if pulls:
pad_pulls.cd()
x_frame_pullhist = x_frame.pullHist()
pad_data.cd()
comb_pdf.plotOn(x_frame,
RooFit.Components(sig_pdf.namePtr().GetName()),
RooFit.LineStyle(ROOT.kDashed))
comb_pdf.plotOn(x_frame,
RooFit.Components(bkg_pdf.namePtr().GetName()),
RooFit.LineStyle(ROOT.kDotted))
# comb_pdf.plotPull(n_sig)
if plot_importance >= 3:
x_frame.Draw()
if pulls:
pad_pulls.cd()
x_frame.SetTitleSize(0.05, 'Y')
x_frame.SetTitleOffset(0.7, 'Y')
x_frame.SetLabelSize(0.04, 'Y')
# c11 = TCanvas("c11", "RooFit\ pulls" + data_name)
# c11.cd()
# frame_tmp = x_frame
frame_tmp = x.frame()
# frame_tmp.SetTitle("significance")
frame_tmp.SetTitle("Roofit\ pulls\ " + data_name)
frame_tmp.addObject(x_frame_pullhist)
frame_tmp.SetMinimum(-5)
frame_tmp.SetMaximum(5)
# frame_tmp.GetYaxis().SetTitle("significance")
frame_tmp.GetYaxis().SetNdivisions(5)
frame_tmp.SetTitleSize(0.1, 'X')
frame_tmp.SetTitleOffset(1, 'X')
frame_tmp.SetLabelSize(0.1, 'X')
frame_tmp.SetTitleSize(0.1, 'Y')
frame_tmp.SetTitleOffset(0.5, 'Y')
frame_tmp.SetLabelSize(0.1, 'Y')
frame_tmp.Draw()
# raw_input("")
if not blind and nll_profile:
# nll_range = RooRealVar("nll_range", "Signal for nLL", n_sig.getVal(),
# -10, 2 * n_sig.getVal())
sframe = n_sig.frame(RooFit.Bins(20), RooFit.Range(1, 1000))
# HACK for best n_cpu
lnL = comb_pdf.createNLL(data, RooFit.NumCPU(8))
# HACK end
lnProfileL = lnL.createProfile(ROOT.RooArgSet(n_sig))
lnProfileL.plotOn(sframe, RooFit.ShiftToZero())
c4 = TCanvas("c4", "NLL Profile")
c4.cd()
# input("press ENTER to show plot")
sframe.Draw()
if plot_importance >= 3:
pass
params = comb_pdf.getVariables()
params.Print("v")
# print bkg_cdf.getVal()
if sPlot:
sPlotData = ROOT.RooStats.SPlot(
"sPlotData",
"sPlotData",
data, # variable fitted to, RooDataSet
comb_pdf, # fitted pdf
ROOT.RooArgList(
n_sig,
n_bkg,
# NSigB0s
))
sweights = np.array([
sPlotData.GetSWeight(i, 'n_sig') for i in range(data.numEntries())
])
return n_sig.getVal(), n_bkg_below_sig, sweights
if blind:
return blind_n_sig.getVal(), n_bkg_below_sig, comb_pdf
else:
return n_sig.getVal(), n_bkg_below_sig, comb_pdf
def readData( filePath, dataSetName, NTuple = False, observables = None, **kwargs ) :
"""reads data from file (RooDataSet or TTree(s))
"""
from ROOT import RooFit
noNAN = ( ' && '.join( '( %s==%s )' % ( obs, obs ) for obs in observables ) ) if hasattr( observables, '__iter__' ) else ''
cuts = kwargs.pop( 'cuts', '' )
tmp_file = None
if observables :
print 'P2VV - INFO: readData: reading data for observables [ %s ]' % ', '.join( obs.GetName() for obs in observables )
dataSetArgs = { }
if 'WeightVar' in kwargs : dataSetArgs['WeightVar'] = kwargs.pop('WeightVar')
if NTuple :
from ROOT import RooDataSet, TFile
assert observables != None, 'P2VV - ERROR: readData: set of observables is required for reading an n-tuple'
# create data set from NTuple file(s)
print 'P2VV - INFO: readData: reading NTuple "%s" from file "%s"' % ( dataSetName, filePath )
ntupleFile = TFile.Open(filePath)
ntupleOrig = ntupleFile.Get(dataSetName)
if not ntupleOrig : raise RuntimeError( 'P2VV - ERROR: could not locate tree "%s" in file "%s"' % ( dataSetName, filePath ) )
if 'ntupleCuts' in kwargs :
ntupleCuts = kwargs.pop( 'ntupleCuts', '' )
print 'P2VV - INFO: readData: applying cuts on n-tuple: %s' % ntupleCuts
import tempfile
import os
from ROOT import TFile, gFile
orig_file = gFile
d = None
if 'TMP' in os.environ:
d = os.environ['TMP']
elif 'TMPDIR' in os.environ:
d = os.environ['TMPDIR']
elif os.access(os.path.dirname(filePath), os.W_OK):
d = os.path.dirname(filePath)
else:
d = '/tmp'
fd, temp_name = tempfile.mkstemp(suffix = '.root', dir = d)
os.close(fd)
os.remove(temp_name)
tmp_file = TFile.Open(temp_name, 'recreate')
ntuple = ntupleOrig.CopyTree(ntupleCuts)
else :
ntuple = ntupleOrig
if cuts : print 'P2VV - INFO: readData: applying cuts on data set: %s' % cuts
data = RooDataSet( dataSetName, dataSetName
, [ obs._var for obs in observables ]
, Import = ntuple
, Cut = noNAN + ' && ' + cuts if cuts else noNAN
, **dataSetArgs
)
else :
from ROOT import TFile
# get data set from file
print 'P2VV - INFO: readData: reading RooDataset "%s" from file "%s"' % ( dataSetName, filePath )
file = TFile.Open( filePath, 'READ' )
assert file, 'P2VV - ERROR: readData: file "%s" could not be opened' % filePath
if cuts : print 'P2VV - INFO: readData: applying cuts: %s' % cuts
# loop over category states
states = tuple( [ [ ( cat[0], ind ) for ind in cat[1] ] for cat in kwargs.pop( 'Categories', [ ( '', [ '' ] ) ] ) ] )
from itertools import product
for it, state in enumerate( product(*states) ) :
# get data set
dsName = '_'.join( str(catSt[0]) + str(catSt[1]) for catSt in state )
dsName = dataSetName + ( ( '_' + dsName ) if dsName else '' )
dataSet = file.Get(dsName)
assert dataSet, 'P2VV - ERROR: data set "%s" not found' % dsName
if it == 0 :
if observables :
from ROOT import RooDataSet
data = RooDataSet( dataSetName, dataSetName
, [ obs._var for obs in observables ]
, Import = dataSet
, Cut = noNAN + ' && ' + cuts if cuts else noNAN
, **dataSetArgs
)
else :
data = dataSet
else :
data.append(dataSet)
file.Close()
print 'P2VV - INFO: read dataset with %s entries (%.1f weighted)' % ( data.numEntries(), data.sumEntries() )
# import data set into current workspace
from P2VV.RooFitWrappers import RooObject
importIntoWS = kwargs.pop( 'ImportIntoWS', True )
rData = None
if importIntoWS :
# import data set into current workspace
from P2VV.RooFitWrappers import RooObject
wsData = RooObject().ws().put( data, **kwargs )
rData = wsData
else :
rData = data
if tmp_file:
tmp_file.Close()
os.remove(tmp_file.GetName())
if orig_file: orig_file.cd()
return rData
theData = theData.reduce("xL > 3.0")
if args.prompt:
theData = theData.reduce("xL < 1.5")
if args.ptcuts is not None:
theData = theData.reduce("trigp_pT > " + str(args.ptcuts))
theData = theData.reduce("trign_pT > " + str(args.ptcuts))
cuts += "P_t_" + str(args.ptcuts) + "_"
#### #### Plotting variables
#### TrakTrak Data
if args.noplot:
print("TrakTrak data plotting . . .")
print("All : " + str(theData.numEntries()))
ttFrame = tt_mass.frame(Title("KK mass"))
theData.plotOn(ttFrame)
ttFrame.Draw()
c.SaveAs(region + "/tt_mass" + cuts + ".png")
#### MuMu Data
print("MuMu data plotting . . .")
print("All : " + str(theData.numEntries()))
mumuFrame = mm_mass.frame(Title("KK mass"))
theData.plotOn(mumuFrame)
mumuFrame.Draw()
c.SaveAs(region + "/mm_mass.png")
#### X Data
for i in range(n_toys % n_p):
n_t[i] += 1
for n in n_t:
c = Calculator(pdf, sigmat_cat, t, st, n)
calculators.append(c)
c.start()
args = RooArgSet(da, dft)
while len(calculators):
for i, calculator in enumerate(calculators):
msg = calculator.queue().get()
if msg == 'done':
calculator.join()
calculators.pop(i)
continue
else:
d_a, d_ft = msg
da.setVal(d_a[0])
dft.setVal(d_ft[0])
dft.setError(d_ft[1])
test_data.add(args)
if test_data.numEntries() % 100 == 0:
print 'completed ', test_data.numEntries()
diff = FormulaVar(Name = 'diff', Formula = '@0 - @1', Arguments = (dft, da), data = test_data)
from ROOT import TFile
f = TFile("dilution.root", "recreate")
f.WriteTObject(test_data, "data")
f.Close()
def make_fit():
adc_bin = 12 #18 for low-m gg, 24 for jpsi
adc_min = 0. #10.
adc_max = 400.
#adc_max = 1200
ptmax = 0.18
#mmin = 1.6
#mmin = 2.1
#mmax = 2.6
#mmin = 1.5
#mmax = 5.
mmin = 2.9
mmax = 3.2
#mmin = 3.4
#mmax = 4.6
#east/west projections and 2D plot
ew = 1
p2d = 2 # 0: single projection by 'ew', 1: 2D plot, 2: both projections
#plot colors
model_col = rt.kMagenta
model_col = rt.kBlue
out = open("out.txt", "w")
lmg = 6
ut.log_results(out, "in " + infile, lmg)
strlog = "adc_bin " + str(adc_bin) + " adc_min " + str(
adc_min) + " adc_max " + str(adc_max)
strlog += " ptmax " + str(ptmax) + " mmin " + str(mmin) + " mmax " + str(
mmax)
ut.log_results(out, strlog, lmg)
#adc distributions
adc_east = RooRealVar("jZDCUnAttEast", "ZDC ADC east", adc_min, adc_max)
adc_west = RooRealVar("jZDCUnAttWest", "ZDC ADC west", adc_min, adc_max)
#kinematics variables
m = RooRealVar("jRecM", "e^{+}e^{-} mass (GeV)", 0., 10.)
y = RooRealVar("jRecY", "rapidity", -1., 1.)
pT = RooRealVar("jRecPt", "pT", 0., 10.)
#adc distributions
#adc_east = RooRealVar("zdce", "ZDC ADC east", adc_min, adc_max)
#adc_west = RooRealVar("zdcw", "ZDC ADC west", adc_min, adc_max)
#kinematics variables
#m = RooRealVar("mee", "e^{+}e^{-} mass (GeV)", 0., 10.)
#y = RooRealVar("rapee", "rapidity", -1., 1.)
#pT = RooRealVar("ptpair", "pT", 0., 10.)
strsel = "jRecPt<{0:.3f} && jRecM>{1:.3f} && jRecM<{2:.3f}".format(
ptmax, mmin, mmax)
#strsel = "ptpair<{0:.3f} && mee>{1:.3f} && mee<{2:.3f}".format(ptmax, mmin, mmax)
data_all = RooDataSet("data", "data", tree,
RooArgSet(adc_east, adc_west, m, y, pT))
print "All input:", data_all.numEntries()
data = data_all.reduce(strsel)
print "Sel input:", data.numEntries()
model = Model2D(adc_east, adc_west)
r1 = model.model.fitTo(data, rf.Save())
ut.log_results(out, ut.log_fit_result(r1), lmg)
ut.log_results(out, "Fit parameters:\n", lmg)
out.write(ut.log_fit_parameters(r1, lmg + 2) + "\n")
#out.write(ut.table_fit_parameters(r1))
#print ut.table_fit_parameters(r1)
#create the plot
if p2d != 2: can = ut.box_canvas()
nbins, adc_max = ut.get_nbins(adc_bin, adc_min, adc_max)
adc_east.setMax(adc_max)
adc_west.setMax(adc_max)
frame_east = adc_east.frame(rf.Bins(nbins), rf.Title(""))
frame_west = adc_west.frame(rf.Bins(nbins), rf.Title(""))
data.plotOn(frame_east, rf.Name("data"))
model.model.plotOn(frame_east, rf.Precision(1e-6), rf.Name("model"),
rf.LineColor(model_col))
data.plotOn(frame_west, rf.Name("data"))
model.model.plotOn(frame_west, rf.Precision(1e-6), rf.Name("model"),
rf.LineColor(model_col))
#reduced chi^2 in east and west projections
ut.log_results(out, "chi2/ndf:\n", lmg)
ut.log_results(
out,
" East chi2/ndf: " + str(frame_east.chiSquare("model", "data", 16)),
lmg)
ut.log_results(
out,
" West chi2/ndf: " + str(frame_west.chiSquare("model", "data", 16)),
lmg)
ut.log_results(out, "", 0)
ytit = "Events / ({0:.0f} ADC units)".format(adc_bin)
frame_east.SetYTitle(ytit)
frame_west.SetYTitle(ytit)
frame_east.SetTitle("")
frame_west.SetTitle("")
frame = [frame_east, frame_west]
if p2d == 0: plot_projection(frame[ew], ew)
plot_pdf = PlotPdf(model, adc_east, adc_west)
if p2d == 1: plot_2d(plot_pdf)
if p2d == 2:
frame2 = ut.prepare_TH1D("frame2", adc_bin, adc_min,
2. * adc_max + 4.1 * adc_bin)
plot_proj_both(frame2, frame_east, frame_west, adc_bin, adc_min,
adc_max, ptmax, mmin, mmax)
lhead = ["east ZDC", "west ZDC"]
if p2d == 1:
leg = ut.prepare_leg(0.003, 0.9, 0.3, 0.1, 0.035)
else:
leg = ut.prepare_leg(0.66, 0.8, 0.32, 0.13, 0.03)
if p2d == 0: leg.AddEntry(None, "#bf{Projection to " + lhead[ew] + "}", "")
leg.SetMargin(0.05)
leg.AddEntry(None,
"#bf{#it{p}_{T} < " + "{0:.2f}".format(ptmax) + " GeV/c}", "")
mmin_fmt = "{0:.1f}".format(mmin)
mmax_fmt = "{0:.1f}".format(mmax)
leg.AddEntry(
None, "#bf{" + mmin_fmt + " < #it{m}_{e^{+}e^{-}} < " + mmax_fmt +
" GeV/c^{2}}", "")
leg.Draw("same")
pleg = ut.prepare_leg(0.99, 0.87, -0.4, 0.11, 0.035)
pleg.SetFillStyle(1001)
#pleg.AddEntry(None, "STAR Preliminary", "")
pleg.AddEntry(None, "AuAu@200 GeV", "")
pleg.AddEntry(None, "UPC sample", "")
#pleg.Draw("same")
#ut.print_pad(gPad)
#b3d = TBuffer3D(0)
#b3d = None
#gPad.GetViewer3D().OpenComposite(b3d)
#print b3d
#print "All input: ", data.numEntries()
#print "All input: 858"
#all input data
nall = float(tree.Draw("", strsel))
print "All input: ", nall
n_1n1n = float(model.num_1n1n.getVal())
print "1n1n events: ", n_1n1n
ratio_1n1n = n_1n1n / nall
sigma_ratio_1n1n = ratio_1n1n * TMath.Sqrt(
(nall - n_1n1n) / (nall * n_1n1n))
print "Ratio 1n1n / all: ", ratio_1n1n, "+/-", sigma_ratio_1n1n
ut.log_results(out, "Fraction of 1n1n events:\n", lmg)
ut.log_results(out, "All input: " + str(nall), lmg)
ut.log_results(out, "1n1n events: " + str(model.num_1n1n.getVal()), lmg)
ratio_str = "Ratio 1n1n / all: " + str(ratio_1n1n) + " +/- " + str(
sigma_ratio_1n1n)
ut.log_results(out, ratio_str, lmg)
if p2d != 2:
#ut.print_pad(gPad)
ut.invert_col(gPad)
can.SaveAs("01fig.pdf")
if interactive == True: start_interactive()
tupleDataSet = RooDataSet("tupleDataSet", "tupleDataSet",
RooArgSet(*varList),
RooFit.Import(tree1))
tupleDataSet.Print();
tmt = RooRealVar('tMinusT','tMinusT',0,-maxV,maxV);
#tDataSet = RooDataSet('tDataSet','tDataSet',RooArgSet(tmt,weightvar), RooFit.WeightVar(weightvar));
from ROOT import TH1D
tHist = TH1D('tDataSet','tDataSet',100,-0.3,0.3)#RooArgSet(tmt,weightvar), RooFit.WeightVar(weightvar));
for i in range(tupleDataSet.numEntries()):
tmt.setVal(tupleDataSet.get(i).find('Bs_ct').getValV()-tupleDataSet.get(i).find('Bs_TRUETAU').getValV()*1000.0);
#tDataSet.add(RooArgSet(tmt,weightvar),tupleDataSet.get(i).find(weightVarName).getValV());
tHist.Fill(tmt.getValV(),tupleDataSet.get(i).find(weightVarName).getValV());
#print tupleDataSet.get(i).find(weightVarName).getValV()," ",tDataSet.weight()
#tupleDataSet.merge(tDataSet);
tHist.Fit('gaus');
from ROOT import TCanvas
aCanvas = TCanvas();
'''
tDataSet.Print();
