from numpy import empty,zeros,max
import matplotlib.pyplot as plt

# Constants
M = 100         # Grid squares on a side
V = 1.0         # Voltage at top wall
target = 1e-6   # Target accuracy
overskip = 0.9  # Overshoot parameter for Gauss-Seidel

# Create arrays to hold potential values
phi = zeros([M+1,M+1],float)
phi[0,:] = V
##phiprime = empty([M+1,M+1],float)

# Main loop
delta = 1.0
while delta>target:
    delta = 0.0

    # Calculate new values of the potential
    for i in range(M+1):
        for j in range(M+1):
            if not (i==0 or i==M or j==0 or j==M): # not at the boundaries
                oldphi = phi[i,j]
                phi[i,j] = (phi[i+1,j] + phi[i-1,j] \
                            + phi[i,j+1] + phi[i,j-1])*\
                            (1.0+overskip)*0.25 - overskip * phi[i,j]
                delt = abs(oldphi - phi[i,j])   # amount of change this step
                if delt > delta:
                    delta = delt    # maximum change this step
                

# Make a plot
plt.imshow(phi,origin="lower",cmap="cool")
plt.show()
plt.contour(phi,origin="lower",cmap="autumn")
plt.show()

print("Done.")