import numpy as np
import matplotlib.pyplot as plt
import math
from scipy.fftpack import fft #paket za FFT
from scipy.fftpack import ifftshift, fftshift
from mpl_toolkits.mplot3d import Axes3D #paket za 3d plot
plt.style.use('dark_background') # izbor boja za plot
plt.rcParams['image.cmap'] = 'plasma'
N=1048 #rezolucija
L=5. # širina zaslona na kojem se javlja difrakcioni uzorak
lamb=632*10**(-6) #talasna dužina svjetlosti u milimetrima (mm)
w=0.05 #širina otvora (mm)
z=100. #udaljenost otvor - zaslon (mm)
k=2*np.pi/lamb #talasni broj
delta_src=L/N #prostorna frekvencija
#definisanje ravni otvora
xv=np.linspace(-N/2,N/2-1,N)
yv=np.linspace(-N/2,N/2-1,N)
[x1,y1]=np.meshgrid(xv,yv)
x1=x1*delta_src #(mm)
y1=y1*delta_src # (mm)
#definisanje kvadratnog otvora
X=np.where(np.abs(x1)<w/2,1,0)
Y=np.where(np.abs(y1)<w/2,1,0)
u1=X*Y #elektricno polje na otvoru ima konstantnu vrijednost
#crtanje otvora
plt.imshow(u1)
plt.ylim(400,650)
plt.xlim(400,650)
plt.show()
# FOURIEROVA TRANSFORMACIJA - izracunavanje integrala
g=u1
G=fftshift(np.fft.fft2(g))
h=np.exp(1j*k*z)/1j*lamb*z*np.exp(1j*np.pi/(lamb*z)*(x1**2+y1**2))
u2=h*G #Vrijednost električnog polja u ravni zaslona
# Crtanje difrakcionog uzorka na zaslonu
fig=plt.figure()
plt.imshow(np.abs(u2))
plt.xlim(480,568)
plt.ylim(480,568)
tick_locs = [0,262,524,786,1048]
tick_lbls = [-2.5,-1.25,0,1.25,2.5]
plt.xticks(tick_locs, tick_lbls)
plt.yticks(tick_locs, tick_lbls)
plt.xlabel('X(mm)')
plt.ylabel('Y(mm)')
plt.show()
# intenzitet zračenja (normirano na jedinicu)
Isquare=1./(lamb*z)**2*u2
Isquare=Isquare/Isquare.max()
# Plotanje difrakcione slike i intenziteta zračenja
fig1 = plt.figure(figsize=plt.figaspect(.3))
ax1=fig1.add_subplot(1,2,1)
plt.imshow(np.abs(u2))
tick_locs = [0,262,524,786,1048]
tick_lbls = [-2.5,-1.25,0,1.25,2.5]
plt.xticks(tick_locs, tick_lbls)
plt.yticks(tick_locs, tick_lbls)
ax1.set_title('Difrakciona slika')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm)')
ax2=fig1.add_subplot(1,2,2,projection='3d')
ax2=fig1.gca(projection='3d')
surf=ax2.plot_surface(x1,y1,np.abs(Isquare),cmap='viridis', edgecolor='none')
ax2.set_title('Relativni intenzitet')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm))')
plt.show()
#DRUGI DIO - KRUŽNI OTOVOR
#Definisati radijus kruznog otvora
Radi=(x1**2+y1**2)**0.5
XX=np.where(np.abs(Radi) <= w/2,1,0)
YY=np.where(np.abs(Radi) <= w/2,1,0)
u3=XX*YY #električno polje ima stalnu vrijednost u svakoj tački otvora
plt.imshow(u3)
plt.show()
#FOURIEROVA TRANSFORMACIJA I IZRAČUNAVANJE INTEGRALA
gg=u3
GG=fftshift(np.fft.fft2(gg))
hh=np.exp(1j*k*z)/1j*lamb*z*np.exp(1j*np.pi/(lamb*z)*(x1**2+y1**2))
HH=fftshift(np.fft.fft2(hh))
u4=hh*GG
# DIFRAKCIONA SLIKA
plt.imshow(np.abs(u4))
tick_locs = [0,262,524,786,1048]
tick_lbls = [-2.5,-1.25,0,1.25,2.5]
plt.xticks(tick_locs, tick_lbls)
plt.yticks(tick_locs, tick_lbls)
plt.show()
# intenzitet zračnje (normiran na jedinicu)
Icirc=1./(lamb*z)**2*u4
Icirc=Icirc/Icirc.max()
# Crtanje difrakcione slike i intenziteta zračenja
fig1 = plt.figure(figsize=plt.figaspect(.3))
ax1=fig1.add_subplot(1,2,1)
plt.imshow(np.abs(u4))
tick_locs = [0,262,524,786,1048]
tick_lbls = [-2.5,-1.25,0,1.25,2.5]
plt.xticks(tick_locs, tick_lbls)
plt.yticks(tick_locs, tick_lbls)
ax1.set_title('Difrakciona slika')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm))')
ax2=fig1.add_subplot(1,2,2,projection='3d')
ax2=fig1.gca(projection='3d')
surf=ax2.plot_surface(x1,y1,np.abs(Icirc),cmap='viridis', edgecolor='none')
ax2.set_title('Relativni intenzitet')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm))')
plt.show()
#REZIME - plotanje difrakcione slike i intenziteta zračenja za kvadratnu i kružnu pukotinu
fig1 = plt.figure(figsize=(11,10))
ax1=fig1.add_subplot(2,2,1)
plt.imshow(np.abs(u4))
tick_locs = [0,262,524,786,1048]
tick_lbls = [-2.5,-1.25,0,1.25,2.5]
plt.xticks(tick_locs, tick_lbls)
plt.yticks(tick_locs, tick_lbls)
ax1.set_title('Difrakciona slika')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm))')
ax2=fig1.add_subplot(2,2,2,projection='3d')
ax2=fig1.gca(projection='3d')
surf=ax2.plot_surface(x1,y1,np.abs(Icirc), edgecolor='none')
ax2.set_title('Relativni intenzitet')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm))')
ax3=fig1.add_subplot(2,2,3)
plt.imshow(np.abs(u2))
tick_locs = [0,262,524,786,1048]
tick_lbls = [-2.5,-1.25,0,1.25,2.5]
plt.xticks(tick_locs, tick_lbls)
plt.yticks(tick_locs, tick_lbls)
ax3.set_title('Difrakciona slika')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm))')
ax4=fig1.add_subplot(2,2,4,projection='3d')
ax4=fig1.gca(projection='3d')
surf=ax4.plot_surface(x1,y1,np.abs(Isquare), edgecolor='none')
ax4.set_title('Relativni intenzitet')
plt.xlabel('X(mm)')
plt.ylabel('Y(mm))')
fig1.savefig('Fraunhofer_2d.png')
plt.show()