Diffusion with Sink using Python and Numba

numba-diffusion.py

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

mu = 0.1
Lx, Ly = 601, 601

@jit('void(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

start = time.time()
diffusionObstacleStep(u,tempU,iterNum)
print(f"This took {time.time()-start:.3f}s")

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

start = time.time()
diffusionObstacleStep2(u,tempU,iterNum)
print(f"This took {time.time()-start:.3f}s")

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