Model.py

from scipy.signal import fftconvolve as convolve
from scipy.interpolate import interp1d
from scipy.stats import poisson, norm, binom
from AdditionalPDFs import gpoisson, borel
from scipy.integrate import simps
from numba import njit 
import numpy as np

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
########################################################################################################
## DETECTOR RESPONSE MODEL
########################################################################################################
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



########################################################################################################
## VARIABLES:
########################################################################################################

########################################################################################################
## Q: bin centres of studied charge distribution in arb. units;
## Q_0: pedestal position in arb. units
## G: gain in arb. units
## mu: mean number of Geiger discharges in n.p.e
## lbda: GP-branching parameter (unitless) 
## sigma_0: electonics noise/pedestal width in arb. units
## sigma_1: gain spread width in arb. units;
## tau_Ap: afterpulse time constant in nanoseconds;
## p_Ap: afterpulse scaling parameter (unitless)
## DCR: dark count rate in gigahertz;
## tau: slow-component time constant of SiPM pulse in nanoseconds
## tau_R: recovery time of SiPM in nanoseconds
## t_gate: gate length in nanoseconds
## t_0: pre-gate integration window in nanoseconds.
## k_max: estimated maximum number of peaks to fit, based on percentile of Poisson distribution with estimated mean in PeakOTron. 
## len_pad: number of bins to pad for convolution. Default value: 100 bins.
##
########################################################################################################





########################################################################################################
## FUNCTION DEFINITIONS:
########################################################################################################



########################################################################################################
## epsilon() 
########################################################################################################
##
## Returns machine epsilon for thresholding zeros
##
########################################################################################################


@njit('float64()')
def epsilon():
    return np.finfo(np.float64).eps

########################################################################################################
## sigmak(k, sigma_0, sigma_1) 
########################################################################################################
##
## Returns width of peak number k, sigma_k(k, sigma_0, sigma_1) = sqrt(sigma_0^2 + k*sigma__1^2)
##
########################################################################################################

@njit
def sigmak(k, sigma_0, sigma_1):
    return np.sqrt(sigma_0**2 + k*sigma_1**2)

########################################################################################################
## dabsQdt(t, tau, t_gate)) 
########################################################################################################
##
## Returns |dQ/dt|(t, tau, t_gate) see Chmill Et Al.
##
########################################################################################################

@njit
def dabsQdt(t, tau, t_gate): 
    return 2*np.exp(-t_gate/(2*tau))*np.abs(np.sinh((t_gate/2 - t)/tau))/tau

########################################################################################################
## dp1dt(t, tau_R, tau_Ap, t_gate)
########################################################################################################
##
## Returns dα/dt(t, tau_R, tau_Ap, t_gate) see Chmill Et Al.
##
########################################################################################################

@njit
def dp1dt(t, tau_R, tau_Ap, t_gate):
    return ((tau_Ap + tau_R)/(tau_Ap*(tau_Ap - np.exp(-t_gate/tau_Ap)*(tau_Ap + tau_R*(1 - np.exp(-t_gate/tau_Ap))))))*(1 - np.exp(-t/tau_R))*np.exp(-t/tau_Ap)

########################################################################################################
## tQ( Q, tau, t_gate)
########################################################################################################
##
## Returns time as a function of charge. 
##
########################################################################################################

@njit
def tQ( Q, tau, t_gate):
    return t_gate/2 - tau*np.arccosh(((1-Q)*np.exp(t_gate/(2*tau)) + np.exp(-t_gate/(2*tau)))/2)

########################################################################################################
## PH_max(t_gate, tau)
########################################################################################################
##
## Returns maximum pulse height of afterpulse
##
########################################################################################################

@njit
def PH_max(t_gate, tau): 
    return (1-np.exp(-t_gate/(2*tau)))**2

########################################################################################################
## Alpha(tau_R, tau_Ap, t_gate, p_Ap)
########################################################################################################
##
## Returns total afterpulse probability, α, occuring between 0 and t_gate.
##
########################################################################################################

@njit
def Alpha(tau_R, tau_Ap, t_gate, p_Ap):
    return (p_Ap/(tau_Ap + tau_R))*(tau_Ap - np.exp(-t_gate/tau_Ap)*(tau_Ap + tau_R*(1 - np.exp(-t_gate/tau_R))))


