# +----------------------------------------------------+
# |                  Lekce 1 - demo 1                  |
# +----------------------------------------------------+
#  FINITE DIFFERENCE METHODS for initial value problem
#   
#  PROBLEM: to solve initial value problem for diffusion equation
#    dU/dt = d^2U/dx^2  on (x,t) in <0,1> x <0,T> 
#    with Gaussian initial condition and Dirichlet/Neumann boundary condition

# Dirichlet or Neumann condition
Neumann=False
import numpy as np
import matplotlib.pyplot as plt

#interval x \in <a,b>
a=0.0
b=1.0

#Exact solution. For t=0 gives initial condition
# (tahle verze dobra pro male t - jinak vice clenu sumy)
def ExactU(x,t):
    if (Neumann):
        sgn=1
    else:
        sgn=-1
    tau=1.0+80*t
    s=np.exp(-20*(x-0.5)**2/tau)
    s=s+sgn*(np.exp(-20*(x+0.5)**2/tau)+np.exp(-20*(x-1.5)**2/tau))
    s=s/np.sqrt(tau)
    return s

#Construction of grigpoints
Nx=25
h=(b-a)/Nx
x=np.linspace(a,b,Nx+1)
v=np.array([ExactU(x[i],0.0) for i in range(Nx+1)])

#Grid for plotting
Nx0=100
x0=np.linspace(a,b,Nx0)
u0=np.array([ExactU(x0[i],0.0) for i in range(Nx0)])

#One step of Euler method
def StepEU(v,dt):
    tmp=dt/h**2
    v[1:-1]=v[1:-1]*(1-2*tmp)+tmp*(v[0:-2]+v[2:])

#Propagace reseni
t=0.0
#dt=0.00075
dt=0.001
for n in range(250):
    StepEU(v,dt)
    t=t+dt
    if (n%10==0):
      ut=np.array([ExactU(x0[i],t) for i in range(Nx0)])
      plt.text(0.3, .1,'t='+str(t)[:7])
      plt.plot(x0,u0,'g')
      plt.plot(x0,ut,'r')
      plt.plot(x,v,'bo')
      plt.plot(x,v,'b')
      plt.show()