prototypes/dml_iv/coverage_experiment.py [54:77]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
class SeparateModel:
    def __init__(self, model0, model1):
        self.model0 = model0
        self.model1 = model1
        return
    def fit(self, XZ, T):
        Z0 = (XZ[:, -1] == 0)
        Z1 = (XZ[:, -1] == 1)
        self.model0.fit(XZ[Z0, :-1], T[Z0])
        self.model1.fit(XZ[Z1, :-1], T[Z1])
        return self
    def predict(self, XZ):
        Z0 = (XZ[:, -1] == 0)
        Z1 = (XZ[:, -1] == 1)
        t_pred = np.zeros(XZ.shape[0])
        if np.sum(Z0) > 0:
                t_pred[Z0] = self.model0.predict(XZ[Z0, :-1])
        if np.sum(Z1) > 0:
                t_pred[Z1] = self.model1.predict(XZ[Z1, :-1])
                
        return t_pred
    @property
    def coef_(self):
        return np.concatenate((self.model0.coef_, self.model1.coef_))
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



prototypes/dml_iv/utilities.py [85:115]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
class SeparateModel:
    """ Splits the data based on the last feature and trains
    a separate model for each subsample. At fit time, it
    uses the last feature to choose which model to use
    to predict.
    """
    def __init__(self, model0, model1):
        self.model0 = model0
        self.model1 = model1
        return

    def fit(self, XZ, T):
        Z0 = (XZ[:, -1] == 0)
        Z1 = (XZ[:, -1] == 1)
        self.model0.fit(XZ[Z0, :-1], T[Z0])
        self.model1.fit(XZ[Z1, :-1], T[Z1])
        return self

    def predict(self, XZ):
        Z0 = (XZ[:, -1] == 0)
        Z1 = (XZ[:, -1] == 1)
        t_pred = np.zeros(XZ.shape[0])
        if np.sum(Z0) > 0:
            t_pred[Z0] = self.model0.predict(XZ[Z0, :-1])
        if np.sum(Z1) > 0:
            t_pred[Z1] = self.model1.predict(XZ[Z1, :-1])
        return t_pred

    @property
    def coef_(self):
        return np.concatenate((self.model0.coef_, self.model1.coef_))
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -