In [ ]:
%matplotlib inline

In [ ]:
import numpy as np
from numpy import newaxis
import scipy.sparse as sps
from scipy.sparse.linalg import spsolve
import matplotlib.pyplot as plt

In [ ]:
from fem import *
eval_phi1d = eval_lagrange_d0

In [ ]:
from poly import eval_P

In [ ]:
speed = -2.0

order = 4
semh = SEMhat(order)

N = 30
n_dofs = (order+1)*N

In [ ]:
L = 1.0

vertices  = np.linspace(0, L, N+1)
EtoV      = np.zeros((N, 2), dtype=np.int)
EtoV[:,0] = np.arange(N)
EtoV[:,1] = np.arange(N)+1

topo  = Interval()
xq = topo.ref_to_phys(vertices[EtoV], semh.xgll)
jacb_det = topo.calc_jacb(vertices[EtoV])[0]

EtoV[-1,-1] = EtoV[0,0]
jacb_det

In [ ]:
# Make elem to dof map
EtoD = np.arange(N*(order+1))
EtoD = EtoD.reshape((N, -1))

dof_phys = xq.ravel()

# Averaging operator
rows = EtoD[:,[0,-1]].ravel()
cols = EtoV.ravel()
vals = np.ones_like(cols)

FtoD = sps.coo_matrix((vals, (rows, cols)),
                       shape=(n_dofs,N)).tocsr()
AVG = FtoD.dot(FtoD.T)/2.0

# Extract face DOFS
vals = np.ones(len(rows))
FD = sps.coo_matrix((vals, (rows, rows))).tocsr()
# Set face signs
vals[::2] = -1
SD = sps.coo_matrix((vals, (rows, rows))).tocsr()

# Jump operator
JUMP = FtoD.dot(SD.dot(FtoD).T)

In [ ]:
# Build Advection operator
C = sps.kron(sps.eye(N), semh.Ch).tocsr()

# Build full elemental mass matrix
x, w = topo.get_quadrature(order+1)
P = eval_phi1d(semh.xgll, x).T
G = sps.dia_matrix((w, 0), shape=(len(x), len(x)))
Bf = P.T.dot(G.dot(P))*jacb_det
Bfinv = np.linalg.inv(Bf)

# Using trick from book
V = eval_P(order, semh.xgll).T
Minv = V.dot(V.T)/jacb_det
Binv = sps.kron(sps.eye(N), Minv).tocsr()

print np.max(np.abs(Minv-Bfinv))
print np.max(np.abs(Minv.dot(Bf)-np.eye(Bf.shape[0])))
print np.max(np.abs(Bfinv.dot(Bf)-np.eye(Bf.shape[0])))

In [ ]:
# Problem setup
a  = 1000.*(dof_phys-.5)**2
u0 = np.exp(-a)

u0 = np.zeros_like(dof_phys)
u0[(dof_phys>=.4) & (dof_phys<=.6)] = 1.0

In [ ]:
CFL = 0.75
xa = np.unique(dof_phys)
dx = np.min(xa[1:]-xa[:-1])
dt = CFL*dx/np.abs(speed)
dt *= 0.25 # safety factor

T = 10.0
nt = int(np.ceil(T/dt))
assert T/nt<=dt
dt = T/nt

nt, "%.2e"%dt, nt*dt, np.abs(speed)*dt/dx

In [ ]:
# Time stepping
def apply_F(u):
    f = speed*u
    c = np.abs(speed)
    flux = AVG.dot(f)+c/2.0*JUMP.dot(u)
    return Binv.dot(-C.dot(f)+SD.dot(FD.dot(f)-flux))

In [ ]:
# Integrate with RK4
ue = u0
u  = u0.copy()

for k in range(nt):
    
    k1 = apply_F(u)
    k2 = apply_F(u+(dt/2.0)*k1)
    k3 = apply_F(u+(dt/2.0)*k2)
    k4 = apply_F(u+dt*k3)
    
    u = u+(dt/6.0)*(k1+2*k2+2*k3+k4)

np.max(np.abs(u-ue))

In [ ]:
plt.plot(dof_phys[:-1], u[:-1])
plt.plot(dof_phys[:-1], u0[:-1], 'g--')

In [ ]: