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

def func(x, y):
    return x**2 - y**2

h_x = 0.1
h_y = 0.1

a_x = 0
b_x = 1

a_y = 0
b_y = 1

eps = 0.001

n_y = int(ceil((b_y - a_y) / h_y))
n_x = int(ceil((b_x - a_x) / h_x))

matrix = zeros((n_x, n_y))
th = a_y

for i in range(n_y):
    matrix[i][0] = th**2
    th += h_y
th = a_y
for i in range(n_y):
    matrix[i][n_y - 1] = th**2 - 1
    th += h_y

while 1:
    tmp = array(matrix)
    matrix[1:(n_x - 1), 1:(n_y - 1)] = (h_y ** 2 * (matrix[2:n_x, 1:(n_y - 1)] + matrix[:(n_x - 2), 1:(n_y - 1)]) + h_x ** 2 * (matrix[1:(n_x - 1), 2:n_y] + matrix[1:(n_x - 1), 0:(n_y - 2)])) / (2 * (h_x ** 2 + h_y ** 2))
    matrix[0, 1:n_y - 1] = matrix[1, 1:n_y - 1]
    th = h_x
    for i in range(1, n_y):
        matrix[n_x - 1][i] = 1 - th**2
        th += h_x
    if linalg.norm(matrix - tmp) < eps:
        break

x = arange(0, 1, 0.1)
y = arange(0, 1, 0.1)
y_grid, x_grid = meshgrid(y, x)
z_grid = func(x_grid, y_grid)

fig = plt.figure()
ax = Axes3D(fig)
ax.view_init(30, 60)
ax.plot_surface(y_grid, x_grid, matrix)
plt.show()

fig = plt.figure()
ax = Axes3D(fig)
ax.view_init(30, 60)
ax.plot_surface(y_grid, x_grid, z_grid)
plt.show()

for i in range(len(matrix)):
    for j in range(len(matrix[i])):
        print('{0:2f}'.format(fabs(z_grid[i][j] - matrix[i][j])), end=' ')
    print()