pandasで条件分岐(case when的な)によるデータ加工を網羅したい

pandasでデータ加工をしつつ分析を進めていると、大抵の処理は問題なくできるのだけれど、SQLのcase whenのような条件分岐の処理をpandasでやろうとするとあまりの面倒くささに死にたくなることがよくある。

しかも毎回やり方を忘れるからここに書いておくことにする。

できるだけSQLと対比して書くことにする。

今回は２０１７年の日経平均株価を例にcase whenのパターンを書いてみたい。

元データはこんな感じ。

dt	stock_id	open	high	low	close	market_id
2017-01-04	1001	19298	19594	19277	19594	1
2017-01-05	1001	19602	19615	19473	19520	1
2017-01-06	1001	19393	19472	19354	19454	1
2017-01-10	1001	19414	19484	19255	19301	1
2017-01-11	1001	19358	19402	19325	19364	1

• dt:日付
• stock_id:株価ID
• open:始値
• high:最高値
• low:最低値
• close:終値
• market_id:市場ID

１変数が条件分岐の対象＆バイナリ化したい場合

条件分岐の対象が１つだけの場合には、pandasで加工するにしてもSQLで加工するにしてもそれほど大変ではない。。
さらに、[0,１]のように２値の値を取るように変換する場合はさらにラク。

例えば、丁度日経平均が2万円台に到達しようとするあたりなので、openが2万円を超えたら１がつくようなflg変数を追加する場合を考える。

SQLでは簡単に書ける。

select
*,
case when open >= 20000 then 1 else 0 end as flg_20t --20000円以上
from
sample_data --仮にテーブル名はsample_dataとする

では、pandasではどうなるかと言うと、まだ１行で書ける。

sample_data['flg_open_20t'] = sample_data['open'].apply(lambda x : 1 if x >= 20000 else 0)
sample_data.loc[sample_data.dt>='2017-06-01',:].head()

dt	stock_id	open	high	low	close	market_id	flg_open_20t
2017-06-01	1001	19692	19887	19686	19860	1	0
2017-06-02	1001	19970	20239	19967	20177	1	0
2017-06-05	1001	20135	20224	20104	20170	1	1
2017-06-06	1001	20122	20152	19948	19979	1	1
2017-06-07	1001	19951	20023	19908	19984	1	0

うん、できてるね。

次。

１変数が条件分岐の対象&複数の状態に分岐させたい（カテゴリ化したい）

次は、case whenの条件分岐の対象になる列は１つのまま変わらないのだけれど、FLGではなくてカテゴリ変数のように複数の状態を表現する場合を考える。

例えば、

1. openが１万円未満→0
2. openが１万円以下ー1万５千円未満→1
3. openが１万5千円以下ー2万円未満→2
4. openが2万円以上→3

を表現したい。
SQLでは、こんな感じ。まだまだ楽勝ですわ。

select
*,
case when open < 10000 then 0
     when open >= 10000 and open < 15000 then 1
     when open >= 15000 and open < 20000 then 2
     else 3 end as cate --カテゴリ変数
from
sample_data --仮にテーブル名はsample_dataとする

さて、pandasではどうかというと

sample_data['cate'] = sample_data['open'].apply(lambda x : 0 if x < 10000 else (1 if x >= 10000 and x < 15000 else(2 if x >= 15000 and x < 20000 else 3)))

だいぶ複雑。これならもうlambdaは使わずに普通に関数にした方が良さそう。

ということで

# 条件分岐の関数を定義
def func_cate(x):
    if  x < 10000:
        return 0
    elif x >= 10000 and x < 15000:
        return 1
    elif x >= 15000 and x < 20000:
        return 2
    else:
        return 3

sample_data['cate_f'] = sample_data['open'].apply(func_cate)
sample_data.head()

dt	stock_id	open	high	low	close	market_id	flg_open_20t	cate	cate_f
2017-01-04	1001	19298	19594	19277	19594	1	0	2	2
2017-01-05	1001	19602	19615	19473	19520	1	0	2	2
2017-01-06	1001	19393	19472	19354	19454	1	0	2	2
2017-01-10	1001	19414	19484	19255	19301	1	0	2	2
2017-01-11	1001	19358	19402	19325	19364	1	0	2	2

でけた。

次。

１変数が条件分岐の対象＆元の値を変換

上の２つは１つの変数を条件分岐の対象としつつ、新たに作成する変数の数値も外側から与えるタイプの処理だったので、まだなんとかなった。

でも、元の値の一部を置換するタイプの処理はだいぶ面倒になる。

