diffusion with random source/sink injections ∂ θ(x, t) / ∂ t = ∇ 2 θ(x, t) + S(x,t)

randomInjection.py

from matplotlib import pyplot as plt
from matplotlib import colors
import numpy as np
from tomIntegrate import randr
from time import sleep

NX = NY = 64
figsize=(16, 16)
ptake = 0.001
pgive = 0.0004
nsteps = 512
takeRange = (2, 10)
giveRange = (20, 1000)
D = 14.0  # diffusivity
norm = colors.normalize(vmin=-.11, vmax=3)
#norm = None

def poisson2d((NX, NY)):
    '''Returns a dense square coefficient matrix for the 2D Poisson equation.
    Fails for NX =/= NY .
    '''
    N = NX * NY
    main = np.eye(N) * -4
    oneup = np.hstack((
    np.zeros((NX * NY, 1)),
    np.vstack((
    np.eye(NX * NY - 1),
    np.zeros((1, NX * NY - 1))
    ))
    ))
    twoup = np.hstack((
    np.zeros((NX * NY, NX)),
    np.vstack((
    np.eye(NX * NY - NX),
    np.zeros((NX, NX * NY - NX))
    ))
    ))
    return main + oneup + twoup + oneup.T + twoup.T

A = poisson2d((NX, NY))

def S():
    S_ = np.zeros((NX, NY))
    for i in range(NX):
        for j in range(NY):
            t = randr(0, 1)
            if t < ptake:
                S_[i,j] = -randr(*takeRange)
            elif t > 1-pgive:
                S_[i,j] = randr(*giveRange)

    return np.array(S_)



boundaries = []
for x in range(NX):
    boundaries.append((x, 0))
    boundaries.append((x, NY-1))
for y in range(NY):
    boundaries.append((0, y))
    boundaries.append((NX-1, y))
mask = np.ones((NX, NY))
for x, y in boundaries:
    mask[x, y] = 0  # enforce a Dirichlet BC because it's easy

def applyBCs(T):
    T = T.reshape((NX, NY))
    #T[NX/2, NY/4] = 0 # and a sink
    #T[NX/2, 3*NY/4] = 1 # and a source
    return (T * np.logical_and(T, mask)).reshape((NX * NY,))

def renormalize(T):
    target = NX * NY
    current = sum(T)
    return T * target / current

T = np.zeros((NX * NY,))
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(1, 1, 1)
plt.show(block=False)
dt = 0.01
for step in range(nsteps):
    print "step", step
    T = applyBCs(T)
    update = np.dot(A, T) * D + S().reshape((NX * NY,))
    #T += update / sum(update) #* NX * NY
    T += update * dt
    #T = renormalize(T)
    #print min(T), max(T)
    T = T.reshape((NX * NY,))
    interpolation = 'none'
    #interpolation = None
    imgplot = ax.imshow(T.reshape((NX, NY)), interpolation=interpolation, norm=norm)
    #imgplot.set_cmap('bone')
    imgplot.set_cmap('hot')
    #imgplot.set_cmap('Blues')
    #print sum(T) / len(T)
    #ax.imshow(S(), interpolation=interpolation)
    plt.draw()
    #sleep(.005)
    fig.savefig('frame%03d' % step)

#plt.show()