Skip to content

Commit

Permalink
Merge pull request #332 from Marcellocosti/kde
Browse files Browse the repository at this point in the history 
Add possibility to include KDE templates in v2 fit procedure
  • Loading branch information
@stefanopolitano
stefanopolitano authored Jan 9, 2025
2 parents 0b1df2f + 1d9f35b commit 6ec7fe6
Show file tree
Hide file tree
Showing 8 changed files with 578 additions and 162 deletions.
52 changes: 52 additions & 0 deletions run3/flow/flow_analysis_utils.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
'''
Analysis utilities for flow analysis
'''

import ROOT
import os
import sys
Expand Down Expand Up @@ -53,6 +54,55 @@ def get_vn_versus_mass(thnSparse, inv_mass_bins, mass_axis, vn_axis, debug=False

return hist_mass_proj

def get_occupancy(thnSparse, occupancy_axis, debug=False):
'''
Project occupancy versus mass
Input:
- thnSparse:
THnSparse, input THnSparse obeject (already projected in centrality and pt)
- occupancy_axis:
int, axis number for occupancy
- debug:
bool, if True, create a debug file with the projections (default: False)
Output:
- hist_occupancy:
TH1D, histogram with vn as a function of mass
'''
hist_occupancy = thnSparse.Projection(occupancy_axis)

if debug:
outfile = ROOT.TFile('debug.root', 'RECREATE')
hist_occupancy.Write()
outfile.Close()

return hist_occupancy

def get_evselbits(thnSparse, evselbits_axis, debug=False):
'''
Project evselbits versus mass
Input:
- thnSparse:
THnSparse, input THnSparse obeject (already projected in centrality and pt)
- evselbits_axis:
int, axis number for evselbits
- debug:
bool, if True, create a debug file with the projections (default: False)
Output:
- hist_evselbits:
TH1D, histogram with vn as a function of mass
'''
hist_evselbits = thnSparse.Projection(evselbits_axis)

if debug:
outfile = ROOT.TFile('debug.root', 'RECREATE')
hist_evselbits.Write()
outfile.Close()

return hist_evselbits

def get_resolution(dets, det_lables, cent_min_max):
'''
Expand Down Expand Up @@ -417,6 +467,8 @@ def get_vnfitter_results(vnFitter, secPeak, useRefl):
vn_results['fBkgFuncMass'] = vnFitter.GetMassBkgFitFunc()
vn_results['fBkgFuncVn'] = vnFitter.GetVnVsMassBkgFitFunc()
vn_results['fSgnFuncMass'] = vnFitter.GetMassSignalFitFunc()
vn_results['fMassTemplFuncts'] = vnFitter.GetMassTemplFuncts()
vn_results['fVnTemplFuncts'] = vnFitter.GetVnTemplFuncts()
bkg, bkgUnc = ctypes.c_double(), ctypes.c_double()
vnFitter.Background(3, bkg, bkgUnc)
vn_results['bkg'] = bkg.value
Expand Down
54 changes: 51 additions & 3 deletions run3/flow/get_vn_vs_mass.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
import numpy as np
import yaml
from ROOT import TLatex, TFile, TCanvas, TLegend, TH1D, TH1F, TDatabasePDG, TGraphAsymmErrors # pylint: disable=import-error,no-name-in-module
from ROOT import gROOT, gPad, gInterpreter, kBlack, kRed, kAzure, kGray, kOrange, kGreen, kFullCircle, kFullSquare, kOpenCircle # pylint: disable=import-error,no-name-in-module
from ROOT import gROOT, gPad, gInterpreter, kBlack, kRed, kAzure, kGray, kOrange, kGreen, kMagenta, kFullCircle, kFullSquare, kOpenCircle # pylint: disable=import-error,no-name-in-module
from flow_analysis_utils import get_centrality_bins, get_vnfitter_results, get_ep_vn, getD0ReflHistos, get_particle_info # pylint: disable=import-error,no-name-in-module
sys.path.append('../../..')
sys.path.append('../..')
Expand All @@ -19,6 +19,7 @@
from ROOT import InvMassFitter, VnVsMassFitter
from utils.StyleFormatter import SetGlobalStyle, SetObjectStyle, DivideCanvas
from utils.FitUtils import SingleGaus, DoubleGaus, DoublePeakSingleGaus, DoublePeakDoubleGaus, RebinHisto
from kde_producer import kde_producer_grid, kde_producer_sim

def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
outputdir, suffix, vn_method, batch):
Expand Down Expand Up @@ -120,6 +121,26 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
print('ERROR: only kGaus, k2Gaus and k2GausSigmaRatioPar signal functions supported! Exit!')
sys.exit()

KDEtemplates = [[None]*len(fitConfig['TemplsFlags']) for _ in range(len(ptMins))]
if fitConfig.get('IncludeKDETempls'):
for iPt in range(len(ptMins)):
for iFlag, flag in enumerate(fitConfig['TemplsFlags']):
if fitConfig.get('FromGrid'):
KDEtemplates[iPt][iFlag], hist_templ = kde_producer_grid(fitConfig['TemplsInputs'][iFlag],
'fM', ptMins[iPt], ptMaxs[iPt], flag)
elif fitConfig.get('FromSim'):
KDEtemplates[iPt][iFlag] = kde_producer_sim(fitConfig['TemplsInputs'][iFlag],
fitConfig['TemplsTreeNames'][iFlag],
ptMins[iPt], ptMaxs[iPt], fitConfig['TemplsFlags'][iFlag])
elif fitConfig.get('FromFile'):
templFile = TFile.Open(f'{fitConfig["FromFile"]}', 'r')
KDEtemplates[iPt][iFlag] = templFile.Get(f'KDE_pt_{ptMins[iPt]}_{ptMaxs[iPt]}_flag{flag}')
templFile.Close()
else:
print(f'ERROR: incorrect setting for including KDEs in fit! Exit!')
sys.exit()
KDEtemplatesFuncts = [[KDE.GetFunction() for KDE in KDEtemplatesPt] for KDEtemplatesPt in KDEtemplates]

# set particle configuration
if particleName == 'Ds':
_, massAxisTit, decay, massForFit = get_particle_info(particleName)
Expand All @@ -137,6 +158,8 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
hMassIns, hMassOuts, hMassInsForFit, hMassOutsForFit = [], [], [], []
fTotFuncMass, fTotFuncVn, fSgnFuncMass, fBkgFuncMass, fMassBkgRflFunc,fBkgFuncVn = [], [], [], [], [], []
hMCSgn, hMCRefl = [], []
fMassTemplFuncts = [[None]*len(fitConfig['TemplsFlags']) for _ in range(len(ptMins))]
fVnTemplFuncts = [[None]*len(fitConfig['TemplsFlags']) for _ in range(len(ptMins))]
hist_reso = infile.Get('hist_reso')
hist_reso.SetDirectory(0)
reso = hist_reso.GetBinContent(1)
Expand Down Expand Up @@ -366,6 +389,11 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
vnFitter[iPt].SetTemplateReflections(hMCRefl[iPt],reflFuncStr,massMin,massMax)
vnFitter[iPt].SetFixReflOverS(SoverR)
vnFitter[iPt].SetReflVnOption(0) # kSameVnSignal
if fitConfig.get('IncludeKDETempls'):
vnFitter[iPt].SetKDETemplates(KDEtemplatesFuncts[iPt], fitConfig['TemplsTreeNames'],
fitConfig['InitWeights'][iPt], fitConfig['MinWeights'][iPt],
fitConfig['MaxWeights'][iPt])

# collect fit results
vnFitter[iPt].SimultaneousFit(False)
vnResults = get_vnfitter_results(vnFitter[iPt], secPeak, useRefl)
Expand All @@ -374,6 +402,8 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
fSgnFuncMass.append(vnResults['fSgnFuncMass'])
fBkgFuncMass.append(vnResults['fBkgFuncMass'])
fBkgFuncVn.append(vnResults['fBkgFuncVn'])
fMassTemplFuncts[iPt] = vnResults['fMassTemplFuncts']
fVnTemplFuncts[iPt] = vnResults['fVnTemplFuncts']

if useRefl:
fMassBkgRflFunc.append(vnResults['fMassBkgRflFunc'])
Expand Down Expand Up @@ -442,6 +472,12 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
latex.DrawLatex(0.18, 0.60, f'Signif. (3#sigma) = {round(vnResults["signif"], 2)}')
if useRefl:
latex.DrawLatex(0.18, 0.20, f'RoverS = {SoverR:.2f}')
for iMassTemplFunct, massTemplFunct in enumerate(fMassTemplFuncts[iPt]):
SetObjectStyle(massTemplFunct, color=kMagenta+iMassTemplFunct*2, linewidth=3)
massTemplFunct.SetLineColor(1)
massTemplFunct.Draw('same')
cSimFit[iCanv].Modified()
cSimFit[iCanv].Update()
cSimFit[iPt].cd(2)
hVnForFit[iPt].GetYaxis().SetRangeUser(-1, 1)
hVnForFit[iPt].GetYaxis().SetTitle(f'#it{{v}}_{{{harmonic}}} ({vn_method})')
Expand All @@ -458,6 +494,13 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
if secPeak:
latex.DrawLatex(0.18, 0.75,
f'#it{{v}}{harmonic}(D^{{+}}) = {vnResults["vnSecPeak"]:.3f} #pm {vnResults["vnSecPeakUnc"]:.3f}')

if fitConfig.get('drawvncomps'):
for iVnTemplFunct, vnTemplFunct in enumerate(fVnTemplFuncts[iPt]):
vnTemplFunct.SetLineColor(iVnTemplFunct+1)
vnTemplFunct.Draw('same')
cSimFit[iCanv].Modified()
cSimFit[iCanv].Update()
cSimFit[iCanv].Modified()
cSimFit[iCanv].Update()

Expand Down Expand Up @@ -701,7 +744,10 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
suffix_pdf = ')'
else:
suffix_pdf = ''
cSimFit[iPt].SaveAs(f'{outputdir}/SimFit{suffix}_{particleName}.pdf{suffix_pdf}')
if len(ptMins)==1:
cSimFit[iPt].SaveAs(f'{outputdir}/SimFit{suffix}_{particleName}.pdf')
else:
cSimFit[iPt].SaveAs(f'{outputdir}/SimFit{suffix}_{particleName}.pdf{suffix_pdf}')
outfile_name = f'{outputdir}/raw_yields{suffix}.root'
outFile = TFile(outfile_name, 'recreate')
if vn_method == 'sp' or vn_method == 'ep':
Expand All @@ -717,6 +763,9 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
fSgnFuncMass[ipt].Write(f'fSgnFuncMass_pt{ptmin*10:.0f}_{ptmax*10:.0f}')
fBkgFuncMass[ipt].Write(f'fBkgFuncMass_pt{ptmin*10:.0f}_{ptmax*10:.0f}')
fBkgFuncVn[ipt].Write(f'fBkgFuncVn_pt{ptmin*10:.0f}_{ptmax*10:.0f}')
for iFlag in range(len(KDEtemplatesFuncts[ipt])):
KDEtemplatesFuncts[ipt][iFlag].Write(f'{fitConfig["TemplsTreeNames"][iFlag]}_pt{ptmin*10:.0f}_{ptmax*10:.0f}_flag{fitConfig["TemplsFlags"][iFlag]}')

hSigmaSimFit.Write()
hMeanSimFit.Write()
hMeanSecPeakFitMass.Write()
Expand Down Expand Up @@ -774,7 +823,6 @@ def get_vn_vs_mass(fitConfigFileName, centClass, inFileName,
if not batch:
input('Press enter to exit')


if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Arguments')
parser.add_argument('fitConfigFileName', metavar='text', default='config_Ds_Fit.yml')
Expand Down
59 changes: 55 additions & 4 deletions run3/flow/invmassfitter/InvMassFitter.cxx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,8 @@ InvMassFitter::InvMassFitter() :
fFixRflOverSig(kFALSE),
fHistoTemplRfl(0x0),
fSmoothRfl(kFALSE),
fTemplates(kFALSE),
fNParsTempls(0),
fRawYieldHelp(0),
fRflFunc(0x0),
fBkRFunc(0x0),
Expand Down Expand Up @@ -157,6 +159,8 @@ InvMassFitter::InvMassFitter(const TH1F *histoToFit, Double_t minvalue, Double_t
fFixRflOverSig(kFALSE),
fHistoTemplRfl(0x0),
fSmoothRfl(kFALSE),
fTemplates(kFALSE),
fNParsTempls(0),
fRawYieldHelp(0),
fRflFunc(0x0),
fBkRFunc(0x0),
Expand Down Expand Up @@ -242,6 +246,9 @@ void InvMassFitter::SetNumberOfParams(){
if(fReflections) fNParsRfl=1;
else fNParsRfl=0;

if(fTemplates) fNParsTempls=this->fTemplatesFuncts.size();
else fNParsTempls=0;

if(fSecondPeak) fNParsSec=3;
else fNParsSec=0;

Expand All @@ -252,7 +259,6 @@ Int_t InvMassFitter::MassFitter(Bool_t draw){
/// returns 0 if the fit fails
/// returns 1 if the fit succeeds
/// returns 2 if there is no signal and the fit is performed with only background

TVirtualFitter::SetDefaultFitter("Minuit");

Double_t integralHisto=fHistoInvMass->Integral(fHistoInvMass->FindBin(fMinMass),fHistoInvMass->FindBin(fMaxMass),"width");
Expand Down Expand Up @@ -314,6 +320,10 @@ Int_t InvMassFitter::MassFitter(Bool_t draw){
fRflFunc = CreateReflectionFunction("freflect");
fBkRFunc = CreateBackgroundPlusReflectionFunction("fbkgrfl");
}
if(fTemplates){
printf(" ---> Final fit includes templates\n");
fTemplFunc = CreateTemplatesFunction("ftempl");
}
fTotFunc = CreateTotalFitFunction("funcmass");

if(doFinalFit){
Expand Down Expand Up @@ -360,6 +370,13 @@ Int_t InvMassFitter::MassFitter(Bool_t draw){
fBkRFunc->SetLineColor(kRed+1);
fBkRFunc->SetLineStyle(7);
}
if(fTemplates){
for(Int_t ipar=0; ipar<fNParsTempls; ipar++){
fTemplFunc->SetParameter(ipar,fTotFunc->GetParameter(ipar+fNParsBkg+fNParsSig+fNParsSec+fNParsRfl));
fTemplFunc->SetParError(ipar,fTotFunc->GetParError(ipar+fNParsBkg+fNParsSig+fNParsSec+fNParsRfl));
}
fTemplFunc->SetLineColor(kGreen+1);
}
fMass=fSigFunc->GetParameter(1);
fMassErr=fSigFunc->GetParError(1);
fSigmaSgn=fSigFunc->GetParameter(2);
Expand Down Expand Up @@ -473,6 +490,21 @@ TF1* InvMassFitter::CreateReflectionFunction(TString fname){
return funcrfl;
}
//______________________________________________________________________________
TF1* InvMassFitter::CreateTemplatesFunction(TString fname){
/// Creates a function for templates in the D+ inv. mass distribution
TF1* functempl = new TF1(fname.Data(),this,&InvMassFitter::FitFunction4Templ,fMinMass,fMaxMass,this->fTemplatesFuncts.size(),"InvMassFitter","FitFunction4Templ");
for(int iPar=0; iPar<functempl->GetNpar(); iPar++) {
functempl->SetParName(iPar, Form("w_%s",this->fTemplatesFuncts[iPar].GetName()));
if(this->fWeightsLowerLims[iPar] >= this->fWeightsUpperLims[iPar]) {
functempl->FixParameter(iPar,this->fInitWeights[iPar]);
} else {
functempl->SetParameter(iPar,this->fInitWeights[iPar]);
functempl->SetParLimits(iPar,this->fWeightsLowerLims[iPar],this->fWeightsUpperLims[iPar]);
}
}
return functempl;
}
//______________________________________________________________________________
TF1* InvMassFitter::CreateBackgroundPlusReflectionFunction(TString fname){
/// Creates the function with sum of background and reflections
///
Expand Down Expand Up @@ -549,11 +581,11 @@ TF1* InvMassFitter::CreateSignalFitFunction(TString fname, Double_t integsig){

//______________________________________________________________________________
TF1* InvMassFitter::CreateTotalFitFunction(TString fname){
/// Creates the total fit fucntion (signal+background+possible second peak)
/// Creates the total fit function (signal+background+possible second peak)
///

SetNumberOfParams();
Int_t totParams=fNParsBkg+fNParsRfl+fNParsSec+fNParsSig;
Int_t totParams=fNParsBkg+fNParsRfl+fNParsTempls+fNParsSec+fNParsSig;
TF1* ftot=new TF1(fname.Data(),this,&InvMassFitter::FitFunction4Mass,fMinMass,fMaxMass,totParams,"InvMassFitter","FitFunction4Mass");
for(Int_t ipar=0; ipar<fNParsBkg; ipar++){
ftot->SetParameter(ipar,fBkgFunc->GetParameter(ipar));
Expand Down Expand Up @@ -587,6 +619,15 @@ TF1* InvMassFitter::CreateTotalFitFunction(TString fname){
ftot->SetParLimits(ipar+fNParsBkg+fNParsSig+fNParsSec,parmin,parmax);
}
}
if(fTemplates && fTemplFunc){
for(Int_t ipar=0; ipar<fNParsTempls; ipar++){
Double_t parmin,parmax;
fTemplFunc->GetParLimits(ipar,parmin,parmax);
ftot->SetParLimits(ipar+fNParsBkg+fNParsSig+fNParsSec+fNParsRfl,parmin,parmax);
ftot->SetParameter(ipar+fNParsBkg+fNParsSig+fNParsSec+fNParsRfl,fTemplFunc->GetParameter(ipar));
ftot->SetParName(ipar+fNParsBkg+fNParsSig+fNParsSec+fNParsRfl,fTemplFunc->GetParName(ipar));
}
}
return ftot;
}
//__________________________________________________________________________
Expand Down Expand Up @@ -773,6 +814,14 @@ Double_t InvMassFitter::FitFunction4SecPeak (Double_t *x, Double_t *par){
return secgaval;
}
//_________________________________________________________________________
Double_t InvMassFitter::FitFunction4Templ(Double_t *x, Double_t *par){
Double_t totalTempl=0;
for(int iFunc=0; iFunc<this->fTemplatesFuncts.size(); iFunc++) {
totalTempl += par[iFunc] * this->fTemplatesFuncts[iFunc].Eval(x[0]);
}
return totalTempl;
}
//_________________________________________________________________________
Double_t InvMassFitter::FitFunction4Mass(Double_t *x, Double_t *par){
/// Total fit function (signal+background+possible second peak)
///
Expand All @@ -783,7 +832,9 @@ Double_t InvMassFitter::FitFunction4Mass(Double_t *x, Double_t *par){
if(fSecondPeak) sec=FitFunction4SecPeak(x,&par[fNParsBkg+fNParsSig]);
Double_t refl=0;
if(fReflections) refl=FitFunction4Refl(x,&par[fNParsBkg+fNParsSig+fNParsSec]);
return bkg+sig+sec+refl;
Double_t templ=0;
if(fTemplates) templ=FitFunction4Templ(x,&par[fNParsBkg+fNParsSig+fNParsSec+fNParsRfl]);
return bkg+sig+sec+refl+templ;
}

//_________________________________________________________________________
Expand Down
Loading

0 comments on commit 6ec7fe6

Please sign in to comment.