Step 1: Import packages

In [1]:
%matplotlib inline
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt

Step 2: Problem description

Here you will set up the problem for $$ u_t + c u_x = 0$$ with periodic BC on the interval [0,1]

In [36]:
a = 1.0
T = 4.0 / a # end time

Step 3: Set up the grid

dx will be the grid spacing in the $x$-direction

x will be the grid coordinates

xx will be really fine grid coordinates

In [44]:
nx = 90
k = 10
x = np.linspace(0, 2*np.pi, nx, endpoint=False)
dx = x[1] - x[0]
xx = np.linspace(0, 2*np.pi, 1000, endpoint=False)

Now define an initial condition

In [45]:
def f(x):
    u = np.sin(k * x)
    return u
In [46]:
plt.plot(xx, f(xx), lw=3, clip_on=False)
Out[46]:
[<matplotlib.lines.Line2D at 0x10e3fde48>]

Step 4: Setting dt

Now we need a time step. Let $$ \Delta t = \Delta x \frac{\lambda}{a}$$

So we need a parameter $\lambda$

What happens when $\lambda>1.0$?

When the `method` changes to FTCS, what is the impact of $\lambda$?

In [47]:
lmbda = 0.6
dt = dx * lmbda / a
nt = int(T/dt)
print('T = %g' % T)
print('tsteps = %d' % nt)
print('    dx = %g' % dx)
print('    dt = %g' % dt)
print('lambda = %g' % lmbda)

T = 4
tsteps = 95
    dx = 0.0698132
    dt = 0.0418879
lambda = 0.6

Now make an index list, called $J$, so that we can access $J+1$ and $J-1$ easily

In [48]:
J = np.arange(0, nx)  # all vertices
Jm1 = np.roll(J, 1)
Jp1 = np.roll(J, -1)

Step 5: Run and Animate

For ipython notebooks be sure to use clear_output. Alternatively, animation from matplotlib may be useful.

In [49]:
method = 'FTBS'
plotit = True

uFTBS = f(x)
uLW = f(x)
if plotit:
    fig = plt.figure(figsize=(10,10))
    ax = fig.add_subplot(111)
    ax.set_title('u vs x')

for n in range(0, nt):
    
    uFTBS[J] = uFTBS[J] - lmbda * (uFTBS[J] - uFTBS[Jm1])  # FTBS
    
    uLW[J] = uLW[J] - lmbda * (1.0 / 2.0) * (uLW[Jp1] - uLW[Jm1]) \
           + (lmbda**2 / 2.0) * (uLW[Jp1] - 2 * uLW[J] + uLW[Jm1])
    
            
    uex = f((xx - c * (n+1) * dt) % 1.0)

if plotit:
    ax.plot(x, uFTBS, '-', lw=3, clip_on=False, label=method)
    ax.plot(x, uLW, '-', lw=3, clip_on=False, label=method)
    ax.plot(xx, uex, 'k-', lw=3, clip_on=False, label='exact')
    ax.legend(frameon=False)
    plt.show()
In [ ]: