def ATE(y, treatment):
    """Average Treatment Effect
    Note! treatment assumed to have values {0, 1}"""
    return (y * treatment).sum() / (treatment.sum()) - (
        y * (1 - treatment)
    ).sum() / ((1 - treatment).sum())