はじめに　

本稿の目的は、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

本編　

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を比較することができる.

EconMLを用いたDoubleMachineLearning