from pydiana import Diana,ROOT
import pandas as pd
import numpy as np
from array import array
pd.options.mode.chained_assignment = None
#ROOT.gROOT.SetBatch(True) #used to turn off ROOT's canvas
#ROOT.gErrorIgnoreLevel = ROOT.kWarning #Used to turn off ROOT Info
###################
# #
# Plot Types #
# #
###################
[docs]
def root_scatter(y,x=None,yerr=None,xerr=None):
"""
Interface Function for making a root scatter plot
Parameters
----------
y : array-like containing the y coordinates of the plot
x : array-like containing the x coordinates (if not provided the index of y will be used)
yerr : array-like containg the errors on the y coordinates (if absent set to 0)
xerr : array-like containg the errors on the x coordinates (if absent set to 0)
Returns
-------
ROOT TGraphErrors instance
"""
y = array("d",y)
if x is not None:
x = array('d',x)
else:
x = array('d',np.arange(len(y)))
if xerr is not None:
xerr = array('d',xerr)
else:
xerr = ROOT.nullptr#array('d',np.zeros(len(y)))
if yerr is not None:
yerr = array('d',yerr)
else:
yerr = ROOT.nullptr#array('d',np.zeros(len(y)))
if len(x) != len(y):
raise Exception(f"X,Y,ErrX,ErrY must have the same length (x[{len(x)}],y[{len(y)}],xerr[{len(xerr)}],yerr[{len(yerr)}]).")
gr = ROOT.TGraphErrors(len(y),x,y,xerr,yerr)
return gr
[docs]
def root_hist(height,bin_edges,name=None,error_y=None):
"""
Interface for converting a numpy/matplotlib histogram in a ROOT histogram
Parameters
----------
height : array-like containing all the heights of the bins
bin_edges : array-like containing all the bin edges (len(bin_edges) = len(height)+1)
name (str) : string containing the name of the histogram
Returns
-------
ROOT TH1D instance
"""
if name is None:
name='hist'
if len(height) != len(bin_edges)-1:
raise Exception(f"Arrays provided are of wrong length (height[{len(height)}], bin_edges[{len(bin_edges)}])")
rh=np.concatenate([[0],height,[0]])
#erh=np.concatenate([[0],bin_errors,[0]])
bin_edges = array('d',bin_edges)
rh = array('d',rh)
h = ROOT.TH1D('h',name,len(height),bin_edges)
h.SetContent(rh)
if error_y is not None:
[h.SetBinError(i+1,error_y[i]) for i in range(len(height))]
return h
[docs]
def root_hist2d(height,xbin_edges,ybin_edges,name=None,title=None):
"""
Interface for converting a 2d numpy/matplotlib histogram in a ROOT histogram
Parameters
----------
height : array-like containing all the heights of the bins
xbin_edges : array-like containing all the bin edges (len(bin_edges) = len(height)+1)
ybin_edges : array-like containing all the bin edges (len(bin_edges) = len(height)+1)
name (str) : string containing the name of the histogram
Returns
-------
ROOT TH1D instance
"""
if name is None:
name='hist2d'
if title is None:
title='hist2d'
height=np.array(height)
if height.shape != (len(ybin_edges)-1 ,len(xbin_edges)-1) :
raise Exception(f"Arrays provided are of wrong length")
xbin_edges = np.sort(xbin_edges)
ybin_edges = np.sort(ybin_edges)
xbin_edges = array('d',xbin_edges)
ybin_edges = array('d',ybin_edges)
hist2d = ROOT.TH2D(name,title,
height.shape[1],xbin_edges,
height.shape[0],ybin_edges)
for ybin in range(height.shape[0]):
for xbin in range(height.shape[1]):
hist2d.SetBinContent(xbin,ybin,height[ybin][xbin])
return hist2d
##################################
# #
# Root functions interface #
# #
##################################
[docs]
def root_func(expr:str,name:str='',parnames:dict=None,
parlims:dict=None,parvals:dict=None,
parfix:list=None,xrange=None):
"""
Interface for defining a ROOT function
Parameters
----------
expr (str) : String containing the expression of the function
name (str) : String containing the name of the function
parnames (dict) : Dictionary containing the names of the function parameters (indexed as in the formula)
parlims (dict) : Dictionary containing the limits of the function parameters (indexed as in the formula)
parvals (dict) : Dictionary containing the values of the function parameters (indexed as in the formula)
parfix (list) : List containing all the parameters to fix (attention the parameter value must be specified in parvals)(indexed as in the formula)
xrange (array-like) : Array containing the domain of the function
Returns
-------
ROOT TF1 Instance
"""
if xrange is None:
xrange=(-np.inf,+np.inf)
if len(xrange) != 2:
raise Exception(f"Must provide 2 values to specify xrange")
func = ROOT.TF1(name,expr,xrange[0],xrange[1])
if parnames is not None:
for par in parnames:
func.SetParName(par,parnames[par])
if parvals is not None:
for par in parvals:
func.SetParameter(par,parvals[par])
if parlims is not None:
for par in parlims:
func.SetParLimits(par,parlims[par][0],parlims[par][1])
if parfix is not None:
for par in parfix:
func.FixParameter(par,parvals[par])
return func
[docs]
def eval_root_func(func,x,params:list=None):
"""
Interface for evaluating root functions
Parameters
----------
func : ROOT.TF1 instance
x : array-like containing the points used for evaluation
params : array-like containing the values of the function parameters
Returns
-------
y : array-like with the values of the evaluated function
"""
if params is not None:
params=list(params)
_=[func.SetParameter(i,p) for i,p in enumerate(params)]
x = array("d",x)
y = []
_ = [y.append(func.Eval(xp)) for xp in x]
y = np.array(y)
return y
#####################
# #
# Standard Fits #
# #
#####################
[docs]
def gaussian_fit(histogram,degree:int=None,
type:str ='Chebyshev',params:list=None,
parfix:list=None,xrange:list=None,
name:str = 'mygaus',unit:str = 'mV',
get_detail_matrices:bool=False,
fit_options:str = 'QSLB',
parlims:list=None,
):
"""
Fitting one gaussian to premade histogram
Parameters
----------
histogram : dictionary containing "counts" with the heights of the histogram bars, and "bin_edges" containing the edges of the histogram bins (left and right)
params : array-like containing the parameters of the gaussian (constant,mean.stddev)
parfix : array-like contatining the index of the parameters to fix (0=constant,1=mean,2=stddev)
name : string with the name of the function
unit : units of mean/stddev
Returns
-------
fit_results : pandas dataframe containing the fit results
gaus : root fitted function
"""
if xrange is None:
xrange = [min(histogram['bin_edges']),max(histogram['bin_edges'])]
elif len(xrange)==2:
xrange = [min(xrange),max(xrange)]
else:
warnings.warn('Expected 2 values for xrange')
if xrange[0]<min(histogram['bin_edges']):
xrange[0]=min(histogram['bin_edges'])
if xrange[1]>max(histogram['bin_edges']):
xrange[1]=max(histogram['bin_edges'])
bkg = make_background_function(degree=degree,start_idx=3,type =type)
#expr = "[0]/(TMath::Sqrt(2*TMath::Pi())*[2])*exp(-0.5*((x-[1])*(x-[1]))/([2]*[2]))"
expr='gaus(0)'
parnames = {0:"Constant",1:f"Mean[{unit}]",2:f"StdDev[{unit}]"}
if params is not None:
parvals = {i:pp for i,pp in enumerate(params)}
else:
bincenters = 0.5*(histogram['bin_edges'][1:]+histogram['bin_edges'][:-1])
mean = np.sum(bincenters*histogram['counts'])/(np.sum(histogram['counts']))
var = np.sum(histogram['counts']*(bincenters-mean)**2)/np.sum(histogram['counts'])
stddev = np.sqrt(var)
const = np.sum(histogram['counts'])/np.sqrt(2*np.pi*var)
parvals = {0:const,1:mean,2:stddev}
if parlims is None and params is not None:
parlims = {}
parlims[0] = (0,parvals[0]*2)
parlims[1] =(parvals[1]-5*parvals[2],parvals[1]+5*parvals[2])
parlims[2] =(0,2*parvals[2])
if bkg != "":
expr +=bkg
for i in range(degree+1):
parnames[i+3]="p"+str(i)
gaus = root_func(expr = expr,
name = name,
parvals = parvals,
parnames = parnames,
xrange = xrange,
parfix=parfix,
parlims=parlims,
)
res = root_hist_fit(heights=histogram['counts'],
bin_edges=histogram['bin_edges'],
func=gaus,
get_detail_matrices=get_detail_matrices,
options=fit_options,
)
#format results
res[0]['Units']['MinX']=unit
res[0]['Units']['MaxX']=unit
return res
[docs]
def root_fitting(graph,func,get_detail_matrices:bool=False,options:str='QSB'):
"""
Root fitting function
Parameters
----------
graph : root graph instance
func : root function instance
get_detail_matrices : bool for enabling the return of the covariance and correlation matrices
Returns
-------
Fit details and fitted function and if get detailed matrices also the covariance and correlation matrices.
"""
if get_detail_matrices and 'S' not in options:
options +='S'
resultpr = graph.Fit(func,f'{options}')
graph.Draw()
if get_detail_matrices:
cov = ROOT.TMatrixD(resultpr.GetCovarianceMatrix())
cor = ROOT.TMatrixD(resultpr.GetCorrelationMatrix())
resmat = {}
for mat,lab in zip((cov,cor),('Covariance','Correlation')):
nrow = mat.GetNrows()
ncol = mat.GetNcols()
cmat = [[mat(r,c) for c in range(ncol)] for r in range(nrow)]
resmat[lab]=np.array(cmat)
return extract_fit_values(func),func,resmat
return extract_fit_values(func),func
[docs]
def root_hist_fit(heights,bin_edges,func,get_detail_matrices:bool=False,options='QSBL'):
"""
Histogram fitting function
Parameters
----------
heights : Array-like containing the heights of the histogram
bin_edges : Array-like containing the edges of the bins
func : root function instance
get_detail_matrices : bool for enabling the return of the covariance and correlation matrices
Returns
-------
Fit details and fitted function
"""
xmin = max((func.GetXmin(),np.min(bin_edges)))
xmax = min((func.GetXmax(),np.max(bin_edges)))
bin_edges = np.array(bin_edges)
mask = np.where((bin_edges >=xmin) & (bin_edges <=xmax) )[0]
bin_edges = np.array(bin_edges)[mask]
heights = np.array(heights)[mask[:-1]]
hist = root_hist(height=heights,bin_edges=bin_edges)
func.SetRange(xmin,xmax)
fit =root_fitting(graph=hist,func = func,
get_detail_matrices=get_detail_matrices,
options=options)
del hist
return fit
[docs]
def make_background_function(degree:int,start_idx:int,type:str ='Chebyshev'):
bkg_expr=''
if degree is None:
return bkg_expr
if degree < 0:
print("Degree less than zero, not accepted. Background will not be used.")
else:
if type == "pol":
bkg_expr+=f"+pol{degree}({start_idx})"
elif type == "Chebyshev":
bkg_expr+=f"+cheb{degree}({start_idx})"
return bkg_expr