import numpy as np
import matplotlib.pyplot as plt

start = np.array([
    [  0,-70,  0,  0,  0],
    [  0,  0,  0,  0,  0],
    [  0,  0,  0,  0,  0],
    [  0,  0,  0,  0,  0],
    [  0,  0,  0,100,  0],
])

# the value of each cell is its number of neighbors
boundary_conditions = np.array([
    [  2,  3,  3,  3,  2],
    [  3,  4,  4,  4,  3],
    [  3,  4,  4,  4,  3],
    [  3,  4,  4,  4,  3],
    [  2,  3,  3,  3,  2],
])

# this needs to be 2x2 bigger to allow pasting the start
# four times - left, right, up, and down
temporary_canvas = np.zeros((7, 7))

for i in range(10_000):
    temporary_canvas[:-2,1:-1] += start # paste left
    temporary_canvas[2:,1:-1] += start # paste right
    temporary_canvas[1:-1,:-2] += start # paste up
    temporary_canvas[1:-1,2:] += start # paste down
    temporary_canvas[1:-1,1:-1] /= boundary_conditions
    start = np.copy(temporary_canvas[1:-1,1:-1])
    # reestablish Dirichlet boundary conditions
    start[0,1] = -70
    start[4,3] = 100
    # reset the canvas for next loop
    temporary_canvas.fill(0)

print(start)
plt.imshow(start, cmap="turbo")
plt.show()