import numpy import scipy.linalg import matplotlib.pyplot as plt def sor_step(A, b, x, omega=1.0): N = len(b) x_new = (1.0-omega)*x for j in range(N): x_new[j] += omega/A[j,j] * \ (b[j] - numpy.dot(A[j,:j], x_new[:j]) - numpy.dot(A[j,j+1:], x[j+1:])) return x_new def sor(A, b, steps, omega=1.0): N = len(b) x = numpy.zeros_like(b) for i in range(steps): x = sor_step(A, b, x, omega) return x # Aufgabe 10.2.1 def solve_poisson2d(rho, h): def linindex(x, y): return (x % N) + N*(y % N) N, N = rho.shape rho = rho.reshape(N*N) A = numpy.zeros((N*N, N*N)) # periodic for x in range(N): for y in range(N): i = linindex(x,y) A[i, i] = -4.0 A[i, linindex(x+1,y)] = 1.0 A[i, linindex(x-1,y)] = 1.0 A[i, linindex(x,y+1)] = 1.0 A[i, linindex(x,y-1)] = 1.0 A *= 1/(h*h) A[0,:] = 0.0 A[0,0] = 1.0 # plt.imshow(A) # plt.show() # phi = sor(A, rho, 200, 1.8) phi = scipy.linalg.solve(A, rho) phi = phi.reshape((N, N)) return phi # Aufgabe 10.2.2 L = 1.0 rho = numpy.load('rho.npy') N, N = rho.shape h = L/N phi = solve_poisson2d(rho, h) plt.subplot(221) plt.imshow(rho) plt.colorbar() plt.subplot(222) plt.imshow(phi) plt.colorbar() plt.subplot(223) ls = [] for y in [0, N/4, N/2, N*3/4]: plt.plot(phi[y,:], 'o-') ls.append('y=%d' % y) plt.legend(ls) plt.subplot(224) ls = [] for x in [0, N/4, N/2, N*3/4]: plt.plot(phi[:,x], 'o-') ls.append('x=%d' % x) plt.legend(ls) plt.show()