# EXEMPLE DE DIFFUSION 2D PAR DIFFERENCES FINIES

from pylab import *

# PARAMETRES PHYSIQUES
C = 1.0   #coefficient de diffusion
LX = 2.0  #taille du domaine
LY = 1.0  #taille du domaine
T = 0.10  #temps d'integration

# PARAMETRES NUMERIQUES
NX = 51   #nombre de points de grille
NY = 51   #nombre de points de grille
NT = 1000 #nombre de pas de temps

dx = LX/(NX-2) #pas de grille (espace)
dy = LY/(NY-2) #pas de grille (espace)
dt = T/NT      #pas de grille (temps)

dd = sqrt(dx**2+dy**2)


def Timestep(U,U_old):
   #Nouveau tableau
   U_new = zeros((NX,NY))
   U_old = zeros((NX,NY))
   #Boucle sur les points de grille, sauf C.L.
   for i in range (2, NX-2):
      for j in range (2, NY-2):
         ddU=(U[i-1,j]-2*U[i,j]+U[i+1,j])/(dx**2)+(U[i-1,j]-2*U[i,j]+U[i+1,j])/(dy**2)
         U_new[i,j]=2*U[i,j]-U_old[i,j]+((C**2)*(dt**2))*ddU

   #Condition aux limites périodiques
   U_new[NX-1,:]=U_new[1,:]
   U_new[0,:]=U_new[NX,:]
   U_new[:,NX-1]=U_new[:,1]
   U_new[:,0]=U_new[:,NX]

   return U_new


### PROGRAMME PRINCIPAL ###

# Maillage
xx = zeros((NX,NY))
yy = zeros((NX,NY))
for i in range (0,NX):
   for j in range (0,NY):
      xx[i,j]=i*dx
      yy[i,j]=j*dy

# Initialisation
U_data = zeros((NX,NY))
U_old = zeros((NX,NY))
x=linspace(0,LX,NX)
y=linspace(0,LY,NY)

a=100.0
U_data =exp(-a*(((x-0.3)**2)+((y-0.3)**2)))
U_zero = U_data.copy()
U_old  = U_data.copy()


# Boucle principale
for n in range(0,NT):

   #Plot tous les 100 pas de temps
   if (n%100==1):
      hold(False)
      plotlabel= "t = " + str(n*dt)
      pcolor(xx,yy,U_data)
      hold(True)
      contour(xx,yy,U_data,colors='k')
      title(plotlabel)
      draw()
      ginput()
   U_new = Timestep(U_data,U_old)
   U_old=U_data
   U_data=U_new