def sor_diffusion_step(data, alpha, deltaX): n = len(data) m = len(data[0]) for j in range(1, n-1): for k in range(1, n-1): nei_avg = (data[j-1][k] + data[j+1][k] + data[j][k-1]+ data[j][k+1])/4 data[j][k] = (1-alpha)*data[j][k] + alpha*nei_avg def adi_diffusion_step(data, alpha, mode): #implicit on rows CORNERS = 0 EDGES = 1 n = len(data) m = len(data[0]) a = [-alpha] * (n-2) + [0] b = [1] + (n-2) * [1+2*alpha] + [1] c = [0] + (n-2) * [-alpha] new_data = [data[0][:]] d = [None for index in range(n)] if mode == CORNERS: for j in range(1, n-1): for k in range(m): d[k] = data[j][k] + alpha*(data[j+1][k] -2*data[j][k] + data[j-1][k]) new_stripe = tridiagonal_solver(a,b,c,d) new_data.append(new_stripe) elif mode == EDGES: for j in range(1, n-1): d[0] = data[j][0] d[m-1] = data[j][m-1] for k in range(1, m-1): d[k] = data[j][k] + alpha*(data[j+1][k] -2*data[j][k] + data[j-1][k]) new_stripe = tridiagonal_solver(a,b,c,d) new_data.append(new_stripe) new_data.append(data[n-1][:]) return new_data def transpose(mat): return map(list, zip(*mat)) def tridiagonal_solver(a, b, c, d): c = c[:] a = [None]+a n = len(b) c[0] = c[0]/b[0] d[0] = d[0]/b[0] for i in range(1,n-1): m = 1/(b[i]-c[i-1]*a[i]) c[i] = m * c[i] d[i] = m * (d[i] - d[i-1]*a[i]) i = n-1 m = 1/(b[i]-c[i-1]*a[i]) d[i] = m * (d[i] - d[i-1]*a[i]) x = [None]*n x[n-1] = d[n-1] for i in range(n-1)[::-1]: x[i] = d[i] - c[i]*x[i+1] return x def finite_difference_diffusion(): data = [0.0]*10 + [5.0] + [0.0]*10 next_data= data[:] deltaT = 1.0 deltaX = 1.0 D = 1.0 steps = 100 history = [data[:]] for step in xrange(steps): for j in xrange(1, len(data)-1): next_data[j] = data[j] + D*deltaT*(data[j+1] - 2.0*data[j]+ data[j-1])/(deltaX**2) data = next_data[:] history.append(data[:]) return history def implicit_finit_difference_diffusion(): n = 103 data = [0.0]*((n-3)/4) + [10.0] + [0.0]*((n-3)/4) + [20.0] +[0.0]*((n-3)/4) + [5.0]+ [0.0]*((n-3)/4) next_data= data[:] deltaT = 1.0 deltaX = 1.0 D = .030 steps = n history = [data[:]] alpha = D*deltaT/deltaX**2 a = [-alpha]*(n-2) + [0] b = [1] + (n-2)*[1+2*alpha] + [1] c = [0] + [-alpha]*(n-2) for step in xrange(steps): data = tridiagonal_solver(a,b,c,data) history.append(data[:]) return history if __name__ == "__main__": n = 5 alpha = 0.5 a = [-alpha]*(n-2) + [0] b = [1] + (n-2)*[1+2*alpha] + [1] c = [0] + [-alpha]*(n-2) d = [1.0] * n print tridiagonal_solver(a,b,c,d)