一维Burgers方程数值解法

# -*- coding: utf-8 -*-
"""
Created on Tue Jan 20 14:32:23 2015
1D burges equation
@author: myjiayan
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation

def u_initial():
    first = np.ones(40)
    second = np.zeros(41)
    result = np.array(list(first)+list(second))
    return result

def computeF(u):
    F = 0.5*u**2
    return F


def maccormack(u,nt,dt,dx):
    un = np.zeros((nt,len(u)))
    un[0] = u.copy()
    ustar = u.copy()
    
    for t in xrange(1,nt):
        F = computeF(u)
        ustar[:-1] = u[:-1] - (F[1:]-F[:-1])*dt/dx
        ustar[-1] = 0.0
        Fstar = computeF(ustar)
        un[t,1:] = 0.5*(u[1:]+ustar[1:]-(Fstar[1:]-Fstar[:-1])*dt/dx)
        un[t,0] = 1.0
        u = un[t].copy()
    
    return un

if __name__ == '__main__':
    nx = 81
    nt = 70
    dx = 4.0/(nx-1)

    def animate(data):
        x = np.linspace(0,4,nx)
        y = data
        line.set_data(x,y)
        return line,

    u = u_initial()
    sigma = .5
    dt = sigma*dx

    un = maccormack(u,nt,dt,dx)
    fig = plt.figure();
    ax = plt.axes(xlim=(0,4),ylim=(-.5,2));
    line, = ax.plot([],[],lw=2);
    anim = animation.FuncAnimation(fig, animate, frames=un, interval=50)
    plt.show()