酒飲みによる機械学習日記

〜ただの酒好きが、なんとなく流行っている機械学習について語っていきます、たまに金融ネタも〜

Pythonでアニメーション2

前回はPythonでsin波を動的にプロットしてみました。

nts-524.hatenablog.jp

今回はその続きとして、sin波とcos波を同時にアニメーションとしてプロットすることに挑戦しました!!

結果は...

想像以上に簡単でした。普段通りのやり方でラベル(凡例)もつけることができました。

f:id:nts-524:20181008194142g:plain

以下ソースコード全文です!!

# =-=-=-=-=-=-=-=-=-=-=-=- module -=-=-=-=-=-=-=-=-=-=-=-= #
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
import seaborn as sns
import math, os, argparse

sns.set_style('whitegrid')

# =-=-=-=-=-=-=-=-=-=-=-=- main -=-=-=-=-=-=-=-=-=-=-=-= #
def main():
    # argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--x_min', type=float, default=0.0)
    parser.add_argument('--x_max', type=float, default=10.0)
    parser.add_argument('--result', type=str, default='result/anime2.gif')
    args = parser.parse_args()

    # make result directory
    if not os.path.exists('result/'):
        os.makedirs('result/')

    # define plot function
    fig = plt.figure()
    x = []
    y_sin = []
    y_cos = []
    def plot(frame, x_min=args.x_min, x_max=args.x_max):
        """update function for graph"""
        # clean graph
        plt.cla()
        # add data
        x.append(frame)
        y_sin.append(math.sin(frame))
        y_cos.append(math.cos(frame))
        # replot
        plt.xlim(x_min, x_max)
        plt.ylim(-1, 1)
        plt.plot(x, y_sin, label='$y=sin(x)$')
        plt.plot(x, y_cos, label='$y=cos(x)$')
        plt.legend()

    # visualize result
    anime = animation.FuncAnimation(
        fig, plot,
        interval=10, # [ms]
        frames=np.arange(args.x_min, args.x_max, 0.1),
        repeat=False,
    )
    anime.save(args.result, writer='pillow')
    plt.show()

if __name__ == '__main__':
    main()