S-Learner Effect

Problem

Return the S-learner CATE at x for a linear model with treatment and treatment-by-x interaction: treatment_coefficient + interaction_coefficient*x.

Hint

The treatment effect varies with x only through the interaction.

case 1

in: [2,3,0.5]

expected: 4

case 2

in: [0,5,2]

expected: 5

case 3

in: [-2,1,-0.5]

expected: 2

case 4

in: [10,0,0]

expected: 0

Starter codePython

def s_learner_effect(x, treatment_coefficient, interaction_coefficient):
    # Your code here
    pass

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