########################################################################################################
## dAlpha(tau_R, tau_Ap, t_gate, p_Ap)
########################################################################################################
##
## Returns error on total afterpulse probability, α, occuring between 0 and t_gate.
##
########################################################################################################


@njit
def dAlpha(tau_R, tau_Ap, t_gate, p_Ap, dtau_Ap, dp_Ap):
    dalphadtauAp = p_Ap*((t_gate*(-tau_Ap - tau_R*(1 - np.exp(-t_gate/tau_R)))*np.exp(-t_gate/tau_Ap))/tau_Ap**2 + 1 - np.exp(-t_gate/tau_Ap))/(tau_Ap+tau_R) - p_Ap*(tau_Ap - (tau_Ap + tau_R*(1 - np.exp(-t_gate/tau_R)))*np.exp(-t_gate/tau_Ap))/(tau_Ap + tau_R)**2
    
    dalphadpAp = (tau_Ap - (tau_Ap + tau_R*(1 - np.exp(-t_gate/tau_R)))*np.exp(-t_gate/tau_Ap))/(tau_Ap + tau_R)
    
    dalpha =  np.sqrt((dalphadtauAp**2)*(dtau_Ap**2) + (dalphadpAp**2)*(dp_Ap**2))
    
    return dalpha


########################################################################################################
## DiracDelta(length, idx_0, dx)
########################################################################################################
##
## Returns Dirac Delta Function at the bin, with probability density 1/dx.
##
########################################################################################################

@njit
def DiracDelta(length, idx_0, dx):
    delta = np.zeros(length)
    delta[idx_0] = 1/dx
    return delta

########################################################################################################
## dp1dQ(Q, tau, tau_R, tau_Ap, t_gate) 
########################################################################################################
##
## Returns afterpulse p.d.f as a function of n.p.e 
## Note: as in Chmill et al. However, the two branches of solutions given by t and t_gate -t are summed together. This demonstrably yields a normalised probability distribution. 
##
########################################################################################################

@njit
def dp1dQ(Q, tau, tau_R, tau_Ap, t_gate):
    _PH_max = PH_max(t_gate, tau)

    tb0 = tQ(Q, tau, t_gate)
    tb1 = t_gate - tb0
    
    N0 = dp1dt(tb0, tau_R, tau_Ap, t_gate)
    N1 = dp1dt(tb1, tau_R, tau_Ap, t_gate)
    
    D0 = dabsQdt(tb0, tau, t_gate)
    D1 = dabsQdt(tb1, tau, t_gate)
    
    D0 = np.where(D0<epsilon(), epsilon(), D0)
    D1 = np.where(D1<epsilon(), epsilon(), D1)
    
    pdf0 = N0/D0
    pdf1 = N1/D1
    
    pdf = pdf0 + pdf1
    
    pdf = np.where((Q>0) & (Q<_PH_max), pdf, 0)
    
    return pdf

########################################################################################################
## dpDCR1dQ(Q, tau, t_gate, t_0)
########################################################################################################
## Returns dark count p.d.f as a function of n.p.e 
## Note: the distribution is normalised by factor (t_gate+t_0)/tau
########################################################################################################

@njit
def dpDCR1dQ(Q, tau, t_gate, t_0):
    
    Q_min =  np.exp(-t_0/tau)*(1 - np.exp(-t_gate/tau))
    Q_max =  (1 - np.exp(-t_gate/tau))
    
    Q_th = np.where(np.abs(Q)<epsilon(), epsilon(), Q)
    
    pdf_bfg = np.where((Q_th>Q_min) & (Q_th<Q_max), 1/(Q_th), 0)
    pdf_dg = np.where((Q_th>0) & (Q_th<Q_max), 1/(1-Q_th), 0)
    
    pdf = (pdf_bfg + pdf_dg)*(tau/(t_gate + t_0))

    return pdf

########################################################################################################
## I_dp1dQ(Q, tau, t_gate, t_0)
########################################################################################################
##
## Returns dp1dQ, corrects for any nans/infs and ensures normalisation by applying Simpson's rule.
########################################################################################################

