update advec dir

Author: sveinlin

src/advec/advec1D.py

@@ -0,0 +1,367 @@
+import numpy as np
+import matplotlib.pyplot as plt
+
+def solver_FECS(I, U0, v, L, dt, C, T, user_action=None):
+    Nt = int(round(T/float(dt)))
+    t = np.linspace(0, Nt*dt, Nt+1)   # Mesh points in time
+    dx = v*dt/C
+    Nx = int(round(L/dx))
+    x = np.linspace(0, L, Nx+1)       # Mesh points in space
+    # Make sure dx and dt are compatible with x and t
+    dx = x[1] - x[0]
+    dt = t[1] - t[0]
+    C = v*dt/dx
+
+    u   = np.zeros(Nx+1)
+    u_n = np.zeros(Nx+1)
+
+    # Set initial condition u(x,0) = I(x)
+    for i in range(0, Nx+1):
+        u_n[i] = I(x[i])
+
+    if user_action is not None:
+        user_action(u_n, x, t, 0)
+
+    for n in range(0, Nt):
+        # Compute u at inner mesh points
+        for i in range(1, Nx):
+            u[i] = u_n[i] - 0.5*C*(u_n[i+1] - u_n[i-1])
+
+        # Insert boundary condition
+        u[0] = U0
+
+        if user_action is not None:
+            user_action(u, x, t, n+1)
+
+        # Switch variables before next step
+        u_n, u = u, u_n
+
+
+def solver(I, U0, v, L, dt, C, T, user_action=None,
+           scheme='FE', periodic_bc=True):
+    Nt = int(round(T/float(dt)))
+    t = np.linspace(0, Nt*dt, Nt+1)   # Mesh points in time
+    dx = v*dt/C
+    Nx = int(round(L/dx))
+    x = np.linspace(0, L, Nx+1)       # Mesh points in space
+    # Make sure dx and dt are compatible with x and t
+    dx = x[1] - x[0]
+    dt = t[1] - t[0]
+    C = v*dt/dx
+    print 'dt=%g, dx=%g, Nx=%d, C=%g' % (dt, dx, Nx, C)
+
+    u   = np.zeros(Nx+1)
+    u_n = np.zeros(Nx+1)
+    u_nm1 = np.zeros(Nx+1)
+    integral = np.zeros(Nt+1)
+
+    # Set initial condition u(x,0) = I(x)
+    for i in range(0, Nx+1):
+        u_n[i] = I(x[i])
+
+    # Insert boundary condition
+    u[0] = U0
+
+    # Compute the integral under the curve
+    integral[0] = dx*(0.5*u_n[0] + 0.5*u_n[Nx] + np.sum(u_n[1:-1]))
+
+    if user_action is not None:
+        user_action(u_n, x, t, 0)
+
+    for n in range(0, Nt):
+        if scheme == 'FE':
+            if periodic_bc:
+                i = 0
+                u[i] = u_n[i] - 0.5*C*(u_n[i+1] - u_n[Nx])
+                u[Nx] = u[0]
+                #u[i] = u_n[i] - 0.5*C*(u_n[1] - u_n[Nx])
+            for i in range(1, Nx):
+                u[i] = u_n[i] - 0.5*C*(u_n[i+1] - u_n[i-1])
+        elif scheme == 'LF':
+            if n == 0:
+                # Use upwind for first step
+                if periodic_bc:
+                    i = 0
+                    #u[i] = u_n[i] - C*(u_n[i] - u_n[Nx-1])
+                    u_n[i] = u_n[Nx]
+                for i in range(1, Nx+1):
+                    u[i] = u_n[i] - C*(u_n[i] - u_n[i-1])
+            else:
+                if periodic_bc:
+                    i = 0
+                    # Must have this,
+                    u[i] = u_nm1[i] - C*(u_n[i+1] - u_n[Nx-1])
+                    # not this:
+                    #u_n[i] = u_n[Nx]
+                for i in range(1, Nx):
+                    u[i] = u_nm1[i] - C*(u_n[i+1] - u_n[i-1])
+                if periodic_bc:
+                    u[Nx] = u[0]
+        elif scheme == 'UP':
+            if periodic_bc:
+                u_n[0] = u_n[Nx]
+            for i in range(1, Nx+1):
+                u[i] = u_n[i] - C*(u_n[i] - u_n[i-1])
+        elif scheme == 'LW':
+            if periodic_bc:
+                i = 0
+                # Must have this,
+                u[i] = u_n[i] - 0.5*C*(u_n[i+1] - u_n[Nx-1]) + \
+                       0.5*C*(u_n[i+1] - 2*u_n[i] + u_n[Nx-1])
+                # not this:
+                #u_n[i] = u_n[Nx]
+            for i in range(1, Nx):
+                u[i] = u_n[i] - 0.5*C*(u_n[i+1] - u_n[i-1]) + \
+                       0.5*C*(u_n[i+1] - 2*u_n[i] + u_n[i-1])
+            if periodic_bc:
+                u[Nx] = u[0]
+        else:
+            raise ValueError('scheme="%s" not implemented' % scheme)
+
+        if not periodic_bc:
+            # Insert boundary condition
+            u[0] = U0
+
+        # Compute the integral under the curve
+        integral[n+1] = dx*(0.5*u[0] + 0.5*u[Nx] + np.sum(u[1:-1]))
+
+        if user_action is not None:
+            user_action(u, x, t, n+1)
+
+        # Switch variables before next step
+        u_nm1, u_n, u = u_n, u, u_nm1
+        print 'I:', integral[n+1]
+    return integral
+
+def run_FECS(case):
+    """Special function for the FECS case."""
+    if case == 'gaussian':
+        def I(x):
+            return np.exp(-0.5*((x-L/10)/sigma)**2)
+    elif case == 'cosinehat':
+        def I(x):
+            return np.cos(np.pi*5/L*(x - L/10)) if x < L/5 else 0
+
+    L = 1.0
+    sigma = 0.02
+    legends = []
+
+    def plot(u, x, t, n):
+        """Animate and plot every m steps in the same figure."""
+        plt.figure(1)
+        if n == 0:
+            lines = plot(x, u)
+        else:
+            lines[0].set_ydata(u)
+            plt.draw()
+            #plt.savefig()
+        plt.figure(2)
+        m = 40
+        if n % m != 0:
+            return
+        print 't=%g, n=%d, u in [%g, %g] w/%d points' % \
+              (t[n], n, u.min(), u.max(), x.size)
+        if np.abs(u).max() > 3:  # Instability?
+            return
+        plt.plot(x, u)
+        legends.append('t=%g' % t[n])
+        if n > 0:
+            plt.hold('on')
+
+    plt.ion()
+    U0 = 0
+    dt = 0.001
+    C = 1
+    T = 1
+    solver(I=I, U0=U0, v=1.0, L=L, dt=dt, C=C, T=T,
+           user_action=plot)
+    plt.legend(legends, loc='lower left')
+    plt.savefig('tmp.png'); plt.savefig('tmp.pdf')
+    plt.axis([0, L, -0.75, 1.1])
+    plt.show()
+
+def run(scheme='UP', case='gaussian', C=1, dt=0.01):
+    """General admin routine for explicit and implicit solvers."""
+
+    if case == 'gaussian':
+        def I(x):
+            return np.exp(-0.5*((x-L/10)/sigma)**2)
+    elif case == 'cosinehat':
+        def I(x):
+            return np.cos(np.pi*5/L*(x - L/10)) \
+                   if 0 < x < L/5 else 0
+
+    L = 1.0
+    sigma = 0.02
+    global lines  # needs to be saved between calls to plot
+
+    def plot(u, x, t, n):
+        """Plot t=0 and t=0.6 in the same figure."""
+        plt.figure(1)
+        global lines
+        if n == 0:
+            lines = plt.plot(x, u)
+            plt.axis([x[0], x[-1], -0.5, 1.5])
+            plt.xlabel('x'); plt.ylabel('u')
+            plt.axes().set_aspect(0.15)
+            plt.savefig('tmp_%04d.png' % n)
+            plt.savefig('tmp_%04d.pdf' % n)
+        else:
+            lines[0].set_ydata(u)
+            plt.axis([x[0], x[-1], -0.5, 1.5])
+            plt.title('C=%g, dt=%g, dx=%g' %
+                      (C, t[1]-t[0], x[1]-x[0]))
+            plt.legend(['t=%.3f' % t[n]])
+            plt.xlabel('x'); plt.ylabel('u')
+            plt.draw()
+            plt.savefig('tmp_%04d.png' % n)
+        plt.figure(2)
+        eps = 1E-14
+        if abs(t[n] - 0.6) > eps and abs(t[n] - 0) > eps:
+            return
+        print 't=%g, n=%d, u in [%g, %g] w/%d points' % \
+              (t[n], n, u.min(), u.max(), x.size)
+        if np.abs(u).max() > 3:  # Instability?
+            return
+        plt.plot(x, u)
+        plt.hold('on')
+        plt.draw()
+        if n > 0:
+            y = [I(x_-v*t[n]) for x_ in x]
+            plt.plot(x, y, 'k--')
+            if abs(t[n] - 0.6) < eps:
+                filename = ('tmp_%s_dt%s_C%s' % \
+                            (scheme, t[1]-t[0], C)).replace('.', '')
+                np.savez(filename, x=x, u=u, u_e=y)
+
+    plt.ion()
+    U0 = 0
+    T = 0.7
+    v = 1
+    # Define video formats and libraries
+    codecs = dict(flv='flv', mp4='libx264', webm='libvpx',
+                  ogg='libtheora')
+    # Remove video files
+    import glob, os
+    for name in glob.glob('tmp_*.png'):
+        os.remove(name)
+    for ext in codecs:
+        name = 'movie.%s' % ext
+        if os.path.isfile(name):
+            os.remove(name)
+
+    if scheme == 'CN':
+        integral = solver_theta(
+            I, v, L, dt, C, T, user_action=plot, FE=False)
+    elif scheme == 'BE':
+        integral = solver_theta(
+            I, v, L, dt, C, T, theta=1, user_action=plot)
+    else:
+        integral = solver(
+            I=I, U0=U0, v=v, L=L, dt=dt, C=C, T=T,
+            scheme=scheme, user_action=plot)
+    # Finish figure(2)
+    plt.figure(2)
+    plt.axis([0, L, -0.5, 1.1])
+    plt.xlabel('$x$');  plt.ylabel('$u$')
+    plt.axes().set_aspect(0.5)  # no effect
+    plt.savefig('tmp1.png'); plt.savefig('tmp1.pdf')
+    plt.show()
+    # Make videos from figure(1) animation files
+    for codec in codecs:
+        cmd = 'ffmpeg -i tmp_%%04d.png -r 25 -vcodec %s movie.%s' % \
+              (codecs[codec], codec)
+        os.system(cmd)
+    print 'Integral of u:', integral.max(), integral.min()
+
+def solver_theta(I, v, L, dt, C, T, theta=0.5, user_action=None, FE=False):
+    """
+    Full solver for the model problem using the theta-rule
+    difference approximation in time (no restriction on F,
+    i.e., the time step when theta >= 0.5).
+    Vectorized implementation and sparse (tridiagonal)
+    coefficient matrix.
+    """
+    import time;  t0 = time.clock()  # for measuring the CPU time
+    Nt = int(round(T/float(dt)))
+    t = np.linspace(0, Nt*dt, Nt+1)   # Mesh points in time
+    dx = v*dt/C
+    Nx = int(round(L/dx))
+    x = np.linspace(0, L, Nx+1)       # Mesh points in space
+    # Make sure dx and dt are compatible with x and t
+    dx = x[1] - x[0]
+    dt = t[1] - t[0]
+    C = v*dt/dx
+    print 'dt=%g, dx=%g, Nx=%d, C=%g' % (dt, dx, Nx, C)
+
+    u   = np.zeros(Nx+1)
+    u_n = np.zeros(Nx+1)
+    u_nm1 = np.zeros(Nx+1)
+    integral = np.zeros(Nt+1)
+
+    # Set initial condition u(x,0) = I(x)
+    for i in range(0, Nx+1):
+        u_n[i] = I(x[i])
+
+    # Compute the integral under the curve
+    integral[0] = dx*(0.5*u_n[0] + 0.5*u_n[Nx] + np.sum(u_n[1:-1]))
+
+    if user_action is not None:
+        user_action(u_n, x, t, 0)
+
+    # Representation of sparse matrix and right-hand side
+    diagonal = np.zeros(Nx+1)
+    lower    = np.zeros(Nx)
+    upper    = np.zeros(Nx)
+    b        = np.zeros(Nx+1)
+
+    # Precompute sparse matrix (scipy format)
+    diagonal[:] = 1
+    lower[:] = -0.5*theta*C
+    upper[:] =  0.5*theta*C
+    if FE:
+        diagonal[:] += 4./6
+        lower[:] += 1./6
+        upper[:] += 1./6
+    # Insert boundary conditions
+    upper[0] = 0
+    lower[-1] = 0
+
+    diags = [0, -1, 1]
+    import scipy.sparse
+    import scipy.sparse.linalg
+    A = scipy.sparse.diags(
+        diagonals=[diagonal, lower, upper],
+        offsets=[0, -1, 1], shape=(Nx+1, Nx+1),
+        format='csr')
+    #print A.todense()
+
+    # Time loop
+    for n in range(0, Nt):
+        b[1:-1] = u_n[1:-1] + 0.5*(1-theta)*C*(u_n[:-2] - u_n[2:])
+        if FE:
+            b[1:-1] += 1./6*u_n[:-2] + 1./6*u_n[:-2] + 4./6*u_n[1:-1]
+        b[0] = u_n[Nx]; b[-1] = u_n[0]  # boundary conditions
+        b[0] = 0; b[-1] = 0  # boundary conditions
+        u[:] = scipy.sparse.linalg.spsolve(A, b)
+
+        if user_action is not None:
+            user_action(u, x, t, n+1)
+
+        # Compute the integral under the curve
+        integral[n+1] = dx*(0.5*u[0] + 0.5*u[Nx] + np.sum(u[1:-1]))
+
+        # Update u_n before next step
+        u_n, u = u, u_n
+
+    t1 = time.clock()
+    return integral
+
+
+if __name__ == '__main__':
+    #run(scheme='LF', case='gaussian', C=1)
+    #run(scheme='UP', case='gaussian', C=0.8, dt=0.01)
+    #run(scheme='LF', case='gaussian', C=0.8, dt=0.001)
+    #run(scheme='LF', case='cosinehat', C=0.8, dt=0.01)
+    #run(scheme='CN', case='gaussian', C=1, dt=0.01)
+    run(scheme='LW', case='gaussian', C=1, dt=0.01)