SQL
pandas
python3

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

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

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

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

今回は2017年の日経平均株価を例に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

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

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

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

SQLでは簡単に書ける。

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

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

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

うん、できてるね。

次。

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

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

例えば、
1. openが1万円未満→0
2. openが1万円以下ー1万5千円未満→1
3. openが1万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

でけた。

次。

1変数が条件分岐の対象&元の値を変換

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

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

例えば、こんな場合。

openの値が
1. 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が私を苦しめ始めるのはこのあたりから。。。

1つの列だけを条件分岐の対象としている間は、列を特定して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

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

おしまい