from numpy import zeros,fromfunction 
from math import sin,pi 
import matplotlib 
from pylab import *
 
M = 101	 
N = 1000	 
a = 1.0/(M-1) 
x = [a*i for i in range(0,M)]
tau = 0.00002	 
Gamma = 0.5	
w=0.0

def Evolve_One_Timestep (dens): 
    dens_new = zeros((M,), float)
    for k in range (1, M-1):
        diff = Gamma*(dens[k+1] - 2*dens[k] + dens[k-1])/a/a - w/2.*(dens[k+1]-dens[k-1])/a
        dens_new[k] = dens[k] + tau*diff
    #Randbedingungen
    dens_new[0] = 0 
    dens_new[M-1] = 1 
    return dens_new

density = zeros((M,), float)
#Stufenfoermiges Anfangsprofil
for i in range(M/2,M):
    density[i] = 1 
for n in range(0,N): 
    density = Evolve_One_Timestep(density)
    if (n%100==0):
        plotlabel= "t = " + str(n*tau)  
        plot(x, density, label=plotlabel)
xlabel("x")
ylabel("n(x)")
legend() 
show()