"""
ACIC2016 dataset
"""
# stdlib
import random
from pathlib import Path
from typing import Any, Tuple
import glob
# third party
import numpy as np
import pandas as pd
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.model_selection import train_test_split
import catenets.logger as log
from .network import download_if_needed
np.random.seed(0)
random.seed(0)
FILE_ID = "0B7pG5PPgj6A3N09ibmFwNWE1djA"
PREPROCESSED_FILE_ID = "1iOfEAk402o3jYBs2Prfiz6oaailwWcR5"
NUMERIC_COLS = [
0,
3,
4,
16,
17,
18,
20,
21,
22,
24,
24,
25,
30,
31,
32,
33,
39,
40,
41,
53,
54,
]
N_NUM_COLS = len(NUMERIC_COLS)
def get_acic_covariates(
fn_csv: Path, keep_categorical: bool = False, preprocessed: bool = True
) -> np.ndarray:
X = pd.read_csv(fn_csv)
if not keep_categorical:
X = X.drop(columns=["x_2", "x_21", "x_24"])
else:
# encode categorical features
feature_list = []
for cols_ in X.columns:
if type(X.loc[X.index[0], cols_]) not in [np.int64, np.float64]:
enc = OneHotEncoder(drop="first")
enc.fit(np.array(X[[cols_]]).reshape((-1, 1)))
for k in range(len(list(enc.get_feature_names()))):
X[cols_ + list(enc.get_feature_names())[k]] = enc.transform(
np.array(X[[cols_]]).reshape((-1, 1))
).toarray()[:, k]
feature_list.append(cols_)
X.drop(feature_list, axis=1, inplace=True)
if preprocessed:
X_t = X.values
else:
scaler = StandardScaler()
X_t = scaler.fit_transform(X)
return X_t
def preprocess_simu(
fn_csv: Path,
n_0: int = 2000,
n_1: int = 200,
n_test: int = 500,
error_sd: float = 1,
sp_lin: float = 0.6,
sp_nonlin: float = 0.3,
prop_gamma: float = 0,
prop_omega: float = 0,
ate_goal: float = 0,
inter: bool = True,
i_exp: int = 0,
keep_categorical: bool = False,
preprocessed: bool = True,
) -> Tuple:
X = get_acic_covariates(
fn_csv, keep_categorical=keep_categorical, preprocessed=preprocessed
)
np.random.seed(i_exp)
# shuffle indices
n_total, n_cov = X.shape
ind = np.arange(n_total)
np.random.shuffle(ind)
ind_test = ind[-n_test:]
ind_1 = ind[n_0: (n_0 + n_1)]
# create treatment indicator (treatment assignment does not matter in test set)
w = np.zeros(n_total).reshape((-1, 1))
w[ind_1] = 1
# create dgp
coeffs_ = [0, 1]
# sample baseline coefficients
beta_0 = np.random.choice(coeffs_, size=n_cov, replace=True, p=[1 - sp_lin, sp_lin])
intercept = np.random.choice([x for x in np.arange(-1, 1.25, 0.25)])
# sample treatment effect coefficients
gamma = np.random.choice(
coeffs_, size=n_cov, replace=True, p=[1 - prop_gamma, prop_gamma]
)
omega = np.random.choice(
[0, 1], replace=True, size=n_cov, p=[prop_omega, 1 - prop_omega]
)
# simulate mu_0 and mu_1
mu_0 = (intercept + np.dot(X, beta_0)).reshape((-1, 1))
mu_1 = (intercept + np.dot(X, gamma + beta_0 * omega)).reshape((-1, 1))
if sp_nonlin > 0:
coefs_sq = [0, 0.1]
beta_sq = np.random.choice(
coefs_sq, size=N_NUM_COLS, replace=True, p=[1 - sp_nonlin, sp_nonlin]
)
omega = np.random.choice(
[0, 1], replace=True, size=N_NUM_COLS, p=[prop_omega, 1 - prop_omega]
)
X_sq = X[:, NUMERIC_COLS] ** 2
mu_0 = mu_0 + np.dot(X_sq, beta_sq).reshape((-1, 1))
mu_1 = mu_1 + np.dot(X_sq, beta_sq * omega).reshape((-1, 1))
if inter:
# randomly add some interactions
ind_c = np.arange(n_cov)
np.random.shuffle(ind_c)
inter_list = list()
for i in range(0, n_cov - 2, 2):
inter_list.append(X[:, ind_c[i]] * X[:, ind_c[i + 1]])
X_inter = np.array(inter_list).T
n_inter = X_inter.shape[1]
beta_inter = np.random.choice(
coefs_sq, size=n_inter, replace=True, p=[1 - sp_nonlin, sp_nonlin]
)
omega = np.random.choice(
[0, 1], replace=True, size=n_inter, p=[prop_omega, 1 - prop_omega]
)
mu_0 = mu_0 + np.dot(X_inter, beta_inter).reshape((-1, 1))
mu_1 = mu_1 + np.dot(X_inter, beta_inter * omega).reshape((-1, 1))
ate = np.mean(mu_1 - mu_0)
mu_1 = mu_1 - ate + ate_goal
y = (
w * mu_1
+ (1 - w) * mu_0
+ np.random.normal(0, error_sd, n_total).reshape((-1, 1))
)
X_train, y_train, w_train, mu_0_train, mu_1_train = (
X[ind[: (n_0 + n_1)], :],
y[ind[: (n_0 + n_1)]],
w[ind[: (n_0 + n_1)]],
mu_0[ind[: (n_0 + n_1)]],
mu_1[ind[: (n_0 + n_1)]],
)
X_test, y_test, w_test, mu_0_t, mu_1_t = (
X[ind_test, :],
y[ind_test],
w[ind_test],
mu_0[ind_test],
mu_1[ind_test],
)
return (
X_train,
w_train,
y_train,
np.asarray([mu_0_train, mu_1_train]).squeeze().T,
X_test,
w_test,
y_test,
np.asarray([mu_0_t, mu_1_t]).squeeze().T,
)
def get_acic_orig_filenames(data_path: Path, simu_num: int) -> list:
return sorted(glob.glob((data_path / ("data_cf_all/" + str(simu_num) +
'/zymu_*.csv')).__str__()))
def get_acic_orig_outcomes(data_path: Path,
simu_num: int,
i_exp: int) -> Tuple:
file_list = get_acic_orig_filenames(data_path=data_path,
simu_num=simu_num)
out = pd.read_csv(file_list[i_exp])
w = out['z']
y = w * out['y1'] + (1 - w) * out['y0']
mu_0, mu_1 = out['mu0'], out['mu1']
return y.values, w.values, mu_0.values, mu_1.values
def preprocess_acic_orig(fn_csv: Path,
data_path: Path,
preprocessed: bool = False,
keep_categorical: bool = True,
simu_num: int = 1,
i_exp: int = 0,
train_size: int = 4000,
random_split: bool = False
)-> Tuple:
X = get_acic_covariates(
fn_csv, keep_categorical=keep_categorical, preprocessed=preprocessed
)
y, w, mu_0, mu_1 = get_acic_orig_outcomes(data_path=data_path, simu_num=simu_num, i_exp=i_exp)
if not random_split:
X_train, y_train, w_train, mu_0_train, mu_1_train = X[:train_size, :], y[:train_size], \
w[:train_size], mu_0[:train_size], \
mu_1[:train_size]
X_test, y_test, w_test, mu_0_test, mu_1_test = X[train_size:, :], y[train_size:], \
w[train_size:], mu_0[train_size:], \
mu_1[train_size:]
else:
X_train, X_test, y_train, y_test, w_train, w_test, \
mu_0_train, mu_0_test, mu_1_train, mu_1_test = train_test_split(X, y, w, mu_0, mu_1,
test_size=1 - train_size,
random_state=i_exp)
return (
X_train,
w_train,
y_train,
np.asarray([mu_0_train, mu_1_train]).squeeze().T,
X_test,
w_test,
y_test,
np.asarray([mu_0_test, mu_1_test]).squeeze().T,
)
def preprocess(fn_csv: Path,
data_path: Path,
preprocessed: bool = True,
original_acic_outcomes: bool = False,
**kwargs: Any,
) -> Tuple:
if not original_acic_outcomes:
return preprocess_simu(fn_csv=fn_csv, preprocessed=preprocessed, **kwargs)
else:
return preprocess_acic_orig(fn_csv=fn_csv, preprocessed=preprocessed,
data_path=data_path, **kwargs)
def load(
data_path: Path,
preprocessed: bool = True,
original_acic_outcomes: bool = False,
**kwargs: Any,
) -> Tuple:
"""
ACIC2016 dataset dataloader.
- Download the dataset if needed.
- Load the dataset.
- Preprocess the data.
- Return train/test split.
Parameters
----------
data_path: Path
Path to the CSV. If it is missing, it will be downloaded.
preprocessed: bool
Switch between the raw and preprocessed versions of the dataset.
original_acic_outcomes: bool
Switch between new simulations (Inductive bias paper) and original acic outcomes
Returns
-------
train_x: array or pd.DataFrame
Features in training data.
train_t: array or pd.DataFrame
Treatments in training data.
train_y: array or pd.DataFrame
Observed outcomes in training data.
train_potential_y: array or pd.DataFrame
Potential outcomes in training data.
test_x: array or pd.DataFrame
Features in testing data.
test_potential_y: array or pd.DataFrame
Potential outcomes in testing data.
"""
if preprocessed:
csv = data_path / "x_trans.csv"
download_if_needed(csv, file_id=PREPROCESSED_FILE_ID)
else:
arch = data_path / "data_cf_all.tar.gz"
download_if_needed(
arch, file_id=FILE_ID, unarchive=True, unarchive_folder=data_path
)
csv = data_path / "data_cf_all/x.csv"
log.debug(f"load dataset {csv}")
return preprocess(csv, data_path=data_path, preprocessed=preprocessed,
original_acic_outcomes=original_acic_outcomes,
**kwargs)
Datasets
Models
