【pandas_flavor】Pandas DataFrameのメソッドを追加する

TL;DR

BEFORE

dataframe_ = dataframe.loc[(dataframe.time == 'pre') & \
                           (dataframe.group == 'exp') & \
                           (dataframe.cond == 'a'), :]
sns.regplot(x='mood', y='score', data=dataframe_)

↓↓↓

AFTER

dataframe.by(time='pre', cond='exp', group='a').regplot(x='trait', y='score')

pandas_flavorを使うことで、好きなメソッドをpandas DataFrame（およびSeries）に追加することができます。

モチベーション

Long形式のデータから条件に合う箇所を抽出するのはめんどい！

例えばこんなデータがあったとします。

被験者50人をふたつの群（group: exp, ctrl）に分け、それぞれになにやら介入を行ったという設定です。
介入の前後（time: pre, post）で課題を行わせ、課題中のふたつの条件（cond: a, b）における成績（score）を測定しました。
同時に、課題をやっているときの気分（mood）も、条件（cond: a, b）ごとに測定しました。¹

測定データは、上の画像のようにlong形式でまとめるとその後の解析がやりやすくなりますね。

さて、いろいろな解析をやる前にひとまず、
preにおけるexp群の課題条件aのときのscoreとmoodとの相関関係をプロットすることにしましょう。

以上の条件に合う行を抽出してくるので、↓こんなコードになりますね。

dataframe_ = dataframe.loc[(dataframe.time == 'pre') & \
                           (dataframe.group == 'exp') & \
                           (dataframe.cond == 'a'), :]
sns.regplot(x='mood', y='score', data=dataframe_)

条件を表したbool型のSeriesを作り、.locに入れてやってます。
うーん、なんか汚い。

.query()メソッドを使えば、↓こんなふうにも書けます。

dataframe_ = dataframe.query('time == "pre" & group == "exp" & cond == "a"')
sns.regplot(x='mood', y='score', data=dataframe_)

こっちのほうがだいぶスッキリしていますが、もうちょっといい感じにならんかなぁ。
.query()使う方法はboolのSeriesを使う方法に比べて遅いそうですし。
やっぱり、Long形式のデータから条件に合う箇所を抽出するのはめんどい！

Pandas DataFrameのメソッドを追加する

ならメソッドを作ればいいじゃない

そこで、DataFrameから条件に合うような行を抽出してくるような新しいメソッドを作ってやりましょう。
↓こんなふうに使える.by()メソッドをDataFrameに新しく追加します。

dataframe.by(time='pre', cond='exp', group='a')

pandas_flavorというパッケージを使うと、簡単にこれを達成することができます。

インストール方法

pipか、

pip install pandas_flavor

condaで一発です。

conda install -c conda-forge pandas_flavor

使用例

import pandas_flavor as pf


@pf.register_dataframe_method
def by(self, **args):
    for key in args.keys():
        self = self.loc[self.loc[:, key] == args[key], :]
    return self

関数を書いて、デコレータとして@pf.register_dataframe_methodを付けてやるだけです。
こちらの例では、**argsとやることで引数を辞書型で受け取ります。
これによって、各引数で指定した行を抽出しています。

さらに、seabornの各種関数をメソッド化させるのもいいですね。

@pf.register_dataframe_method
def regplot(self, **args):
    return sns.regplot(data=self, **args)

と、こんな感じです。
pandas.Seriesにメソッドを追加してやりたければ、@pf.register_series_methodで同じことができるようです。

今回の例は…、まぁ.query()使えばいいような気もしますが、色々と応用ができそうです。

---

1. 言うまでもありませんが、すべてでっち上げの心理学実験です。数値はrandomモジュールで生成しています。 ↩