%matplotlib inline

import numpy
from scipy import integrate
from matplotlib import pyplot
numpy.set_printoptions(precision=3)

L = 1.
J = 100
dX = float(L)/float(J)
X_grid = numpy.array([float(dX)/2. + j*dX for j in range(J)])
x_grid = numpy.array([j*dX for j in range(J+1)])

T = 200
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_vec = lambda U, V: 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))

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)])
g = numpy.array([1. for j in range(J)])
a = numpy.array([0. for j in range(J)])
s = lambda U: numpy.array([0.001*X_grid[j]*U[j] for j in range(J)])

pyplot.ylim((0., 2.1))
pyplot.xlabel('X'); pyplot.ylabel('concentration')
pyplot.plot(X_grid, U)
pyplot.plot(X_grid, V)
pyplot.show()

pyplot.plot(X_grid, g)
pyplot.xlabel('X'); pyplot.ylabel('g')
pyplot.show()

pyplot.plot(X_grid, a)
pyplot.plot(X_grid, s(U))
pyplot.xlabel('X'); pyplot.ylabel('g')
pyplot.ylim((-.0001, 0.0021))
pyplot.show()

# syntax to get one element in an array: http://stackoverflow.com/a/7332880/3078529
a_left = lambda a: numpy.concatenate((a[1:J], a[J-1:J]))
a_right = lambda a: numpy.concatenate((a[0:1], a[0:J-1]))

f_vec_u = lambda U, V, g, a: numpy.subtract(f_vec(U, V),
                                         numpy.multiply(numpy.subtract(a_left(a), a_right(a)),
                                                        numpy.multiply(float(dt)/(float(2.*dX)),
                                                                       numpy.divide(U,g))))

f_vec_v = lambda U, V, g, a: numpy.subtract(numpy.multiply(-1., f_vec(U, V)),
                                         numpy.multiply(numpy.subtract(a_left(a), a_right(a)),
                                                        numpy.multiply(float(dt)/(float(2.*dX)),
                                                                       numpy.divide(V,g))))

sigma_u_func = lambda g: numpy.divide(float(D_u*dt)/float(2.*dX*dX),
                                      numpy.multiply(g, g))
sigma_v_func = lambda g: numpy.divide(float(D_v*dt)/float(2.*dX*dX),
                                      numpy.multiply(g, g))

# syntax to get one element in an array: http://stackoverflow.com/a/7332880/3078529
g_left = lambda g: numpy.concatenate((g[1:J], g[J-1:J]))
g_right = lambda g: numpy.concatenate((g[0:1], g[0:J-1]))

rho_u_func = lambda g, a: numpy.multiply(float(-dt)/float(4.*dX),
                                      numpy.add(numpy.divide(a, g),
                                                numpy.multiply(numpy.subtract(g_left(g), g_right(g)),
                                                               numpy.divide(float(D_u)/(2.*dX),
                                                                            numpy.power(g, 3)))))
rho_v_func = lambda g, a: numpy.multiply(float(-dt)/float(4.*dX),
                                      numpy.add(numpy.divide(a, g),
                                                numpy.multiply(numpy.subtract(g_left(g), g_right(g)),
                                                               numpy.divide(float(D_v)/(2.*dX),
                                                                            numpy.power(g, 3)))))

g_rhs = lambda t, g, a: numpy.divide(numpy.subtract(a_left(a), a_right(a)),
                                     numpy.array([dX]+[2.*dX for j in range(J-2)]+[dX]))

g_stepper = integrate.ode(g_rhs)
g_stepper = g_stepper.set_integrator('dopri5', nsteps=10, max_step=dt)
g_stepper = g_stepper.set_initial_value(g, 0.)
g_stepper = g_stepper.set_f_params(a)

a_rhs = lambda t, a, s: s

a_stepper = integrate.ode(a_rhs)
a_stepper = a_stepper.set_integrator('dopri5', nsteps=10, max_step=dt)
a_stepper = a_stepper.set_initial_value(a, 0.)
a_stepper = a_stepper.set_f_params(s(U))

x_grid_rhs = lambda t, x_grid, a: numpy.concatenate(([0.],
                                                     numpy.divide(numpy.add(a[0:J-1], a[1:J]), 2.),
                                                     a[J-1:J]))

x_stepper = integrate.ode(x_grid_rhs)
x_stepper = x_stepper.set_integrator('dopri5', nsteps=10, max_step=dt)
x_stepper = x_stepper.set_initial_value(x_grid, 0.)
x_stepper = x_stepper.set_f_params(a)

U_record = []
V_record = []
g_record = []
a_record = []
x_record = []

U_record.append(U)
V_record.append(V)
g_record.append(g)
a_record.append(a)
x_record.append(x_grid)

