In [1]:
import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import interact  

Backward Euler for unconditional stability

In [2]:
D = 4.
L = 1
M = 50
T = 0.1
N = 1000  # critical value for Euler is 2000 (sigma=0.5)

h = L/M  # spatial spacing
k = T/N  # time step

sigma = D*k/h**2
# matrix for backward Euler
A = np.zeros((M-1,M-1))
I = np.arange(M-1)
A[I,I] = 1 + 2*sigma     # set the main diagonal
A[I[1:],I[:-1]] = -sigma # set the subdiagonal
A[I[:-1],I[1:]] = -sigma # set the superdiagonal

def f(x):
    fvals = np.zeros_like(x)
    fvals[ x>0.5 ] = 1
    return fvals

def l(t):
    return 0*t

def r(t):
    return np.sin(100*t)
    #return 0*t

x = np.linspace(0,L,M+1)

w = f(x)  # approx to u(x,0)
wa = np.zeros((M+1,N+1))  # place to store values entire grid
wa[:,0] = w  

for j in range(N):
    # update w
    newt = k*(j+1)
    b = w[1:-1].copy()
    b[ 0] += sigma*l(newt)
    b[-1] += sigma*r(newt)
    w[1:-1] = np.linalg.solve(A,b)  # linear solve (could use tridiagonal solve for greater efficiency)
    w[0]  = l(newt)
    w[-1] = r(newt)
    wa[:,j+1] = w
    if np.abs(w).max()> 10: break
    
def plotframe(j):
    plt.figure()
    plt.plot(x,wa[:,j])
    plt.xlim(0,L)
    plt.ylim(-2,2)
    plt.plot(1.01,-2+4*j/N,'mo',clip_on=False)
    plt.title('Backward Euler, t ='+str(round(k*j,3)))
    
jmax = j
interact(plotframe,j=(0,jmax));    
    

Crank-Nicolson

Done quickly, not carefully checked.

In [7]:
D = 4.
L = 1
M = 50
T = 0.1
N = 500

h = L/M  # spatial spacing
k = T/N  # time step

sigma = D*k/h**2
# matrix for backward Euler
A = np.zeros((M-1,M-1))
I = np.arange(M-1)
A[I,I] = 2 + 2*sigma     # set the main diagonal
A[I[1:],I[:-1]] = -sigma # set the subdiagonal
A[I[:-1],I[1:]] = -sigma # set the superdiagonal

def f(x):
    fvals = np.zeros_like(x)
    fvals[ x>0.5 ] = 1
    return fvals

def l(t):
    return 0*t

def r(t):
    return 0*t

x = np.linspace(0,L,M+1)

w = f(x)  # approx to u(x,0)
wa = np.zeros((M+1,N+1))  # place to store values entire grid
wa[:,0] = w  

for j in range(N):
    # update w
    oldt = k*j
    newt = k*(j+1)
    b = (2-2*sigma)*w[1:-1] 
    b[:-1] += sigma*w[2:-1]
    b[1:]  += sigma*w[1:-2]
    b[ 0] += sigma*(l(oldt)+l(newt))
    b[-1] += sigma*(r(oldt)+r(newt))
    #display(b.shape)
    w[1:-1] = np.linalg.solve(A,b)
    w[0]  = l(newt)
    w[-1] = r(newt)
    wa[:,j+1] = w
    if np.abs(w).max()> 10: break
    
def plotframe(j):
    plt.figure()
    plt.plot(x,wa[:,j])
    plt.xlim(0,L)
    plt.ylim(-2,2)
    plt.plot(1.01,-2+4*j/N,'mo',clip_on=False)
    plt.title('Backward Euler, t ='+str(round(k*j,3)))
    
jmax = j
interact(plotframe,j=(0,jmax)); 

Using a banded solver in Backward Euler for efficiency

In [26]:
import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import solve_banded  ###########

D = 4.
L = 1.
M = 50 
T = .01
N = 20

h = L/M  # spatial grid spacing
k = T/N  # time step

sigma = D*k/h**2

def f(x): 
    fvals = 0*x
    fvals[ (x>.5) & (x<0.7) ] = 1
    return fvals
def l(t): return 0.
def r(t): return 0.

x = np.linspace(0,L,M+1)

w = f(x)  # set initial values

# 3 by M-1 matrix instead of M-1 by M-1
A = np.zeros((3,M-1))  # just make space for the 3 diagonals
A[0,1:] = -sigma       # subdiagonal
A[1,:  ] = 1.+2.*sigma  # diagonal
A[2,:-1] = -sigma       # superdiagonal

#display(A)

plt.plot(x,w,'k')

for j in range(N):
        # first update the boundary values
        newt = k*(j+1)
        w[ 0] = l(newt)
        w[-1] = r(newt)
        rhs = np.array(w[1:-1])
        rhs[ 0] += sigma*w[ 0]
        rhs[-1] += sigma*w[-1]
        w[1:-1] = solve_banded( (1,1), A, rhs )

        plt.plot(x,w) #p.set_ydata(w)  # update the curve in the plot
plt.title('j='+str(j).zfill(4))
display( w )
array([0.        , 0.00778376, 0.01562752, 0.02358957, 0.03172473,
       0.04008273, 0.04870651, 0.05763069, 0.06688004, 0.0764682 ,
       0.08639641, 0.09665253, 0.10721031, 0.11802888, 0.12905262,
       0.14021134, 0.15142093, 0.1625843 , 0.17359285, 0.18432826,
       0.1946647 , 0.20447135, 0.21361518, 0.22196396, 0.22938928,
       0.23576967, 0.2409935 , 0.24496172, 0.24759021, 0.24881177,
       0.24857758, 0.24685806, 0.24364322, 0.23894235, 0.23278322,
       0.22521066, 0.21628481, 0.20607895, 0.19467712, 0.18217158,
       0.16866033, 0.15424473, 0.13902726, 0.12310973, 0.10659172,
       0.08956959, 0.07213582, 0.05437883, 0.03638328, 0.01823062,
       0.        ])
In [ ]: