pythonでheatmap(color map)のgif動画を作成 #Python

カラーマップのgif動画を作成するときにはまったので投稿．

ポイント

pythonで動画を作成する際には大きく分けて2つの方法がある．
1つ目が画像を生成・配列格納を行いつつ，最終的にレンダリングする方法．2つ目が最初に画像を保存し，それらの画像をレンダリングする方法．

前者は計算のろかったので，後者を採用．その際にはまった点を重点的に記述．

ちなみにレンダリングの方法はここを参照した．

カラーバーが変になる問題

最初に取り組んだ際に生じたのが，カラーバーが複数個画像に載ってしまう現象．更に実行速度が極端に落ちた．

前のプロットをそのまま引き継いだがためのメモリ消費以外ありえないと感じたので，まずは前のプロットを消去する方法を考えた．結論から言うと，「plt.figure()」を挟むことによって解決した．

また，カラーバーの上限と下限を固定したかったので，あらかじめカラーバーの設定を入れることにした．

def plot_heatmap(x, y, dt, level):
  # x, yのメッシュとdt, color barの設定level
  # は不変なので，最初に関数内に保持し，
  # 後に変数のtだけ代入して画像を生成する関数_impを
  # 返す．

  def _imp(t):
    # 前のplotの消去
    plt.figure()

    # カラーマップ取得
    contourf = plt.contourf(x, y, f(x, y, t * dt), level)

    # カラーバーの生成
    plt.colorbar(contourf)

    # 画像の保存
    plt.savefig("test/figure{}.png".format(t))

  return _imp

コード全容

その他の部分は特に問題なし．
上のgif動画は下のコードを実行することで得ることが可能．
ちなみに今回は単純な関数で動画生成を行ったが，実際にはあらかじめcsvにデータをとりためて，それらを行列データとしてimportしたのちにメッシュに区切ったx,yと対応付けた関数として扱った．

その際に使用したコードはこちら．

import numpy as np
import matplotlib.pyplot as plt
import os
from PIL import Image, ImageDraw

def rendering():
  path = os.getcwd() + "/test/"
  files = [path + "figure{}.png".format(t) for t in range(len(os.listdir(path)))]
  images = list(map(lambda file: Image.open(file), files))
  images.pop(0).save("test.gif" ,save_all = True, append_images = images, duration = 100, optimize = False, loop = 0)

def f(x, y, t):
  return np.cos(np.sqrt(x ** 2 + y ** 2) + t)

def plot_heatmap(x, y, dt, level):
  def _imp(t):
    plt.figure()
    contourf = plt.contourf(x, y, f(x, y, t * dt), level)
    plt.colorbar(contourf)
    plt.savefig("test/figure{}.png".format(t))

  return _imp

def make_savedir():
  if not os.path.isdir("test"):
    os.mkdir("test")

if __name__ == "__main__":
  # 定義域を決定
  x_min, x_max, y_min, y_max = -5, 5, -5, 5

  # 各軸の分割数を決定
  num = 50

  # 時間ステップの進む速さと終了ステップ数を決定
  dt = 0.1
  T = 20

  # 画像を保存するディレクトリを作成
  make_savedir()

  # x, yのグリッドを作成
  x, y = np.meshgrid(np.linspace(x_min, x_max, num), np.linspace(y_min, y_max, num))

  # 関数値の取る最大，最小値を設定
  f_min, f_max = -1, 1

  # 等高線の上限・下限と分割数を設定
  level = np.linspace(f_min, f_max, 10)

  # タイムステップを代入することによって，そのときの画像を
  # 生成する関数が返ってくる．
  img_creation = plot_heatmap(x, y, t_array, level)

  for t in range(T):
    img_creation(t)

  # できた画像をレンダリングして動画にする
  rendering()