熱拡散の計算01

import matplotlib.pyplot as plt

# グリッド分割数
DIVISION_COUNT = 256

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

def generate_initial_grid(division_count):
    """
    初期状態のグリッドを作成
    """
    grid = []
    grid .extend([0.0]*int(DIVISION_COUNT/4))
    grid .extend([1.0]*int(DIVISION_COUNT/2))
    grid .extend([0.0]*int(DIVISION_COUNT/4))
    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(1000*2**i):
            grid = calcurate_step(grid)
        plt.plot(grid)

if __name__ == '__main__':
    main()