import numpy as np

#-------------------------------------------------------------------------------
#------------------- Initialisations -------------------------------------------
#-------------------------------------------------------------------------------

Nx = 30     # entier
Ny = 30     # entier

#------- Potentiels imposés sur les 4 bords du rectangle -----------------------

Vh, Vb, Vg, Vd = 200, 0, 100, -100
 
M0 = np.zeros( (Nx+1, Ny+1), dtype = float)

Px = np.ones( Nx+1, dtype = float)  # Tableau unidimensionnel de taille Nx + 1
Py = np.ones( Ny+1, dtype = float)  # Tableau unidimensionnel de taille Ny + 1
  
#------------------- Remplissage des bords de M0 -------------------------------

M0[0:Nx+1, 0] = Vg * Px       # première colonne
M0[0:Nx + 1, Ny] = Vd * Px    # dernière colonne
M0[0, 0:Ny+1] = Vh * Py       # pemière ligne
M0[Nx, 0:Ny+1] = Vb * Py      # dernière ligne

M1 = M0.copy()
erreur = 1
maxi = 1e-3

#-------------------------------------------------------------------------------
#---------------------    Début du programme -----------------------------------
#-------------------------------------------------------------------------------

while erreur > maxi :
    for i in range(1, Nx) :
        for j in range(1, Ny) :
            M1[i, j] = 0.25* ( M0[i-1, j] + M0[i + 1, j] \
            + M0[i, j - 1] + M0[i, j+1] )
            
    E = M1 - M0
    erreur = np.amax( abs(E) )
    M0 = M1.copy()

#-------------------------------------------------------------------------------
#---------------------- Représentation 3D du potentiel -------------------------
#-------------------------------------------------------------------------------

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure(1)
ax = Axes3D(fig)
h = 0.01
x = h*np.arange(Nx+1)      # h est le pas du quadrillage
y = h*np.arange(Ny+1)
X, Y = np.meshgrid(x, y)
ax.plot_surface(X,Y,M1, rstride=2, cstride=2, color = 'y')

