ENH: effect_group (ATET/ATC) for TreatmentEffect.aipw and aipw_wls

Is your feature request related to a problem? Please describe

TreatmentEffect.ra, .ipw and .ipw_ra take effect_group ("all"/treated/control), but .aipw and .aipw_wls don't they only return ATE. The aipw_wls docstring says "Option for effect on treated or on untreated is not available" and test_teffects.py skips the ATET block with # not available for aipw. Stata added atet to teffects aipw in version 18, which is presumably why it was left out originally.

Describe the solution you'd like

Add effect_group to aipw and aipw_wls with the same accepted values as the other estimators ("all", 1/"treated", 0/"untreated"/"control").

ATET needs its own effect estimating equations rather than a restriction of the ATE moments to the treated. For the treated, the counterfactual control mean combines the outcome-model prediction averaged over treated observations with an augmentation term from the control residuals weighted by the propensity odds $p/(1-p)$:

$$ \hat\mu_{0|T=1} = \frac{1}{n_1}\sum_i \Big[ T_i,\hat\mu_0(x_i) + (1-T_i),\frac{\hat p_i}{1-\hat p_i},{y_i-\hat\mu_0(x_i)} \Big], \qquad \hat\mu_{1|T=1} = \frac{1}{n_1}\sum_i T_i,y_i $$

ATC is symmetric. The GMM effect moments keep both components, stacked with the outcome and treatment model scores as now, so the standard errors account for the first-step estimation. This is the standard doubly robust ATT estimating equation (Hahn 1998; Moodie, Saarela & Stephens 2018) and corresponds to what Stata 18 added to teffects aipw, atet.

For aipw_wls the WNLS outcome-model weights s

Describe alternatives you have considered

No response

References

• Hahn, J. (1998). "On the Role of the Propensity Score in Efficient Semiparametric Estimation of Average Treatment Effects." Econometrica 66(2), 315–331. doi:10.2307/2998560
• Moodie, E.E.M., Saarela, O. & Stephens, D.A. (2018). "A doubly robust weighting estimator of the average treatment effect on the treated." Stat 7(1), e205. doi:10.1002/sta4.205
• StataCorp (2023), Stata 18 [CAUSAL] teffects aipw, “Methods and formulas”

Additional context

No response