import matplotlib.pyplot as plt

def inittab():  # Creation d'un tableau vierge avec NxN éléments et conditions aux limites
    tab = []
    for i in range(N):  # Création  d'un tableau avec NxN éléments nuls
        tab.append(N*[0.])
    for i in range(a,b):  # Ajout des conditions initiales
        tab[a][i]  = 100.
        tab[b][i]  = -100.
    return tab

def iteration(tab): #Effectue une itération et renvoie le nouveau tableau
    Vnew = inittab()
    for i in range(1,N-1):
        for j in range(1,N-1):
            if (i==a or i==b) and (j>a-1 and j<b):
                pass
            else:
                Vnew[i][j] = (tab[i-1][j]+tab[i+1][j]+tab[i][j-1]+tab[i][j+1])/4.
    return Vnew

def compare(tab1, tab2):    # Retourne le plus grand écart constaté
        ecart = 0
        for i in range(1,N-1):
            for j in range(1,N-1):
                diff = abs(tab1[i][j]-tab2[i][j])
                if diff>ecart:
                    ecart = diff
        return ecart


# Initialisation
N = 101     # dimension du tableau
e = 30       # ecart entre les plans
d = 30       # dimension des plans
a = (N-e)//2
b = (N+e)//2
eps = 0.01  # ecart toléré
Npas = 0    # nombre d'itérations
V = inittab()   # premiere distribution

#Iterations
while compare(V, iteration(V)) > eps :  # Tant que l'écart entre le nouveau potentiel et l'ancien est plus grand la tolérence, itération
    V = iteration(V)
    Npas += 1                           # On incrémente aussi le nombre de pas


# Affichage
fig = plt.figure(0)
ax = fig.add_subplot(1,1,1)
ax.set_aspect('equal')
plt.imshow(V, interpolation='nearest', cmap=plt.cm.seismic) #RdBu)    #hot)  #ocean)
plt.colorbar()
plt.show()