例えば、こんな場合。

openの値が

1. １万9000円未満ならばNULLに置換
2. 1万9000円以上2万円未満ならば19500に置換
3. 2万円以上ならopenの値そのまま

データに外れ値があって特定の値に置換して使いたい場合にやる必要があることが多い。

SQLなら別に大したことない。

select
*,
case when open < 19000 then NULL
     when open >= 19000 and open < 20000 then 19500
     else open end as open_2 --加工済みopen値
from
sample_data --仮にテーブル名はsample_dataとする

pandasなら

import numpy as np
def func_open_2(x):
    if  x < 19000:
        return np.NaN
    elif x >= 19000 and x < 20000:
        return 19500
    else:
        return x

sample_data['open_2'] = sample_data['open'].apply(func_open_2)

| dt | stock_id | open | high | low | close | market_id | flg_open_20t | cate | cate_f | open_2 |
|--:|:--|--:|--:|--:|--:|--:|--:|--:|--:|--:|--:|
| 2017-01-04 | 1001 | 19298 | 19594 | 19277 | 19594 | 1 | 0 | 2 | 2 | 19500.0 |
| 2017-01-05 | 1001 | 19602 | 19615 | 19473 | 19520 | 1 | 0 | 2 | 2 | 19500.0 |
| 2017-01-06 | 1001 | 19393 | 19472 | 19354 | 19454 | 1 | 0 | 2 |2|19500.0 |
| 2017-01-06 | 1001 | 19393 | 19472 | 19354 | 19454 | 1 | 0 | 2 |2|19500.0 |
| 2017-01-18 | 1001 | 18753 | 18941 | 18650 | 18894 | 1 | 0 | 2 | 2 | NaN |
| 2017-11-24 | 1001 | 22390 | 22567 | 22381 | 22550 | 1 | 1 | 3 | 3 | 22390.0 |

うん。ええやろ。

次。

複数変数が条件分岐の対象＆カテゴリ化したい

ここからが本丸。

本当にpandasが私を苦しめ始めるのはこのあたりから。。。

１つの列だけを条件分岐の対象としている間は、列を特定してapplyを使っておけば良かったのだが、複数列が条件分岐に必要となると途端にpandasは扱いが面倒になる。

こういう場合を考えてみる。

1. openがcloseよりも低い（値上がり）→1
2. openがcloseよりも高い（値下がり）→2
3. openがcloseと同じ→0

SQLだったら簡単に実現できる。

select
*,
case when open < close then 1
     when open > close then 2
     else 0 end as cate_change --価格変動カテゴリ
from
sample_data --仮にテーブル名はsample_dataとする

pandasなら

def func_cate_change(x):
    if x.open < x.close:
        return 1
    elif x.open > x.close:
        return 2
    else:
        return 0

sample_data['cate_change'] = sample_data.apply(lambda x:func_cate_change(x),axis=1)

applyの処理がpandas.Seriesからpandas.DataFrameに変わっているのと処理の方向が列方向に対するものに変わっているのがポイント

このaxisの処理を忘れてよく何をやっているかわからなくなる orz

axis=0 →行方向（デフォルト）
axis=1 →列方向

dt	stock_id	open	high	low	close	market_id	flg_open_20t	cate	cate_f	open_2	cate_change
2017-01-04	1001	19298	19594	19277	19594	1	0	2	2	19500.0	1
2017-01-05	1001	19602	19615	19473	19520	1	0	2	2	19500.0	2
2017-01-06	1001	19393	19472	19354	19454	1	0	2	2	19500.0	1
2017-01-10	1001	19414	19484	19255	19301	1	0	2	2	19500.0	2

いいね！！

ちなみに

# 別にこっちでも通る
sample_data['cate_change_2'] = sample_data.apply(func_cate_change,axis=1)

複数変数が条件分岐の対象＆値を加工する

条件が複数列にまたがるだけならまだしも、新たに作る列に用いる値が複数にまたがることもある。

例えば、

1. highとlowが100円以上離れているなら、openとcloseの平均を取る
2. highとlowが100円未満なら、highとlowの平均を取る

SQLなら

select
*,
case when high - low >= 100 then (open + close) /2
     else (high + low) /2 end as avg_100 --変則平均値
from
sample_data --仮にテーブル名はsample_dataとする

pandasなら

def func_avg_100(x):
    if x.high - x.low <= 100:
        return (x.high + x.low) /2
    else:
        return (x.open + x.close) /2
    
