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]