# script gof-test.py
''' 
   goodness-of-fit for Poission statistics in histogram bins


.. author:: Guenter Quast <g.quast@kit.edu> 
     fuer den Kurs Rechnernutzung in der Physik

Modification: May-21 saturated model for with definiton Nt = No
'''
# --------------------------------------------------------------

import numpy as np
import matplotlib.pyplot as plt, scipy.special as sp
from scipy import stats

# standard definition of Poisson goodness-of-fit
def gof_Poisson(No, Nt):
  return np.sum( (Nt-No) + No*(np.log(No+.005)-np.log(Nt)) )

# alternative definition of Poisson goodness-of-fit
def gof_Poisson2(No, Nt):
  return np.inner( (Nt-No), np.log(Nt) ) +\
         np.sum( np.log(sp.gamma(No+1.)) - np.log(sp.gamma(Nt+1.)) )

def chi2(No, Nt):
   return np.sum( (Nt-No)**2 / Nt )

##### ---- main Program starts here -----
if __name__ == "__main__":

  Nbins = 20
  lam = 10

  nexp = 100000
  Nt = np.ones(Nbins) * lam
#  g= [ gof_Poisson(np.random.poisson(lam=lam,size=Nbins), Nt) \
#        for i in range(nexp) ]
  g= [ gof_Poisson(np.random.poisson(lam=lam,size=Nbins), Nt) \
        for i in range(nexp) ]
  # normalize g (factor 2!):
  g = 2*np.array(g)

  # print  
  print('*==* goodness-of-fit Test for histogram')
  print('  %i bins, true values = %.2g'%(Nbins, lam) )
  print('    mean of gof = %.3g'%(g.mean() ) )
  print('    std of gof  = %.3g'%(g.std() ) )
 
  # histogram goodness-of-fit
  fig=plt.figure('goodness-of-fit', figsize=(10, 4.))
  fig.subplots_adjust(left=0.14, bottom=0.1, right=0.97, top=0.93,
                    wspace=None, hspace=.25)#
  ax1=fig.add_subplot(1, 1, 1)
  nbins=100
  bc, be, _ = ax1.hist(g, nbins, density=False, alpha=0.5, label='g.o.f. Poisson')
  bw = be[1]-be[0]
  ax1.set_ylabel('freqeuncy')
  ax1.set_xlabel('Value of Tgof')
# compare with chi2-distribution
  xes=np.linspace(0.95*min(g), max(g), 100)
  ax1.plot(xes, nexp*bw*stats.chi2.pdf(xes, Nbins), 'r--',
           label='$\chi^2$ distribution')
  ax1.text(0.05, 0.85,  'uniform distribution',
     transform=ax1.transAxes )
  ax1.text(0.05, 0.75,  '  %i bins, lam=%.3g '%(Nbins,lam),
     transform=ax1.transAxes )
  ax1.legend(loc='best')

  plt.show()