Spectrum.py 1.88 KB
from numpy import histogram, log10, shape, logspace
from matplotlib.pyplot import figure, show, ylim, setp
from matplotlib import gridspec
from os.path import isfile
from Read import ReadEnergy, ReadWeight, ReadElmag, ReadNbLeptonsProd, resultDirectory, ElmagFile
from Analytic import Analytic_flux

def spectrum(fileName):
   energy = ReadEnergy(fileName)
   weight = ReadWeight(fileName)
   print "file", fileName, "->", shape(energy)[0], "events" 
   nbBins=100
   energy =log10(energy)
   flux,ener=histogram(energy,nbBins,weights=weight)
   ener = 10**ener
   enercenter=(ener[1:nbBins+1]+ener[0:nbBins])/2
   binSize=ener[1:nbBins+1]-ener[0:nbBins]
   flux = enercenter**2 * flux/binSize 
   return enercenter,flux

def drawSpectrum(files):
   fig = figure()
   gs = gridspec.GridSpec(2, 1, height_ratios=[4,1]) 
   fig.subplots_adjust(hspace=0.001)
   ax1 = fig.add_subplot(gs[0])
   ax2 = fig.add_subplot(gs[1],sharex=ax1)

   for fileId in files:
      energy,flux = spectrum(fileId)
      maxflux = ReadNbLeptonsProd(fileId)
      flux /= maxflux
      p=ax1.plot(energy,flux,drawstyle='steps-mid',label=fileId)

      yfit = Analytic_flux(energy)
      ax1.plot(energy,yfit,'--m',linewidth=2)

      error=flux/yfit-1
      ax2.plot(energy,error,'+',color=p[0].get_color())

      Elmag=resultDirectory+fileId+ElmagFile
      if isfile(Elmag):
         energy, flux = ReadElmag(fileId)
         ax1.plot(energy,flux,"-",color=p[0].get_color(),label="Elmag - "+fileId)

   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.grid(b=True,which='major')
   ax1.legend(loc="best")
   ax1.set_ylabel("$E^2 \\frac{dN}{dE}$ [GeV]")
   ax2.set_xscale('log')
   ax2.grid(b=True,which='major')
   ax2.set_xlabel("energy [GeV]")
   ax2.set_ylabel("error")
   step=0.2*(max(error)-min(error))
   ylim(min(error)-step, max(error)+step)
   xticklabels = ax1.get_xticklabels()
   setp(xticklabels, visible=False)

   show()