def fit(self, X, y, w, p=None):
"""
Fit method for a CATENet. Takes covariates, outcome variable and treatment indicator as
input
Parameters
----------
X: pd.DataFrame or np.array
Covariate matrix
y: np.array
Outcome vector
w: np.array
Treatment indicator
p: np.array
Vector of (known) treatment propensities. Currently only supported for TwoStepNets.
"""
# some quick input checks
if p is not None:
raise NotImplementedError('Only two-step-nets take p as input. ')
X = check_X_is_np(X)
self._check_inputs(w, p)
train_func = self._get_train_function()
train_params = self.get_params()
self._params, self._predict_funs = train_func(X, y, w, **train_params)
def predict(self, X, return_po: bool = False, return_prop: bool = False):
"""
Predict treatment effect estimates using a CATENet. Depending on method, can also return
potential outcome estimate and propensity score estimate.
Parameters
----------
X: pd.DataFrame or np.array
Covariate matrix
return_po: bool, default False
Whether to return potential outcome estimate
return_prop: bool, default False
Whether to return propensity estimate
Returns
-------
array of CATE estimates, optionally also potential outcomes and propensity
"""
X = check_X_is_np(X)
predict_func = self._get_predict_function()
return predict_func(X,
trained_params=self._params,
predict_funs=self._predict_funs,
return_po=return_po,
return_prop=return_prop)
def fit_and_select_params(self, X, y, w, p=None, param_grid: dict = None):
# some quick input checks
if param_grid is None:
raise ValueError("No param_grid to evaluate. ")
X = check_X_is_np(X)
self._check_inputs(w, p)
param_grid = ParameterGrid(param_grid)
self_param_dict = self.get_params()
train_function = self._get_train_function()
models = []
losses = []
param_settings = []
for param_setting in param_grid:
if self.verbose > 0:
print('Testing parameter setting: ' + ' '.join([
key + ': ' + str(value)
for key, value in param_setting.items()
]))
# replace params
train_param_dict = {
key: (val if not key in param_setting.keys() else
param_setting[key])
for key, val in self_param_dict.items()
}
if p is not None:
params, funs, val_loss = train_function(X,
y,
w,
p=p,
return_val_loss=True,
**train_param_dict)
else:
params, funs, val_loss = train_function(X,
y,
w,
return_val_loss=True,
**train_param_dict)
models.append((params, funs))
losses.append(val_loss)
# save results
param_settings.extend(param_grid)
self._selection_results = {
'param_settings': param_settings,
'val_losses': losses
}
# find lowest loss and set params
best_idx = jnp.array(losses).argmin()
self._params, self._predict_funs = models[best_idx]
self.set_params(**param_settings[best_idx])
def fit(self, X, y, w, p=None):
# overwrite super so we can pass p as extra param
# some quick input checks
X = check_X_is_np(X)
self._check_inputs(w, p)
train_func = self._get_train_function()
train_params = self.get_params()
self._params, self._predict_funs = train_func(X, y, w, p,
**train_params)
def fit(self, X, y, w, p=None):
# overwrite super so we can pass p as extra param
# some quick input checks
X = check_X_is_np(X)
self._check_inputs(w, p)
train_func = self._get_train_function()
train_params = self.get_params()
if self.rescale_transformation:
self._params, self._predict_funs, self._scale_factor = train_func(X, y, w, p,
**train_params)
else:
self._params, self._predict_funs = train_func(X, y, w, p, **train_params)