sample_data['avg_100'] = sample_data.apply(lambda x:func_avg_100(x),axis=1)

dt	stock_id	open	high	low	close	market_id	flg_open_20t	cate	cate_f	open_2	cate_change	cate_change_2	avg_100
2017-01-04	1001	19298	19594	19277	19594	1	0	2	2	19500.0	1	1	19446.0
2017-01-05	1001	19602	19615	19473	19520	1	0	2	2	19500.0	2	2	19561.0
2017-01-06	1001	19393	19472	19354	19454	1	0	2	2	19500.0	1	1	19423.5
2017-01-10	1001	19414	19484	19255	19301	1	0	2	2	19500.0	2	2	19357.5
2017-01-11	1001	19358	19402	19325	19364	1	0	2	2	19500.0	1	1	19361.0

面倒だけど、まぁセーフ

複数変数が条件分岐の対象＆外側から引数を与える

ここまでは、条件や値が複数になることはあっても、どのような条件で分岐するかやどのような値を取るかということは所与の値を決めて作っていた。

しかし、このままだと条件分岐のパターンを作るたびに関数を作成する必要があり非常に面倒。
そこで、引数を加えて任意の条件や値でも分岐できることを考える。

想定するパターンとしては、下記を想定。

1. openがcloseに対して＋X円以上ならYを代入
2. openがcloseに対して＋X円未満ならZを代入

SQLでは、プロシージャや関数を使わないと基本的にはかけないので飛ばす。

pandasなら

def func_xyz(row,x,y,z):
    if row.open >= x:
        return y
    else:
        return z

# x,y,zを事前に定義
x = 19500
y = 0
z = 1
sample_data['xyz'] = sample_data.apply(lambda row:func_xyz(row,x,y,z),axis=1)

引数を調整すれば、これで任意の変数を自由に条件分岐に使える

listとかdict使って制御した方が、頭良さそうだけど、まぁええやろ

複数変数が条件分岐の対象＆欠損値置換

最後に欠損値の置換を条件に応じて変える場合を考える。

1. 値下がりする日(cate_change:1)のopen_2が欠損値ならば、期間全体の平均値を置換する
2. 値上がりする日(cate_change:2)のopen_2が欠損値ならば、当日のopenを置換する
3. 変動がない日(cate_change:0)のopen_2が欠損値ならば、そのままopen_2を採用する
4. 欠損値でないならそのままopen_2を採用する

まずは、SQLでかくと

select
a.*,
case when open_2 is NULL
     case when cate_change = 1 then b.avg_open
          when cate_change = 2 then a.open
          else a.open_2 end
     else open_2 end as open_3
       
from
    sample_data as a,
    --単純結合 
    (
     select
         avg(open) as avg_open
    from
        sample_data
    ) as b
;

SQLで書いてもかなり複雑だわ。。

pandasなら

import math

def func_open_3(row,avg_open):
    # スカラーでnp.NaNを評価したい場合は、mathで判定するのがいいっぽい
    if math.isnan(row.open_2):
        if row.cate_change == 1:
            return avg_open
        elif row.cate_change == 2:
            return row.open
        else:
            return row.open_2
    else:
        return row.open_3

# 行に対する処理と列に対する処理を同時に行うのは困難なので、まず欠損値処理を行う
sample_data['open_3'] = sample_data.open_2.fillna(method='bfill')

# 列に対する処理を実施する前に、open_2の行方向の平均値を事前に計算しておく
avg_open = sample_data['open'].mean()

# 列方向に対する処理を行う
sample_data['open_3'] = sample_data.apply(lambda x: func_open_3(x,avg_open),axis=1)

行方向の処理（平均や最大値の取得）と列方向の処理(条件分岐・置換)が込み入ってるとかなり複雑になってくるイメージ。

dt	stock_id	open	high	low	close	market_id	flg_open_20t	cate	cate_f	open_2	cate_change	cate_change_2	avg_100	xyz	open_3
3931	2017-01-17	1001	19038	19043	18812	18813	1	0	2	2	19500.0	2	2	18925.5	1
3932	2017-01-18	1001	18753	18941	18650	18894	1	0	2	2	NaN	1	1	18823.5	1
3933	2017-01-19	1001	19082	19122	18982	19072	1	0	2	2	19500.0	2	2	19077.0	1
3935	2017-01-23	1001	18938	19024	18879	18891	1	0	2	2	NaN	2	2	18914.5	1

ということで、色々やったから忘れた時、ここ見直せば今後自分を助けてくれるはず。。。

おしまい