How to use the causallib.evaluation.weight_evaluator.WeightEvaluatorPredictions2 function in causallib

# #################### #


class PropensityEvaluatorPredictions(WeightEvaluatorPredictions):
    """Data structure to hold propensity-model predictions"""

    def __init__(self, weight_by_treatment_assignment, weight_for_being_treated, treatment_assignment_pred,
                 propensity, propensity_by_treatment_assignment):
        super(PropensityEvaluatorPredictions, self).__init__(weight_by_treatment_assignment,
                                                             weight_for_being_treated,
                                                             treatment_assignment_pred)
        self.propensity = propensity
        self.propensity_by_treatment_assignment = propensity_by_treatment_assignment


class PropensityEvaluatorPredictions2(WeightEvaluatorPredictions2):
    """Data structure to hold propensity-model predictions"""

    def __init__(self, weight_matrix, propensity_matrix, treatment_assignment, treatment_assignment_prediction=None):
        super(PropensityEvaluatorPredictions2, self).__init__(weight_matrix, treatment_assignment,
                                                              treatment_assignment_prediction)
        self.propensity_matrix = propensity_matrix

    @property
    def propensity(self):
        propensity = self._extract_vector_from_matrix(self.propensity_matrix,
                                                      self._treatment_assignment)
        return propensity

    @property
    def propensity_by_treatment_assignment(self):
        # TODO: remove propensity_by_treatment if expected-ROC is not to be used.

def _estimator_predict2(self, X, a):
        """Predict on data"""
        weight_matrix = self.estimator.compute_weight_matrix(X, a, use_stabilized=False)
        treatment_assignment_pred = self.estimator.learner.predict(
            X)  # TODO: maybe add predict_label to interface instead
        treatment_assignment_pred = pd.Series(treatment_assignment_pred, index=X.index)

        fold_prediction = WeightEvaluatorPredictions2(weight_matrix, a, treatment_assignment_pred)
        return fold_prediction