Diffusion with Sink using Python and Numba

numba-diffusion.py

import numpy as np
from numba.decorators import jit
from numba import *

mu = 0.1
Lx, Ly = 101, 101

@jit(arg_types=[double[:,:],double[:,:],int32])
def diffusionObstacleStep(u,tempU,iterNum):
  for n in range(iterNum):
    a = Lx/2-3
    b = Lx/2-3
    for i in range(Lx-1):
      for j in range(Ly-1):
        if (i-a)**2+(j-b)**2 <  3:
          tempU[i,j] = u[i,j]  
        else:
          tempU[i,j] = u[i,j] + mu * (u[i+1,j]-2*u[i,j]+u[i-1,j] + \
                                  u[i,j+1]-2*u[i,j]+u[i,j-1] )
    for i in range(Lx-1):
      for j in range(Ly-1):
        u[i,j] = tempU[i,j]
        tempU[i,j] = 0.0
    

u = np.zeros([Lx, Ly], dtype=np.float64)
tempU = np.zeros([Lx, Ly], dtype=np.float64)
u[Lx/2, Ly/2] = 1000.0
print 'With u[Lx/2,Ly/2]= 1000 set iterNum to 10 and iterate 10 times'

iterNum = 10 
for i in range(10):
  diffusionObstacleStep(u,tempU,iterNum)
  print i, ': Value in center: ', u[Lx/2,Ly/2]

print 'reset lattice, set iterNum to 100 and iterate once\n' 

u = np.zeros([Lx,Ly],dtype=np.float64)
u[Lx/2,Ly/2] = 1000
tempU = np.zeros([Lx,Ly],dtype=np.float64)
iterNum = 20

print 'With u[Lx/2,Ly/2]= 1000 set iterNum to 20 and iterate 5 times'

for i in range(5):
  diffusionObstacleStep(u,tempU,iterNum)
  print i, ': Value in center: ', u[Lx/2,Ly/2]

print 'reset lattice, set iterNum to 100 and iterate once\n' 

u = np.zeros([Lx,Ly],dtype=np.float64)
u[Lx/2,Ly/2] = 1000
tempU = np.zeros([Lx,Ly],dtype=np.float64)
iterNum = 100


diffusionObstacleStep(u,tempU,iterNum)

print 'Value in center: ', u[Lx/2,Ly/2]

print 'reset lattice, run 100 iterations without numba\n' 

def diffusionObstacleStep2(u,tempU,iterNum):
  for n in range(iterNum):
    a = Lx/2-3
    b = Lx/2-3
    for i in range(Lx-1):
      for j in range(Ly-1):
        if (i-a)**2+(j-b)**2 <  3:
          tempU[i,j] = u[i,j]  
        else:
          tempU[i,j] = u[i,j] + mu * (u[i+1,j]-2*u[i,j]+u[i-1,j] + \
                                  u[i,j+1]-2*u[i,j]+u[i,j-1] )
    u[:,:] = np.copy(tempU)

u = np.zeros([Lx,Ly],dtype=np.float64)
u[Lx/2,Ly/2] = 1000
tempU = np.zeros([Lx,Ly],dtype=np.float64)
iterNum = 100


diffusionObstacleStep2(u,tempU,iterNum)

print 'Value in center: ', u[Lx/2,Ly/2]