def test_inverse_onehot(self):
T = np.random.randint(4, size=100)
T_oh = OneHotEncoder(categories='auto',
sparse=False).fit_transform(T.reshape(-1, 1))[:,
1:]
T_inv = inverse_onehot(T_oh)
np.testing.assert_array_equal(T, T_inv)
def fit(self, Y, T, X=None, W=None, *, sample_weight=None):
# TODO Allow for non-vector y, i.e. of shape (n, 1)
assert np.ndim(Y) == 1, "Can only accept single dimensional outcomes Y! Use Y.ravel()."
if (X is None) and (W is None):
raise AttributeError("At least one of X or W has to not be None!")
if np.any(np.all(T == 0, axis=0)) or (not np.any(np.all(T == 0, axis=1))):
raise AttributeError("Provided crossfit folds contain training splits that " +
"don't contain all treatments")
XW = self._combine(X, W)
filtered_kwargs = _filter_none_kwargs(sample_weight=sample_weight)
self._model_propensity.fit(XW, inverse_onehot(T), **filtered_kwargs)
self._model_regression.fit(np.hstack([XW, T]), Y, **filtered_kwargs)
return self
def model_cate(self, T=1):
"""
Get the fitted final CATE model.
Parameters
----------
T: alphanumeric
The treatment with respect to which we want the fitted CATE model.
Returns
-------
model_cate: object of type(model_final)
An instance of the model_final object that was fitted after calling fit which corresponds
to the CATE model for treatment T=t, compared to baseline. Available when multitask_model_final=False.
"""
if self._multitask_model_final:
raise AttributeError("A single multitask model was fitted for all treatments! Use multitask_model_cate.")
_, T = self._expand_treatments(None, T)
ind = inverse_onehot(T).item() - 1
assert ind >= 0, "No model was fitted for the control"
return super().model_final.models_cate[ind]