зеркало из https://github.com/py-why/EconML.git
454 строки
26 KiB
Python
454 строки
26 KiB
Python
import numpy as np
|
|
from econml.dml import LinearDML
|
|
from sklearn.linear_model import LinearRegression, MultiTaskLassoCV, MultiTaskLasso, Lasso
|
|
from econml.inference import StatsModelsInference
|
|
from econml.tests.test_statsmodels import _summarize
|
|
from econml.sklearn_extensions.linear_model import WeightedLasso, WeightedMultiTaskLassoCV, WeightedMultiTaskLasso
|
|
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
|
|
from statsmodels.tools.tools import add_constant
|
|
from econml.utilities import cross_product
|
|
import matplotlib.pyplot as plt
|
|
import os
|
|
import time
|
|
import argparse
|
|
import warnings
|
|
import joblib
|
|
from sklearn.model_selection import GridSearchCV
|
|
from sklearn.multioutput import MultiOutputRegressor
|
|
|
|
|
|
class GridSearchCVList:
|
|
|
|
def __init__(self, estimator_list, param_grid_list, scoring=None,
|
|
n_jobs=None, iid='warn', refit=True, cv='warn', verbose=0, pre_dispatch='2*n_jobs',
|
|
error_score='raise-deprecating', return_train_score=False):
|
|
self._gcv_list = [GridSearchCV(estimator, param_grid, scoring=scoring,
|
|
n_jobs=n_jobs, iid=iid, refit=refit, cv=cv, verbose=verbose,
|
|
pre_dispatch=pre_dispatch, error_score=error_score,
|
|
return_train_score=return_train_score)
|
|
for estimator, param_grid in zip(estimator_list, param_grid_list)]
|
|
return
|
|
|
|
def fit(self, X, y, **fit_params):
|
|
self.best_ind_ = np.argmax([gcv.fit(X, y, **fit_params).best_score_ for gcv in self._gcv_list])
|
|
self.best_estimator_ = self._gcv_list[self.best_ind_].best_estimator_
|
|
self.best_score_ = self._gcv_list[self.best_ind_].best_score_
|
|
self.best_params_ = self._gcv_list[self.best_ind_].best_params_
|
|
return self
|
|
|
|
def predict(self, X):
|
|
return self.best_estimator_.predict(X)
|
|
|
|
|
|
def _coverage_profile(est, X_test, alpha, true_coef, true_effect):
|
|
cov = {}
|
|
d_t = true_coef.shape[1] // (X_test.shape[1] + 1)
|
|
d_y = true_coef.shape[0]
|
|
coef_interval = est.coef__interval(alpha=alpha)
|
|
intercept_interval = est.intercept__interval(alpha=alpha)
|
|
true_coef = true_coef.flatten()
|
|
est_coef = np.concatenate((est.intercept_[..., np.newaxis], est.coef_), axis=-1).flatten()
|
|
est_coef_lb = np.concatenate((intercept_interval[0][..., np.newaxis], coef_interval[0]), axis=-1).flatten()
|
|
est_coef_ub = np.concatenate((intercept_interval[1][..., np.newaxis], coef_interval[1]), axis=-1).flatten()
|
|
cov['coef'] = est_coef
|
|
cov['coef_lower'] = est_coef_lb
|
|
cov['coef_upper'] = est_coef_ub
|
|
cov['true_coef'] = true_coef
|
|
cov['coef_stderr'] = est.model_final.coef_stderr_.flatten()
|
|
cov['coef_sqerror'] = (est_coef - true_coef)**2
|
|
cov['coef_cov'] = ((true_coef >= est_coef_lb) & (true_coef <= est_coef_ub))
|
|
cov['coef_length'] = est_coef_ub - est_coef_lb
|
|
effect_interval = est.effect_interval(X_test, T0=np.zeros(
|
|
(X_test.shape[0], d_t)), T1=np.ones((X_test.shape[0], d_t)), alpha=alpha)
|
|
true_eff = true_effect(X_test, np.ones((X_test.shape[0], d_t))).reshape(effect_interval[0].shape)
|
|
est_effect = est.effect(X_test, T0=np.zeros((X_test.shape[0], d_t)), T1=np.ones((X_test.shape[0], d_t)))
|
|
cov['x_test'] = np.repeat(X_test, d_y, axis=0)
|
|
cov['effect'] = est_effect.flatten()
|
|
cov['effect_lower'] = effect_interval[0].flatten()
|
|
cov['effect_upper'] = effect_interval[1].flatten()
|
|
cov['true_effect'] = true_eff.flatten()
|
|
cov['effect_sqerror'] = ((est_effect - true_eff)**2).flatten()
|
|
cov['effect_stderr'] = est.model_final.prediction_stderr(
|
|
cross_product(add_constant(X_test), np.ones((X_test.shape[0], d_t)))).flatten()
|
|
cov['effect_cov'] = ((true_eff >= effect_interval[0]) & (true_eff <= effect_interval[1])).flatten()
|
|
cov['effect_length'] = (effect_interval[1] - effect_interval[0]).flatten()
|
|
return cov
|
|
|
|
|
|
def _append_coverage(key, coverage, est, X_test, alpha, true_coef, true_effect):
|
|
cov = _coverage_profile(est, X_test, alpha, true_coef, true_effect)
|
|
if key not in coverage:
|
|
coverage[key] = {}
|
|
for cov_key, value in cov.items():
|
|
coverage[key][cov_key] = [value]
|
|
else:
|
|
for cov_key, value in cov.items():
|
|
coverage[key][cov_key].append(value)
|
|
|
|
|
|
def _agg_coverage(coverage, qs=np.array([.005, .025, .1, .9, .975, .995])):
|
|
mean_coverage_est = {}
|
|
std_coverage_est = {}
|
|
q_coverage_est = {}
|
|
for key, cov_dict in coverage.items():
|
|
mean_coverage_est[key] = {}
|
|
std_coverage_est[key] = {}
|
|
q_coverage_est[key] = {}
|
|
for cov_key, cov_list in cov_dict.items():
|
|
mean_coverage_est[key][cov_key] = np.mean(cov_list, axis=0)
|
|
std_coverage_est[key][cov_key] = np.std(cov_list, axis=0)
|
|
q_coverage_est[key][cov_key] = np.percentile(cov_list, qs * 100, axis=0)
|
|
return mean_coverage_est, std_coverage_est, q_coverage_est
|
|
|
|
|
|
def plot_coverage(coverage, cov_key, n, n_exp, hetero_coef_list, d_list, d_x_list, p_list, t_list,
|
|
cov_type_list, alpha_list, prefix="", folder="", print_matrix=False):
|
|
if not os.path.exists('figures'):
|
|
os.makedirs('figures')
|
|
if not os.path.exists(os.path.join("figures", folder)):
|
|
os.makedirs(os.path.join("figures", folder))
|
|
joblib.dump(coverage, os.path.join("figures", folder, "{}data.jbl".format(prefix)))
|
|
for hetero_coef in hetero_coef_list:
|
|
for d in d_list:
|
|
for d_x in d_x_list:
|
|
if d_x > d:
|
|
continue
|
|
for p in p_list:
|
|
for d_t in t_list:
|
|
for cov_type in cov_type_list:
|
|
for alpha in alpha_list:
|
|
key = "n_{}_n_exp_{}_hetero_{}_d_{}_d_x_{}_p_{}_d_t_{}_cov_type_{}_alpha_{}".format(
|
|
n, n_exp, hetero_coef, d, d_x, p, d_t, cov_type, alpha)
|
|
if print_matrix:
|
|
print(coverage[key][cov_key])
|
|
plt.figure()
|
|
plt.title("{}{}_{}".format(prefix, key, cov_key))
|
|
plt.hist(coverage[key][cov_key].flatten())
|
|
plt.savefig(os.path.join("figures", folder, "{}{}_{}.png".format(prefix,
|
|
key,
|
|
cov_key)))
|
|
plt.close()
|
|
|
|
|
|
def print_aggregate(mean_coverage, std_coverage, q_coverage, file_gen=lambda: None):
|
|
for key, cov in mean_coverage.items():
|
|
print(key, file=file_gen())
|
|
with np.printoptions(formatter={'float': '{:.2f}'.format}, suppress=True):
|
|
print("Mean Coef (True, RMSE) \t (Coverage, Mean Length) \t Mean StdErr (True, RMSE)\t "
|
|
"[Mean Lower (std), Mean Upper (std)]\t (True Quantiles)\r\n{}".format(
|
|
"\r\n".join(["{:.2f} ({:.2f}, {:.2f}) \t ({:.2f}, {:.2f}) \t {:.2f} ({:.2f}, {:.2f}) "
|
|
" \t [{:.2f} ({:.2f}), {:.2f} ({:.2f})] \t {}".format(est,
|
|
true,
|
|
np.sqrt(
|
|
sqerr),
|
|
coverage,
|
|
length,
|
|
stderr,
|
|
true_stderr,
|
|
np.sqrt(
|
|
std_stderr**2 +
|
|
(stderr -
|
|
true_stderr)**2),
|
|
lower,
|
|
std_lower,
|
|
upper,
|
|
std_upper,
|
|
true_qs)
|
|
for (est, true, sqerr, coverage, length, stderr, std_stderr,
|
|
true_stderr, lower, std_lower, upper, std_upper, true_qs)
|
|
in zip(cov['coef'],
|
|
cov['true_coef'],
|
|
cov['coef_sqerror'],
|
|
cov['coef_cov'],
|
|
cov['coef_length'],
|
|
cov['coef_stderr'],
|
|
std_coverage[key]['coef_stderr'],
|
|
std_coverage[key]['coef'],
|
|
cov['coef_lower'],
|
|
std_coverage[key]['coef_lower'],
|
|
cov['coef_upper'],
|
|
std_coverage[key]['coef_upper'],
|
|
q_coverage[key]['coef'].T)])), file=file_gen())
|
|
print("Effect SqError: {}".format(np.mean(cov['effect_sqerror'])), file=file_gen())
|
|
|
|
print("Point\t Mean Coef (True, RMSE) \t (Coverage, Mean Length)\t Mean StdErr (True, RMSE)\t "
|
|
" [Mean Lower (std), Mean Upper (std)]\t (True Quantiles)\r\n{}".format(
|
|
"\r\n".join(["{}\t {:.2f} ({:.2f}, {:.2f}) \t ({:.2f}, {:.2f}) \t {:.2f} ({:.2f}, {:.2f}) \t "
|
|
"[{:.2f} ({:.2f}), {:.2f} ({:.2f})] \t "
|
|
"{}".format(','.join(x.astype(int).astype(str)),
|
|
est,
|
|
true,
|
|
np.sqrt(
|
|
sqerr),
|
|
coverage,
|
|
length,
|
|
stderr,
|
|
true_stderr,
|
|
np.sqrt(
|
|
std_stderr**2 + (stderr - true_stderr)**2),
|
|
lower,
|
|
std_lower,
|
|
upper,
|
|
std_upper,
|
|
true_qs)
|
|
for (x, est, true, sqerr, coverage, length, stderr, std_stderr,
|
|
true_stderr, lower, std_lower, upper, std_upper, true_qs)
|
|
in zip(cov['x_test'],
|
|
cov['effect'],
|
|
cov['true_effect'],
|
|
cov['effect_sqerror'],
|
|
cov['effect_cov'],
|
|
cov['effect_length'],
|
|
cov['effect_stderr'],
|
|
std_coverage[key]['effect_stderr'],
|
|
std_coverage[key]['effect'],
|
|
cov['effect_lower'],
|
|
std_coverage[key]['effect_lower'],
|
|
cov['effect_upper'],
|
|
std_coverage[key]['effect_upper'],
|
|
q_coverage[key]['effect'].T)])), file=file_gen())
|
|
print("Effect SqError: {}".format(np.mean(cov['effect_sqerror'])), file=file_gen())
|
|
|
|
|
|
def run_all_mc(first_stage, folder, n_list, n_exp, hetero_coef_list, d_list,
|
|
d_x_list, p_list, t_list, cov_type_list, alpha_list):
|
|
|
|
if not os.path.exists("results"):
|
|
os.makedirs('results')
|
|
results_filename = os.path.join("results", "{}.txt".format(folder))
|
|
|
|
np.random.seed(123)
|
|
coverage_est = {}
|
|
coverage_lr = {}
|
|
n_tests = 0
|
|
n_failed_coef = 0
|
|
n_failed_effect = 0
|
|
cov_tol = .04
|
|
for n in n_list:
|
|
for hetero_coef in hetero_coef_list:
|
|
for d in d_list:
|
|
for d_x in d_x_list:
|
|
if d_x > d:
|
|
continue
|
|
for p in p_list:
|
|
for d_t in t_list:
|
|
X_test = np.unique(np.random.binomial(1, .5, size=(20, d_x)), axis=0)
|
|
t0 = time.time()
|
|
for it in range(n_exp):
|
|
X = np.random.binomial(1, .8, size=(n, d))
|
|
T = np.hstack([np.random.binomial(1, .5 * X[:, 0] + .25,
|
|
size=(n,)).reshape(-1, 1) for _ in range(d_t)])
|
|
true_coef = np.hstack([np.hstack([it + np.arange(p).reshape(-1, 1),
|
|
it + np.ones((p, 1)), np.zeros((p, d_x - 1))])
|
|
for it in range(d_t)])
|
|
|
|
def true_effect(x, t):
|
|
return cross_product(
|
|
np.hstack([np.ones((x.shape[0], 1)), x[:, :d_x]]), t) @ true_coef.T
|
|
y = true_effect(X, T) + X[:, [0] * p] +\
|
|
(hetero_coef * X[:, [0]] + 1) * np.random.normal(0, 1, size=(n, p))
|
|
|
|
XT = np.hstack([X, T])
|
|
X1, X2, y1, y2, X_final_first, X_final_sec, y_sum_first, y_sum_sec, \
|
|
n_sum_first, n_sum_sec, var_first, var_sec = _summarize(XT, y)
|
|
X = np.vstack([X1, X2])
|
|
y = np.concatenate((y1, y2))
|
|
X_final = np.vstack([X_final_first, X_final_sec])
|
|
y_sum = np.concatenate((y_sum_first, y_sum_sec))
|
|
n_sum = np.concatenate((n_sum_first, n_sum_sec))
|
|
var_sum = np.concatenate((var_first, var_sec))
|
|
first_half_sum = len(y_sum_first)
|
|
first_half = len(y1)
|
|
for cov_type in cov_type_list:
|
|
class SplitterSum:
|
|
def __init__(self):
|
|
return
|
|
|
|
def split(self, X, T):
|
|
return [(np.arange(0, first_half_sum),
|
|
np.arange(first_half_sum, X.shape[0])),
|
|
(np.arange(first_half_sum, X.shape[0]),
|
|
np.arange(0, first_half_sum))]
|
|
|
|
est = LinearDML(model_y=first_stage(),
|
|
model_t=first_stage(),
|
|
cv=SplitterSum(),
|
|
discrete_treatment=False)
|
|
est.fit(y_sum,
|
|
X_final[:, -d_t:],
|
|
X_final[:, :d_x],
|
|
X_final[:, d_x:-d_t],
|
|
freq_weight=n_sum,
|
|
sample_var=var_sum,
|
|
inference=StatsModelsInference(cov_type=cov_type))
|
|
|
|
class Splitter:
|
|
def __init__(self):
|
|
return
|
|
|
|
def split(self, X, T):
|
|
return [(np.arange(0, first_half), np.arange(first_half, X.shape[0])),
|
|
(np.arange(first_half, X.shape[0]), np.arange(0, first_half))]
|
|
|
|
lr = LinearDML(model_y=first_stage(),
|
|
model_t=first_stage(),
|
|
cv=Splitter(),
|
|
discrete_treatment=False)
|
|
lr.fit(y, X[:, -d_t:], X=X[:, :d_x], W=X[:, d_x:-d_t],
|
|
inference=StatsModelsInference(cov_type=cov_type))
|
|
for alpha in alpha_list:
|
|
key = ("n_{}_n_exp_{}_hetero_{}_d_{}_d_x_"
|
|
"{}_p_{}_d_t_{}_cov_type_{}_alpha_{}").format(
|
|
n, n_exp, hetero_coef, d, d_x, p, d_t, cov_type, alpha)
|
|
_append_coverage(key, coverage_est, est, X_test,
|
|
alpha, true_coef, true_effect)
|
|
_append_coverage(key, coverage_lr, lr, X_test,
|
|
alpha, true_coef, true_effect)
|
|
if it == n_exp - 1:
|
|
n_tests += 1
|
|
mean_coef_cov = np.mean(coverage_est[key]['coef_cov'])
|
|
mean_eff_cov = np.mean(coverage_est[key]['effect_cov'])
|
|
mean_coef_cov_lr = np.mean(coverage_lr[key]['coef_cov'])
|
|
mean_eff_cov_lr = np.mean(coverage_lr[key]['effect_cov'])
|
|
[print("{}. Time: {:.2f}, Mean Coef Cov: ({:.4f}, {:.4f}), "
|
|
"Mean Effect Cov: ({:.4f}, {:.4f})".format(key,
|
|
time.time() - t0,
|
|
mean_coef_cov,
|
|
mean_coef_cov_lr,
|
|
mean_eff_cov,
|
|
mean_eff_cov_lr),
|
|
file=f)
|
|
for f in [None, open(results_filename, "a")]]
|
|
coef_cov_dev = mean_coef_cov - (1 - alpha)
|
|
if np.abs(coef_cov_dev) >= cov_tol:
|
|
n_failed_coef += 1
|
|
[print("BAD coef coverage on "
|
|
"average: deviation = {:.4f}".format(coef_cov_dev), file=f)
|
|
for f in [None, open(results_filename, "a")]]
|
|
eff_cov_dev = mean_eff_cov - (1 - alpha)
|
|
if np.abs(eff_cov_dev) >= cov_tol:
|
|
n_failed_effect += 1
|
|
[print("BAD effect coverage on "
|
|
"average: deviation = {:.4f}".format(eff_cov_dev), file=f)
|
|
for f in [None, open(results_filename, "a")]]
|
|
|
|
[print("Finished {} Monte Carlo Tests. Failed Coef Coverage Tests: {}/{}."
|
|
"Failed Effect Coverage Tests: {}/{}. (Coverage Tolerance={})".format(n_tests,
|
|
n_failed_coef,
|
|
n_tests,
|
|
n_failed_effect,
|
|
n_tests,
|
|
cov_tol),
|
|
file=f) for f in [None, open(results_filename, "a")]]
|
|
|
|
agg_coverage_est, std_coverage_est, q_coverage_est = _agg_coverage(coverage_est)
|
|
agg_coverage_lr, std_coverage_lr, q_coverage_lr = _agg_coverage(coverage_lr)
|
|
|
|
[print("\nResults for: {}\n--------------------------\n".format(folder), file=f)
|
|
for f in [None, open(results_filename, "a")]]
|
|
|
|
plot_coverage(agg_coverage_est, 'coef_cov', n, n_exp, hetero_coef_list, d_list, d_x_list,
|
|
p_list, t_list, cov_type_list, alpha_list, prefix="sum_", folder=folder)
|
|
plot_coverage(agg_coverage_lr, 'coef_cov', n, n_exp, hetero_coef_list, d_list, d_x_list,
|
|
p_list, t_list, cov_type_list, alpha_list, prefix="orig_", folder=folder)
|
|
plot_coverage(agg_coverage_est, 'effect_cov', n, n_exp, hetero_coef_list, d_list, d_x_list,
|
|
p_list, t_list, cov_type_list, alpha_list, prefix="sum_", folder=folder)
|
|
plot_coverage(agg_coverage_lr, 'effect_cov', n, n_exp, hetero_coef_list, d_list, d_x_list,
|
|
p_list, t_list, cov_type_list, alpha_list, prefix="orig_", folder=folder)
|
|
|
|
[print("Summarized Data\n----------------", file=f) for f in [None, open(results_filename, "a")]]
|
|
print_aggregate(agg_coverage_est, std_coverage_est, q_coverage_est)
|
|
print_aggregate(agg_coverage_est, std_coverage_est, q_coverage_est, lambda: open(results_filename, "a"))
|
|
[print("\nUn-Summarized Data\n-----------------", file=f) for f in [None, open(results_filename, "a")]]
|
|
print_aggregate(agg_coverage_lr, std_coverage_lr, q_coverage_lr)
|
|
print_aggregate(agg_coverage_lr, std_coverage_lr, q_coverage_lr, lambda: open(results_filename, "a"))
|
|
|
|
|
|
def monte_carlo(first_stage=lambda: LinearRegression(), folder='lr'):
|
|
n_exp = 1000
|
|
n_list = [500]
|
|
hetero_coef_list = [0, 1]
|
|
d_list = [1, 10]
|
|
d_x_list = [1, 5]
|
|
p_list = [1, 5]
|
|
t_list = [1, 2]
|
|
cov_type_list = ['HC1']
|
|
alpha_list = [.01, .05, .2]
|
|
run_all_mc(first_stage, folder, n_list, n_exp, hetero_coef_list,
|
|
d_list, d_x_list, p_list, t_list, cov_type_list, alpha_list)
|
|
|
|
|
|
def monte_carlo_lasso(first_stage=lambda: WeightedLasso(alpha=0.01,
|
|
fit_intercept=True,
|
|
tol=1e-6, random_state=123), folder='lasso'):
|
|
n_exp = 1000
|
|
n_list = [500]
|
|
hetero_coef_list = [1]
|
|
d_list = [20]
|
|
d_x_list = [5]
|
|
p_list = [1, 2]
|
|
t_list = [1, 3]
|
|
cov_type_list = ['HC1']
|
|
alpha_list = [.01, .05, .2]
|
|
run_all_mc(first_stage, folder, n_list, n_exp, hetero_coef_list,
|
|
d_list, d_x_list, p_list, t_list, cov_type_list, alpha_list)
|
|
|
|
|
|
def monte_carlo_rf(first_stage=lambda: RandomForestRegressor(n_estimators=100,
|
|
max_depth=3, min_samples_leaf=10), folder='rf'):
|
|
n_exp = 1000
|
|
n_list = [500, 5000]
|
|
hetero_coef_list = [1]
|
|
d_list = [20]
|
|
d_x_list = [5]
|
|
p_list = [1]
|
|
t_list = [2]
|
|
cov_type_list = ['HC1']
|
|
alpha_list = [.01, .05, .2]
|
|
run_all_mc(first_stage, folder, n_list, n_exp, hetero_coef_list,
|
|
d_list, d_x_list, p_list, t_list, cov_type_list, alpha_list)
|
|
|
|
|
|
def monte_carlo_gcv(folder='gcv'):
|
|
def first_stage():
|
|
return GridSearchCVList([LinearRegression(),
|
|
WeightedMultiTaskLasso(alpha=0.05, fit_intercept=True,
|
|
tol=1e-6, random_state=123),
|
|
RandomForestRegressor(n_estimators=100, max_depth=3,
|
|
min_samples_leaf=10, random_state=123),
|
|
MultiOutputRegressor(GradientBoostingRegressor(n_estimators=20,
|
|
max_depth=3,
|
|
min_samples_leaf=10,
|
|
random_state=123))],
|
|
param_grid_list=[{},
|
|
{},
|
|
{},
|
|
{}],
|
|
cv=3,
|
|
iid=True)
|
|
n_exp = 1000
|
|
n_list = [1000, 5000]
|
|
hetero_coef_list = [1]
|
|
d_list = [20]
|
|
d_x_list = [5]
|
|
p_list = [1]
|
|
t_list = [2]
|
|
cov_type_list = ['HC1']
|
|
alpha_list = [.01, .05, .2]
|
|
run_all_mc(first_stage, folder, n_list, n_exp, hetero_coef_list,
|
|
d_list, d_x_list, p_list, t_list, cov_type_list, alpha_list)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description='Description of your program')
|
|
parser.add_argument('-e', '--exp', help='What experiment (default=all)', required=False, default='all')
|
|
args = vars(parser.parse_args())
|
|
if args['exp'] in ['lr', 'all']:
|
|
monte_carlo()
|
|
if args['exp'] in ['lasso', 'all']:
|
|
monte_carlo_lasso()
|
|
if args['exp'] in ['rf', 'all']:
|
|
monte_carlo_rf()
|
|
if args['exp'] in ['gcv', 'all']:
|
|
monte_carlo_gcv()
|