How to use the causallib.estimation.base_weight.WeightEstimator function in causallib

for treatment_value in treatment_values:
            subgroup_mask = stratify_by == treatment_value
            aggregated_value = np.average(y[subgroup_mask], weights=sample_weight[subgroup_mask])
            res[treatment_value] = aggregated_value
        res = pd.Series(res)
        return res

    def evaluate_balancing(self, X, a, y, w):
        pass  # TODO: implement: (1) table one with smd (2) gather lots of metric (ks, kl, smd) (3) plot CDF of each feature

    def __repr__(self):
        repr_string = create_repr_string(self)
        return repr_string


class PropensityEstimator(WeightEstimator):
    """
    Interface for causal estimators balancing datasets through propensity (i.e. treatment probability) estimation
    (e.g. inverse probability weighting).
    """

    def __init__(self, learner, use_stabilized=False, *args, **kwargs):
        """

        Args:
            learner: Initialized sklearn model.
            use_stabilized (bool): Whether to re-weigh the learned weights with the prevalence of the treatment.
                                   See Also: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4351790/#S6title
        """
        super(PropensityEstimator, self).__init__(learner, use_stabilized=use_stabilized)
        if not hasattr(self.learner, "predict_proba"):
            raise AttributeError("Propensity Estimator must use a machine learning that can predict probabilities"

def __init__(self, estimator):
        """
        Args:
            estimator (WeightEstimator):
        """
        if not isinstance(estimator, WeightEstimator):
            raise TypeError("WeightEvaluator should be initialized with WeightEstimator, got ({}) instead."
                            .format(type(estimator)))
        super(WeightEvaluator, self).__init__(estimator)
        self._plot_functions.update({"weight_distribution": plot_propensity_score_distribution_folds,
                                     "covariate_balance_love": plot_mean_features_imbalance_love_folds,
                                     "covariate_balance_slope": plot_mean_features_imbalance_slope_folds})