for ti in range(1,N):
    sigma_u = sigma_u_func(g)
    sigma_v = sigma_v_func(g)
    rho_u = rho_u_func(g, a)
    rho_v = rho_v_func(g, a)
    
    A_u = numpy.diagflat([-sigma_u[j]+rho_u[j] for j in range(1,J)], -1) +\
          numpy.diagflat([1.+sigma_u[0]+rho_u[0]]+[1.+2.*sigma_u[j] for j in range(1,J-1)]+[1.+sigma_u[J-1]-rho_u[J-1]]) +\
          numpy.diagflat([-(sigma_u[j]+rho_u[j]) for j in range(0,J-1)], 1)
        
    B_u = numpy.diagflat([sigma_u[j]-rho_u[j] for j in range(1,J)], -1) +\
          numpy.diagflat([1.-sigma_u[0]-rho_u[0]]+[1.-2.*sigma_u[j] for j in range(1,J-1)]+[1.-sigma_u[J-1]+rho_u[J-1]]) +\
          numpy.diagflat([sigma_u[j]+rho_u[j] for j in range(0,J-1)], 1)
        
    A_v = numpy.diagflat([-sigma_v[j]+rho_v[j] for j in range(1,J)], -1) +\
          numpy.diagflat([1.+sigma_v[0]+rho_v[0]]+[1.+2.*sigma_v[j] for j in range(1,J-1)]+[1.+sigma_v[J-1]-rho_v[J-1]]) +\
          numpy.diagflat([-(sigma_v[j]+rho_v[j]) for j in range(0,J-1)], 1)
        
    B_v = numpy.diagflat([sigma_v[j]-rho_v[j] for j in range(1,J)], -1) +\
          numpy.diagflat([1.-sigma_v[0]-rho_v[0]]+[1.-2.*sigma_v[j] for j in range(1,J-1)]+[1.-sigma_v[J-1]+rho_v[J-1]]) +\
          numpy.diagflat([sigma_v[j]+rho_v[j] for j in range(0,J-1)], 1)
            
    U_new = numpy.linalg.solve(A_u, B_u.dot(U) + f_vec_u(U, V, g, a))
    V_new = numpy.linalg.solve(A_v, B_v.dot(V) + f_vec_v(U, V, g, a))
    
    while a_stepper.successful() and a_stepper.t + dt < ti*dt:
        a_stepper.integrate(a_stepper.t + dt)
    while g_stepper.successful() and g_stepper.t + dt < ti*dt:
        g_stepper.integrate(g_stepper.t + dt)
    while x_stepper.successful() and x_stepper.t + dt < ti*dt:
        x_stepper.integrate(x_stepper.t + dt)
        
    g_stepper = g_stepper.set_f_params(a)
    x_stepper = x_stepper.set_f_params(a)
    a_stepper = a_stepper.set_f_params(s(U))
    
    U = U_new
    V = V_new
    
    # these are the correct "y" values to save for the current time step since
    # we integrate only up to t < ti*dt
    g = g_stepper.y
    a = a_stepper.y
    x_grid = x_stepper.y
    
    U_record.append(U)
    V_record.append(V)
    g_record.append(g)
    a_record.append(a)
    x_record.append(x_grid)

sum(numpy.multiply(numpy.diff(x_record[0]),U_record[0]) + numpy.multiply(numpy.diff(x_record[0]),V_record[0]))

sum(numpy.multiply(numpy.diff(x_record[-1]),U_record[-1]) + numpy.multiply(numpy.diff(x_record[-1]),V_record[-1]))

U_record = numpy.array(U_record)
V_record = numpy.array(V_record)

fig, ax = pyplot.subplots()
pyplot.xlabel('X'); pyplot.ylabel('t')
heatmap = ax.pcolormesh(x_record[0], t_grid, U_record, vmin=0., vmax=1.2)
colorbar = pyplot.colorbar(heatmap)
colorbar.set_label('concentration U')

a_record = numpy.array(a_record)

fig, ax = pyplot.subplots()
pyplot.xlabel('X'); pyplot.ylabel('t')
heatmap = ax.pcolormesh(X_grid, t_grid, a_record)
colorbar = pyplot.colorbar(heatmap)
colorbar.set_label('local growth velocity a(X,t)')

fig, ax = pyplot.subplots()
pyplot.xlabel('X'); pyplot.ylabel('t')
s_record = numpy.empty(shape=(N, J))
for ti in range(N):
    s_record[ti] = s(U_record[ti])
heatmap = ax.pcolormesh(X_grid, t_grid, s_record)
colorbar = pyplot.colorbar(heatmap)
colorbar.set_label('local growth acceleration s')