Doubly Robust ATE

Problem

Return the augmented IPW (doubly robust) ATE: the mean of (mu1 - mu0) plus the weighted treated residual (y - mu1)/e minus the weighted control residual (y - mu0)/(1-e). The estimate is consistent if either the outcome model or the propensity model is correct. Return 0.0 for empty or mismatched input, a treated unit with e <= 0, or a control with e >= 1.

Hint

Model predictions are corrected by the weighted residuals.

Test cases

case 1

in: [[1,0,1,0],[10,5,8,4],[0.5,0.5,0.4,0.6],[9,6,7,5],[8,4,6,3]]

expected: 1.5

case 2

in: [[1],[7],[0.7],[6],[5]]

expected: 2.428571428571429

case 3

in: [[0],[2],[0.3],[3],[1]]

expected: 0.5714285714285714

case 4

in: [[],[],[],[],[]]

expected: 0

case 5

in: [[1],[5],[0],[1],[1]]

expected: 0

Starter codePython

def doubly_robust_ate(treatment, outcome, propensity, mu1, mu0):
    # Your code here
    pass

5 test cases · run your code against them in the editor