%matplotlib inline

import numpy
from matplotlib import pyplot

numpy.set_printoptions(precision=3)

L = 1.
J = 200
dx = float(L)/float(J-1)
x_grid = numpy.array([j*dx for j in range(J)])

T = 500
N = 1000
dt = float(T)/float(N-1)
t_grid = numpy.array([n*dt for n in range(N)])

D_v = float(10.)/float(100.)
D_u = 0.01 * D_v

k0 = 0.067
f = lambda u, v: dt*(v*(k0 + float(u*u)/float(1. + u*u)) - u)
g = lambda u, v: -f(u,v)
 
sigma_u = float(D_u*dt)/float((2.*dx*dx))
sigma_v = float(D_v*dt)/float((2.*dx*dx))

total_protein = 2.26

no_high = 10
U =  numpy.array([0.1 for i in range(no_high,J)] + [2. for i in range(0,no_high)])
V = numpy.array([float(total_protein-dx*sum(U))/float(J*dx) for i in range(0,J)])

pyplot.ylim((0., 2.1))
pyplot.xlabel('x')
pyplot.ylabel('concentration')
pyplot.plot(x_grid, U)
pyplot.plot(x_grid, V)
pyplot.show()

A_u = numpy.diagflat([-sigma_u for i in range(J-1)], -1) +\
      numpy.diagflat([1.+sigma_u]+[1.+2.*sigma_u for i in range(J-2)]+[1.+sigma_u]) +\
      numpy.diagflat([-sigma_u for i in range(J-1)], 1)
        
B_u = numpy.diagflat([sigma_u for i in range(J-1)], -1) +\
      numpy.diagflat([1.-sigma_u]+[1.-2.*sigma_u for i in range(J-2)]+[1.-sigma_u]) +\
      numpy.diagflat([sigma_u for i in range(J-1)], 1)
        
A_v = numpy.diagflat([-sigma_v for i in range(J-1)], -1) +\
      numpy.diagflat([1.+sigma_v]+[1.+2.*sigma_v for i in range(J-2)]+[1.+sigma_v]) +\
      numpy.diagflat([-sigma_v for i in range(J-1)], 1)
        
B_v = numpy.diagflat([sigma_v for i in range(J-1)], -1) +\
      numpy.diagflat([1.-sigma_v]+[1.-2.*sigma_v for i in range(J-2)]+[1.-sigma_v]) +\
      numpy.diagflat([sigma_v for i in range(J-1)], 1)

def f_vec_ee(U,V):
    return numpy.multiply(dt, numpy.subtract(numpy.multiply(V, 
           numpy.add(k0, numpy.divide(numpy.multiply(U,U), numpy.add(1., numpy.multiply(U,U))))), U))

def f_vec_rk(U, V):
    f_vec = lambda U, V: numpy.subtract(numpy.multiply(V, 
                         numpy.add(k0, numpy.divide(numpy.multiply(U,U), numpy.add(1., numpy.multiply(U,U))))), U)
    k1 = f_vec(U, V)
    k2 = f_vec(U + numpy.multiply(dt/2., k1), V - numpy.multiply(dt/2., k1))
    k3 = f_vec(U + numpy.multiply(dt/2., k2), V - numpy.multiply(dt/2., k2))
    k4 = f_vec(U + numpy.multiply(dt, k3), V - numpy.multiply(dt, k3))
    
    return numpy.multiply(dt/6., k1 + numpy.multiply(2., k2) + numpy.multiply(2., k3) + k4)

U_record_ee = numpy.empty(shape=(N,J))
V_record_ee = numpy.empty(shape=(N,J))

U_record_rk = numpy.empty(shape=(N,J))
V_record_rk = numpy.empty(shape=(N,J))

U_record_ee[0][:] = U[:]
V_record_ee[0][:] = V[:]

U_record_rk[0][:] = U[:]
V_record_rk[0][:] = V[:]

for ti in range(1,N):
    U_record_ee[ti][:] = numpy.linalg.solve(A_u, B_u.dot(U_record_ee[ti-1][:]) +
                                            f_vec_ee(U_record_ee[ti-1][:],V_record_ee[ti-1][:]))
    V_record_ee[ti][:] = numpy.linalg.solve(A_v, B_v.dot(V_record_ee[ti-1][:]) - 
                                            f_vec_ee(U_record_ee[ti-1][:],V_record_ee[ti-1][:]))
    
    U_record_rk[ti][:] = numpy.linalg.solve(A_u, B_u.dot(U_record_rk[ti-1][:]) +
                                            f_vec_rk(U_record_rk[ti-1][:],V_record_rk[ti-1][:]))
    V_record_rk[ti][:] = numpy.linalg.solve(A_v, B_v.dot(V_record_rk[ti-1][:]) - 
                                            f_vec_rk(U_record_rk[ti-1][:],V_record_rk[ti-1][:]))

print 'Explicit Euler', numpy.sum(numpy.multiply(dx, U_record_ee[0]) + numpy.multiply(dx, V_record_ee[0]))
print 'Runge-Kutta 4 ', numpy.sum(numpy.multiply(dx, U_record_ee[0]) + numpy.multiply(dx, V_record_ee[0]))

print 'Explicit Euler %.14f' % numpy.sum(numpy.multiply(dx, U_record_ee[-1]) + numpy.multiply(dx, V_record_ee[-1]))
print 'Runge-Kutta 4  %.14f' % numpy.sum(numpy.multiply(dx, U_record_rk[-1]) + numpy.multiply(dx, V_record_rk[-1]))

pyplot.ylim((0., 2.1))
pyplot.xlabel('x')
pyplot.ylabel('concentration')
pyplot.plot(x_grid, U_record_ee[-1])
pyplot.plot(x_grid, V_record_ee[-1])
pyplot.plot(x_grid, U_record_rk[-1])
pyplot.plot(x_grid, V_record_rk[-1])
pyplot.show()

fig, ax = pyplot.subplots()
pyplot.xlabel('x')
pyplot.ylabel('t')
pyplot.ylim((0., T))
heatmap = ax.pcolormesh(x_grid, t_grid, U_record_ee, vmin=0., vmax=1.2)
colorbar = pyplot.colorbar(heatmap)
colorbar.set_label('concentration U')

fig, ax = pyplot.subplots()
pyplot.xlabel('x')
pyplot.ylabel('t')
pyplot.ylim((0., T))
heatmap = ax.pcolormesh(x_grid, t_grid, U_record_rk, vmin=0., vmax=1.2)
colorbar = pyplot.colorbar(heatmap)
colorbar.set_label('concentration U')