#!/usr/bin/python3
# 2D Laplace equation on a square grid, Jacobi method

import numpy as np
import matplotlib.pyplot as plt

N = 100        # lattice = (N+1) x (N+1)
delta = 1.E-5  # desired accuracy for testing convergence
U = 1.0        # potential at the top boundary

phi = np.zeros((N+1, N+1))
phi[0, :] = U   # boundary condition
phiprime = np.copy(phi)
eps = delta + 1 # error, initially unimportant but must be > delta

while eps > delta:
   for i in range(1, N):
      for j in range(1, N):
         phiprime[i, j] = (phi[i+1, j] + phi[i-1, j] + phi[i, j+1] + phi[i, j-1]) / 4   
                   
   # error = maximal difference between new and old iteration
   eps = np.max(np.abs(phi - phiprime)) 
   # exchange phi' <-> phi
   phiprime, phi = phi, phiprime
   
plt.imshow(phi)
plt.show()