def I_dp1dQ(Q, dx, tau, tau_R, tau_Ap, t_gate):
    
    I = dp1dQ(Q, tau, tau_R, tau_Ap, t_gate)
    I = np.nan_to_num(I, nan=0.0, posinf=0.0, neginf=0.0)
    I = I/max(epsilon(), simps(I, dx=dx))

    return I

########################################################################################################
## I_dpDCR1dQ(Q, dx, tau, t_gate, t_0)
########################################################################################################
##
## Returns dpDCR1dQ, corrects for any nans/infs and ensures normalisation by applying Simpson's rule.
########################################################################################################

def I_dpDCR1dQ(Q, dx, tau, t_gate, t_0):
    I = dpDCR1dQ(Q, tau, t_gate, t_0)
    I = np.nan_to_num(I, nan=0.0, posinf=0.0, neginf=0.0)
    I = I/max(epsilon(), simps(I, dx=dx))

    return I

########################################################################################################
## Convolve(A, B, idx_0, length, dx, norm=True) 
########################################################################################################
##
## Calculates convolution of A and B using scipy fftconvolve, selects the convolved range relative to the pedestal index due to convolution shift, 
## and multiplies by dx to renormalise the distribution. Then, corrects for any nans/infs and ensures normalisation by applying Simpson's rule.
##
########################################################################################################

def Convolve(A, B, idx_0, length, dx, norm=True):
    A_o_B = convolve(A, B)[idx_0:length+idx_0]*dx
    A_o_B = np.nan_to_num(A_o_B, nan=0.0, posinf=0.0, neginf=0.0)
    if(norm):
        A_o_B = A_o_B/max(epsilon(), simps(A_o_B, dx=dx))
    return A_o_B


########################################################################################################
## DRM(Q, Q_0, G, mu, lbda, sigma_0, sigma_1, tau_Ap, p_Ap, DCR, tau, tau_R, t_gate, t_0, k_max, len_pad) 
########################################################################################################
##
## Calculates full detector response model. 
## Notes: 
## - The model is fitted in n.p.e scale. k = (Q - Q_0)/G. 
## - To rescale to the input parameter unit, the final pdf is scaled by a factor of 1/G -
##   this is because dp/dQ = (dp/dk)*(dk/dQ),  dk/dQ = d/dQ (Q - Q_0)/G = 1/G
##
########################################################################################################

