import numpy as np

import pylab as pl

dx = 0.01 # Step size in x-direction

dy = 0.01 # Step size in y-direction

nx = int(1/dx) # Number of bins over the x-axis

ny = int(1/dy) # Number of bins over the y-axis

a=1 # diffusion coefficient

dt=(dx*dy)**2/(dx**2+dy**2)/2/a/2

total_time=200000

# Compute the distance of (x, y) from (0.5, 0.5)

def dist(x,y): return np.sqrt((x-0.5)**2 + (y-0.5)**2)

# Initialize unow and u as zero matrices

un = np.zeros([nx,ny])

u = np.zeros([nx,ny])

# Store the initial condition in unow

for i in range(nx):

 for j in range(ny):

    if dist(i*dx,j*dy)<=0.1:

       un[i,j] = 1

for k in range(total_time):

 u[1:-1,1:-1]=un[1:-1,1:-1]+dt*a*(

 (un[2:,1:-1]-2*un[1:-1,1:-1]+un[:-2,1:-1])/dx/dx+

 (un[1:-1,2:]-2*un[1:-1,1:-1]+un[1:-1,:-2])/dy/dy)

 un=u

pl.imshow(un,cmap='hot')

pl.colorbar()

pl.axis('off')

pl.title(str(total_time)+'steps'+'='+str(total_time*dt)+'s')