# +----------------------------------------------------+
# |                  Lekce 1 - demo 2                  |
# +----------------------------------------------------+
#  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
#  Test of implicit Euler method

# 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
#dt=0.00075
dt=0.002

#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)
    for i in range(10):
      s=s+sgn**(i+1)*(np.exp(-20*(x+0.5+i)**2/tau)+np.exp(-20*(x-1.5-i)**2/tau))
    s=s/np.sqrt(tau)
    return s

#Construction of grigpoints
Nx=50
h=(b-a)/Nx
alpha=dt/h**2
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):
    v[1:-1]=v[1:-1]*(1-2*alpha)+alpha*(v[0:-2]+v[2:])

#Implicit Euler method, matrix for inversion defined first
M=np.zeros([Nx-1,Nx-1])
for i in range(Nx-2):
    M[i,i]=1.0+alpha*2
    M[i,i+1]=-alpha
    M[i+1,i]=-alpha
M[Nx-2,Nx-2]=1.0+alpha*2
def StepIE(v,dt):
    v[1:-1]=np.linalg.solve(M,v[1:-1])

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