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