matplotlibのtips

さくっとグラフを作成したいときにmatplotlibは非常に便利. ipython notebookをサーバー上に立てておけば,リモートで作図したグラフなどをscpなどしなくても描画とコピペができるのでいろいろ捗る.

しかし個人的に棒グラフを描画するbar関数の使い勝手が悪いと思っている. 具体的には

• 棒グラフを複数並べたい場合,棒グラフ間のマージンの設定
• グラフのラベルをそれぞれの真ん中に持ってくる時の幅の計算

あたりが割とﾒﾝﾄﾞｸｻｲ.

なので,スニペット的なものを書いてみた.

#! /usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def bars(target, size, labels=None, width=0.25, margin=0.5,
         output=None,title=None, xlabel=None, ylabel=None, loc="best"):
    """
    plot bars

    required args
    --------------

    target: type dict,key:label,value:points(array)
    size: type int, number of points

    optional args
    ---------------

    labels: type array, group label
    width: type float, bar's width
    margin: type float, length between bars
    output: type str,output file name
    title: type str, title name
    xlabel: type str, xlabel name
    ylabel: type str, ylabel name
    loc: type str, legend location

    """

    block = len(target) * width + margin
    ind = np.arange(size) * block

    index = 0
    for l, points in target.iteritems():
        plt.bar(ind + width * index,
                points, width,
                color=cm.cool(float(index) / len(target)),
                label=l)

        index += 1

    # centring
    plt.xlim(-margin, ind[-1] + width*len(target)+margin)

    if title:
        plt.title(title)

    if xlabel:
        plt.xlabel(xlabel)

    if ylabel:
        plt.ylabel(ylabel)

    if labels:
        plt.xticks(ind+(len(target)/2)*width, labels)

    plt.legend(loc=loc)

    # draw bars
    if output:
        plt.savefig(output)
    else:
        plt.show()

引数はいろいろあるが最低限必要なのは以下の2つだけ

• target : keyがラベル,valueがプロットするy座標の配列となっている辞書
• size: y座標の配列の長さ

これ以外の引数は

• labels :　各棒グラフグループの名前
• output :その場に表示するかファイルに落とすか
• width : 各棒グラフの幅
• margin : 棒グラフグループ間の距離
• title : グラフのタイトル
• xlabel : 横軸の名前
• ylabel : 縦軸の名前
• loc : 凡例の位置

を必要に応じて設定する.

サンプルコードと実行例

サンプルコード : gist

実行例