Heat Transfer mini program

heat_grid.py

import numpy as np
import matplotlib.pyplot as plt

s = [10, 10]
grid = np.zeros(s)
esp = 1e-3


def init():
    grid[:, :] = 650/4


def setBound():
    grid[0, :] = 50
    grid[:, -1] = 100
    grid[-1, :] = 200
    grid[:, 0] = 300


# main
init()
setBound()
it = 0

while True:
    it += 1
    print(it)
    oldgrid = grid.copy()
    grid = np.mean([
        np.roll(grid, 1, axis=0),
        np.roll(grid, 1, axis=1),
        np.roll(grid, -1, axis=0),
        np.roll(grid, -1, axis=1)], axis=0)
    setBound()
    # print(grid)
    if not (np.abs(grid - oldgrid) > esp).any():
        break


X, Y = np.meshgrid(np.arange(s[0]), np.arange(s[1]))
# plt.contourf(X, Y, grid, 10)
cont = plt.contour(X, Y, grid, 10)
plt.clabel(cont)
plt.xticks([])
plt.yticks([])
plt.subplots_adjust(top=1, bottom=0, right=1, left=0,
                    hspace=0, wspace=0)
plt.show()