############################################################################
#
#   Name:           Grundlagen_1D_2.py
#   Python:         3.3.2 on win32
#   Autor:          Adrian Haas
#   
############################################################################
#1D wave equation
############################################################################
#
# The script animates a gaussian puls hitting a media in one dimension,
# using the explicit FDM (Finite Difference Method) for the wave equation.
# The output is the reflected/transmitted power in % of the source.
# A short mathematical description of the wave model for electromagnetic
# waves is in the file http://www.engineering-tools.ch/wellengleichung.pdf
#
############################################################################

import numpy
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import Rectangle

############################################################################
#define parameters and grid
############################################################################


#constants
eps0=8.854e-12
mu0=1.256e-6
c0=1/numpy.sqrt(eps0*mu0)
c=c0

#define grid
delta_x=0.005
delta_t=delta_x/c #stability criteria S=1
Nx=400
Nt=600

x_0 = numpy.arange(0,Nx*delta_x,delta_x) 
t_0 = numpy.arange(0,Nt*delta_t,delta_t)

Ez=numpy.zeros(Nx*Nt).reshape(Nx,Nt)

#source
func0=lambda t,a,mean,var: a*numpy.exp(-(t-mean)**2/2/var)
lq=numpy.arange(0,50,1)

#layers
eps=numpy.ones(Nx)*eps0
mu=numpy.ones(Nx)*mu0

#add obstacle
epsr=2
eps[200:250]=eps[200:250]*epsr



############################################################################
#Iteration 
############################################################################

for n in range(1,Nt):

    #wave equation
    Ez[1:Nx-1,n]=(1/(eps[1:Nx-1]*mu[1:Nx-1])*delta_t**2/delta_x**2
                  *(Ez[2:Nx,n-1]-2*Ez[1:Nx-1,n-1]+Ez[0:Nx-2,n-1])
                  +2*Ez[1:Nx-1,n-1]-Ez[1:Nx-1,n-2])

    #inserting source at timestep 10
    if n==10:
     Ez[69:119,n-1]= Ez[69:119,n-1]+func0(lq,1,25,40)
     Ez[70:120,n]=Ez[70:120,n]+func0(lq,1,25,40)

    #border condition, erease reflection  by copying (x-1,t-1) 
    Ez[0,n]=Ez[1,n-1]
    Ez[len(x_0)-1,n]=Ez[len(x_0)-2,n-1]


############################################################################
#Power calculation
############################################################################


#relative ~E**2 instead E**2/Z

sp=numpy.sum((Ez[:,11])**2)
rp=numpy.cumsum((Ez[10,:])**2)
tp=numpy.cumsum((Ez[len(x_0)-10,:])**2)


############################################################################
#Animation
############################################################################

#animation function
def ani_wave(j):

  wave=Ez[:,j]
  s_w.set_ydata(wave)

  time="Time: "+str(round(j*numpy.max(t_0)*1e9/len(t_0),1))+" ns"
  reflected_power="Reflected power="+str(numpy.round(rp[j]/sp*100,3))+"%"
  transmitted_power="Transmitted power="+str(numpy.round(tp[j]/sp*100,3))+"%"

  textz.set_text(time)
  textr.set_text(reflected_power)
  textt.set_text(transmitted_power)

  return s_w,textz,textr,textt


#init parameters
wave=Ez[:,0]
zeit="Time: 0 ns"
source_power="Source power=100%"
reflected_power="Reflected="+str(numpy.round(rp[0]/sp*100,3))+"%"
transmitted_power="Transmitted="+str(numpy.round(tp[0]/sp*100,3))+"%"


#create plot
fig=plt.figure("Reflexion und Transmission")
plt.title("$u_{tt}-c^2 u_{xx}=0$")
ax=fig.add_subplot(111)

s_w,=ax.plot(x_0,wave, linewidth=1)

textz=fig.text(0.2, 0.2, zeit, color='green')
texts=fig.text(0.2, 0.25, source_power, color='green')
textr=fig.text(0.2, 0.8, reflected_power, color='red')
textt=fig.text(0.2, 0.85, transmitted_power, color='blue')

ax.set_ylim(-1.5,1.5)

currentAxis = plt.gca()
currentAxis.add_patch(Rectangle((200/Nx*x_0[Nx-1], -1.5), 50/Nx*x_0[Nx-1], 3,
                      alpha=0.2, facecolor='gray'))
currentAxis.add_patch(Rectangle((0, -1.5), 0.05, 3,
                      alpha=0.2, facecolor='red'))
currentAxis.add_patch(Rectangle((1.95, -1.5), 0.1, 3,
                      alpha=0.2, facecolor='blue'))

#animate
wave_ani = animation.FuncAnimation(fig, ani_wave, frames=len(t_0),interval=1,
                                   blit=False)

plt.show()