%matplotlib inline
from pylab import *
from numpy import *

def diffusion_init(maxx, maxt, D, nx, SC):
    # initialize variables for solving the diffusion equation
    # on 2D square domain
    #choose time step according to stability condition
    dx = maxx/(nx+1)
    dt = SC*dx**2/(4*D)
    nt = int(maxt/dt)+1
    
    #define array for storing the solution
    U = zeros((nx+2, nx+2))
    
    x = arange(nx+2)*dx
    t = arange(nt)*dt
    return U, dx, dt, x, t

# constants taken from http://en.wikipedia.org/wiki/File:Brusselator_space.gif
#initial field values
U0 = 0.
V0 = 1.
Vmax = 5.0

noise = 2.

#diffusion coefficients
DU = 0.2
DV = 0.02

#constant concentrations
A = 1
B = 3

def RHS(U, V):
    fV = B*U - U**2*V
    fU = A - fV - U
    return fU, fV


from scipy.sparse import dia_matrix, eye, kron
from scipy.sparse.linalg import factorized
def d2matrix(nelem):
    elements = ones((3,nelem))
    elements[1,:] *= -2
    return dia_matrix((elements, [-1,0,1]), shape=(nelem,nelem)).tocsc()

def brusselator_explicit(Uold, Vold, dx, dt, nx, nt, DU, DV):
    alphaU = DU*dt/dx**2
    alphaV = DV*dt/dx**2
    
    # matrix of second differences with Neumann BC in 1D
    d2x = d2matrix(nx+2)
    d2x[0, 1] += 1
    d2x[-1, -2] += 1
    
    # obtain 2D difference operator as Kronecker product of 1D operators
    # and construct the diffusion operators
    mat2D = kron(d2x, eye(nx+2)) + kron(eye(nx+2), d2x)
    M2U = eye((nx+2)**2)+mat2D*alphaU
    M2V = eye((nx+2)**2)+mat2D*alphaV
    for it in range(0,nt-1):
        # U[it, boundary] -= alpha*dx*(NBC[it]+NBC[it+1])
        # explicit diffusion operator
        U = M2U.dot(Uold.ravel()).reshape(nx+2, nx+2)
        V = M2V.dot(Vold.ravel()).reshape(nx+2, nx+2)
        
        # Euler integration of the reaction part
        fU, fV = RHS(U, V)
        U += fU*dt
        V += fV*dt
        
        Uold, Vold = U, V
        
        # store the animation frames
        if it in frames:
            i = where(frames==it)
            Uframes[i] = U
            Vframes[i] = V

maxt = 3000.
maxx = 150.

nx = 200   # number of unknown grid points in spatial direction

U, dx, dt, x, t = diffusion_init(maxx, maxt, max((DU, DV)), nx, SC=.05)
V, dx, dt, x, t = diffusion_init(maxx, maxt, max((DU, DV)), nx, SC=.05)

random.seed(123)
Uold = ones((nx+2, nx+2))*0
Vold = V0 + 1.5*noise*randn(nx+2, nx+2)
Vold[0, Vold[0] > Vmax] = Vmax

frames = arange(int(0.9*len(t)),len(t),4)
Uframes = zeros((len(frames), nx+2, nx+2))
Vframes = zeros((len(frames), nx+2, nx+2))
#print frames

brusselator_explicit(Uold, Vold, dx, dt, nx, len(t), DU, DV)

figure(figsize=(12,12))
subplot(221)
pcolormesh(Uframes[:, 30, :], rasterized=True, vmin=0, vmax=5)
subplot(223)
pcolormesh(Vframes[:, 30, :], rasterized=True, vmin=0, vmax=5)
subplot(222)
pcolormesh(Uframes[2001, :-2, :-2], rasterized=True, vmin=0, vmax=5)
subplot(224)
pcolormesh(Vframes[2000, :-2, :-2], rasterized=True, vmin=0, vmax=5)

fig = figure(figsize=(8,4))
ax1 = fig.add_axes([0, 0, 0.5, 1.0])
ax2 = fig.add_axes([0.5, 0, 0.5, 1.0])
setp(ax2.get_yticklabels(), visible=False)
#ax2 = fig.add_subplot(122)

from matplotlib import animation
from JSAnimation.IPython_display import display_animation
#anim = animation.ArtistAnimation(fig, frames, interval=50, repeat_delay=3000,
#    blit=True)
def animate(data):
    mesh = ax1.pcolormesh(Vframes[data, :-2, :-2], rasterized=True, vmin=0, vmax=5)
    mesh = ax2.pcolormesh(Uframes[data, :-2, :-2], rasterized=True, vmin=0, vmax=5)
    return mesh
anim = animation.FuncAnimation(fig, animate, frames=arange(2000, 2046, 4, dtype=int), interval=100,blit=True)
display_animation(anim, default_mode='loop')

anim.save("anim.gif", writer="imagemagick")