def DRM(Q, Q_0, G, mu, lbda, sigma_0, sigma_1, tau_Ap, p_Ap, DCR, tau, tau_R, t_gate, t_0, k_max, len_pad):
    
    
    ########################################################################################################
    # Primary Geiger Model
    #########################################################################################################
    ##
    ## 1. Correct negative t_0 and ensures len_pad is an integer.
    ## 2. Calculate padded Q and indices of the original Q so that these can be selected at the end of the function. 
    ## 3. Convert Q to  n.p.e units for the main fit.
    ##
    ########################################################################################################
    
    t_0 = abs(t_0)
    len_pad = int(len_pad)
    

    
    len_Q = len(Q)
    len_Q_pad = len_Q + 2*len_pad
    Q_idx = np.arange(0, len_Q)
    Q_idx_pad =  np.arange(-len_pad, len_Q+len_pad)
    
    
    f_lin = interp1d(Q, Q_idx, fill_value='extrapolate')
    f_lin_inv = interp1d(Q_idx, Q, fill_value='extrapolate')
    
    Q_pad = f_lin_inv(Q_idx_pad)
    k_pad = (Q_pad - Q_0)/G
    dx = abs(f_lin_inv(1) - f_lin_inv(0))
    
    dx_k = dx/G
    
    idx_k_0_pad = np.argmin(abs(k_pad))


    Q_idx_orig = np.where(np.in1d(Q_idx_pad,
                                    np.intersect1d(
                                    Q_idx_pad,
                                    Q_idx
                                )
                        ), True, False)

        
    k_pg = np.arange(0, max(k_max, 3)).astype(int)
    
 
    ########################################################################################################
    # Primary Geiger Model
    ########################################################################################################
    ## 1. Calculate the probabilities:
    ##  - Primary Geiger, GP(k;μ,λ) 
    ##  - Afterpulse, Binom(k, i; alpha)
    ##
    ## 2. Calculate the p.d.fs:
    ##  - dp0dQs(Q; sigma_k(k; sigma_0/G, sigma_1/G) - Gaus(Q; μ=k, σ = sigma_k)
    ##  - Note: sigma_0 and sigma_1 are in bin units. To convert to n.p.e, divide by G. 
    ##
    ##  - dpndQs(k; tau, tau_R, tau_Ap, t_gate) - Afterpulse distributions. 
    ##     dp2dQ = dp1dQ o dp1dQ, dp3dQ = dp1dQ o dp1dQ o dp1dQ ... etc, where o is convolution. 
    ##
    ## 3. Calculate primary Geiger p.d.f:
    ##   - For zeroth order afterpulse, return Binom(0; k, α) x GP(k;μ,λ) x Gaus(Q; k; sigma_k)
    ##   - For nth order afterpulse, return Binom(n, k, α) x GP(k;μ,λ) x (dpndQ o Gaus(Q; k; sigma_k))   (Afterpulse p.d.f convolved with noise peak)
    ##   - The above statement is valid because of convolution distributivity rules: f o (g+h) = f o g + f o h
    ##   - Sum across all k to k_max;
    ## 
    ## RESULT: pdf_pg - Primary Geiger Discharge p.d.f 
    ########################################################################################################

        
    pdf_pg = np.zeros(len_Q_pad)
    
    Ppgs = np.expand_dims(gpoisson.pmf(k_pg, mu, lbda), -1)   
    
    Pns = np.zeros((int(k_max), int(k_max)))
    
    alpha = Alpha(tau_R, tau_Ap, t_gate, p_Ap)
    
    dp0dQs = np.vstack([norm.pdf(k_pad, loc=_k, scale=sigmak(_k, sigma_0/G, sigma_1/G)) for _k in k_pg])
    
    dpndQs = []
        

    for _k in k_pg:
        if(_k==0):
            _dpndQ = np.zeros(len(k_pad))
            Pns[0,0] = binom.pmf(0, 0, alpha)
        elif(_k==1):
            _dpndQ = I_dp1dQ(k_pad, dx_k, tau, tau_R, tau_Ap, t_gate)
            
            
            Pns[1,0:2] = np.array(binom.pmf(np.arange(0, 2).astype(int), _k, alpha))
        else:
            _dpndQ = Convolve(dpndQs[1], dpndQs[-1], idx_k_0_pad, len_Q_pad, dx_k, norm=True)
            Pns[_k,0:_k+1] = np.array(binom.pmf(np.arange(0, _k+1).astype(int), _k, alpha))

        dpndQs.append(_dpndQ)
            
    dpndQs = np.vstack(dpndQs)
  
    for _k in k_pg:

        _dp0dQ = float(Pns[_k, 0])*dp0dQs[_k, :]
        
        if(_k>0):
            
            
            _dpndQ_sum = np.sum(np.vstack([
                                    float(Pns[_k, _j])*Convolve(dpndQs[_j, :],
                                                             dp0dQs[_k, :],
                                                             idx_k_0_pad,
                                                             len_Q_pad,
                                                             dx_k,
                                                             norm=True   
                                                            )
                                    for _j in range(1, _k+1)
                                ]),
                               axis=0)

        else:
            _dpndQ_sum = np.zeros(len_Q_pad)
            
        _dp0dQ = np.nan_to_num(_dp0dQ, nan=0.0, posinf=0.0, neginf=0.0)
        _dpndQ_sum = np.nan_to_num(_dpndQ_sum, nan=0.0, posinf=0.0, neginf=0.0)
            
  
        pdf_pg += float(Ppgs[_k])*(_dp0dQ + _dpndQ_sum)
        

         
    ########################################################################################################
    # Dark Model
    ########################################################################################################
    ##
    ## 1. Calculate the probabilities:
    ##  - μDCR = DCR x (t_0 + t_gate)
    ##  - Primary Geiger -  Poisson(k;μDCR) 
    ##  - XTalk - Borel(k;λ) 
    ##
    ## 2. Calculate the p.d.fs:
    ##  - Notation as in Chmill et al:
    ##      - f_1 is 1st order dark distribution; f_2 = f_1 o f_1, f_3 = f_1 o f_1 o f_1 ... etc
    ##      - h_1 is 1st order stretched XT distribution; h_1 = 1/2 x f_1(Q/2; pars), h_2 = 1/3 x f_1(Q/3; pars)
    ##      - Convolutions of f_n and h_m shown as f_n_o_h_m. 
    ##  
    ## 3. Sum probabilities and pdfs as in Chmill et al. 
    ##    - Note: The 0th term is incorrect in Chmill et al. The probability should be Poisson(0; μDCR); the distribution should be
    ##      Dirac Delta distribution at pedestal 𝛿(Q - Q_0) (i.e. no discharge.)
    ##
    ## 4. Renormalise distribution
    ##  - Since p.d.f only goes up to 4th order, the distribution is not always correctly normalised. 
    ##    We correct for the normalisation by appling Simpson's rule. The effect of poor normalisation is only relevant 
    ##    for λ > 0.4: 5% off at this value. 
    ## 
    ## RESULT: pdf_dark - Dark Count p.d.f
    ########################################################################################################
    

    mu_DCR = DCR*(t_0 + t_gate)
    
    Pp0 = float(poisson.pmf(0, mu_DCR))
    Pp1 = float(poisson.pmf(1, mu_DCR))
    Pp2 = float(poisson.pmf(2, mu_DCR))
    Pp3 = float(poisson.pmf(3, mu_DCR))
    Pp4 = float(poisson.pmf(4, mu_DCR))
    
    Pb0 = float(borel.pmf(0, lbda))
    Pb1 = float(borel.pmf(1, lbda))
    Pb2 = float(borel.pmf(2, lbda))
    Pb3 = float(borel.pmf(3, lbda))
    
    
    f_0 = np.zeros(len_Q_pad)
    f_0[idx_k_0_pad] = 1/dx_k
    
    f_1 = I_dpDCR1dQ(k_pad, dx_k, tau, t_gate, t_0)
    
    f_2 = Convolve(f_1, f_1, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    f_3 = Convolve(f_1, f_2, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    f_4 = Convolve(f_1, f_3, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    
    h_1 = I_dpDCR1dQ(k_pad/(1+1), dx_k, tau, t_gate, t_0)/(1+1)
    h_2 = I_dpDCR1dQ(k_pad/(2+1), dx_k, tau, t_gate, t_0)/(2+1)
    h_3 = I_dpDCR1dQ(k_pad/(3+1), dx_k, tau, t_gate, t_0)/(3+1)
    

    f_1_o_h_1 = Convolve(f_1, h_1, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    f_1_o_h_2 = Convolve(f_1, h_2, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    h_1_o_h_1 = Convolve(h_1, h_1, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    f_2_o_h_1 = Convolve(f_2, h_1, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    

    
    
    dp0DCRdQ = Pp0*f_0
    dp1DCRdQ = Pp1*Pb0*f_1 
    dp2DCRdQ = Pp1*Pb1*h_1 + Pp2*(Pb0**2)*f_2
    dp3DCRdQ = Pp1*Pb2*h_2 + 2*Pp2*Pb0*Pb1*f_1_o_h_1 + Pp3*(Pb0**3)*f_3
    dp4DCRdQ = Pp1*Pb3*h_3 + 2*Pp2*Pb0*Pb2*f_1_o_h_2 + Pp2*(Pb1**2)*h_1_o_h_1 + 3*Pp3*(Pb0**2)*Pb1*f_2_o_h_1 + Pp4*(Pb0**4)*f_4
    
    pdf_dark = dp0DCRdQ + dp1DCRdQ + dp2DCRdQ + dp3DCRdQ + dp4DCRdQ
    pdf_dark = pdf_dark/max(epsilon(), simps(pdf_dark, dx=dx_k))
    

    
    ########################################################################################################
    # Convolve Primary Geiger and Dark Models
    ########################################################################################################
    ##
    ## 1. Convolve and normalise pdf_pg and pdf_dark
    ## pdf_drm = pdf_pg o pdf_dark 
    ## 
    ## RESULT: pdf_drm, final detector response model p.d.f
    ########################################################################################################
    
        
    pdf_drm =  Convolve(pdf_pg, pdf_dark, idx_k_0_pad, len_Q_pad, dx_k, norm=True)
    pdf_drm = pdf_drm[Q_idx_orig]/G
        
    return pdf_drm