熱拡散の計算03(数値解析)

import matplotlib.pyplot as plt

# グリッド分割数
DIVISION_COUNT = 256

# 境界条件
BOUNDARY_CONDITION = {
    0 : 0.0,
    int(DIVISION_COUNT-1) : 0.0
}

def generate_initial_grid(division_count):
    """
    初期状態のグリッドを作成
    """
    grid = []
    for i in range(int(DIVISION_COUNT/2)):
        grid.append(i/DIVISION_COUNT*2)
    for i in range(int(DIVISION_COUNT/2)):
        grid.append(-i/DIVISION_COUNT*2+1)
    return grid
    
def calcurate_step(grid):
    """
    Δtステップの計算を行う
    """
    next_grid = [0]*DIVISION_COUNT
    for x in range(DIVISION_COUNT):
        if x in BOUNDARY_CONDITION:
            next_grid[x] = BOUNDARY_CONDITION[x]
        else:
            next_grid[x]=grid[x]+(grid[x+1]-2*grid[x]+grid[x-1])/DIVISION_COUNT
    return next_grid

def main():
    grid = generate_initial_grid(DIVISION_COUNT)
    plt.plot(grid)
    for i in range(8):
        for j in range(2000*2**i):
            grid = calcurate_step(grid)
        plt.plot(grid)

if __name__ == '__main__':
    main()