ATT Propensity Weights

Problem

Compute ATT propensity weights: treated units get weight 1.0 and control units get e / (1 - e) with the propensity e clipped to [1e-6, 1 - 1e-6]. Return the weights in input order; empty input returns [].

Hint

ATT weights reweight controls to look like the treated group instead of the full population.

Test cases

case 1

in: [[1,0,0,1],[0.8,0.4,0.9,0.2]]

expected: [1,0.6666666666666667,9.000000000000002,1]

case 2

in: [[1,1,0,0],[0.5,0.5,0.5,0.5]]

expected: [1,1,1,1]

case 3

in: [[],[]]

expected: []

case 4

in: [[0],[0.5]]

expected: [1]

Starter codePython

def att_weights(treatment, propensity):
    # Your code here
    pass

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