def score()

in econml/iv/dr/_dr.py [0:0]

• 51 lines of code
• 25 McCabe index (conditional complexity)

    def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None):
        # T and Z only allow single continuous or binary, keep the shape of (n,) for continuous and (n,1) for binary
        T = T.ravel() if not self._discrete_treatment else T
        Z = Z.ravel() if not self._discrete_instrument else Z

        if hasattr(self._model_y_xw, 'score'):
            y_xw_score = self._model_y_xw.score(X=X, W=W, Target=Y, sample_weight=sample_weight)
        else:
            y_xw_score = None
        if hasattr(self._model_t_xw, 'score'):
            t_xw_score = self._model_t_xw.score(X=X, W=W, Target=T, sample_weight=sample_weight)
        else:
            t_xw_score = None

        if hasattr(self._prel_model_effect, 'score'):
            # we need to undo the one-hot encoding for calling effect,
            # since it expects raw values
            raw_T = inverse_onehot(T) if self._discrete_treatment else T
            raw_Z = inverse_onehot(Z) if self._discrete_instrument else Z
            effect_score = self._prel_model_effect.score(Y, raw_T,
                                                         Z=raw_Z, X=X, W=W, sample_weight=sample_weight)
        else:
            effect_score = None

        if self._projection:
            if hasattr(self._model_t_xwz, 'score'):
                WZ = self._combine(W, Z, Y.shape[0])
                t_xwz_score = self._model_t_xwz.score(X=X, W=WZ, Target=T, sample_weight=sample_weight)
            else:
                t_xwz_score = None

            if hasattr(self._model_tz_xw, 'score'):
                T_proj = self._model_t_xwz.predict(X, WZ).reshape(T.shape)
                # if discrete, return shape (n,1); if continuous return shape (n,)
                target = (T * T_proj).reshape(T.shape[0],)
                tz_xw_score = self._model_tz_xw.score(X=X, W=W, Target=target, sample_weight=sample_weight)
            else:
                tz_xw_score = None
            return y_xw_score, t_xw_score, t_xwz_score, tz_xw_score, effect_score

        else:
            if hasattr(self._model_z_xw, 'score'):
                z_xw_score = self._model_z_xw.score(X=X, W=W, Target=Z, sample_weight=sample_weight)
            else:
                z_xw_score = None

            if hasattr(self._model_tz_xw, 'score'):
                if self._discrete_treatment:
                    if self._discrete_instrument:
                        # target will be discrete and will be inversed from FirstStageWrapper
                        target = T * Z
                    else:
                        target = inverse_onehot(T) * Z
                else:
                    if self._discrete_instrument:
                        target = T * inverse_onehot(Z)
                    else:
                        target = T * Z
                tz_xw_score = self._model_tz_xw.score(X=X, W=W, Target=target, sample_weight=sample_weight)
            else:
                tz_xw_score = None

            return y_xw_score, t_xw_score, z_xw_score, tz_xw_score, effect_score