はじめに
本稿の目的は、Pythonのライブラリであるeconmlを用いたDouble Machine Learningに関する理解を深めることです.
使用データはEconML公式ドキュメント内のコード例(以下コード例)内にあるオレンジジュースの売上データです.
https://github.com/py-why/EconML/blob/main/notebooks/Double%20Machine%20Learning%20Examples.ipynb
今回のテーマは因果推論で、問題を抽象化すると以下の類型に該当します.
- Continuous outcome
- Continuous treatment
- Observational data
環境など
- Google Colab
EconMLのインストール
以下のようにインストールします.
pip install econml
本編
関連ライブラリのインポート
関連ライブラリをインポートするためのコードは次のとおりです. 結果として使用しなかったインポート文も含まれます. コード例とほとんど同じです.
# Ignore warnings
import warnings
warnings.filterwarnings("ignore")
# Main imports
# econmlモジュールから特定のクラスをインポート
from econml.dml import (
DML, # Double Machine Learning (DML)
LinearDML, # 線形なDML
SparseLinearDML, # スパースな線形DML
CausalForestDML # Causal Forest DM
)
# Helper imports
import numpy as np
from itertools import product
from sklearn.linear_model import (
Lasso,
LassoCV,
LogisticRegression,
LogisticRegressionCV,
LinearRegression,
MultiTaskElasticNet,
MultiTaskElasticNetCV
)
from sklearn.ensemble import (
RandomForestRegressor,
RandomForestClassifier,
GradientBoostingRegressor,
GradientBoostingRegressor
)
from sklearn.multioutput import(
MultiOutputRegressor
)
from sklearn.preprocessing import (
PolynomialFeatures
)
import matplotlib.pyplot as plt
import matplotlib
from sklearn.model_selection import (
train_test_split
)
# A few more imports
import os
import pandas as pd
import urllib.request
from sklearn.preprocessing import StandardScaler
import seaborn as sns
Double Machine Learning(以下DML)について
DMLのモデルはデータ生成過程に以下のような構造的な仮定を置きます.
\begin{align*}
Y &= \theta(X) \cdot T + g(X, W) + \epsilon & E[\epsilon | X, W] &= 0 \\
T &= f(X, W) + \eta & E[\eta \mid X, W] &= 0 \\
& & E[\eta \cdot \epsilon | X, W] &= 0
\end{align*}
推定したいのは$\theta(X)$で、$T$は介入(Treatment)です.
$X$と$W$はいずれも特徴量ですが、$T$の係数になるかどうかが異なります.
$Y$のモデル化にあたり、$T$について線形となっていることから、部分線形モデルと呼ばれるようです.
先ほどインポートしたLinearDML, SparseLinearDML, CausalForestDMLは$\theta(X) \cdot T$をどのようにモデル化するかに応じて使い分けます.
上式を次のように変形できるようです.
Y - E[Y | X, W]
= \theta(X) \cdot (T - E[T | X, W]) + \epsilon
さらに, $\tilde{Y}$, $\tilde{T}$をそれぞれ次のように定義すると,
\begin{split}\tilde{Y} =~& Y - q(X, W)\\
\tilde{T} =~& T - f(X, W) = \eta\end{split}
以下のように変形できます.
\tilde{Y} = \theta(X) \cdot \tilde{T} + \epsilon
従って、まず初めに$f$と$q$($g$でないことに留意)を予測し、そのあと$\theta(X)$を予測します. $f$と$q$のモデル化にあたっては、既製品の機械学習モデルを使用して問題なく、残差を残差で回帰する手法はFWL定理というよく知らた性質を用いた手法のようです.
なお、EconMLの公式ドキュメントで言及されている文献（Chernozhukov2016）を確認すると、学習時に使用するデータを工夫したり、一般的なスコア関数（と文献内で呼ばれている、推定値を得るための偏微分=0などといった条件式）に変形を加えるなど、$\theta(X)$の推定量(推定値？)が一致性(consistency)や漸近正規性(asymptotic normality)を担保できるような仕掛けが施されているようです.
データ
以下ではオレンジジュースの売上データを読み込みます.
# Import the data
file_name = "oj_large.csv"
if not os.path.isfile(file_name):
print("Downloading file (this might take a few seconds)")
urllib.request.urlretrieve("https://msalicedatapublic.z5.web.core.windows.net/datasets/OrangeJuice/oj_large.csv", file_name)
oj_data = pd.read_csv(file_name)
以下のようなカラム名のデータで、今回は$Y$がlogmove, $T$がpriceの対数をとったものになります.
# Check the data types and the number of missing values in each column
print(oj_data.info())
"""
RangeIndex: 28947 entries, 0 to 28946
Data columns (total 17 columns):
# Column Non-Null Count Dtype
--- ------ -------------- -----
0 store 28947 non-null int64
1 brand 28947 non-null object
2 week 28947 non-null int64
3 logmove 28947 non-null float64
4 feat 28947 non-null int64
5 price 28947 non-null float64
6 AGE60 28947 non-null float64
7 EDUC 28947 non-null float64
8 ETHNIC 28947 non-null float64
9 INCOME 28947 non-null float64
10 HHLARGE 28947 non-null float64
11 WORKWOM 28947 non-null float64
12 HVAL150 28947 non-null float64
13 SSTRDIST 28947 non-null float64
14 SSTRVOL 28947 non-null float64
15 CPDIST5 28947 non-null float64
16 CPWVOL5 28947 non-null float64
"""
ちなみに、$Y = \theta(X) \cdot T + g(X, W) + \epsilon$において、$Y$が$log(売上)$, $T$が$log(価格)$となるとき、$T$の係数$\theta(X)$は価格弾力性に相当するようです.
$X$はモデル上は特徴量を指定せず、interceptだけとしても良いのですが、コード例では$X$をINCOMEとし、収入により価格弾力性に違いがあるのかを検証していました.
DMLを実施する前にこれら3つの変数の関係性を確認しましょう.
以下はそのためのコードとグラフ出力結果です.
# Plot scatter plots.
plt.scatter(
np.log(oj_data['price']),
oj_data['logmove'],
c=oj_data['INCOME'], # 色を変数によって分ける
cmap='viridis', # 色を変数によって分ける
alpha=0.5 # 透明度を調整
)
plt.xlabel('logprice')
plt.ylabel('logmove')
plt.title('Scatter Plot between logprice and logmove')
# カラーバーを追加する（必要に応じて）
cbar = plt.colorbar()
cbar.set_label('INCOME')
plt.show()
上図を見ると、確かにlogpriceとlogmoveには負の相関がありそうですが、それ以上のことは直感的には分かりませんね.
DMLのための学習用データ準備
学習用のデータを以下のように準備します.
# Prepare data for DML
Y = oj_data['logmove'].values.reshape(-1, 1)
T = np.log(oj_data['price']).values.reshape(-1, 1)
T_df = pd.DataFrame()
T_df['logprice'] = np.log(oj_data['price'])
scaler = StandardScaler()
W1 = scaler.fit_transform(
oj_data[[
c for c in oj_data.columns if c not in ['price', 'logmove', 'brand', 'week', 'store','INCOME']
]]
)
W2 = pd.get_dummies(oj_data[['brand']]).values
W = np.concatenate([W1, W2], axis=1)
X = scaler.fit_transform(oj_data[['INCOME']].values)
X_df = pd.DataFrame()
X_df[scaler.get_feature_names_out()] = X
$T$, $X$は作業の都合上、データフレーム版も準備しました.
また、特徴量はスケーリングします.
LinearDML
LinearDMLの学習は次のように実施します.
# Train estimator
est_linear = LinearDML(
model_y = GradientBoostingRegressor(),
model_t = GradientBoostingRegressor(),
featurizer=None,
treatment_featurizer=None,
linear_first_stages=False,
discrete_treatment=False,
cv=2
)
est_linear.fit(Y, T, X=X, W=W)
est_linear.summary()
"""
Coefficient Results
point_estimate stderr zstat pvalue ci_lower ci_upper
X0 0.147 0.022 6.636 0.0 0.104 0.191
CATE Intercept Results
point_estimate stderr zstat pvalue ci_lower ci_upper
cate_intercept -2.636 0.022 -117.858 0.0 -2.68 -2.592
"""
上図から$\theta(X)$を線形モデルと仮定していることが分かります. 初見では切片の数値の大きさが目に止まりますね.
先ほどの散布図と合わせて可視化します. 以下はそのコードと実行結果です.
# Plot scatter plots with synchronized y-axes.
fig, axes = plt.subplots(1, 2, figsize=[10, 5], sharey=True)
# First subplot
scatter1 = axes[0].scatter(
T,
Y,
c=X, # 色を変数によって分ける
cmap='viridis', # 色を変数によって分ける
alpha=0.5 # 透明度を調整
)
# 体裁の調整
axes[0].set_xlabel('T')
axes[0].set_ylabel('Y')
axes[0].set_title('T and Y')
# Second subplot
scatter2 = axes[1].scatter(
T,
T * est_linear.const_marginal_effect(X),
c=X, # 色を変数によって分ける
cmap='viridis', # 色を変数によって分ける
alpha=0.5 # 透明度を調整
)
# 体裁の調整
axes[1].set_xlabel('T')
axes[1].set_ylabel('theta(X) * T')
axes[1].set_title('T and theta(X) * T')
# カラーバーを追加する（必要に応じて）
cbar1 = plt.colorbar(scatter1, ax=axes[0])
cbar1.set_label('X')
cbar2 = plt.colorbar(scatter2, ax=axes[1])
cbar2.set_label('X')
plt.show()
上図のとおり、同じ$T:=log(price)$の水準であれば、$X:=INCOME$が大きい方が、$\theta(X)$も大きくなることが分かります. また、全体として切片の影響が強く、$T$と$\theta(X) \cdot T$の関係は線形に近いことも見て取れます.
LinearDMLの引数'treatment_featurizer'にPolynomialFeatures()を使用することで、以下のようなモデルにも対応することが可能です.
$$Y = \theta_{1}(X) \cdot T + \theta_{2}(X) \cdot T^2 + g(X, W) + \epsilon$$
今回は引数'treatment_featurizer'を使用せずに試して見ました.
# Train estimator
poly = PolynomialFeatures(degree=2, interaction_only=False, include_bias=False)
T_composite = poly.fit_transform(T_df)
col_name = poly.get_feature_names_out()
T_composite_df = pd.DataFrame()
T_composite_df[col_name] = T_composite
est_linear2 = LinearDML(
model_y = GradientBoostingRegressor(),
model_t = MultiOutputRegressor(GradientBoostingRegressor()),
featurizer=None,
treatment_featurizer=None,
linear_first_stages=False,
discrete_treatment=False,
cv=2
)
est_linear2.fit(Y, T=T_composite_df, X=X, W=W)
est_linear2.summary()
"""
Coefficient Results
point_estimate stderr zstat pvalue ci_lower ci_upper
X0|logprice 0.099 0.073 1.361 0.173 -0.043 0.241
X0|logprice^2 0.042 0.042 1.006 0.315 -0.04 0.124
CATE Intercept Results
point_estimate stderr zstat pvalue ci_lower ci_upper
cate_intercept|logprice -3.316 0.077 -43.197 0.0 -3.467 -3.166
cate_intercept|logprice^2 0.473 0.046 10.316 0.0 0.383 0.563
"""
上図のケースでは$T$を予測するための$f$が$T$と$T^2$の項の計2種類必要になるので、LinearDMLの引数'model_t'をMultiOutputRegressor()としました. 先ほどと同様$T$に係る切片の影響が強いようです. これを再びグラフにしてみたいと思います. 以下がそのためのコードと出力結果です.
estimator = est_linear2
# Plot scatter plots with synchronized y-axes.
fig, axes = plt.subplots(1, 2, figsize=[10, 5], sharey=True)
# First subplot
scatter1 = axes[0].scatter(
T,
Y,
c=X, # 色を変数によって分ける
cmap='viridis', # 色を変数によって分ける
alpha=0.5 # 透明度を調整
)
# 体裁の調整
axes[0].set_xlabel('T')
axes[0].set_ylabel('Y')
axes[0].set_title('T and Y')
# Second subplot
scatter2 = axes[1].scatter(
T,
(T_composite * estimator.const_marginal_effect(X)[:, 0, :]).sum(axis=1),
c=X, # 色を変数によって分ける
cmap='viridis', # 色を変数によって分ける
alpha=0.5 # 透明度を調整
)
# 体裁の調整
axes[1].set_xlabel('T')
axes[1].set_ylabel('theta(X) * T')
axes[1].set_title('T and theta(X) * T')
# カラーバーを追加する（必要に応じて）
cbar1 = plt.colorbar(scatter1, ax=axes[0])
cbar1.set_label('X')
cbar2 = plt.colorbar(scatter2, ax=axes[1])
cbar2.set_label('X')
plt.show()
グラフの形状から、$T^2$の係数内の切片の推定値が正であることが見て取れ、実際そのようになっています.
$T^2$の項をモデルに追加しても同じような右肩下がりの形状となりました. 左の散布図から右の散布図を引いた残りが$g$ということになりますね.
元の散布図からこのような関係性が抽出できるのは興味深いです.
SparseLinearDML
SparseLinearDMLの学習は次のように実施します.
# Train estimator
est_sparse_linear = SparseLinearDML(
model_y = GradientBoostingRegressor(),
model_t = MultiOutputRegressor(GradientBoostingRegressor()),
featurizer=PolynomialFeatures(degree=2, include_bias=False),
treatment_featurizer=PolynomialFeatures(degree=2, interaction_only=False, include_bias=False),
discrete_treatment=False,
linear_first_stages=False)
est_sparse_linear.fit(Y, T=T_df, X=X_df, W=W)
est_sparse_linear.summary()
"""
Coefficient Results
point_estimate stderr zstat pvalue ci_lower ci_upper
x0|logprice 0.054 0.039 1.388 0.165 -0.022 0.131
x0|logprice^2 0.065 0.03 2.169 0.03 0.006 0.123
x0^2|logprice -0.101 0.037 -2.704 0.007 -0.174 -0.028
x0^2|logprice^2 0.054 0.024 2.222 0.026 0.006 0.102
CATE Intercept Results
point_estimate stderr zstat pvalue ci_lower ci_upper
cate_intercept|logprice -3.205 0.069 -46.506 0.0 -3.34 -3.07
cate_intercept|logprice^2 0.414 0.045 9.158 0.0 0.325 0.502
"""
LinearDMLと異なり、正則化をかけているからか、控えめな係数が出力されている気がします. ただ、出来上がりのグラフは先ほどまでとあまり変わり映えしません.
CausalForestDML
CausalForestDMLの学習は次のように実施します. 線形回帰モデルではないので先ほ度までと異なりsummary()メソッドはありません.
# Train estimator
est_causal_forest = CausalForestDML(
model_y = GradientBoostingRegressor(),
model_t = MultiOutputRegressor(GradientBoostingRegressor()),
discrete_treatment=False,
featurizer=PolynomialFeatures(degree=2, include_bias=False),
treatment_featurizer=PolynomialFeatures(degree=2, interaction_only=False, include_bias=False),
n_estimators=1000,
min_impurity_decrease=0.001,
verbose=0,
cv=6,
random_state=42
)
est_causal_forest.fit(Y, T=T_df, X=X_df, W=W)
グラフは下図のようになりました. あまり変わり映えしませんね. 名前からしてtree系のモデルなのでfeaturizerやtreatment_featrurizerは不要なのかもしれません.
まとめ
- EconMLのDMLを使用することで、観察データ,連続値による介入, 連続値のOutcomeにおけるCATEを推定することができる.
- EconMLのDMLを使用することで介入$T$に関する複数のDMLを比較することができる.