Spectrum.py 1.43 KB
from numpy import histogram, log10, shape, logspace
from matplotlib.pyplot import figure, show
from os.path import isfile
from Read import ReadEnergy, ReadWeight, ReadElmag, 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()
   ax = fig.add_subplot(111)
   for fileId in files:
      energy,flux = spectrum(fileId)
      maxflux = 10 #shape(energy)[0]
      flux /= maxflux
      p=ax.plot(energy,flux,drawstyle='steps-mid',label=fileId)

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

   xfit = logspace(log10(10**(-3)),log10(2*10**4),200)
   yfit = Analytic_flux(xfit)
   ax.plot(xfit,yfit,'--m',linewidth=2,label="Analytic ($\\sqrt{E}$)")

   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.legend(loc="best")
   ax.set_xlabel("energy [GeV]")
   ax.set_ylabel("$E^2 \\frac{dN}{dE}$ [GeV]")

   show()