Map.py 4.8 KB
from numpy import shape, histogram2d, arange, newaxis, array, nditer, zeros, histogram
from numpy import sqrt, pi, exp, log, log10, sin, cos, where
from matplotlib.pyplot import figure, matshow, savefig, show
from matplotlib.colors import LogNorm
from Read import  resultDirectory, ReadRadial, ReadAngleVsEnergy
from Analytic import Analytic_theta
from Constants import degre

thetalim = 180
philim = 180
step = 0.25 #deg
bining = array([int(2*philim/step)+1,int(2*thetalim/step)+1])
philim=float(philim)
thetalim=float(thetalim)
limits = [[-thetalim,thetalim],[-philim,philim]]

def computeMap(theta,phi,weight,energy,dir_source,saveId,jet_opening_angle=180,Elim=1e-3,nbBins=50,borne=[180,180]):
   '''
      Plot 2D histogram of arrival direction of the photons with energy upper than Elim
      Input: list of directories
      Input (optional): lower energy limit (in GeV (default = 100MeV)
      Input (optional): nature of the source (default isotrop) 
      Input (optional): limits on theta, phi 
      Output: 2D histogram of theta cos(phi), theta sin(phi)
   '''

   # apply selection
   cond = (dir_source*degre <= jet_opening_angle) & (energy>Elim) #& (abs(theta)<thetalim) & (abs(phi)<thetalim)
   theta    = theta[cond]
   phi      = phi[cond]
   weight   = weight[cond]
   energy   = energy[cond]

   fig = figure()

   #===============================================================================================#
   #  IMAGE
   #===============================================================================================#
   ax = fig.add_subplot(111)
   H,xedges,yedges,im = ax.hist2d(theta*cos(phi),theta*sin(phi),bining,weights=weight,
         norm=LogNorm(),range=limits)
   cbar=fig.colorbar(im, ax=ax)
   cbar.ax.set_ylabel("counts")
   ax.set_xlim((-borne[0],borne[0]))
   ax.set_ylim((-borne[1],borne[1]))
   ax.set_xlabel("$\\theta$ $\\cos(\\phi)$ [deg]")
   ax.set_ylabel("$\\theta$ $\\sin(\\phi)$ [deg]")
   ax.grid(b=True,which='major')
   if jet_opening_angle == 180:
      ax.set_title(saveId+" - isotrop")
      savefig(resultDirectory+saveId+"/map_isotrop.png")
   else:
      ax.set_title(saveId+" - jet opening angle: %d degre"%(int(jet_opening_angle)))
      savefig(resultDirectory+saveId+"/map_"+str(int(jet_opening_angle))+".png")

   if borne != [180,180]: 
      H,xedges,yedges,im = ax.hist2d(theta*cos(phi),theta*sin(phi),bining,weights=weight,
            norm=LogNorm(),range=limits)
   #===============================================================================================#
   #  RADIAL DISTRIBUTION AND ENERGY VERSUS ANGLE
   #===============================================================================================#
   theta2 = zeros(shape(H)[0]*shape(H)[1])
   N_evts = zeros(shape(H)[0]*shape(H)[1])
   theta_cos_phi = xedges[::-1]
   theta_sin_phi = yedges[::-1]
   it = nditer(H,flags=["multi_index"])
   k = 0
   while not it.finished:
      i = int(it.multi_index[1])
      j = int(it.multi_index[0])
      theta2[k]= sqrt(theta_cos_phi[i]**2+theta_sin_phi[j]**2)
      N_evts[k]= it[0]
      k += 1
      it.iternext()

   theta2=log10(theta2)
   thetamax=max(borne)
   dn,angle = histogram(theta2,nbBins,range=[log10(1e-2),log10(thetamax)],weights=N_evts)
   angle=10**angle
   thetacenter = (angle[1:nbBins+1]+angle[0:nbBins])/2
   dtheta = angle[1:nbBins+1]-angle[0:nbBins]
   dndtheta = dn/dtheta*thetacenter

   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 thetacenter, dndtheta , enercenter, angle

def drawRadial(files):
   '''
      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',label=fileId)

   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("$\\theta*dN/d\\theta$")
   
   show()

def drawAngle_vs_energy(files):
   '''
      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',label=fileId)

      yfit = Analytic_theta(ener,fileId)
      ax1.plot(ener,yfit,'--',color="m",linewidth=2,label="analytic")

   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$ [rad]")
   ax1.set_xlabel("energy [GeV]")

   show()