def bootstrap()

in causalml/inference/iv/drivlearner.py [0:0]

• 26 lines of code
• 5 McCabe index (conditional complexity)

    def bootstrap(self, X, assignment, treatment, y, p, pZ, size=10000, seed=None):
        """Runs a single bootstrap. Fits on bootstrapped sample, then predicts on whole population."""
        idxs = np.random.choice(np.arange(0, X.shape[0]), size=size)
        X_b = X[idxs]

        if isinstance(p[0], (np.ndarray, pd.Series)):
            p0_b = {self.t_groups[0]: convert_pd_to_np(p[0][idxs])}
        else:
            p0_b = {g: prop[idxs] for g, prop in p[0].items()}
        if isinstance(p[1], (np.ndarray, pd.Series)):
            p1_b = {self.t_groups[0]: convert_pd_to_np(p[1][idxs])}
        else:
            p1_b = {g: prop[idxs] for g, prop in p[1].items()}

        pZ_b = pZ[idxs]
        assignment_b = assignment[idxs]
        treatment_b = treatment[idxs]
        y_b = y[idxs]
        self.fit(
            X=X_b,
            assignment=assignment_b,
            treatment=treatment_b,
            y=y_b,
            p=(p0_b, p1_b),
            pZ=pZ_b,
            seed=seed,
        )
        te_b = self.predict(X=X)
        return te_b