Angle.py 3.22 KB
from numpy import log10, histogram, arange
from matplotlib.pyplot import figure, show
from matplotlib import gridspec
from Read import ReadRadial, ReadAngleVsEnergy, ReadGeneration
from Analytic import Analytic_theta
from Constants import degre

def angle_vs_energy(theta,energy,weight,nbBins=50):
   '''
      Compute arrival angle versus energy
      Input: directory name
      Output: energy, arrival angle
   '''
   energy =log10(energy)
   angle,ener = histogram(energy,nbBins,weights=weight*theta)
   nb,   ener = histogram(energy,nbBins,weights=weight)
   ener = 10**ener
   enercenter = (ener[1:nbBins+1]+ener[0:nbBins])/2
   angle /= nb

   return enercenter, angle

def radial(theta,weight,nbBins=50):
   cond = theta>0
   theta = log10(theta[cond])
   weight = weight[cond]
   dn,angle2 = histogram(theta,nbBins,range=[log10(1e-3),log10(25)],weights=weight)
   angle2=10**angle2
   #theta = theta**2
   #dn,angle2 = histogram(theta,nbBins,range=[0,25],weights=weight)
   thetacenter = (angle2[1:nbBins+1]+angle2[0:nbBins])/2
   dtheta = angle2[1:nbBins+1]-angle2[0:nbBins]
   dndtheta = dn/dtheta
   dndtheta /= max(dndtheta)

   return thetacenter, dndtheta
 
def drawRadial(files,plot_gen=False,plot_energy_range=False):
   '''
      Plot flux versus arrival angle
      Input: list of directories
      Output: graph of flux versus angle
   '''
   fig = figure()
   ax = fig.add_subplot(111)

   for fileId in files:
      theta,dndtheta2 = ReadRadial(fileId,[0,1])
      ax.plot(theta,dndtheta2,drawstyle='steps-mid')

      if plot_energy_range:
         labels = ["MeV band","GeV band","TeV band"]
         i = 2
         for n in arange(0,3,1):
            theta,dndtheta2 = ReadRadial(fileId,[0,i])
            ax.plot(theta,dndtheta2,drawstyle='steps-mid',label=labels[n])
            i+=1

      if plot_gen:
         generation = ReadGeneration(fileId,[0])
         i=5
         for gen in generation:
            theta,dndtheta2 = ReadRadial(fileId,[0,i])
            ax.plot(theta,dndtheta2,drawstyle='steps-mid',label="gen=%.0f"%gen)
            i+=1

   ax.legend(loc="best")
   ax.set_xscale('log')
   ax.set_yscale('log')
   ax.grid(b=True,which='major')
   ax.set_xlabel("$\\theta$ [deg]")
   ax.set_ylabel("$dN/d\\theta$ [arbitrary]")
   
   show()

def drawAngle_vs_energy(files,plot_gen=False,plot_analytic=False):
   '''
      Plot arrival angle versus energy + analytic expression
      Input: list of directories
      Output: graph of arrival angle versus energy
   '''
   fig = figure()
   ax1 = fig.add_subplot(111)

   for fileId in files:
      ener,angle = ReadAngleVsEnergy(fileId,[0,1])
      p=ax1.plot(ener,angle,drawstyle='steps-mid')

      if plot_gen:
         generation = ReadGeneration(fileId,[0])
         i=2
         for gen in generation:
            ener,angle = ReadAngleVsEnergy(fileId,[0,i])
            ax1.plot(ener,angle,drawstyle='steps-mid',label="gen=%.0f"%gen)
            i+=1

      if plot_analytic:
         yfit = Analytic_theta(ener,fileId)
         ax1.plot(ener,yfit,'--',color=p[0].get_color(),linewidth=2)

   ax1.legend(loc="best")
   ax1.set_xscale('log')
   ax1.set_yscale('log')
   ax1.set_ylim([0,180])
   ax1.grid(b=True,which='major')
   ax1.set_ylabel("$\\theta$ [deg]")
   ax1.set_xlabel("energy [GeV]")

   show()