flame.py

# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import seaborn as sns

def draw_flame(N=10, M=10):
    X = np.zeros((N, M))

    # i=0
    X[0, :] = np.random.randint(low=0, high=255, size=M)

    # i=1
    for j in range(M):
        idx = np.linspace(j-1, j+1, 3, dtype=np.int)
        idx = idx[np.where(idx >= 0)[0]]
        idx = idx[np.where(idx < M)[0]]
        X[1, j] = int(np.mean(X[0, idx]))

    # others
    for i in range(2, N):
        for j in range(M):
            noise_min = np.random.randint(-5, 5)
            noise_max = noise_min + np.random.randint(0, 5)
            idx = np.linspace(j-noise_min, j+noise_max, 3, dtype=np.int)
            idx = idx[np.where(idx >= 0)[0]]
            idx = idx[np.where(idx < M)[0]]
            mean = int((np.sum(X[i-1, idx]) + X[i-2, j]) / 4.0)
            X[i, j] = mean

    X = np.flipud(X)
    return X

def flame(n, N, M):
    X = draw_flame(N, M)
    fig = plt.figure()
    im = plt.imshow(X, interpolation="nearest", cmap="hot")

    def update(t):
        X = draw_flame(N, M)
        im.set_array(X)

    anim = animation.FuncAnimation(fig, func=update, frames=n, repeat=False, interval=10)
    # plt.show()
    anim.save("output.gif", writer="imagemagick")

if __name__ == '__main__':
    flame(100, 15, 50)