In [1]:
import numpy as np
import matplotlib.pyplot as plt
from IPython import display
import math as mt

Advection equation

I show some schemes to solve a PDE. In this notebook, a simple advection equation is solved. $$\frac{\partial u}{\partial t}+c\frac{\partial u}{\partial x}=0$$ Here $c$ is a constant. The CFL number is fixed by $0.1$.

FTCS

A time-discretization is forward and a space-discretization is central. $$u_j^{k+1}=u_j^k-\frac{c\Delta t}{2\Delta x}(u_{j+1}^k-u_{j-1}^k)$$

In [2]:
def FTCS(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
        
    jj=np.arange(nx); 
    jp=np.r_[jj[0:nx-1]+1,0];
    jm=np.r_[nx-1,jj[1:nx]-1] 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
    
    cf=c*dt/2/dx
    udata=np.zeros((nt,nx))
    udata[0,:]=u    
    for it in range(nt):

        u=u-cf*(u[jp]-u[jm])
        
        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    #plt.ylim(-0.2,1.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show()             
In [3]:
%%time
FTCS(1,128,0.1,1280)
CPU times: user 8.73 s, sys: 375 ms, total: 9.11 s
Wall time: 9.9 s

Lax-Friedrichs

$$u_j^{k+1}=\frac{u_{j+1}^k+u_{j-1}^k}{2}-\frac{c\Delta t}{2\Delta x}(u_{j+1}^k-u_{j-1}^k)$$

In [4]:
def LF(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    
    jj=np.arange(nx); 
    jp=np.r_[jj[0:nx-1]+1,0];
    jm=np.r_[nx-1,jj[1:nx]-1] 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
    
    cf=c*dt/2/dx
    udata=np.zeros((nt,nx))
    udata[0,:]=u    
    for it in range(nt):

        u=(u[jp]+u[jm])/2-cf*(u[jp]-u[jm])
        
        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    #plt.ylim(-0.2,1.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show()                         
In [5]:
%%time
LF(1,128,0.1,1280)
CPU times: user 8.58 s, sys: 359 ms, total: 8.94 s
Wall time: 9.76 s

Lax-Wendroff

$$u_j^{k+1}=u_{j}^k-\frac{c\Delta t}{2\Delta x}(u_{j+1}^k-u_{j-1}^k)+\frac{c\Delta t^2}{2\Delta x^2}(u_{j+1}^k-2u_j^k+u_{j-1}^k)$$

In [6]:
def LW(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    
    jj=np.arange(nx); 
    jp=np.r_[jj[0:nx-1]+1,0];
    jm=np.r_[nx-1,jj[1:nx]-1] 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
    
    cf=c*dt/2/dx
    udata=np.zeros((nt,nx))
    udata[0,:]=u    
    for it in range(nt):

        u=u-cf*(u[jp]-u[jm])+2*cf*cf*(u[jp]-2*u+u[jm])
        
        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    #plt.ylim(-0.2,1.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show()                          
In [7]:
%time LW(1,128,0.1,1280)
CPU times: user 8.62 s, sys: 484 ms, total: 9.11 s
Wall time: 9.87 s

Up-wind (1st order)

$$u_j^{k+1}=u_{j}^k-\frac{c\Delta t}{\Delta x}(u_{j}^k-u_{j-1}^k)$$

In [8]:
def uw1(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    
    jj=np.arange(nx); 
    jp=np.r_[jj[0:nx-1]+1,0];
    jm=np.r_[nx-1,jj[1:nx]-1] 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
    
    cf=c*dt/2/dx
    udata=np.zeros((nt,nx))
    udata[0,:]=u    
    for it in range(nt):

        u=u-2*cf*(u[jj]-u[jm])
        
        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    #plt.ylim(-0.2,1.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show()               
In [9]:
%%time 
uw1(1,128,0.1,1280)
CPU times: user 8.5 s, sys: 500 ms, total: 9 s
Wall time: 9.58 s

Up-wind (2nd order)

$$u_j^{k+1}=u_{j}^k-\frac{c\Delta t}{2\Delta x}(3u_{j}^k-4u_{j-1}^k+u_{j-2}^k)$$

In [10]:
def uw2(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    
    jj=np.arange(nx); 
    jp=np.r_[jj[0:nx-1]+1,0];
    jm=np.r_[nx-1,jj[1:nx]-1] 
    jm2=np.r_[nx-2,nx-1,jj[1:nx-1]-1]
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
    
    cf=c*dt/2/dx
    udata=np.zeros((nt,nx))
    udata[0,:]=u    
    for it in range(nt):

        u=u-cf*(3*u[jj]-4*u[jm]+u[jm2])
        
        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    #plt.ylim(-0.2,1.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show()               
In [11]:
%%time 
uw2(1,128,0.1,1280)
CPU times: user 7.66 s, sys: 297 ms, total: 7.95 s
Wall time: 8.51 s

Central differences (4th order)+Runge-Kutta (4th order)

In [12]:
def Cent4RG4(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    
    jj=np.arange(nx); 
    jp2=np.r_[jj[0:nx-2]+2,0,1];     jp=np.r_[jj[0:nx-1]+1,0]
    jm2=np.r_[nx-2,nx-1,jj[2:nx]-2]; jm=np.r_[nx-1,jj[1:nx]-1]
 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
    #plt.plot(xx,u,'-o')
  
    udata=np.zeros((nt,nx))
    udata[0,:]=u
    
    cf=c*dt/12/dx
    for it in range(nt):
        kk1=-cf*(-u[jp2]+8*u[jp]-8*u[jm]+u[jm2])
        v=u+0.5*kk1
        kk2=-cf*(-v[jp2]+8*v[jp]-8*v[jm]+v[jm2])
        v=u+0.5*kk2
        kk3=-cf*(-v[jp2]+8*v[jp]-8*v[jm]+v[jm2])
        v=u+    kk3
        kk4=-cf*(-v[jp2]+8*v[jp]-8*v[jm]+v[jm2])    
        u =u+(kk1+2*kk2+2*kk3+kk4)/6
        #u =tenthlowpassfilter(u,nx)

        #display.clear_output(True)
        #plt.plot(u)
        #plt.ylim(-0.2,1.2)
        #plt.show()
        udata[it,:]=u

    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show()  
In [13]:
%time Cent4RG4(1,128,0.1,1280)
CPU times: user 5.72 s, sys: 516 ms, total: 6.23 s
Wall time: 6.4 s

6th order compact difference and 3rd order TVD Runge-Kutta

In [14]:
def Com6TRG3(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    
    jj=np.arange(nx); 
    jp2=np.r_[jj[0:nx-2]+2,0,1];     jp=np.r_[jj[0:nx-1]+1,0]
    jm2=np.r_[nx-2,nx-1,jj[2:nx]-2]; jm=np.r_[nx-1,jj[1:nx]-1]
 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
  
    udata=np.zeros((nt,nx))
    udata[0,:]=u
    
    cf=c*dt/dx
    for it in range(nt):
        kk1=u     -cf*cmpct(u,nx)
        kk2=3/4*u+1/4*(kk1-cf*cmpct(kk1,nx))
        u  =1/3*u+2/3*(kk2-cf*cmpct(kk2,nx))
        #u  =tenthlowpassfilter(u,nx)

        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    plt.ylim(-0.2,2.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show()  
    
def cmpct(f,nx):
    jj=np.arange(nx)
    jp1=np.r_[jj[0:nx-1]+1,0]
    jm1=np.r_[nx-1,jj[1:nx]-1]
    jp2=np.r_[jj[0:nx-2]+2,0,1]
    jm2=np.r_[nx-2,nx-1,jj[1:nx-1]-1]
    A=np.zeros((nx,nx))
    for i in range(nx):
        A[i,i]=3
    for i in range(nx-1):
        A[i,i+1]=1
        A[i+1,i]=1

    A[0,nx-1]=1
    A[nx-1,0]=1

    cm=np.zeros(nx)
    cm=7/3*(f[jp1]-f[jm1])+1/12*(f[jp2]-f[jm2])

    x =np.zeros((nx,nx+1))
    x[:,0:nx]=A
    x[:,nx]=cm

    return np.linalg.solve(A, cm)

def tenthlowpassfilter(f,nx):
    jj=np.arange(nx)
    jp1=np.r_[jj[0:nx-1]+1,0]
    jm1=np.r_[nx-1,jj[1:nx]-1]
    jp2=np.r_[jj[0:nx-2]+2,0,1]
    jm2=np.r_[nx-2,nx-1,jj[1:nx-1]-1]
    jp3=np.r_[jj[0:nx-3]+3,0,1,2]
    jm3=np.r_[nx-3,nx-2,nx-1,jj[1:nx-2]-1]
    jp4=np.r_[jj[0:nx-4]+4,0,1,2,3]
    jm4=np.r_[nx-4,nx-3,nx-2,nx-1,jj[1:nx-3]-1]
    jp5=np.r_[jj[0:nx-5]+5,0,1,2,3,4]
    jm5=np.r_[nx-5,nx-4,nx-3,nx-2,nx-1,jj[1:nx-4]-1]
    jp6=np.r_[jj[0:nx-6]+6,0,1,2,3,4,5]
    jm6=np.r_[nx-6,nx-5,nx-4,nx-3,nx-2,nx-1,jj[1:nx-5]-1]
    af=0.495
    a0=(193+126*af)/256
    a1=(105+302*af)/256
    a2=15*(-1+2*af)/64
    a3=45*(1-2*af)/512
    a4=5*(-1+2*af)/256
    a5=(1-2*af)/512
    
    A=np.zeros((nx,nx))
    for i in range(nx):
        A[i,i]=1
    for i in range(nx-1):
        A[i,i+1]=af
        A[i+1,i]=af

    A[0,nx-1]=af
    A[nx-1,0]=af

    cm=np.zeros(nx)
    cm=a5/2*(f[jp5]+f[jm5])+a4/2*(f[jp4]+f[jm4])+a3/2*(f[jp3]+f[jm3])+a2/2*(f[jp2]+f[jm2])+a1/2*(f[jp1]+f[jm1])+a0*(f[jj])
    x =np.zeros((nx,nx+1))
    x[:,0:nx]=A
    x[:,nx]=cm
    
    return np.linalg.solve(A, cm)
In [15]:
Com6TRG3(1,128,0.1,1280)#(dx,nx,dt,nt)

6th order compact difference and 4th order Runge-Kutta

In [16]:
def Com6RG4(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    
    jj=np.arange(nx); 
    jp2=np.r_[jj[0:nx-2]+2,0,1];     jp=np.r_[jj[0:nx-1]+1,0]
    jm2=np.r_[nx-2,nx-1,jj[2:nx]-2]; jm=np.r_[nx-1,jj[1:nx]-1]
 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
  
    udata=np.zeros((nt,nx))
    udata[0,:]=u
    
    cf=c*dt/dx
    for it in range(nt):
        g1=u 
        k1=-cf*cmpct(g1,nx)
        g2=u+0.5*k1
        k2=-cf*cmpct(g2,nx)
        g3=u+0.5*k2
        k3=-cf*cmpct(g3,nx)
        g4=u+    k3
        k4=-cf*cmpct(g4,nx)
        u =u+(k1+2*k2+2*k3+k4)/6
        #u =tenthlowpassfilter(u,nx)

        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    plt.ylim(-0.2,2.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show() 
    
def cmpct(f,nx):
    jj=np.arange(nx)
    jp1=np.r_[jj[0:nx-1]+1,0]
    jm1=np.r_[nx-1,jj[1:nx]-1]
    jp2=np.r_[jj[0:nx-2]+2,0,1]
    jm2=np.r_[nx-2,nx-1,jj[1:nx-1]-1]
    A=np.zeros((nx,nx))
    for i in range(nx):
        A[i,i]=3
    for i in range(nx-1):
        A[i,i+1]=1
        A[i+1,i]=1

    A[0,nx-1]=1
    A[nx-1,0]=1

    cm=np.zeros(nx)
    cm=7/3*(f[jp1]-f[jm1])+1/12*(f[jp2]-f[jm2])

    x =np.zeros((nx,nx+1))
    x[:,0:nx]=A
    x[:,nx]=cm
    #plt.plot(gaussian_elimination(x,nx)[:,nx])
    #plt.show()
    return np.linalg.solve(A, cm)

def tenthlowpassfilter(f,nx):
    jj=np.arange(nx)
    jp1=np.r_[jj[0:nx-1]+1,0]
    jm1=np.r_[nx-1,jj[1:nx]-1]
    jp2=np.r_[jj[0:nx-2]+2,0,1]
    jm2=np.r_[nx-2,nx-1,jj[1:nx-1]-1]
    jp3=np.r_[jj[0:nx-3]+3,0,1,2]
    jm3=np.r_[nx-3,nx-2,nx-1,jj[1:nx-2]-1]
    jp4=np.r_[jj[0:nx-4]+4,0,1,2,3]
    jm4=np.r_[nx-4,nx-3,nx-2,nx-1,jj[1:nx-3]-1]
    jp5=np.r_[jj[0:nx-5]+5,0,1,2,3,4]
    jm5=np.r_[nx-5,nx-4,nx-3,nx-2,nx-1,jj[1:nx-4]-1]

    af=0.495
    a0=(193+126*af)/256
    a1=(105+302*af)/256
    a2=15*(-1+2*af)/64
    a3=45*(1-2*af)/512
    a4=5*(-1+2*af)/256
    a5=(1-2*af)/512
    
    A=np.zeros((nx,nx))
    for i in range(nx):
        A[i,i]=1
    for i in range(nx-1):
        A[i,i+1]=af
        A[i+1,i]=af

    A[0,nx-1]=af
    A[nx-1,0]=af

    cm=np.zeros(nx)
    cm=a5/2*(f[jp5]+f[jm5])+a4/2*(f[jp4]+f[jm4])+a3/2*(f[jp3]+f[jm3])+a2/2*(f[jp2]+f[jm2])+a1/2*(f[jp1]+f[jm1])+a0*(f[jj])
    x =np.zeros((nx,nx+1))
    x[:,0:nx]=A
    x[:,nx]=cm
    
    return np.linalg.solve(A, cm)
In [17]:
Com6RG4(1,128,0.1,1280)#(dx,nx,dt,nt)

FFT and 4th order Runge-Kutta

In [18]:
def rk4fft(dx,nx,dt,nt):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    jj=np.arange(nx)
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
    #kk=2*np.pi/xmax*np.r_[np.arange(nx/2),np.arange(-nx/2+1,0,1)]
    kk=2*mt.pi/xmax*np.r_[np.arange(nx/2),np.arange(-nx/2,0)]
  
    udata=np.zeros((nt,nx))
    udata[0,:]=u
    
    for it in range(nt):
        #clear_output(wait=True)
        kk1=-c*dt*np.real(np.fft.ifft(1j*kk*np.fft.fft(u)))
        kk2=-c*dt*np.real(np.fft.ifft(1j*kk*np.fft.fft(u+0.5*kk1)))
        kk3=-c*dt*np.real(np.fft.ifft(1j*kk*np.fft.fft(u+0.5*kk2)))
        kk4=-c*dt*np.real(np.fft.ifft(1j*kk*np.fft.fft(u+    kk3)))
        u  =u+(kk1+2*kk2+2*kk3+kk4)/6
        
         #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    plt.ylim(-0.2,2.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show() 
In [19]:
%%time
rk4fft(1,128,0.1,1280)
CPU times: user 6.39 s, sys: 141 ms, total: 6.53 s
Wall time: 6.99 s

Implicit sheme

$$\frac{u_j^{k+1}-u_j^k}{\Delta t}=-(1-\alpha)c\frac{u_{j+1}^k-u_{j-1}^k}{2\Delta x}-\alpha c\frac{u_{j+1}^{k+1}-u_{j-1}^{k+1}}{2\Delta x}$$ This equation can be re-written as follows, $$u_j^{k+1}+\frac{\alpha\xi}{2}(u_{j+1}^{k+1}-u_{j-1}^{k+1})=u_j^k-\frac{(1-\alpha)\xi}{2}(u_{j+1}^k-u_{j-1}^k)$$ or in an matrix form, $$\left[ {\begin{array}{*{20}{c}} 1&{{h^ + }}&{}&{}&{ - {h^ + }}\\ { - {h^ + }}&1&{{h^ + }}&{}&{}\\ {}&{}& \ddots &{}&{}\\ {}&{}&{ - {h^ + }}&1&{{h^ + }}\\ {{h^ + }}&{}&{}&{ - {h^ + }}&1 \end{array}} \right]\left[ {\begin{array}{*{20}{c}} {u_1^{k + 1}}\\ {}\\ \vdots \\ {}\\ {u_N^{k + 1}} \end{array}} \right] = \left[ {\begin{array}{*{20}{c}} 1&{{h^0}}&{}&{}&{ - {h^0}}\\ { - {h^0}}&1&{{h^0}}&{}&{}\\ {}&{}& \ddots &{}&{}\\ {}&{}&{ - {h^0}}&1&{{h^0}}\\ {{h^0}}&{}&{}&{ - {h^0}}&1 \end{array}} \right]\left[ {\begin{array}{*{20}{c}} {u_1^k}\\ {}\\ \vdots \\ {}\\ {u_N^k} \end{array}} \right]$$ where $h^0=(1-\alpha)\xi/2$, $h^+=\alpha\xi/2$ and $\xi=c\Delta t/\Delta x$.

In [20]:
def implct(dx,nx,dt,nt,alph):
    #u=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    jj=np.arange(nx); 
    xx=dx*jj; xmax=dx*nx; c=1
    tt=dt*np.arange(nt)
  
    udata=np.zeros((nt,nx))
    udata[0,:]=u
    
    #alph=1
    cf=c*dt/dx
    hm=(1-alph)*cf/2
    hp=alph*cf/2
    A=np.zeros((nx,nx))
    for i in range(nx):
        A[i,i]=1
    for i in range(nx-1):
        A[i,i+1]= hp
        A[i+1,i]=-hp

    A[0,nx-1]=-hp
    A[nx-1,0]= hp

    B=np.zeros((nx,nx))
    for i in range(nx):
        B[i,i]=1
    for i in range(nx-1):
        B[i,i+1]=-hm
        B[i+1,i]= hm

    B[0,nx-1]= hm
    B[nx-1,0]=-hm
    
    for it in range(nt):
        u =np.linalg.solve(A, np.dot(B,u))
        #u =tenthlowpassfilter(u,nx)

        #if it%10==0:
        #    display.clear_output(True)
        #    plt.plot(u)
        #    plt.ylim(-0.2,2.2)
        #    plt.show()
        udata[it,:]=u
        
    #u0=np.r_[np.zeros(int(nx/4)), np.ones(int(nx/2)), np.zeros(int(nx/4))]
    u0=np.exp(-(nx/2*dx-np.arange(nx)*dx)**2/20)
    plt.title('at the last time')
    plt.plot(u0,'k--')
    plt.plot(u)
    plt.show()
    plt.pcolor(xx,tt,udata)
    plt.show() 
In [21]:
%%time
implct(1,128,0.1,1280,1)
CPU times: user 7.41 s, sys: 391 ms, total: 7.8 s
Wall time: 7.72 s